1  // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
2  
3  use crate::call::server::{RequestContext, UnaryRequestContext};
4  use crate::call::{BatchContext, Call};
5  use crate::cq::CompletionQueue;
6  use crate::server::{self, RequestCallContext};
7  
8  pub struct Request {
9      ctx: RequestContext,
10  }
11  
12  impl Request {
new(rc: RequestCallContext) -> Request13      pub fn new(rc: RequestCallContext) -> Request {
14          let ctx = RequestContext::new(rc);
15          Request { ctx }
16      }
17  
context(&self) -> &RequestContext18      pub fn context(&self) -> &RequestContext {
19          &self.ctx
20      }
21  
resolve(mut self, cq: &CompletionQueue, success: bool)22      pub fn resolve(mut self, cq: &CompletionQueue, success: bool) {
23          let mut rc = self.ctx.take_request_call_context().unwrap();
24          if !success {
25              server::request_call(rc, cq);
26              return;
27          }
28  
29          match self.ctx.handle_stream_req(cq, &mut rc) {
30              Ok(_) => server::request_call(rc, cq),
31              Err(ctx) => ctx.handle_unary_req(rc, cq),
32          }
33      }
34  }
35  
36  pub struct UnaryRequest {
37      ctx: UnaryRequestContext,
38  }
39  
40  impl UnaryRequest {
new(ctx: RequestContext, rc: RequestCallContext) -> UnaryRequest41      pub fn new(ctx: RequestContext, rc: RequestCallContext) -> UnaryRequest {
42          let ctx = UnaryRequestContext::new(ctx, rc);
43          UnaryRequest { ctx }
44      }
45  
batch_ctx(&self) -> &BatchContext46      pub fn batch_ctx(&self) -> &BatchContext {
47          self.ctx.batch_ctx()
48      }
49  
request_ctx(&self) -> &RequestContext50      pub fn request_ctx(&self) -> &RequestContext {
51          self.ctx.request_ctx()
52      }
53  
resolve(mut self, cq: &CompletionQueue, success: bool)54      pub fn resolve(mut self, cq: &CompletionQueue, success: bool) {
55          let mut rc = self.ctx.take_request_call_context().unwrap();
56          if !success {
57              server::request_call(rc, cq);
58              return;
59          }
60  
61          let reader = self.ctx.batch_ctx_mut().recv_message();
62          self.ctx.handle(&mut rc, cq, reader);
63          server::request_call(rc, cq);
64      }
65  }
66  
67  /// A callback to wait for status for the aborted rpc call to be sent.
68  pub struct Abort {
69      ctx: BatchContext,
70      _call: Call,
71  }
72  
73  impl Abort {
new(call: Call) -> Abort74      pub fn new(call: Call) -> Abort {
75          Abort {
76              ctx: BatchContext::new(),
77              _call: call,
78          }
79      }
80  
batch_ctx(&self) -> &BatchContext81      pub fn batch_ctx(&self) -> &BatchContext {
82          &self.ctx
83      }
84  }
85