refactor: WIP, add add more methods

Author: sjfhsjfh

binaries/cli/src/command/build/distributed.rs

@@ -2,7 +2,7 @@ use communication_layer_request_reply::{TcpConnection, TcpRequestReplyConnection
 use dora_core::descriptor::Descriptor;
 use dora_message::{
     BuildId,
-    cli_to_coordinator::ControlRequest,
+    cli_to_coordinator::{CliToCoordinatorClient, ControlRequest},
     common::{GitSource, LogMessage},
     coordinator_to_cli::ControlRequestReply,
     id::NodeId,
@@ -23,33 +23,17 @@ pub fn build_distributed_dataflow(
     local_working_dir: Option<std::path::PathBuf>,
     uv: bool,
 ) -> eyre::Result<BuildId> {
-    let build_id = {
-        let reply_raw = session
-            .request(
-                &serde_json::to_vec(&ControlRequest::Build {
-                    session_id: dataflow_session.session_id,
-                    dataflow,
-                    git_sources: git_sources.clone(),
-                    prev_git_sources: dataflow_session.git_sources.clone(),
-                    local_working_dir,
-                    uv,
-                })
-                .unwrap(),
-            )
-            .wrap_err("failed to send start dataflow message")?;
-
-        let result: ControlRequestReply =
-            serde_json::from_slice(&reply_raw).wrap_err("failed to parse reply")?;
-        match result {
-            ControlRequestReply::DataflowBuildTriggered { build_id } => {
-                eprintln!("dataflow build triggered: {build_id}");
-                build_id
-            }
-            ControlRequestReply::Error(err) => bail!("{err}"),
-            other => bail!("unexpected start dataflow reply: {other:?}"),
-        }
-    };
-    Ok(build_id)
+    let mut client: CliToCoordinatorClient<_, Vec<u8>, Vec<u8>, std::io::Error> =
+        CliToCoordinatorClient::new(session);
+    let build_id = client.build(dora_message::cli_to_coordinator::BuildReq {
+        session_id: dataflow_session.session_id,
+        dataflow,
+        git_sources: git_sources.clone(),
+        prev_git_sources: dataflow_session.git_sources.clone(),
+        local_working_dir,
+        uv,
+    })?;
+    Ok(build_id.build_id)
 }
 
 pub fn wait_until_dataflow_built(

libraries/communication-layer/request-reply/src/lib.rs

@@ -75,3 +75,13 @@ pub trait RequestReplyConnection: Send + Sync {
 
     fn request(&mut self, request: &Self::RequestData) -> Result<Self::ReplyData, Self::Error>;
 }
+
+impl<T: RequestReplyConnection + ?Sized> RequestReplyConnection for &mut T {
+    type RequestData = T::RequestData;
+    type ReplyData = T::ReplyData;
+    type Error = T::Error;
+    fn request(&mut self, request: &Self::RequestData) -> Result<Self::ReplyData, Self::Error> {
+        (**self).request(request)
+    }
+}
+

libraries/dora-schema-macro/src/client.rs

@@ -37,6 +37,7 @@ pub fn generate_client(schema: &SchemaInput) -> proc_macro2::TokenStream {
 
                     match resp_enum {
                         #response_enum::#variant(resp) => Ok(resp),
+                        #response_enum::Error(err) => Err(::eyre::eyre!("Server returned error: {}", err.msg)),
                         _ => ::eyre::bail!("Unexpected response type"),
                     }
                 }
@@ -71,7 +72,7 @@ pub fn generate_client(schema: &SchemaInput) -> proc_macro2::TokenStream {
                 ReplyData = std::vec::Vec<u8>,
                 Error = E,
             >,
-            E: std::marker::Send + std::marker::Sync + std::error::Error + 'static,
+            E: std::marker::Send + std::marker::Sync + std::error::Error,
         {
             #(#client_methods)*
         }

libraries/dora-schema-macro/src/lib.rs

@@ -5,7 +5,7 @@ mod syntax;
 
 // dora-schema-macro/src/lib.rs
 use proc_macro::TokenStream;
-use quote::{format_ident, quote};
+use quote::quote;
 
 use crate::syntax::SchemaInput;
 
@@ -15,8 +15,7 @@ pub fn dora_schema(input: TokenStream) -> TokenStream {
 
     let protocol_code = protocol::generate_protocol(&schema);
     let client_code = client::generate_client(&schema);
-    // let server_trait_code = server::generate_server_trait(&schema);
-    let server_trait_code = quote! {};
+    let server_trait_code = server::generate_server_trait(&schema);
 
     let expanded = quote! {
         #protocol_code
@@ -26,116 +25,3 @@ pub fn dora_schema(input: TokenStream) -> TokenStream {
 
     TokenStream::from(expanded)
 }
-
-fn generate_server(schema: &SchemaInput) -> proc_macro2::TokenStream {
-    let client_name = &schema.client_name;
-    let server_name = &schema.server_name;
-    let handler_trait_name = format_ident!("{}Handler", server_name);
-    let request_enum = format_ident!("{}To{}Request", client_name, server_name);
-    let response_enum = format_ident!("{}To{}Response", client_name, server_name);
-
-    let trait_methods: Vec<_> = schema
-        .methods
-        .iter()
-        .map(|m| {
-            let handler_name = format_ident!("{}_handler", m.name);
-            let request_type = &m.request;
-            let response_type = &m.response;
-
-            // TODO: better
-            quote! {
-                fn #handler_name(
-                    &self,
-                    request: #request_type
-                ) -> std::pin::Pin<
-                    Box<dyn std::future::Future<Output = ::eyre::Result<#response_type>> + Send>
-                >;
-            }
-        })
-        .collect();
-
-    let dispatch_arms: Vec<_> = schema
-        .methods
-        .iter()
-        .map(|m| {
-            let handler_name = format_ident!("{}_handler", m.name);
-            let request_variant = format_ident!("{}", capitalize(&m.name.to_string()));
-            let response_variant = request_variant.clone();
-
-            quote! {
-                #request_enum::#request_variant(req) => {
-                    match self.#handler_name(req).await {
-                        Ok(resp) => {
-                            let resp_enum = #response_enum::#response_variant(resp);
-                            ::bincode::serialize(&resp_enum)
-                                .wrap_err("Failed to serialize response")
-                        }
-                        Err(e) => {
-                            ::eyre::bail!("Handler error: {:?}", e)
-                        }
-                    }
-                }
-            }
-        })
-        .collect();
-
-    quote! {
-        pub trait #handler_trait_name: Send + Sync {
-            #(#trait_methods)*
-        }
-
-        pub struct #server_name<H, L> {
-            handler: H,
-            listener: L,
-        }
-
-        impl<H, L> #server_name<H, L>
-        where
-            H: #handler_trait_name + 'static,
-            L: ListenConnection,
-        {
-            pub fn new(handler: H, listener: L) -> Self {
-                Self { handler, listener }
-            }
-
-            pub async fn serve(self) -> ::eyre::Result<()> {
-                loop {
-                    let req_bytes = self.listener.receive()
-                        .await
-                        .wrap_err("Failed to receive request")?;
-
-                    let request: #request_enum = match ::bincode::deserialize(&req_bytes) {
-                        Ok(req) => req,
-                        Err(e) => {
-                            eprintln!("Failed to deserialize request: {:?}", e);
-                            continue;
-                        }
-                    };
-
-                    let response_bytes = match request {
-                        #(#dispatch_arms,)*
-                    };
-
-                    match response_bytes {
-                        Ok(bytes) => {
-                            if let Err(e) = self.listener.send(bytes).await {
-                                eprintln!("Failed to send response: {:?}", e);
-                            }
-                        }
-                        Err(e) => {
-                            eprintln!("Handler error: {:?}", e);
-                        }
-                    }
-                }
-            }
-        }
-    }
-}
-
-fn capitalize(s: &str) -> String {
-    let mut chars = s.chars();
-    match chars.next() {
-        Some(first) => first.to_uppercase().chain(chars).collect(),
-        None => String::new(),
-    }
-}

