xref: /tonic/tonic-build/src/server.rs (revision ea7fe66b)
1 use super::{Method, Service};
2 use crate::{generate_doc_comment, generate_doc_comments, naive_snake_case};
3 use proc_macro2::{Span, TokenStream};
4 use quote::quote;
5 use syn::{Ident, Lit, LitStr};
6 
7 /// Generate service for Server.
8 ///
9 /// This takes some `Service` and will generate a `TokenStream` that contains
10 /// a public module containing the server service and handler trait.
11 pub fn generate<T: Service>(service: &T, proto_path: &str) -> TokenStream {
12     let methods = generate_methods(service, proto_path);
13 
14     let server_service = quote::format_ident!("{}Server", service.name());
15     let server_trait = quote::format_ident!("{}", service.name());
16     let server_mod = quote::format_ident!("{}_server", naive_snake_case(&service.name()));
17     let generated_trait = generate_trait(service, proto_path, server_trait.clone());
18     let service_doc = generate_doc_comments(service.comment());
19 
20     // Transport based implementations
21     let path = format!("{}.{}", service.package(), service.identifier());
22     let transport = generate_transport(&server_service, &server_trait, &path);
23 
24     quote! {
25         /// Generated server implementations.
26         pub mod #server_mod {
27             #![allow(unused_variables, dead_code, missing_docs)]
28             use tonic::codegen::*;
29 
30             #generated_trait
31 
32             #service_doc
33             #[derive(Debug)]
34             pub struct #server_service<T: #server_trait> {
35                 inner: _Inner<T>,
36             }
37 
38             struct _Inner<T>(Arc<T>, Option<tonic::Interceptor>);
39 
40             impl<T: #server_trait> #server_service<T> {
41                 pub fn new(inner: T) -> Self {
42                     let inner = Arc::new(inner);
43                     let inner = _Inner(inner, None);
44                     Self { inner }
45                 }
46 
47                 pub fn with_interceptor(inner: T, interceptor: impl Into<tonic::Interceptor>) -> Self {
48                     let inner = Arc::new(inner);
49                     let inner = _Inner(inner, Some(interceptor.into()));
50                     Self { inner }
51                 }
52             }
53 
54             impl<T, B> Service<http::Request<B>> for #server_service<T>
55                 where
56                     T: #server_trait,
57                     B: HttpBody + Send + Sync + 'static,
58                     B::Error: Into<StdError> + Send + 'static,
59             {
60                 type Response = http::Response<tonic::body::BoxBody>;
61                 type Error = Never;
62                 type Future = BoxFuture<Self::Response, Self::Error>;
63 
64                 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
65                     Poll::Ready(Ok(()))
66                 }
67 
68                 fn call(&mut self, req: http::Request<B>) -> Self::Future {
69                     let inner = self.inner.clone();
70 
71                     match req.uri().path() {
72                         #methods
73 
74                         _ => Box::pin(async move {
75                             Ok(http::Response::builder()
76                                .status(200)
77                                .header("grpc-status", "12")
78                                .body(tonic::body::BoxBody::empty())
79                                .unwrap())
80                         }),
81                     }
82                 }
83             }
84 
85             impl<T: #server_trait> Clone for #server_service<T> {
86                 fn clone(&self) -> Self {
87                     let inner = self.inner.clone();
88                     Self { inner }
89                 }
90             }
91 
92             impl<T: #server_trait> Clone for _Inner<T> {
93                 fn clone(&self) -> Self {
94                     Self(self.0.clone(), self.1.clone())
95                 }
96             }
97 
98             impl<T: std::fmt::Debug> std::fmt::Debug for _Inner<T> {
99                 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100                    write!(f, "{:?}", self.0)
101                 }
102             }
103 
104             #transport
105         }
106     }
107 }
108 
109 fn generate_trait<T: Service>(service: &T, proto_path: &str, server_trait: Ident) -> TokenStream {
110     let methods = generate_trait_methods(service, proto_path);
111     let trait_doc = generate_doc_comment(&format!(
112         "Generated trait containing gRPC methods that should be implemented for use with {}Server.",
113         service.name()
114     ));
115 
116     quote! {
117         #trait_doc
118         #[async_trait]
119         pub trait #server_trait : Send + Sync + 'static {
120             #methods
121         }
122     }
123 }
124 
125 fn generate_trait_methods<T: Service>(service: &T, proto_path: &str) -> TokenStream {
126     let mut stream = TokenStream::new();
127 
128     for method in service.methods() {
129         let name = quote::format_ident!("{}", method.name());
130 
131         let (req_message, res_message) = method.request_response_name(proto_path);
132 
133         let method_doc = generate_doc_comments(method.comment());
134 
135         let method = match (method.client_streaming(), method.server_streaming()) {
136             (false, false) => {
137                 quote! {
138                     #method_doc
139                     async fn #name(&self, request: tonic::Request<#req_message>)
140                         -> Result<tonic::Response<#res_message>, tonic::Status>;
141                 }
142             }
143             (true, false) => {
144                 quote! {
145                     #method_doc
146                     async fn #name(&self, request: tonic::Request<tonic::Streaming<#req_message>>)
147                         -> Result<tonic::Response<#res_message>, tonic::Status>;
148                 }
149             }
150             (false, true) => {
151                 let stream = quote::format_ident!("{}Stream", method.identifier());
152                 let stream_doc = generate_doc_comment(&format!(
153                     "Server streaming response type for the {} method.",
154                     method.identifier()
155                 ));
156 
157                 quote! {
158                     #stream_doc
159                     type #stream: Stream<Item = Result<#res_message, tonic::Status>> + Send + Sync + 'static;
160 
161                     #method_doc
162                     async fn #name(&self, request: tonic::Request<#req_message>)
163                         -> Result<tonic::Response<Self::#stream>, tonic::Status>;
164                 }
165             }
166             (true, true) => {
167                 let stream = quote::format_ident!("{}Stream", method.identifier());
168                 let stream_doc = generate_doc_comment(&format!(
169                     "Server streaming response type for the {} method.",
170                     method.identifier()
171                 ));
172 
173                 quote! {
174                     #stream_doc
175                     type #stream: Stream<Item = Result<#res_message, tonic::Status>> + Send + Sync + 'static;
176 
177                     #method_doc
178                     async fn #name(&self, request: tonic::Request<tonic::Streaming<#req_message>>)
179                         -> Result<tonic::Response<Self::#stream>, tonic::Status>;
180                 }
181             }
182         };
183 
184         stream.extend(method);
185     }
186 
187     stream
188 }
189 
190 #[cfg(feature = "transport")]
191 fn generate_transport(
192     server_service: &syn::Ident,
193     server_trait: &syn::Ident,
194     service_name: &str,
195 ) -> TokenStream {
196     let service_name = syn::LitStr::new(service_name, proc_macro2::Span::call_site());
197 
198     quote! {
199         impl<T: #server_trait> tonic::transport::NamedService for #server_service<T> {
200             const NAME: &'static str = #service_name;
201         }
202     }
203 }
204 
205 #[cfg(not(feature = "transport"))]
206 fn generate_transport(
207     _server_service: &syn::Ident,
208     _server_trait: &syn::Ident,
209     _service_name: &str,
210 ) -> TokenStream {
211     TokenStream::new()
212 }
213 
214 fn generate_methods<T: Service>(service: &T, proto_path: &str) -> TokenStream {
215     let mut stream = TokenStream::new();
216 
217     for method in service.methods() {
218         let path = format!(
219             "/{}{}{}/{}",
220             service.package(),
221             if service.package().is_empty() {
222                 ""
223             } else {
224                 "."
225             },
226             service.identifier(),
227             method.identifier()
228         );
229         let method_path = Lit::Str(LitStr::new(&path, Span::call_site()));
230         let ident = quote::format_ident!("{}", method.name());
231         let server_trait = quote::format_ident!("{}", service.name());
232 
233         let method_stream = match (method.client_streaming(), method.server_streaming()) {
234             (false, false) => generate_unary(method, proto_path, ident, server_trait),
235 
236             (false, true) => {
237                 generate_server_streaming(method, proto_path, ident.clone(), server_trait)
238             }
239             (true, false) => {
240                 generate_client_streaming(method, proto_path, ident.clone(), server_trait)
241             }
242 
243             (true, true) => generate_streaming(method, proto_path, ident.clone(), server_trait),
244         };
245 
246         let method = quote! {
247             #method_path => {
248                 #method_stream
249             }
250         };
251         stream.extend(method);
252     }
253 
254     stream
255 }
256 
257 fn generate_unary<T: Method>(
258     method: &T,
259     proto_path: &str,
260     method_ident: Ident,
261     server_trait: Ident,
262 ) -> TokenStream {
263     let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
264 
265     let service_ident = quote::format_ident!("{}Svc", method.identifier());
266 
267     let (request, response) = method.request_response_name(proto_path);
268 
269     quote! {
270         #[allow(non_camel_case_types)]
271         struct #service_ident<T: #server_trait >(pub Arc<T>);
272 
273         impl<T: #server_trait> tonic::server::UnaryService<#request> for #service_ident<T> {
274             type Response = #response;
275             type Future = BoxFuture<tonic::Response<Self::Response>, tonic::Status>;
276 
277             fn call(&mut self, request: tonic::Request<#request>) -> Self::Future {
278                 let inner = self.0.clone();
279                 let fut = async move {
280                     (*inner).#method_ident(request).await
281                 };
282                 Box::pin(fut)
283             }
284         }
285 
286         let inner = self.inner.clone();
287         let fut = async move {
288             let interceptor = inner.1.clone();
289             let inner = inner.0;
290             let method = #service_ident(inner);
291             let codec = #codec_name::default();
292 
293             let mut grpc = if let Some(interceptor) = interceptor {
294                 tonic::server::Grpc::with_interceptor(codec, interceptor)
295             } else {
296                 tonic::server::Grpc::new(codec)
297             };
298 
299             let res = grpc.unary(method, req).await;
300             Ok(res)
301         };
302 
303         Box::pin(fut)
304     }
305 }
306 
307 fn generate_server_streaming<T: Method>(
308     method: &T,
309     proto_path: &str,
310     method_ident: Ident,
311     server_trait: Ident,
312 ) -> TokenStream {
313     let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
314 
315     let service_ident = quote::format_ident!("{}Svc", method.identifier());
316 
317     let (request, response) = method.request_response_name(proto_path);
318 
319     let response_stream = quote::format_ident!("{}Stream", method.identifier());
320 
321     quote! {
322         #[allow(non_camel_case_types)]
323         struct #service_ident<T: #server_trait >(pub Arc<T>);
324 
325         impl<T: #server_trait> tonic::server::ServerStreamingService<#request> for #service_ident<T> {
326             type Response = #response;
327             type ResponseStream = T::#response_stream;
328             type Future = BoxFuture<tonic::Response<Self::ResponseStream>, tonic::Status>;
329 
330             fn call(&mut self, request: tonic::Request<#request>) -> Self::Future {
331                 let inner = self.0.clone();
332                 let fut = async move {
333                     (*inner).#method_ident(request).await
334 
335                 };
336                 Box::pin(fut)
337             }
338         }
339 
340         let inner = self.inner.clone();
341         let fut = async move {
342             let interceptor = inner.1;
343             let inner = inner.0;
344             let method = #service_ident(inner);
345             let codec = #codec_name::default();
346 
347             let mut grpc = if let Some(interceptor) = interceptor {
348                 tonic::server::Grpc::with_interceptor(codec, interceptor)
349             } else {
350                 tonic::server::Grpc::new(codec)
351             };
352 
353             let res = grpc.server_streaming(method, req).await;
354             Ok(res)
355         };
356 
357         Box::pin(fut)
358     }
359 }
360 
361 fn generate_client_streaming<T: Method>(
362     method: &T,
363     proto_path: &str,
364     method_ident: Ident,
365     server_trait: Ident,
366 ) -> TokenStream {
367     let service_ident = quote::format_ident!("{}Svc", method.identifier());
368 
369     let (request, response) = method.request_response_name(proto_path);
370     let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
371 
372     quote! {
373         #[allow(non_camel_case_types)]
374         struct #service_ident<T: #server_trait >(pub Arc<T>);
375 
376         impl<T: #server_trait> tonic::server::ClientStreamingService<#request> for #service_ident<T>
377         {
378             type Response = #response;
379             type Future = BoxFuture<tonic::Response<Self::Response>, tonic::Status>;
380 
381             fn call(&mut self, request: tonic::Request<tonic::Streaming<#request>>) -> Self::Future {
382                 let inner = self.0.clone();
383                 let fut = async move {
384                     (*inner).#method_ident(request).await
385 
386                 };
387                 Box::pin(fut)
388             }
389         }
390 
391         let inner = self.inner.clone();
392         let fut = async move {
393             let interceptor = inner.1;
394             let inner = inner.0;
395             let method = #service_ident(inner);
396             let codec = #codec_name::default();
397 
398             let mut grpc = if let Some(interceptor) = interceptor {
399                 tonic::server::Grpc::with_interceptor(codec, interceptor)
400             } else {
401                 tonic::server::Grpc::new(codec)
402             };
403 
404             let res = grpc.client_streaming(method, req).await;
405             Ok(res)
406         };
407 
408         Box::pin(fut)
409     }
410 }
411 
412 fn generate_streaming<T: Method>(
413     method: &T,
414     proto_path: &str,
415     method_ident: Ident,
416     server_trait: Ident,
417 ) -> TokenStream {
418     let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
419 
420     let service_ident = quote::format_ident!("{}Svc", method.identifier());
421 
422     let (request, response) = method.request_response_name(proto_path);
423 
424     let response_stream = quote::format_ident!("{}Stream", method.identifier());
425 
426     quote! {
427         #[allow(non_camel_case_types)]
428         struct #service_ident<T: #server_trait>(pub Arc<T>);
429 
430         impl<T: #server_trait> tonic::server::StreamingService<#request> for #service_ident<T>
431         {
432             type Response = #response;
433             type ResponseStream = T::#response_stream;
434             type Future = BoxFuture<tonic::Response<Self::ResponseStream>, tonic::Status>;
435 
436             fn call(&mut self, request: tonic::Request<tonic::Streaming<#request>>) -> Self::Future {
437                 let inner = self.0.clone();
438                 let fut = async move {
439                     (*inner).#method_ident(request).await
440                 };
441                 Box::pin(fut)
442             }
443         }
444 
445         let inner = self.inner.clone();
446         let fut = async move {
447             let interceptor = inner.1;
448             let inner = inner.0;
449             let method = #service_ident(inner);
450             let codec = #codec_name::default();
451 
452             let mut grpc = if let Some(interceptor) = interceptor {
453                 tonic::server::Grpc::with_interceptor(codec, interceptor)
454             } else {
455                 tonic::server::Grpc::new(codec)
456             };
457 
458             let res = grpc.streaming(method, req).await;
459             Ok(res)
460         };
461 
462         Box::pin(fut)
463     }
464 }
465