dynamo_runtime/component/
component.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use anyhow::Context;
5use async_trait::async_trait;
6use futures::stream::StreamExt;
7use futures::{Stream, TryStreamExt};
8
9use super::*;
10
11use crate::traits::events::{EventPublisher, EventSubscriber};
12
13#[async_trait]
14impl EventPublisher for Component {
15    fn subject(&self) -> String {
16        format!("namespace.{}.component.{}", self.namespace.name, self.name)
17    }
18
19    async fn publish(
20        &self,
21        event_name: impl AsRef<str> + Send + Sync,
22        event: &(impl Serialize + Send + Sync),
23    ) -> Result<()> {
24        let bytes = serde_json::to_vec(event)?;
25        self.publish_bytes(event_name, bytes).await
26    }
27
28    async fn publish_bytes(
29        &self,
30        event_name: impl AsRef<str> + Send + Sync,
31        bytes: Vec<u8>,
32    ) -> Result<()> {
33        let subject = format!("{}.{}", self.subject(), event_name.as_ref());
34        let Some(nats_client) = self.drt().nats_client() else {
35            anyhow::bail!("KV router's EventPublisher requires NATS");
36        };
37        nats_client.client().publish(subject, bytes.into()).await?;
38        Ok(())
39    }
40}
41
42#[async_trait]
43impl EventSubscriber for Component {
44    async fn subscribe(
45        &self,
46        event_name: impl AsRef<str> + Send + Sync,
47    ) -> Result<async_nats::Subscriber> {
48        let subject = format!("{}.{}", self.subject(), event_name.as_ref());
49        let Some(nats_client) = self.drt().nats_client() else {
50            anyhow::bail!("KV router's EventSubscriber requires NATS");
51        };
52        Ok(nats_client.client().subscribe(subject).await?)
53    }
54
55    async fn subscribe_with_type<T: for<'de> Deserialize<'de> + Send + 'static>(
56        &self,
57        event_name: impl AsRef<str> + Send + Sync,
58    ) -> Result<impl Stream<Item = Result<T>> + Send> {
59        let subscriber = self.subscribe(event_name).await?;
60
61        // Transform the subscriber into a stream of deserialized events
62        let stream = subscriber.map(move |msg| {
63            serde_json::from_slice::<T>(&msg.payload)
64                .with_context(|| format!("Failed to deserialize event payload: {:?}", msg.payload))
65        });
66
67        Ok(stream)
68    }
69}
70
71#[cfg(feature = "integration")]
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    // todo - make a distributed runtime fixture
77    // todo - two options - fully mocked or integration test
78    #[tokio::test]
79    async fn test_publish_and_subscribe() {
80        let rt = Runtime::from_current().unwrap();
81        let dtr = DistributedRuntime::from_settings(rt.clone()).await.unwrap();
82        let ns = dtr.namespace("test_component".to_string()).unwrap();
83        let cp = ns.component("test_component".to_string()).unwrap();
84
85        // Create a subscriber on the component
86        let mut subscriber = cp.subscribe("test_event").await.unwrap();
87
88        // Publish a message from the component
89        cp.publish("test_event", &"test_message".to_string())
90            .await
91            .unwrap();
92
93        // Receive the message
94        if let Some(msg) = subscriber.next().await {
95            let received = String::from_utf8(msg.payload.to_vec()).unwrap();
96            assert_eq!(received, "\"test_message\"");
97        }
98
99        rt.shutdown();
100    }
101}
dynamo_runtime/component/component.rs

dynamo_runtime/component/
component.rs