libraries/dora-schema-macro/src/protocol.rs

@@ -60,12 +60,18 @@ pub fn generate_protocol(schema: &SchemaInput) -> proc_macro2::TokenStream {
         #[derive(Debug, ::serde::Serialize, ::serde::Deserialize)]
         pub enum #response_enum {
             #(#response_variants,)*
+            Error(#error_struct),
         }
 
         #[derive(Debug, ::serde::Serialize, ::serde::Deserialize)]
         pub struct #error_struct {
-            pub message: String,
-            pub source: Option<String>,
+            pub msg: String,
+        }
+
+        impl #error_struct {
+            pub fn new(msg: String) -> Self {
+                Self { msg }
+            }
         }
     }
 }

libraries/dora-schema-macro/src/server.rs

@@ -1,5 +1,100 @@
-use crate::SchemaInput;
+use convert_case::{Case, Casing};
+use proc_macro2::Ident;
+use quote::{format_ident, quote};
+
+use crate::{
+    SchemaInput, protocol::{enum_variant_ident, error_struct_ident, request_enum_ident, response_enum_ident}
+};
+
+pub fn server_trait_ident(schema: &SchemaInput) -> Ident {
+    format_ident!("{}Handler", schema.protocol_name())
+}
+
+pub fn handle_func_ident(schema: &SchemaInput) -> Ident {
+    let snake_case_protocol_name =
+        Casing::from_case(&schema.protocol_name(), Case::Pascal).to_case(Case::Snake);
+    format_ident!("handle_{}", snake_case_protocol_name)
+}
+
+pub fn method_handler_ident(method: &crate::syntax::MethodDef) -> Ident {
+    format_ident!("{}_handler", method.name)
+}
 
 pub fn generate_server_trait(schema: &SchemaInput) -> proc_macro2::TokenStream {
-    todo!()
+    let trait_name = server_trait_ident(schema);
+    let request_enum = request_enum_ident(schema);
+    let response_enum = response_enum_ident(schema);
+    let error_type = error_struct_ident(schema);
+
+    let trait_methods: Vec<_> = schema
+        .methods
+        .iter()
+        .map(|m| {
+            let handler_name = method_handler_ident(m);
+            let request_type = &m.request;
+            let response_type = &m.response;
+
+            // TODO: better
+            quote! {
+                fn #handler_name(
+                    &self,
+                    request: #request_type
+                ) -> ::std::pin::Pin<
+                    Box<dyn ::std::future::Future<Output = ::std::result::Result<#response_type, #error_type>> + Send>
+                >;
+            }
+        })
+        .collect();
+
+    let dispatch_arms = schema
+        .methods
+        .iter()
+        .map(|m| {
+            let handler_name = method_handler_ident(m);
+            let method_variant = enum_variant_ident(m);
+            quote! {
+                #request_enum::#method_variant(request) => {
+                    let response = handler.#handler_name(request).await;
+                    match response {
+                        Ok(resp) => #response_enum::#method_variant(resp),
+                        Err(err) => #response_enum::Error(err),
+                    }
+                }
+            }
+        })
+        .collect::<Vec<_>>();
+
+    let handle_func_name = handle_func_ident(schema);
+
+    // TODO: remove unwraps and handle errors properly
+    let handle_func = quote! {
+        pub async fn #handle_func_name(
+            handler: impl #trait_name,
+            request: ::std::vec::Vec<u8>
+        ) -> ::std::vec::Vec<u8> {
+            let request_enum: #request_enum = match ::serde_json::from_slice(&request) {
+                Ok(req) => req,
+                Err(err) => {
+                    let error_response = #response_enum::Error(#error_type {
+                        msg: format!("Failed to deserialize request: {}", err),
+                    });
+                    return ::serde_json::to_vec(&error_response).unwrap();
+                }
+            };
+
+            let response_enum = match request_enum {
+                #(#dispatch_arms)*
+            };
+
+            ::serde_json::to_vec(&response_enum).unwrap()
+        }
+    };
+
+    quote! {
+        pub trait #trait_name: std::marker::Send + std::marker::Sync {
+            #(#trait_methods)*
+        }
+
+        #handle_func
+    }
 }

