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

dynamo_runtime/component/
component.rs
