Add ability to register ADP as a remote agent to the Datadog Agent

Cargo.toml

@@ -143,6 +143,7 @@ chrono-tz = { version = "0.10", default-features = false }
 iana-time-zone = { version = "0.1", default-features = false }
 backon = { version = "1", default-features = false }
 http-serde-ext = { version = "1", default-features = false }
+uuid = { version = "1.11.0", default-features = false }
 
 [profile.release]
 lto = "thin"

LICENSE-3rdparty.csv

@@ -313,6 +313,7 @@ utf16_iter,https://github.com/hsivonen/utf16_iter,Apache-2.0 OR MIT,Henri Sivone
 utf8-width,https://github.com/magiclen/utf8-width,MIT,Magic Len <[email protected]>
 utf8_iter,https://github.com/hsivonen/utf8_iter,Apache-2.0 OR MIT,Henri Sivonen <[email protected]>
 utf8parse,https://github.com/alacritty/vte,Apache-2.0 OR MIT,"Joe Wilm <[email protected]>, Christian Duerr <[email protected]>"
+uuid,https://github.com/uuid-rs/uuid,Apache-2.0 OR MIT,"Ashley Mannix<[email protected]>, Dylan DPC<[email protected]>, Hunar Roop Kahlon<[email protected]>"
 valuable,https://github.com/tokio-rs/valuable,MIT,The valuable Authors
 walkdir,https://github.com/BurntSushi/walkdir,Unlicense OR MIT,Andrew Gallant <[email protected]>
 want,https://github.com/seanmonstar/want,MIT,Sean McArthur <[email protected]>

bin/agent-data-plane/src/main.rs

