feat: implement StreamData trait for local client engine registry

ryanolson · ryanolson · commit f3d24a9eea1e · 2025-09-18T03:43:24.000Z
This change addresses trait bound conflicts in the local client system by introducing
a new StreamData trait that only requires Send (not Sync) for response types.

Key changes:
- Add StreamData trait for response types that don't need Sync
- Create StreamAnyEngineWrapper for type erasure with StreamData responses
- Update LocalClient to use StreamData constraint for response types
- Add downcast_stream() method for StreamData-compatible downcasting
- Fix clippy warnings about complex types and unused variables
- Rewrite local_client_demo.rs to use new endpoint builder pattern

This allows ManyOut&lt;T&gt; types (Pin&lt;Box&lt;dyn AsyncEngineStream&lt;T&gt;&gt;&gt;) to work
properly with the local engine registry, since streams are consumed
sequentially and don't need to be shared between threads.

Signed-off-by: Ryan Olson &lt;rolson@nvidia.com&gt;
diff --git a/lib/runtime/examples/local_client_demo.rs b/lib/runtime/examples/local_client_demo.rs
@@ -3,19 +3,38 @@
 
 //! Example demonstrating LocalClient functionality
 
-use dynamo_runtime::DistributedRuntime;
-use dynamo_runtime::component::register_local_engine;
-use dynamo_runtime::engine::{AsyncEngine, async_trait};
+use dynamo_runtime::component::LocalClient;
+use dynamo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, async_trait};
+use dynamo_runtime::pipeline::network::Ingress;
+use dynamo_runtime::pipeline::{Context, ManyOut, ResponseStream, SingleIn};
+use dynamo_runtime::protocols::annotated::Annotated;
+use dynamo_runtime::{DistributedRuntime, Runtime, distributed::DistributedConfig};
+use futures::StreamExt;
 use std::sync::Arc;
 
 /// Simple test engine that echoes strings
 struct SimpleEchoEngine;
 
 #[async_trait]