libraries/message/src/cli_to_coordinator.rs

@@ -14,10 +14,7 @@ dora_schema_macro::dora_schema! {
     Cli => Coordinator:
 
     build: BuildReq => BuildResp;
-//     wait_for_build: ControlRequestWaitForBuild => ControlRequestWaitForBuildReply;
-//     start: ControlRequestStart => ControlRequestStartReply;
-//     wait_for_spawn: ControlRequestWaitForSpawn => ControlRequestWaitForSpawnReply;
-//     reload: ControlRequestReload => ControlRequestReloadReply;
+    start: StartReq => StartResp;
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -42,7 +39,30 @@ pub struct BuildResp {
     pub build_id: BuildId,
 }
 
-#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StartReq {
+    pub build_id: Option<BuildId>,
+    pub session_id: SessionId,
+    pub dataflow: Descriptor,
+    pub name: Option<String>,
+    /// Allows overwriting the base working dir when CLI and daemon are
+    /// running on the same machine.
+    ///
+    /// Must not be used for multi-machine dataflows.
+    ///
+    /// Note that nodes with git sources still use a subdirectory of
+    /// the base working dir.
+    pub local_working_dir: Option<PathBuf>,
+    pub uv: bool,
+    pub write_events_to: Option<PathBuf>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StartResp {
+    pub uuid: Uuid,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum ControlRequest {
     Build {
         session_id: SessionId,