datadog_agent_runtime/env/
autodiscovery.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use async_trait::async_trait;
5use datadog_agent_commons::ipc::{client::RemoteAgentClient, config::RemoteAgentClientConfiguration};
6use futures::StreamExt;
7use saluki_common::sync::shutdown::ShutdownHandle;
8use saluki_core::runtime::{InitializationError, Supervisable, Supervisor, SupervisorFuture};
9use saluki_env::autodiscovery::AutodiscoveryEvent;
10use saluki_env::AutodiscoveryProvider;
11use saluki_error::GenericError;
12use tokio::select;
13use tokio::sync::mpsc::{self, Receiver, Sender};
14use tokio::sync::Mutex;
15use tokio::time::sleep;
16use tracing::{debug, warn};
17
18/// Datadog Agent-based autodiscovery provider.
19///
20/// This provider is based primarily on the remote autodiscovery API exposed by the Datadog Agent, which handles the
21/// bulk of the work by detecting changes to underlying environment (such as containers starting and stopping), and
22/// determining if they have services associated with them that require checks to be run against them. The remote
23/// autodiscovery API operates in a streaming fashion, which the provider uses to then broadcast updates to subscribers.
24#[derive(Clone)]
25pub struct RemoteAgentAutodiscoveryProvider {
26    subscribers: AutodiscoverySubscribers,
27}
28
29type AutodiscoverySubscribers = Arc<Mutex<Vec<Sender<AutodiscoveryEvent>>>>;
30
31impl RemoteAgentAutodiscoveryProvider {
32    /// Creates a new `RemoteAgentAutodiscoveryProvider` based on the given configuration, along with a [`Supervisor`] that
33    /// drives the collection and broadcasting of autodiscovery events.
34    ///
35    /// # Errors
36    ///
37    /// If the remote agent client couldn't be created, an error is returned.
38    pub async fn new(client_config: &RemoteAgentClientConfiguration) -> Result<(Self, Supervisor), GenericError> {
39        let client = RemoteAgentClient::connect(client_config).await?;
40        let subscribers = Arc::new(Mutex::new(Vec::new()));
41
42        let provider = Self {
43            subscribers: subscribers.clone(),
44        };
45
46        let mut supervisor = Supervisor::new("autodiscovery")?;
47        supervisor.add_worker(AutodiscoveryEventBroadcaster { client, subscribers });
48
49        Ok((provider, supervisor))
50    }
51}
52
53#[async_trait]
54impl AutodiscoveryProvider for RemoteAgentAutodiscoveryProvider {
55    async fn subscribe(&self) -> Option<Receiver<AutodiscoveryEvent>> {
56        let (sender, receiver) = mpsc::channel::<AutodiscoveryEvent>(16);
57        self.subscribers.lock().await.push(sender);
58        Some(receiver)
59    }
60}
61
62struct AutodiscoveryEventBroadcaster {
63    client: RemoteAgentClient,
64    subscribers: AutodiscoverySubscribers,
65}
66
67#[async_trait]
68impl Supervisable for AutodiscoveryEventBroadcaster {
69    fn name(&self) -> &str {
70        "ad-event-broadcaster"
71    }
72
73    async fn initialize(&self, process_shutdown: ShutdownHandle) -> Result<SupervisorFuture, InitializationError> {
74        let client = self.client.clone();
75        let subscribers = self.subscribers.clone();
76
77        Ok(Box::pin(async move {
78            select! {
79                _ = process_shutdown => {},
80                _ = run_ad_event_broadcaster(client, subscribers) => {},
81            }
82            Ok(())
83        }))
84    }
85}
86
87async fn run_ad_event_broadcaster(mut client: RemoteAgentClient, subscribers: AutodiscoverySubscribers) {
88    debug!("Listening to autodiscovery events from remote agent.");
89
90    loop {
91        let mut autodiscovery_stream = client.get_autodiscovery_stream();
92
93        debug!("Polling autodiscovery event stream.");
94
95        while let Some(result) = autodiscovery_stream.next().await {
96            match result {
97                Ok(response) => {
98                    for proto_config in response.configs {
99                        let event = AutodiscoveryEvent::from(proto_config);
100                        send_to_subscribers(&subscribers, event).await;
101                    }
102                }
103                Err(status) => {
104                    warn!(
105                        ?status,
106                        "Encountered error while listening for autodiscovery events. Retrying in 1 second...",
107                    );
108                    sleep(Duration::from_secs(1)).await;
109                    break;
110                }
111            }
112        }
113
114        debug!("Autodiscovery event stream ended. Reconnecting...");
115    }
116}
117
118async fn send_to_subscribers(subscribers: &AutodiscoverySubscribers, event: AutodiscoveryEvent) {
119    let mut subscribers = subscribers.lock().await;
120    let mut active_subscribers = Vec::with_capacity(subscribers.len());
121
122    for sender in subscribers.drain(..) {
123        if sender.send(event.clone()).await.is_ok() {
124            active_subscribers.push(sender);
125        }
126    }
127
128    *subscribers = active_subscribers;
129}
datadog_agent_runtime/env/autodiscovery.rs

datadog_agent_runtime/env/
autodiscovery.rs