-impl AsyncEngine<String, String, String> for SimpleEchoEngine {
-    async fn generate(&self, request: String) -> Result<String, String> {
-        println!("Engine received: {}", request);
-        Ok(format!("Echo: {}", request))
+impl AsyncEngine<SingleIn<String>, ManyOut<Annotated<String>>, anyhow::Error> for SimpleEchoEngine {
+    async fn generate(
+        &self,
+        request: SingleIn<String>,
+    ) -> Result<ManyOut<Annotated<String>>, anyhow::Error> {
+        println!("Engine received: {}", *request);
+
+        let response = Annotated {
+            data: Some(format!("Echo: {}", *request)),
+            id: None,
+            event: None,
+            comment: None,
+        };
+
+        let context = request.context();
+
+        // Create a simple stream that yields the response once
+        let stream = futures::stream::once(async move { response });
+        Ok(ResponseStream::new(Box::pin(stream), context))
     }
 }
 
@@ -28,16 +47,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     // Create runtime and DRT
     println!("1. Creating runtime...");
-    let runtime = dynamo_runtime::Runtime::builder()
-        .app_name("local-client-demo")
-        .build()
-        .await?;
+    let runtime = Runtime::from_current()?;
 
-    let config = dynamo_runtime::DistributedConfig::builder()
-        .etcd_config(Default::default())
-        .nats_config(Default::default())
-        .is_static(true)
-        .build()?;
+    let config = DistributedConfig {
+        etcd_config: Default::default(),
+        nats_config: Default::default(),
+        is_static: true,
+    };
 
     let drt = DistributedRuntime::new(runtime, config).await?;
     println!("   ✓ Runtime created\n");
@@ -52,40 +68,51 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     // Create and register an engine
     println!("3. Creating and registering engine...");
-    let engine: Arc<dyn AsyncEngine<String, String, String>> = Arc::new(SimpleEchoEngine);
+    let engine: Arc<dyn AsyncEngine<SingleIn<String>, ManyOut<Annotated<String>>, anyhow::Error>> =
+        Arc::new(SimpleEchoEngine);
+
+    // Wrap the engine in an Ingress to make it a PushWorkHandler
+    let ingress = Ingress::for_engine(engine)?;
 
-    // Register the engine for local access
-    let key = register_local_engine(&endpoint, engine.clone()).await?;
-    println!("   ✓ Registered engine with key: {}\n", key);
+    // Create the endpoint instance with the ingress as handler (setup phase)
+    let _endpoint_instance = endpoint
+        .endpoint_builder()
+        .handler(ingress)
+        .create()
+        .await?;
+    println!("   ✓ Engine registered automatically during endpoint creation\n");
 
-    // Demonstrate direct local invocation
-    println!("4. Testing direct local invocation...");
+    // Create a LocalClient using the endpoint's convenience method
+    println!("4. Creating LocalClient...");
+    let local_client: LocalClient<SingleIn<String>, ManyOut<Annotated<String>>, anyhow::Error> =
+        endpoint.local_client().await?;
+    println!("   ✓ LocalClient created successfully\n");
+
+    // Demonstrate local client usage
+    println!("5. Testing LocalClient invocation...");
     println!("   (This bypasses all network layers and invokes the engine directly)");
 
-    // Direct invocation through the registered engine
-    let response = engine
-        .generate("Hello from direct call!".to_string())
-        .await?;
-    println!("   Response: {}\n", response);
-
-    // Show what LocalClient would do
-    println!("5. LocalClient usage (conceptual):");
-    println!("   - LocalClient::from_endpoint(&endpoint) would retrieve the engine");
-    println!("   - local_client.generate(request) would call the engine directly");
-    println!("   - No network overhead, no etcd watching, no instance discovery\n");
-
-    // Show the registered engines
-    println!("6. Registry information:");
-    println!("   - Key format: namespace/component/endpoint");
-    println!("   - Registered key: {}", key);
-    println!("   - Engine is type-erased as AnyAsyncEngine");
-    println!("   - LocalClient downcasts back to specific types\n");
-
-    // Cleanup
-    println!("7. Cleanup...");
-    drt.unregister_local_engine(&key).await;
-    println!("   ✓ Unregistered engine\n");
-
-    println!("=== Demo Complete ===");
+    // Create a request with context
+    let request = Context::new("Hello from LocalClient!".to_string());
+
+    // Generate response using the local client
+    let mut response_stream = local_client.generate(request).await?;
+    let response = response_stream.next().await.expect("Expected response");
+
+    println!("   Request: Hello from LocalClient!");
+    if let Some(data) = &response.data {
+        println!("   Response: {}", data);
+    }
+    println!();
+
+    // Show the benefits
+    println!("6. LocalClient Benefits:");
+    println!("   ✓ No network overhead");
+    println!("   ✓ No etcd watching required");
+    println!("   ✓ No instance discovery needed");
+    println!("   ✓ Direct in-process engine invocation");
+    println!("   ✓ Perfect for testing and local development\n");
+
+    println!("Demo completed successfully!");
     Ok(())
 }
diff --git a/lib/runtime/src/component.rs b/lib/runtime/src/component.rs
@@ -443,7 +443,7 @@ impl Endpoint {
     ) -> Result<local_client::LocalClient<Req, Resp, E>>
     where
         Req: crate::engine::Data,
-        Resp: crate::engine::Data + crate::engine::AsyncEngineContextProvider,
+        Resp: crate::engine::StreamData + crate::engine::AsyncEngineContextProvider,
         E: crate::engine::Data,
     {
         local_client::LocalClient::from_endpoint(self).await
diff --git a/lib/runtime/src/component/local_client.rs b/lib/runtime/src/component/local_client.rs
@@ -5,6 +5,7 @@
 
 use crate::engine::{
     AnyAsyncEngine, AsyncEngine, AsyncEngineContextProvider, Data, DowncastAnyAsyncEngine,
+    StreamData,
 };
 use crate::traits::DistributedRuntimeProvider;
 use crate::v2::entity::{ComponentDescriptor, EndpointDescriptor, NamespaceDescriptor};
@@ -18,7 +19,7 @@ use super::Endpoint;
 pub struct LocalClient<Req, Resp, E>
 where
     Req: Data,
-    Resp: Data + AsyncEngineContextProvider,
+    Resp: StreamData + AsyncEngineContextProvider,
     E: Data,
 {
     engine: Arc<dyn AsyncEngine<Req, Resp, E>>,
@@ -28,7 +29,7 @@ where
 impl<Req, Resp, E> LocalClient<Req, Resp, E>
 where
     Req: Data,
-    Resp: Data + AsyncEngineContextProvider,
+    Resp: StreamData + AsyncEngineContextProvider,
     E: Data,
 {
     /// Create a LocalClient from an endpoint descriptor
@@ -47,7 +48,7 @@ where
 
         // Downcast to the specific types
         let engine = any_engine
-            .downcast::<Req, Resp, E>()
+            .downcast_stream::<Req, Resp, E>()
             .ok_or_else(|| error!("Type mismatch when downcasting local engine for: {}", key))?;
 
         Ok(Self { engine, descriptor })
diff --git a/lib/runtime/src/engine.rs b/lib/runtime/src/engine.rs
@@ -90,6 +90,12 @@ use futures::stream::Stream;
 pub trait Data: Send + Sync + 'static {}
 impl<T: Send + Sync + 'static> Data for T {}
 
+/// A trait for stream data that only requires Send, not Sync.
+/// This is used for response types that are consumed sequentially, not shared between threads.
+/// **Do not manually implement this trait** - the blanket implementation covers all valid types.
+pub trait StreamData: Send + 'static {}
+impl<T: Send + 'static> StreamData for T {}
+
 /// [`DataStream`] is a type alias for a stream of [`Data`] items. This can be adapted to a [`ResponseStream`]
 /// by associating it with a [`AsyncEngineContext`].
 pub type DataUnary<T> = Pin<Box<dyn Future<Output = T> + Send>>;
@@ -350,6 +356,41 @@ where
     }
 }
 
+/// A variant of AnyEngineWrapper that uses StreamData for the response type.
+/// This allows non-Sync response types (like streams) to be stored in the registry.
+struct StreamAnyEngineWrapper<Req, Resp, E>
+where
+    Req: Data,
+    Resp: StreamData + AsyncEngineContextProvider,
+    E: Data,
+{
+    engine: Arc<dyn AsyncEngine<Req, Resp, E>>,
+    _phantom: PhantomData<fn(Req, Resp, E)>,
+}
+
+impl<Req, Resp, E> AnyAsyncEngine for StreamAnyEngineWrapper<Req, Resp, E>
+where
+    Req: Data,
+    Resp: StreamData + AsyncEngineContextProvider,
+    E: Data,
+{
+    fn request_type_id(&self) -> TypeId {
+        TypeId::of::<Req>()
+    }
+
+    fn response_type_id(&self) -> TypeId {
+        TypeId::of::<Resp>()
+    }
+
+    fn error_type_id(&self) -> TypeId {
+        TypeId::of::<E>()
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        &self.engine
+    }
+}
+
 /// An extension trait that provides a convenient way to type-erase an `AsyncEngine`.
 ///
 /// This trait provides the `.into_any_engine()` method on any `Arc<dyn AsyncEngine<...>>`,
@@ -381,6 +422,38 @@ where
     }
 }
 
