dcl_rpc/
service_module_definition.rs

1//! It contains all the types related to the auto-codegeneration for a defined RPC service when a `.proto` is compiled.
2//!
3//! A [`ServiceModuleDefinition`] is auto-generated when a `.proto` is compiled. And it's filled with the defined procedures.
4//! Actually, you probably don't need to use this module or their types.
5//!
6use crate::{rpc_protocol::RemoteError, stream_protocol::Generator};
7use core::future::Future;
8use std::{collections::HashMap, pin::Pin, sync::Arc};
9
10/// The context type received by every procedure
11pub struct ProcedureContext<Context> {
12    /// The context given to the [`RpcServer`](`crate::server::RpcServer`)
13    pub server_context: Arc<Context>,
14    /// The ID given to the Transport which is calling the procedure
15    pub transport_id: u32,
16}
17
18/// General type returned by every procedure
19pub type Response<T> = Pin<Box<dyn Future<Output = Result<T, RemoteError>> + Send>>;
20
21/// Payload for procedures which don't receive a stream
22pub type CommonPayload = Vec<u8>;
23
24/// Response type returned by a unary procedure
25pub type UnaryResponse = Response<Vec<u8>>;
26
27/// Handler type for a unary procedure.
28pub type UnaryRequestHandler<Context> =
29    dyn Fn(CommonPayload, ProcedureContext<Context>) -> UnaryResponse + Send + Sync;
30
31/// Response type returned by a server streams procedure
32pub type ServerStreamsResponse = Response<Generator<Vec<u8>>>;
33
34/// Handler type for a server streams procedure
35pub type ServerStreamsRequestHandler<Context> =
36    dyn Fn(CommonPayload, ProcedureContext<Context>) -> ServerStreamsResponse + Send + Sync;
37
38/// Payload type that a client streams procedure receives
39pub type ClientStreamsPayload = Generator<Vec<u8>>;
40
41/// Response type returned by a client streams procedure
42pub type ClientStreamsResponse = Response<Vec<u8>>;
43
44/// Handler type for a client streams procedure
45pub type ClientStreamsRequestHandler<Context> =
46    dyn Fn(ClientStreamsPayload, ProcedureContext<Context>) -> ClientStreamsResponse + Send + Sync;
47
48/// Payload type that a bidirection streams procedure rceives
49pub type BiStreamsPayload = Generator<Vec<u8>>;
50
51/// Response type returned by a bidirectional streams procedure
52pub type BiStreamsResponse = Response<Generator<Vec<u8>>>;
53
54/// Handler type for a bidirectional streams procedure
55pub type BiStreamsRequestHandler<Context> =
56    dyn Fn(BiStreamsPayload, ProcedureContext<Context>) -> BiStreamsResponse + Send + Sync;
57
58/// Type used for storing procedure definitions given by the codegeneration for the RPC service
59pub enum ProcedureDefinition<Context> {
60    /// Stores unary procedure definitions. Unary Procedure means a basic request<>response
61    Unary(Arc<UnaryRequestHandler<Context>>),
62    /// Stores server streams procedure definitions. [`crate::client::RpcClient`] sends a request and waits for the [`crate::server::RpcServer`] to send all the data that it has and close the stream opened
63    ServerStreams(Arc<ServerStreamsRequestHandler<Context>>),
64    /// Stores client strems procedure definitions. [`crate::client::RpcClient`] sends a request and opens a stream in the [`crate::server::RpcServer`], then [`crate::server::RpcServer`] waits for [`crate::client::RpcClient`] to send all the payloads
65    ClientStreams(Arc<ClientStreamsRequestHandler<Context>>),
66    /// Stores bidirectional streams procedure definitions. A stream is opened on both sides (client and server)
67    BiStreams(Arc<BiStreamsRequestHandler<Context>>),
68}
69
70impl<Context> Clone for ProcedureDefinition<Context> {
71    fn clone(&self) -> Self {
72        match self {
73            Self::Unary(procedure) => Self::Unary(procedure.clone()),
74            Self::ServerStreams(procedure) => Self::ServerStreams(procedure.clone()),
75            Self::ClientStreams(procedure) => Self::ClientStreams(procedure.clone()),
76            Self::BiStreams(procedure) => Self::BiStreams(procedure.clone()),
77        }
78    }
79}
80
81/// It stores all procedures defined for a RPC service
82pub struct ServiceModuleDefinition<Context> {
83    /// Map that stores all procedures for the service
84    ///
85    /// the key is the procedure's name and the value is the handler for the procedure
86    ///
87    procedure_definitions: HashMap<String, ProcedureDefinition<Context>>,
88}
89
90impl<Context> ServiceModuleDefinition<Context> {
91    pub fn new() -> Self {
92        Self {
93            procedure_definitions: HashMap::new(),
94        }
95    }
96
97    /// Add an unary procedure handler to the service definition
98    pub fn add_unary<
99        H: Fn(CommonPayload, ProcedureContext<Context>) -> UnaryResponse + Send + Sync + 'static,
100    >(
101        &mut self,
102        name: &str,
103        handler: H,
104    ) {
105        self.add_definition(name, ProcedureDefinition::Unary(Arc::new(handler)));
106    }
107
108    /// Add a server streams procedure handler to the service definition
109    pub fn add_server_streams<
110        H: Fn(CommonPayload, ProcedureContext<Context>) -> ServerStreamsResponse
111            + Send
112            + Sync
113            + 'static,
114    >(
115        &mut self,
116        name: &str,
117        handler: H,
118    ) {
119        self.add_definition(name, ProcedureDefinition::ServerStreams(Arc::new(handler)));
120    }
121
122    /// Add a client streams procedure handler to the service definition
123    pub fn add_client_streams<
124        H: Fn(ClientStreamsPayload, ProcedureContext<Context>) -> ClientStreamsResponse
125            + Send
126            + Sync
127            + 'static,
128    >(
129        &mut self,
130        name: &str,
131        handler: H,
132    ) {
133        self.add_definition(name, ProcedureDefinition::ClientStreams(Arc::new(handler)));
134    }
135
136    /// Add a bidirectional streams procedure handler to the service definition
137    pub fn add_bidir_streams<
138        H: Fn(BiStreamsPayload, ProcedureContext<Context>) -> BiStreamsResponse
139            + Send
140            + Sync
141            + 'static,
142    >(
143        &mut self,
144        name: &str,
145        handler: H,
146    ) {
147        self.add_definition(name, ProcedureDefinition::BiStreams(Arc::new(handler)));
148    }
149
150    fn add_definition(&mut self, name: &str, definition: ProcedureDefinition<Context>) {
151        self.procedure_definitions
152            .insert(name.to_string(), definition);
153    }
154
155    pub fn get_definitions(&self) -> &HashMap<String, ProcedureDefinition<Context>> {
156        &self.procedure_definitions
157    }
158}
159
160impl<Context> Default for ServiceModuleDefinition<Context> {
161    fn default() -> Self {
162        Self::new()
163    }
164}