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