+/// An extension trait for converting AsyncEngine with StreamData response to AnyAsyncEngine.
+/// This is used for engines whose response types don't need to be Sync.
+pub trait AsStreamAnyAsyncEngine {
+    /// Converts a typed `AsyncEngine` with StreamData response into a type-erased `AnyAsyncEngine`.
+    fn into_stream_any_engine(self) -> Arc<dyn AnyAsyncEngine>;
+
+    /// Converts a typed `AsyncEngine` with StreamData response into a type-erased `AnyAsyncEngine`.
+    /// This is the non-consuming version of into_stream_any_engine.
+    fn as_stream_any(&self) -> Option<Arc<dyn AnyAsyncEngine>>;
+}
+
+impl<Req, Resp, E> AsStreamAnyAsyncEngine for Arc<dyn AsyncEngine<Req, Resp, E>>
+where
+    Req: Data,
+    Resp: StreamData + AsyncEngineContextProvider,
+    E: Data,
+{
+    fn into_stream_any_engine(self) -> Arc<dyn AnyAsyncEngine> {
+        Arc::new(StreamAnyEngineWrapper {
+            engine: self,
+            _phantom: PhantomData,
+        })
+    }
+
+    fn as_stream_any(&self) -> Option<Arc<dyn AnyAsyncEngine>> {
+        Some(Arc::new(StreamAnyEngineWrapper {
+            engine: self.clone(),
+            _phantom: PhantomData,
+        }))
+    }
+}
+
 /// An extension trait that provides a convenient method to downcast an `AnyAsyncEngine`.
 ///
 /// This trait provides the `.downcast<Req, Resp, E>()` method on `Arc<dyn AnyAsyncEngine>`,
@@ -410,6 +483,16 @@ pub trait DowncastAnyAsyncEngine {
         Req: Data,
         Resp: Data + AsyncEngineContextProvider,
         E: Data;
+
+    /// Attempts to downcast an `AnyAsyncEngine` to a specific `AsyncEngine` type with StreamData response.
+    ///
+    /// Returns `Some(engine)` if the type parameters match the original engine,
+    /// or `None` if the types don't match.
+    fn downcast_stream<Req, Resp, E>(&self) -> Option<Arc<dyn AsyncEngine<Req, Resp, E>>>
+    where
+        Req: Data,
+        Resp: StreamData + AsyncEngineContextProvider,
+        E: Data;
 }
 
 impl DowncastAnyAsyncEngine for Arc<dyn AnyAsyncEngine> {
@@ -430,6 +513,24 @@ impl DowncastAnyAsyncEngine for Arc<dyn AnyAsyncEngine> {
             None
         }
     }
