1use syn::{Attribute, Ident, ReturnType, PatType, Receiver, FnArg, ImplItem};
2use proc_macro::TokenStream;
3use proc_macro2::TokenStream as TokenStream2;
4use quote::ToTokens;
5
6#[allow(dead_code)]
7struct RpcMethod {
8 attrs: Vec<Attribute>,
9 ident: Ident,
10 arg: PatType,
11 receiver: Receiver,
12 output: ReturnType,
13}
14
15impl RpcMethod {
16 fn ident_web_api(&self) -> Ident {
17 quote::format_ident!(
18 "{}_web_api",
19 self.ident
20 )
21 }
22
23 fn ident_web_api_compressed(&self) -> Ident {
24 quote::format_ident!(
25 "{}_web_api_compressed",
26 self.ident
27 )
28 }
29
30 fn req_type(&self) -> TokenStream2 {
31 let arg_type = &self.arg.ty;
32 quote::quote! {
33 #arg_type
34 }
35 }
36
37 fn resp_type(&self) -> TokenStream2 {
38 let output = &self.output;
39 let output = match output {
40 ReturnType::Default => syn::Type::Verbatim(quote::quote! {()}),
41 ReturnType::Type(_, t) => (**t).clone(),
42 };
43 quote::quote! {
44 #output
45 }
46 }
47
48 fn client_method(&self, client_field: &Ident) -> TokenStream2 {
49 let ident = &self.ident;
50 let ident_compressed = quote::format_ident!("{}_compressed", ident);
51 let req_arg = &self.arg.pat;
52 let req_arg_type = &self.arg.ty;
53 let resp_type = self.resp_type();
54 let web_api_string_name = self.ident_web_api().to_string();
55 let web_api_compressed_string_name = self.ident_web_api_compressed().to_string();
56
57 quote::quote! {
58 pub async fn #ident(&self, #req_arg: &#req_arg_type) -> Result<#resp_type, ::persia_rpc::PersiaRpcError> {
59 self.#client_field.call_async(#web_api_string_name, #req_arg, false).await
60 }
61
62 pub async fn #ident_compressed(&self, #req_arg: &#req_arg_type) -> Result<#resp_type, ::persia_rpc::PersiaRpcError> {
63 self.#client_field.call_async(#web_api_compressed_string_name, #req_arg, true).await
64 }
65 }
66 }
67
68 fn service_method(&self) -> TokenStream2 {
69 let method_ident = &self.ident;
70 let web_api_ident = self.ident_web_api();
71 let web_api_ident_string = web_api_ident.to_string();
72 let web_api_ident_compressed = self.ident_web_api_compressed();
73 let web_api_ident_compressed_string = web_api_ident_compressed.to_string();
74 let req_type = self.req_type();
76 let call_line = quote::quote! {
77 use bytes::Buf;
78 let input: #req_type = ::tokio::task::block_in_place(|| ::persia_speedy::Readable::read_from_stream_unbuffered(body.reader()))
79 .context(::persia_rpc::SerializationFailure {})?;
80 let output = self.#method_ident(input).await;
81 };
82 quote::quote! {
83 pub async fn #web_api_ident(&self, req: ::hyper::Request<::hyper::Body>) -> Result<::hyper::Response<hyper::Body>, ::hyper::Error> {
84 let result = async move {
85 let body =
86 ::hyper::body::aggregate(req.into_body())
87 .await
88 .context(::persia_rpc::TransportError {
89 msg: format!("hyper read body error: {}", #web_api_ident_string),
90 })?;
91 #call_line
92 let output = ::tokio::task::block_in_place(|| output.write_to_vec())
93 .context(::persia_rpc::SerializationFailure {})?;
94 Ok::<_, ::persia_rpc::PersiaRpcError>(output)
95 }
96 .await;
97 match result {
98 Ok(x) => Ok(::hyper::Response::new(::hyper::body::Body::from(x))),
99 Err(e) => {
100 ::tracing::error!("server side error {:?}", e);
101 let mut resp = ::hyper::Response::default();
102 *resp.status_mut() = ::hyper::StatusCode::INTERNAL_SERVER_ERROR;
103 *resp.body_mut() = ::hyper::body::Body::from(format!("{:#?}", e));
104 Ok(resp)
105 }
106 }
107 }
108
109
110 pub async fn #web_api_ident_compressed(&self, req: ::hyper::Request<::hyper::Body>) -> Result<::hyper::Response<hyper::Body>, ::hyper::Error> {
111 use ::hyper::body::Buf;
112 let result = async move {
113 let body =
114 ::hyper::body::to_bytes(req.into_body())
115 .await
116 .context(::persia_rpc::TransportError {
117 msg: format!("hyper read body error: {}", #web_api_ident_compressed_string),
118 })?;
119 let body = if body.len() > 0 {
120 ::tokio::task::block_in_place(|| {
121 ::lz4::block::decompress(body.as_ref(), None)
122 }).context(::persia_rpc::IOFailure {})?.into()
123 } else {
124 body
125 };
126 #call_line
127 let output = ::tokio::task::block_in_place(|| output.write_to_vec())
128 .context(::persia_rpc::SerializationFailure {})?;
129 let output = if output.len() > 0 {
130 ::tokio::task::block_in_place(|| ::lz4::block::compress(&output, Some(lz4::block::CompressionMode::FAST(3)), true)).context(::persia_rpc::IOFailure {})?
131 } else {
132 output
133 };
134 Ok::<_, ::persia_rpc::PersiaRpcError>(output)
135 }
136 .await;
137 match result {
138 Ok(x) => Ok(::hyper::Response::new(::hyper::body::Body::from(x))),
139 Err(e) => {
140 ::tracing::error!("server side error {:?}", e);
141 let mut resp = ::hyper::Response::default();
142 *resp.status_mut() = ::hyper::StatusCode::INTERNAL_SERVER_ERROR;
143 *resp.body_mut() = ::hyper::body::Body::from(format!("{:#?}", e));
144 Ok(resp)
145 }
146 }
147 }
148 }
149 }
150}
151
152struct Service {
153 rpcs: Vec<RpcMethod>,
154 service_type_name: String,
155}
156
157impl Service {
158 fn client_impl(&self) -> TokenStream2 {
159 let client_type_name = quote::format_ident!("{}Client", self.service_type_name);
160 let client_field = quote::format_ident!("inner_rpc_client");
161 let client_methods: Vec<_> = self.rpcs.iter().map(|x| x.client_method(&client_field)).collect();
162 quote::quote! {
163 pub struct #client_type_name {
164 #client_field: ::persia_rpc::RpcClient
165 }
166
167 impl #client_type_name {
168 pub fn new(client: ::persia_rpc::RpcClient) -> Self {
169 Self {
170 #client_field: client
171 }
172 }
173
174 #( #client_methods )*
175 }
176 }
177 }
178
179 fn service_impl(&self) -> TokenStream2 {
180 let service_ident = quote::format_ident!("{}", self.service_type_name);
181 let service_web_api_impl: Vec<_> = self.rpcs.iter().map(|x| x.service_method()).collect();
182 let match_arms: Vec<_> = self.rpcs.iter().map(|x| {
183 let web_api_string_name = "/".to_string() + x.ident_web_api().to_string().as_str();
184 let web_api_compressed_string_name = "/".to_string() + x.ident_web_api_compressed().to_string().as_str();
185 let ident_web_api = x.ident_web_api();
186 let ident_web_api_compressed = x.ident_web_api_compressed();
187 quote::quote! {
188 #web_api_string_name => {
189 Box::pin(async move { server.#ident_web_api(req).await })
190 }
191
192 #web_api_compressed_string_name => {
193 Box::pin(async move { server.#ident_web_api_compressed(req).await })
194 }
195 }
196 }).collect();
197 quote::quote! {
198 use std::task::{Context, Poll};
199
200 impl #service_ident {
201 #( #service_web_api_impl )*
202 }
203
204 impl ::hyper::service::Service<::hyper::Request<::hyper::Body>> for #service_ident {
205 type Response = ::hyper::Response<::hyper::Body>;
206 type Error = ::hyper::Error;
207 type Future = std::pin::Pin<Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>>;
208
209 fn poll_ready(&mut self, _ctx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
210 Poll::Ready(Ok(()))
211 }
212
213 fn call(&mut self, req: ::hyper::Request<::hyper::Body>) -> Self::Future {
214 let server = self.clone();
215 match req.uri().path() {
216 #( #match_arms )*
217 _ => {
218 Box::pin(
219 async move {
220 let mut not_found = hyper::Response::default();
221 *not_found.status_mut() = hyper::StatusCode::NOT_FOUND;
222 Ok(not_found)
223 }
224 )
225 }
226 }
227 }
228 }
229 }
230 }
231}
232
233#[proc_macro_attribute]
234pub fn service(_attr: TokenStream, tokens: TokenStream) -> TokenStream {
235 let item = syn::parse_macro_input!(tokens as syn::ItemImpl);
236 let service_name = item.self_ty.clone().into_token_stream().to_string();
237 let rpc_methods: Vec<_> = item.items.iter()
238 .filter_map(|item| match item {
239 ImplItem::Method(m) => Some(m),
240 _ => None,
241 })
242 .map(|m| {
243 let mut args = vec![];
244 let mut receiver = None;
245 for arg in &m.sig.inputs {
246 match arg {
247 FnArg::Receiver(captures) => {
248 if captures.mutability.is_some() {
249 panic!("should be &self");
250 } else if receiver.is_some() {
251 panic!("duplicated self");
252 } else {
253 receiver = Some(captures.clone());
254 }
255 }
256 FnArg::Typed(captures) => match *captures.pat {
257 syn::Pat::Ident(_) => args.push(captures.clone()),
258 _ => panic!("patterns aren't allowd in RPC args"),
259 },
260 }
261 }
262 assert_eq!(args.len(), 1, "only support single argument");
263 RpcMethod {
264 attrs: m.attrs.clone(),
265 ident: m.sig.ident.clone(),
266 arg: args[0].clone(),
267 receiver: receiver.unwrap(),
268 output: m.sig.output.clone(),
269 }
270 })
271 .collect();
272
273 let s = Service {
274 rpcs: rpc_methods,
275 service_type_name: service_name,
276 };
277
278 let client_impl = s.client_impl();
279 let service_impl = s.service_impl();
280
281 (quote::quote! {
282 #item
283
284 #service_impl
285
286 #client_impl
287 }).into()
288}