Context clients + more integration tests (#10)

* Heavily refactor the context(s) to implement specific traits, makes the code accessing context features more reusable for users.

Also re-shaped the call/send to a builder-like approach, to prepare it for codegen.

* Macro generated context clients

Author: slinkydeveloper

macros/src/ast.rs

@@ -16,8 +16,8 @@ use syn::parse::{Parse, ParseStream};
 use syn::spanned::Spanned;
 use syn::token::Comma;
 use syn::{
-    braced, parenthesized, Attribute, Error, Expr, ExprLit, FnArg, GenericArgument, Ident, Lit,
-    Pat, PatType, Path, PathArguments, Result, ReturnType, Token, Type, Visibility,
+    braced, parenthesized, parse_quote, Attribute, Error, Expr, ExprLit, FnArg, GenericArgument,
+    Ident, Lit, Pat, PatType, Path, PathArguments, Result, ReturnType, Token, Type, Visibility,
 };
 
 /// Accumulates multiple errors into a result.
@@ -145,7 +145,7 @@ pub(crate) struct Handler {
     pub(crate) restate_name: String,
     pub(crate) ident: Ident,
     pub(crate) arg: Option<PatType>,
-    pub(crate) output: ReturnType,
+    pub(crate) output: Type,
 }
 
 impl Parse for Handler {
@@ -189,20 +189,24 @@ impl Parse for Handler {
         errors?;
 
         // Parse return type
-        let output: ReturnType = input.parse()?;
+        let return_type: ReturnType = input.parse()?;
         input.parse::<Token![;]>()?;
 
-        match &output {
-            ReturnType::Default => {}
+        let output: Type = match &return_type {
+            ReturnType::Default => {
+                parse_quote!(())
+            }
             ReturnType::Type(_, ty) => {
-                if handler_result_parameter(ty).is_none() {
+                if let Some(ty) = extract_handler_result_parameter(ty) {
+                    ty
+                } else {
                     return Err(Error::new(
-                        output.span(),
+                        return_type.span(),
                         "Only restate_sdk::prelude::HandlerResult is supported as return type",
                     ));
                 }
             }
-        }
+        };
 
         // Process attributes
         let mut is_shared = false;
@@ -259,7 +263,7 @@ fn read_literal_attribute_name(attr: &Attribute) -> Result<Option<String>> {
         .transpose()
 }
 
-fn handler_result_parameter(ty: &Type) -> Option<&Type> {
+fn extract_handler_result_parameter(ty: &Type) -> Option<Type> {
     let path = match ty {
         Type::Path(ty) => &ty.path,
         _ => return None,
@@ -280,7 +284,7 @@ fn handler_result_parameter(ty: &Type) -> Option<&Type> {
     }
 
     match &bracketed.args[0] {
-        GenericArgument::Type(arg) => Some(arg),
+        GenericArgument::Type(arg) => Some(arg.clone()),
         _ => None,
     }
 }

macros/src/gen.rs

@@ -2,12 +2,13 @@ use crate::ast::{Handler, Object, Service, ServiceInner, ServiceType, Workflow};
 use proc_macro2::TokenStream as TokenStream2;
 use proc_macro2::{Ident, Literal};
 use quote::{format_ident, quote, ToTokens};
-use syn::{parse_quote, Attribute, ReturnType, Type, Visibility};
+use syn::{Attribute, PatType, Visibility};
 
 pub(crate) struct ServiceGenerator<'a> {
     pub(crate) service_ty: ServiceType,
     pub(crate) restate_name: &'a str,
     pub(crate) service_ident: &'a Ident,
+    pub(crate) client_ident: Ident,
     pub(crate) serve_ident: Ident,
     pub(crate) vis: &'a Visibility,
     pub(crate) attrs: &'a [Attribute],
@@ -20,6 +21,7 @@ impl<'a> ServiceGenerator<'a> {
             service_ty,
             restate_name: &s.restate_name,
             service_ident: &s.ident,
+            client_ident: format_ident!("{}Client", s.ident),
             serve_ident: format_ident!("Serve{}", s.ident),
             vis: &s.vis,
             attrs: &s.attrs,
@@ -50,8 +52,6 @@ impl<'a> ServiceGenerator<'a> {
             ..
         } = self;
 
-        let unit_type: &Type = &parse_quote!(());
-
         let handler_fns = handlers
             .iter()
             .map(
@@ -66,13 +66,9 @@ impl<'a> ServiceGenerator<'a> {
                         (ServiceType::Workflow, false) => quote! { ::restate_sdk::prelude::WorkflowContext },
                     };
 
-                    let output = match output {
-                        ReturnType::Type(_, ref ty) => ty.as_ref(),
-                        ReturnType::Default => unit_type,
-                    };
                     quote! {
                         #( #attrs )*
-                        fn #ident(&self, context: #ctx, #( #args ),*) -> impl std::future::Future<Output=#output> + ::core::marker::Send;
+                        fn #ident(&self, context: #ctx, #( #args ),*) -> impl std::future::Future<Output=::restate_sdk::prelude::HandlerResult<#output>> + ::core::marker::Send;
                     }
                 },
             );
@@ -223,6 +219,123 @@ impl<'a> ServiceGenerator<'a> {
             }
         }
     }
+
+    fn struct_client(&self) -> TokenStream2 {
+        let &Self {
+            vis,
+            ref client_ident,
+            // service_ident,
+            ref service_ty,
+            ..
+        } = self;
+
+        let key_field = match service_ty {
+            ServiceType::Service => quote! {},
+            ServiceType::Object | ServiceType::Workflow => quote! {
+                key: String,
+            },
+        };
+
+        let into_client_impl = match service_ty {
+            ServiceType::Service => {
+                quote! {
+                    impl<'ctx> ::restate_sdk::context::IntoServiceClient<'ctx> for #client_ident<'ctx> {
+                        fn create_client(ctx: &'ctx ::restate_sdk::endpoint::ContextInternal) -> Self {
+                            Self { ctx }
+                        }
+                    }
+                }
+            }
+            ServiceType::Object => quote! {
+                impl<'ctx> ::restate_sdk::context::IntoObjectClient<'ctx> for #client_ident<'ctx> {
+                    fn create_client(ctx: &'ctx ::restate_sdk::endpoint::ContextInternal, key: String) -> Self {
+                        Self { ctx, key }
+                    }
+                }
+            },
+            ServiceType::Workflow => quote! {
+                impl<'ctx> ::restate_sdk::context::IntoWorkflowClient<'ctx> for #client_ident<'ctx> {
+                    fn create_client(ctx: &'ctx ::restate_sdk::endpoint::ContextInternal, key: String) -> Self {
+                        Self { ctx, key }
+                    }
+                }
+            },
+        };
+
+        quote! {
+            /// Struct exposing the client to invoke [#service_ident] from another service.
+            #vis struct #client_ident<'ctx> {
+                ctx: &'ctx ::restate_sdk::endpoint::ContextInternal,
+                #key_field
+            }
+
+            #into_client_impl
+        }
+    }
+
+    fn impl_client(&self) -> TokenStream2 {
+        let &Self {
+            vis,
+            ref client_ident,
+            service_ident,
+            handlers,
+            restate_name,
+            service_ty,
+            ..
+        } = self;
+
+        let service_literal = Literal::string(restate_name);
+
+        let handlers_fns = handlers.iter().map(|handler| {
+            let handler_ident = &handler.ident;
+            let handler_literal = Literal::string(&handler.restate_name);
+
+            let argument = match &handler.arg {
+                None => quote! {},
+                Some(PatType {
+                         ty, ..
+                     }) => quote! { req: #ty }
+            };
+            let argument_ty = match &handler.arg {
+                None => quote! { () },
+                Some(PatType {
+                         ty, ..
+                     }) => quote! { #ty }
+            };
+            let res_ty = &handler.output;
+            let input =  match &handler.arg {
+                None => quote! { () },
+                Some(_) => quote! { req }
+            };
+            let request_target = match service_ty {
+                ServiceType::Service => quote! {
+                    ::restate_sdk::context::RequestTarget::service(#service_literal, #handler_literal)
+                },
+                ServiceType::Object => quote! {
+                    ::restate_sdk::context::RequestTarget::object(#service_literal, &self.key, #handler_literal)
+                },
+                ServiceType::Workflow => quote! {
+                    ::restate_sdk::context::RequestTarget::workflow(#service_literal, &self.key, #handler_literal)
+                }
+            };
+
+            quote! {
+                #vis fn #handler_ident(&self, #argument) -> ::restate_sdk::context::Request<'ctx, #argument_ty, #res_ty> {
+                    self.ctx.request(#request_target, #input)
+                }
+            }
+        });
+
+        let doc_msg = format!(
+            "Struct exposing the client to invoke [`{service_ident}`] from another service."
+        );
+        quote! {
+            #[doc = #doc_msg]
+            impl<'ctx> #client_ident<'ctx> {
+                #( #handlers_fns )*
+            }
+        }
+    }
 }
 
 impl<'a> ToTokens for ServiceGenerator<'a> {
@@ -232,6 +345,8 @@ impl<'a> ToTokens for ServiceGenerator<'a> {
             self.struct_serve(),
             self.impl_service_for_serve(),
             self.impl_discoverable(),
+            self.struct_client(),
+            self.impl_client(),
         ]);
     }
 }