+
+    fn downcast_stream<Req, Resp, E>(&self) -> Option<Arc<dyn AsyncEngine<Req, Resp, E>>>
+    where
+        Req: Data,
+        Resp: StreamData + AsyncEngineContextProvider,
+        E: Data,
+    {
+        if self.request_type_id() == TypeId::of::<Req>()
+            && self.response_type_id() == TypeId::of::<Resp>()
+            && self.error_type_id() == TypeId::of::<E>()
+        {
+            self.as_any()
+                .downcast_ref::<Arc<dyn AsyncEngine<Req, Resp, E>>>()
+                .cloned()
+        } else {
+            None
+        }
+    }
 }
 
 #[cfg(test)]
@@ -482,7 +583,7 @@ mod tests {
         let typed_engine: Arc<dyn AsyncEngine<Req1, Resp1, Err1>> = Arc::new(MockEngine);
 
         // 4. Use the extension trait to erase the type
-        let any_engine = typed_engine.into_any_engine();
+        let any_engine = AsAnyAsyncEngine::into_any_engine(typed_engine);
 
         // Check type IDs are preserved
         assert_eq!(any_engine.request_type_id(), TypeId::of::<Req1>());
diff --git a/lib/runtime/src/pipeline/network/ingress/push_handler.rs b/lib/runtime/src/pipeline/network/ingress/push_handler.rs
@@ -139,9 +139,23 @@ where
     }
 
     fn as_any_engine(&self) -> Option<Arc<dyn crate::engine::AnyAsyncEngine>> {
+        use crate::engine::AsStreamAnyAsyncEngine;
         use crate::pipeline::network::Ingress;
-        // Return the pre-stored type-erased engine
-        Ingress::any_engine(self)
+
+        // First try to get the pre-stored type-erased engine
+        if let Some(engine) = Ingress::any_engine(self) {
+            return Some(engine);
+        }
+
+        // If no pre-stored engine, try to create one from the underlying engine
+        if let Some(service_engine) = Ingress::engine(self) {
+            // Try to convert using StreamData-compatible wrapper
+            if let Some(stream_engine) = service_engine.as_stream_any() {
+                return Some(stream_engine);
+            }
+        }
+
+        None
     }
 
     async fn handle_payload(&self, payload: Bytes) -> Result<(), PipelineError> {
diff --git a/lib/runtime/tests/local_client.rs b/lib/runtime/tests/local_client.rs
@@ -78,7 +78,7 @@ async fn test_local_client_registration_and_retrieval() -> Result<(), Box<dyn st
     let response = response_stream.next().await.expect("Expected response");
 
     assert_eq!(response.data, Some("Hello, LocalClient!".to_string()));
-    println!("LocalClient test passed: received '{}'", response);
+    println!("LocalClient test passed: received '{:?}'", response);
 
     // Note: We don't need to start the endpoint for local client testing
     // The engine is registered during create() and available for local access
@@ -128,7 +128,7 @@ async fn test_local_client_with_ingress() -> Result<(), Box<dyn std::error::Erro
 
     assert_eq!(response.data, Some("Test with Ingress".to_string()));
     println!(
-        "LocalClient with Ingress test passed: received '{}'",
+        "LocalClient with Ingress test passed: received '{:?}'",
         response
     );
 
@@ -166,8 +166,8 @@ async fn test_local_client_type_mismatch() -> Result<(), Box<dyn std::error::Err
     println!("Created endpoint with String engine");
 
     // Try to create a LocalClient with different types (this should fail)
-    let result: Result<LocalClient<SingleIn<i32>, ManyOut<Annotated<i32>>, anyhow::Error>, _> =
-        endpoint.local_client().await;
+    type TestLocalClient = LocalClient<SingleIn<i32>, ManyOut<Annotated<i32>>, anyhow::Error>;
+    let result: Result<TestLocalClient, _> = endpoint.local_client().await;
 
     assert!(result.is_err(), "Expected type mismatch error");
     if let Err(e) = result {
diff --git a/lib/runtime/tests/local_client_simple.rs b/lib/runtime/tests/local_client_simple.rs

Original file line number	Diff line number	Diff line change
`@@ -443,7 +443,7 @@ impl Endpoint {`
`443`	`443`	`) -> Result<local_client::LocalClient<Req, Resp, E>>`
`444`	`444`	`where`
`445`	`445`	`Req: crate::engine::Data,`
`446`		`- Resp: crate::engine::Data + crate::engine::AsyncEngineContextProvider,`
	`446`	`+ Resp: crate::engine::StreamData + crate::engine::AsyncEngineContextProvider,`
`447`	`447`	`E: crate::engine::Data,`
`448`	`448`	`{`
`449`	`449`	`local_client::LocalClient::from_endpoint(self).await`