feat(sidecar)!: Allow for targeted telemetry flush (#1898)

Specifically allowing to chose what to flush.

Co-authored-by: bob.weinand <bob.weinand@datadoghq.com>

Author: bwoebi

datadog-sidecar-ffi/src/lib.rs

@@ -28,7 +28,7 @@ use datadog_sidecar::service::telemetry::InternalTelemetryAction;
 use datadog_sidecar::service::{
     blocking::{self, SidecarTransport},
     DynamicInstrumentationConfigState, InstanceId, QueueId, RuntimeMetadata,
-    SerializedTracerHeaderTags, SessionConfig, SidecarAction,
+    SerializedTracerHeaderTags, SessionConfig, SidecarAction, SidecarFlushOptions,
 };
 use datadog_sidecar::service::{get_telemetry_action_sender, InternalTelemetryActions};
 use datadog_sidecar::shm_remote_config::{path_for_remote_config, RemoteConfigReader};
@@ -361,8 +361,11 @@ pub extern "C" fn ddog_sidecar_ping(transport: &mut Box<SidecarTransport>) -> Ma
 }
 
 #[no_mangle]
-pub extern "C" fn ddog_sidecar_flush_traces(transport: &mut Box<SidecarTransport>) -> MaybeError {
-    try_c!(blocking::flush_traces(transport));
+pub extern "C" fn ddog_sidecar_flush(
+    transport: &mut Box<SidecarTransport>,
+    options: SidecarFlushOptions,
+) -> MaybeError {
+    try_c!(blocking::flush(transport, options));
 
     MaybeError::None
 }

datadog-sidecar/src/service/blocking.rs

@@ -3,7 +3,7 @@
 
 use super::{
     DynamicInstrumentationConfigState, InstanceId, QueueId, SerializedTracerHeaderTags,
-    SessionConfig, SidecarAction,
+    SessionConfig, SidecarAction, SidecarFlushOptions,
 };
 use crate::service::sender::SidecarSender;
 use crate::service::sidecar_interface::SidecarInterfaceChannel;
@@ -440,9 +440,9 @@ pub fn stats(transport: &mut SidecarTransport) -> io::Result<String> {
     transport.with_retry(|s| s.stats().map_err(|e| io::Error::other(e.to_string())))
 }
 
-/// Flushes the outstanding traces.
-pub fn flush_traces(transport: &mut SidecarTransport) -> io::Result<()> {
-    transport.with_retry(|s| s.flush_traces())
+/// Flushes traces/stats and/or telemetry, as specified by options.
+pub fn flush(transport: &mut SidecarTransport, options: SidecarFlushOptions) -> io::Result<()> {
+    transport.with_retry(|s| s.flush(options))
 }
 
 /// Sends a ping to the service.

datadog-sidecar/src/service/mod.rs

@@ -44,7 +44,7 @@ pub(crate) mod tracing;
 
 #[cfg(windows)]
 pub use remote_configs::RemoteConfigNotifyFunction;
-pub use sidecar_interface::DynamicInstrumentationConfigState;
+pub use sidecar_interface::{DynamicInstrumentationConfigState, SidecarFlushOptions};
 pub use telemetry::{get_telemetry_action_sender, InternalTelemetryActions};
 pub(crate) use telemetry::{init_telemetry_sender, telemetry_action_receiver_task};
 

datadog-sidecar/src/service/sender.rs

@@ -13,7 +13,8 @@
 
 use crate::service::{
     sidecar_interface::{
-        DynamicInstrumentationConfigState, SidecarInterfaceChannel, SidecarInterfaceRequest,
+        DynamicInstrumentationConfigState, SidecarFlushOptions, SidecarInterfaceChannel,
+        SidecarInterfaceRequest,
     },
     InstanceId, QueueId, SerializedTracerHeaderTags, SessionConfig, SidecarAction,
 };
@@ -456,9 +457,9 @@ impl SidecarSender {
         self.channel.0.set_write_timeout(d)
     }
 
-    pub fn flush_traces(&mut self) -> io::Result<()> {
+    pub fn flush(&mut self, options: SidecarFlushOptions) -> io::Result<()> {
         self.drain_outbox_blocking();
-        self.channel.call_flush_traces()
+        self.channel.call_flush(options)
     }
 
     pub fn ping(&mut self) -> io::Result<()> {

datadog-sidecar/src/service/sidecar_interface.rs

@@ -22,6 +22,13 @@ pub enum DynamicInstrumentationConfigState {
     NotSet,
 }
 
+#[repr(C)]
+#[derive(Debug, Default, Copy, Clone, Serialize, Deserialize)]
+pub struct SidecarFlushOptions {
+    pub traces_and_stats: bool,
+    pub telemetry: bool,
+}
+
 /// The `SidecarInterface` trait defines the necessary methods for the sidecar service.
 ///
 /// These methods include operations such as enqueueing actions, registering services, setting
@@ -201,9 +208,9 @@ pub trait SidecarInterface {
     /// * `actions` - The DogStatsD actions to send.
     async fn send_dogstatsd_actions(instance_id: InstanceId, actions: Vec<DogStatsDActionOwned>);
 
-    /// Flushes any outstanding traces queued for sending.
+    /// Flushes outstanding traces/stats and/or telemetry, as specified by options.
     #[blocking]
-    async fn flush_traces();
+    async fn flush(options: SidecarFlushOptions);
 
     /// Sets x-datadog-test-session-token on all requests for the given session.
     ///

datadog-sidecar/src/service/sidecar_server.rs

@@ -7,7 +7,8 @@ use crate::service::{
     telemetry::{TelemetryCachedClient, TelemetryCachedClientSet},
     tracing::TraceFlusher,
     DynamicInstrumentationConfigState, InstanceId, QueueId, RuntimeInfo, RuntimeMetadata,
-    SerializedTracerHeaderTags, SessionConfig, SessionInfo, SidecarAction, SidecarInterface,
+    SerializedTracerHeaderTags, SessionConfig, SessionInfo, SidecarAction, SidecarFlushOptions,
+    SidecarInterface,
 };
 use datadog_ipc::platform::{FileBackedHandle, ShmHandle};
 use datadog_ipc::{PeerCredentials, SeqpacketConn};
@@ -991,12 +992,44 @@ impl SidecarInterface for ConnectionSidecarHandler {
         }
     }
 
-    async fn flush_traces(&self, _peer: PeerCredentials) {
-        let flusher = self.server.trace_flusher.clone();
-        if let Err(e) = tokio::spawn(async move { flusher.flush().await }).await {
-            error!("Failed flushing traces: {e:?}");
+    async fn flush(&self, _peer: PeerCredentials, options: SidecarFlushOptions) {
+        if options.traces_and_stats {
+            let flusher = self.server.trace_flusher.clone();
+            if let Err(e) = tokio::spawn(async move { flusher.flush().await }).await {
+                error!("Failed flushing traces: {e:?}");
+            }
+            flush_all_stats_now(&self.server.span_concentrators).await;
+        }
+        if options.telemetry {
+            let workers: Vec<_> = {
+                let clients = self.server.telemetry_clients.inner.lock_or_panic();
+                clients
+                    .values()
+                    .filter_map(|entry| {
+                        entry
+                            .client
+                            .lock_or_panic()
+                            .as_ref()
+                            .map(|c| c.worker.clone())
+                    })
+                    .collect()
+            };
+            futures::future::join_all(workers.into_iter().map(|worker| async move {
+                let _ = worker
+                    .send_msg(TelemetryActions::Lifecycle(
+                        LifecycleAction::FlushMetricAggr,
+                    ))
+                    .await;
+                let _ = worker
+                    .send_msg(TelemetryActions::Lifecycle(LifecycleAction::FlushData))
+                    .await;
+                // now await completion
+                let (tx, rx) = futures::channel::oneshot::channel();
+                let _ = worker.send_msg(TelemetryActions::CollectStats(tx)).await;
+                let _ = rx.await;
+            }))
+            .await;
         }
-        flush_all_stats_now(&self.server.span_concentrators).await;
     }
 
     async fn set_test_session_token(&self, _peer: PeerCredentials, token: String) {