@@ -13,7 +13,10 @@ use std::{
 use memory_accounting::ComponentRegistry;
 use saluki_app::{api::APIBuilder, prelude::*};
 use saluki_components::{
-    destinations::{DatadogEventsServiceChecksConfiguration, DatadogMetricsConfiguration, PrometheusConfiguration},
+    destinations::{
+        DatadogEventsServiceChecksConfiguration, DatadogMetricsConfiguration, DatadogStatusFlareConfiguration,
+        PrometheusConfiguration,
+    },
     sources::{DogStatsDConfiguration, InternalMetricsConfiguration},
     transforms::{
         AggregateConfiguration, ChainedConfiguration, HostEnrichmentConfiguration, OriginEnrichmentConfiguration,
@@ -93,9 +96,11 @@ async fn run(started: Instant) -> Result<(), GenericError> {
     let env_provider =
         ADPEnvironmentProvider::from_configuration(&configuration, &component_registry, &health_registry).await?;
 
+    let status_configuration = DatadogStatusFlareConfiguration::from_configuration(&configuration).await?;
+
     // Create a simple pipeline that runs a DogStatsD source, an aggregation transform to bucket into 10 second windows,
     // and a Datadog Metrics destination that forwards aggregated buckets to the Datadog Platform.
-    let blueprint = create_topology(&configuration, env_provider, &component_registry)?;
+    let blueprint = create_topology(&configuration, env_provider, &component_registry, status_configuration)?;
 
     // Build our administrative API server.
     let primary_api_listen_address = configuration
@@ -158,6 +163,7 @@ async fn run(started: Instant) -> Result<(), GenericError> {
 
 fn create_topology(
     configuration: &GenericConfiguration, env_provider: ADPEnvironmentProvider, component_registry: &ComponentRegistry,
+    status_configuration: DatadogStatusFlareConfiguration,
 ) -> Result<TopologyBlueprint, GenericError> {
     // Create a simple pipeline that runs a DogStatsD source, an aggregation transform to bucket into 10 second windows,
     // and a Datadog Metrics destination that forwards aggregated buckets to the Datadog Platform.
@@ -192,30 +198,32 @@ fn create_topology(
     let events_service_checks_config = DatadogEventsServiceChecksConfiguration::from_configuration(configuration)
         .error_context("Failed to configure Datadog Events/Service Checks destination.")?;
 
+    let int_metrics_config = InternalMetricsConfiguration;
+    let int_metrics_remap_config = AgentTelemetryRemapperConfiguration::new();
+
     let topology_registry = component_registry.get_or_create("topology");
     let mut blueprint = TopologyBlueprint::from_component_registry(topology_registry);
     blueprint
         .add_source("dsd_in", dsd_config)?
+        .add_source("internal_metrics_in", int_metrics_config)?
+        .add_transform("internal_metrics_remap", int_metrics_remap_config)?
         .add_transform("dsd_agg", dsd_agg_config)?
         .add_transform("enrich", enrich_config)?
         .add_destination("dd_metrics_out", dd_metrics_config)?
         .add_destination("dd_events_sc_out", events_service_checks_config)?
+        .add_destination("dd_status_flare_out", status_configuration)?
         .connect_component("dsd_agg", ["dsd_in.metrics"])?
         .connect_component("enrich", ["dsd_agg"])?
         .connect_component("dd_metrics_out", ["enrich"])?
-        .connect_component("dd_events_sc_out", ["dsd_in.events", "dsd_in.service_checks"])?;
+        .connect_component("dd_events_sc_out", ["dsd_in.events", "dsd_in.service_checks"])?
+        .connect_component("internal_metrics_remap", ["internal_metrics_in"])?
+        .connect_component("dd_status_flare_out", ["internal_metrics_remap"])?;
 
     // Insert a Prometheus scrape destination if we've been instructed to enable internal telemetry.
     if configuration.get_typed_or_default::<bool>("telemetry_enabled") {
-        let int_metrics_config = InternalMetricsConfiguration;
-        let int_metrics_remap_config = AgentTelemetryRemapperConfiguration::new();
         let prometheus_config = PrometheusConfiguration::from_configuration(configuration)?;
-
         blueprint
-            .add_source("internal_metrics_in", int_metrics_config)?
-            .add_transform("internal_metrics_remap", int_metrics_remap_config)?
             .add_destination("internal_metrics_out", prometheus_config)?
-            .connect_component("internal_metrics_remap", ["internal_metrics_in"])?
             .connect_component("internal_metrics_out", ["internal_metrics_remap"])?;
     }
 

lib/datadog-protos/build.rs

@@ -31,21 +31,15 @@ fn main() {
 
     // Handle code generation for gRPC service definitions.
     tonic_build::configure()
-        .build_server(false)
+        .build_server(true)
         .include_file("api.mod.rs")
         .compile_protos(
             &[
                 "proto/datadog/api/v1/api.proto",
                 "proto/datadog/workloadmeta/workloadmeta.proto",
+                "proto/datadog/remoteagent/remoteagent.proto",
             ],
             &["proto"],
         )
         .expect("failed to build gRPC service definitions for DCA");
-
-    // Handle code generation for gRPC service definitions.
-    tonic_build::configure()
-        .build_server(true)
-        .include_file("remoteagent.mod.rs")
-        .compile_protos(&["proto/datadog/remoteagent/remoteagent.proto"], &["proto"])
-        .expect("failed to build gRPC service definitions for RemoteAgent")
 }

lib/datadog-protos/src/lib.rs

@@ -12,10 +12,6 @@ mod agent_include {
     include!(concat!(env!("OUT_DIR"), "/api.mod.rs"));
 }
 
-mod remoteagent_include {
-    include!(concat!(env!("OUT_DIR"), "/remoteagent.mod.rs"));
-}
-
 /// Metrics-related definitions.
 pub mod metrics {
     pub use super::include::dd_metric::metric_payload::*;
@@ -33,11 +29,9 @@ pub mod traces {
 pub mod agent {
     pub use super::agent_include::datadog::api::v1::agent_client::AgentClient;
     pub use super::agent_include::datadog::api::v1::agent_secure_client::AgentSecureClient;
+    pub use super::agent_include::datadog::api::v1::remote_agent_server::RemoteAgent;
+    pub use super::agent_include::datadog::api::v1::remote_agent_server::RemoteAgentServer;
     pub use super::agent_include::datadog::model::v1::*;
+    pub use super::agent_include::datadog::remoteagent::*;
     pub use super::agent_include::datadog::workloadmeta::*;
 }
-
-// RemoteAgent definitions.
-pub mod remoteagent {
-    pub use super::remoteagent_include::datadog::remoteagent::*;
-}

lib/saluki-components/Cargo.toml

@@ -33,6 +33,7 @@ paste = { workspace = true }
 pin-project = { workspace = true }
 protobuf = { workspace = true }
 quanta = { workspace = true }
+rand = { workspace = true }
 regex = { workspace = true }
 rustls = { workspace = true }
 saluki-config = { workspace = true }
@@ -54,6 +55,8 @@ snafu = { workspace = true }
 stringtheory = { workspace = true }
 tokio = { workspace = true, features = ["io-util", "macros", "net", "rt", "rt-multi-thread", "signal", "sync"] }
 tokio-util = { workspace = true }
+tonic = { workspace = true }
 tower = { workspace = true, features = ["retry", "timeout", "util"] }
 tracing = { workspace = true }
 url = { workspace = true }
+uuid = { workspace = true, features = ["std", "v7"] }

lib/saluki-components/src/destinations/datadog/mod.rs

@@ -5,3 +5,6 @@ pub use self::events_service_checks::DatadogEventsServiceChecksConfiguration;
 
 mod metrics;
 pub use self::metrics::DatadogMetricsConfiguration;
+
+mod status_flare;
+pub use self::status_flare::DatadogStatusFlareConfiguration;

lib/saluki-components/src/destinations/datadog/status_flare/mod.rs

@@ -0,0 +1,150 @@
+use std::num::NonZeroUsize;
+use std::time::Duration;
+
+use async_trait::async_trait;
+use memory_accounting::{MemoryBounds, MemoryBoundsBuilder};
+use rand::{distributions::Alphanumeric, thread_rng, Rng};
+use saluki_config::GenericConfiguration;
+use saluki_core::components::{
+    destinations::{Destination, DestinationBuilder, DestinationContext},
+    ComponentContext,
+};
+use saluki_env::helpers::remote_agent::RemoteAgentClient;
+use saluki_error::GenericError;
+use saluki_event::DataType;
+use tokio::select;
+use tokio::time::{interval, MissedTickBehavior};
+use tracing::debug;
+use uuid::Uuid;
+
+const DEFAULT_API_LISTEN_PORT: NonZeroUsize = unsafe { NonZeroUsize::new_unchecked(5102) };
+
+/// Datadog Status and Flare Destination
+///
+/// Registers ADP as a remote agent to the Core Agent.
+///
+/// ## Missing
+///
+/// - grpc server to respond to Core Agent
+pub struct DatadogStatusFlareConfiguration {
+    id: String,
+
+    display_name: String,
+
+    api_listen_port: NonZeroUsize,
+
+    client: RemoteAgentClient,
+}
+
+impl DatadogStatusFlareConfiguration {
+    /// Creates a new `DatadogStatusFlareConfiguration` from the given configuration.
+    pub async fn from_configuration(config: &GenericConfiguration) -> Result<Self, GenericError> {
+        let app_details = saluki_metadata::get_app_details();
+        let formatted_full_name = app_details
+            .full_name()
+            .replace(" ", "-")
+            .replace("_", "-")
+            .to_lowercase();
+        let api_listen_port = config
+            .try_get_typed::<NonZeroUsize>("remote_agent_api_listen_port")?
+            .unwrap_or(DEFAULT_API_LISTEN_PORT);
+        let client = RemoteAgentClient::from_configuration(config).await?;
+
+        Ok(Self {
+            id: format!("{}-{}", formatted_full_name, Uuid::now_v7()),
+            display_name: formatted_full_name,
+            api_listen_port,
+            client,
+        })
+    }
+}
+
+#[async_trait]
+impl DestinationBuilder for DatadogStatusFlareConfiguration {
+    fn input_data_type(&self) -> DataType {
+        DataType::Metric | DataType::EventD | DataType::ServiceCheck
+    }
+
+    async fn build(&self, _context: ComponentContext) -> Result<Box<dyn Destination + Send>, GenericError> {
+        Ok(Box::new(DatadogStatusFlare {
+            id: self.id.clone(),
+            display_name: self.display_name.clone(),
+            api_listen_port: self.api_listen_port,
+            client: self.client.clone(),
+        }))
+    }
+}
+
+impl MemoryBounds for DatadogStatusFlareConfiguration {
+    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder) {
+        builder
+            .minimum()
+            // Capture the size of the heap allocation when the component is built.
+            .with_single_value::<DatadogStatusFlare>();
+    }
+}
+
+pub struct DatadogStatusFlare {
+    id: String,
+
+    display_name: String,
+
+    api_listen_port: NonZeroUsize,
+
+    client: RemoteAgentClient,
+}
+
+#[async_trait]
+impl Destination for DatadogStatusFlare {
+    async fn run(mut self: Box<Self>, mut context: DestinationContext) -> Result<(), GenericError> {
+        let Self {
+            id,
+            display_name,
+            api_listen_port,
+            mut client,
+        } = *self;
+
+        let api_endpoint = format!("0.0.0.0:{}", api_listen_port);
+        let auth_token: String = thread_rng()
+            .sample_iter(&Alphanumeric)
+            .take(64)
+            .map(char::from)
+            .collect();
+
+        let mut register_agent = interval(Duration::from_secs(10));
+        register_agent.set_missed_tick_behavior(MissedTickBehavior::Delay);
+
+        let mut health = context.take_health_handle();
+        health.mark_ready();
+        debug!("Datadog Status and Flare destination started.");
+
+        loop {
+            select! {
+                _ = health.live() => continue,
+
+                result = context.events().next() => match result {
+                    Some(_events) => {
+                    },
+                    None => break,
+                },
+
+                // Time to (re)register with the Core Agent.
+                _ = register_agent.tick() => {
+                    match client.register_remote_agent_request(&id, &display_name, &api_endpoint, &auth_token).await {
+                        Ok(resp) => {
+                            let new_refresh_interval = resp.into_inner().recommended_refresh_interval_secs;
+                            register_agent.reset_after(Duration::from_secs(new_refresh_interval as u64));
+                            debug!("Refreshed registration with Core Agent");
+                        }
+                        Err(_) => {
+                            debug!("Failed to refresh registration with Core Agent.");
+                        }
+                    }
+                }
+            }
+        }
+
+        debug!("Datadog Status Flare destination stopped.");
+        Ok(())
+    }
+}

lib/saluki-components/src/destinations/mod.rs

@@ -4,6 +4,8 @@ mod blackhole;
 pub use self::blackhole::BlackholeConfiguration;
 
 mod datadog;
-pub use self::datadog::{DatadogEventsServiceChecksConfiguration, DatadogMetricsConfiguration};
+pub use self::datadog::{
+    DatadogEventsServiceChecksConfiguration, DatadogMetricsConfiguration, DatadogStatusFlareConfiguration,
+};
 mod prometheus;
 pub use self::prometheus::PrometheusConfiguration;

lib/saluki-env/src/helpers/remote_agent/client.rs

@@ -8,9 +8,9 @@ use std::{
 
 use backon::{BackoffBuilder, ConstantBuilder, Retryable as _};
 use datadog_protos::agent::{
-    AgentClient, AgentSecureClient, EntityId, FetchEntityRequest, HostnameRequest, StreamTagsRequest,
-    StreamTagsResponse, TagCardinality, WorkloadmetaEventType, WorkloadmetaFilter, WorkloadmetaKind,
-    WorkloadmetaSource, WorkloadmetaStreamRequest, WorkloadmetaStreamResponse,
+    AgentClient, AgentSecureClient, EntityId, FetchEntityRequest, HostnameRequest, RegisterRemoteAgentRequest,
+    RegisterRemoteAgentResponse, StreamTagsRequest, StreamTagsResponse, TagCardinality, WorkloadmetaEventType,
+    WorkloadmetaFilter, WorkloadmetaKind, WorkloadmetaSource, WorkloadmetaStreamRequest, WorkloadmetaStreamResponse,
 };
 use futures::Stream;
 use pin_project_lite::pin_project;
@@ -209,6 +209,22 @@ impl RemoteAgentClient {
                 .await
         })
     }
+
+    /// TODO
+    pub async fn register_remote_agent_request(
+        &mut self, id: &str, display_name: &str, api_endpoint: &str, auth_token: &str,
+    ) -> Result<Response<RegisterRemoteAgentResponse>, GenericError> {
+        let mut client = self.secure_client.clone();
+        let response = client
+            .register_remote_agent(RegisterRemoteAgentRequest {
+                id: id.to_string(),
+                display_name: display_name.to_string(),
+                api_endpoint: api_endpoint.to_string(),
+                auth_token: auth_token.to_string(),
+            })
+            .await?;
+        Ok(response)
+    }
 }
 
 pin_project! {