chore(components, core, io): cutover Datadog Traces encoder and Checks IPC source to full supervision

lib/saluki-components/src/encoders/datadog/traces/mod.rs

@@ -10,7 +10,6 @@ use http::{uri::PathAndQuery, HeaderName, HeaderValue, Method, Uri};
 use piecemeal::{ScratchBuffer, ScratchWriter};
 use saluki_common::collections::FastHashMap;
 use saluki_common::strings::StringBuilder;
-use saluki_common::task::HandleExt as _;
 use saluki_context::tags::TagSet;
 use saluki_core::accounting::{MemoryBounds, MemoryBoundsBuilder};
 use saluki_core::data_model::event::trace::AttributeValue;
@@ -172,7 +171,6 @@ impl EncoderBuilder for DatadogTraceConfiguration {
         let default_hostname = MetaString::from(default_hostname);
 
         // Create request builder for traces which is used to generate HTTP requests.
-
         let mut trace_rb = RequestBuilder::new(
             TraceEndpointEncoder::new(
                 default_hostname,
@@ -239,24 +237,32 @@ impl Encoder for DatadogTrace {
 
         let mut health = context.take_health_handle();
 
-        // The encoder runs two async loops, the main encoder loop and the request builder loop,
-        // this channel is used to send events from the main encoder loop to the request builder loop safely.
+        // Run our request builder on a background worker pool task.
+        //
+        // We'll shuttle event buffers to the task, and get back encoded payloads ready to be sent. We do this to
+        // isolate the heavy encoding/compression work to dedicated threads to avoid causing schedulinmg latency spikes
+        // on the main async runtime.
+        //
+        // We also ignore shutdown in the request builder task since we want the task to drain the incoming event
+        // buffers channel until it's empty and all payloads have been sent out.
         let (events_tx, events_rx) = mpsc::channel(8);
-        // adds a channel to send payloads to the dispatcher and a channel to receive them.
         let (payloads_tx, mut payloads_rx) = mpsc::channel(8);
         let request_builder_fut = run_request_builder(trace_rb, telemetry, events_rx, payloads_tx, flush_timeout);
-        // Spawn the request builder task on the global thread pool, this task is responsible for encoding traces and flushing requests.
-        let request_builder_handle = context
-            .topology_context()
-            .global_thread_pool() // Use the shared Tokio runtime thread pool.
-            .spawn_traced_named("dd-traces-request-builder", request_builder_fut);
+
+        context
+            .spawner()
+            .noninterruptible("request_builder", |_shutdown| request_builder_fut)
+            .on_worker_pool()
+            .spawn()
+            .await
+            .error_context("Failed to spawn request builder task.")?;
 
         health.mark_ready();
         debug!("Datadog Trace encoder started.");
 
         loop {
             select! {
-                biased; // makes the branches of the select statement be evaluated in order.
+                biased;
 
                 _ = health.live() => continue,
                 maybe_payload = payloads_rx.recv() => match maybe_payload {
@@ -286,13 +292,6 @@ impl Encoder for DatadogTrace {
             }
         }
 
-        // Request build task should now be stopped.
-        match request_builder_handle.await {
-            Ok(Ok(())) => debug!("Request builder task stopped."),
-            Ok(Err(e)) => error!(error = %e, "Request builder task failed."),
-            Err(e) => error!(error = %e, "Request builder task panicked."),
-        }
-
         debug!("Datadog Trace encoder stopped.");
 
         Ok(())

lib/saluki-components/src/sources/checks_ipc/mod.rs

@@ -11,7 +11,6 @@ use datadog_protos::checks::{
     service_check::{ServiceCheck as ProtoServiceCheck, Status as ServiceCheckStatus},
     SendCheckPayloadRequest, SendCheckPayloadResponse,
 };
-use saluki_common::task::HandleExt as _;
 use saluki_config::GenericConfiguration;
 use saluki_context::tags::{Tag, TagSet};
 use saluki_context::Context;
@@ -26,13 +25,12 @@ use saluki_core::{
     components::{sources::*, ComponentContext},
     data_model::event::log::LogStatus,
 };
-use saluki_error::{generic_error, GenericError};
-use saluki_io::net::ListenAddress;
+use saluki_error::{generic_error, ErrorContext as _, GenericError};
+use saluki_io::net::{server::grpc::GrpcServer, ListenAddress};
 use serde::Deserialize;
 use stringtheory::MetaString;
 use tokio::sync::mpsc;
 use tokio::{pin, select};
-use tonic::transport::Server;
 use tonic::{Response, Status};
 use tracing::{debug, trace, warn};
 
@@ -113,19 +111,21 @@ impl Source for ChecksIPC {
 
         let (events_tx, mut events_rx) = mpsc::channel(16);
 
-        let grpc_server = Server::builder().add_service(ChecksServer::new(ChecksService {
+        let ListenAddress::Tcp(grpc_socket_addr) = grpc_endpoint else {
+            return Err(generic_error!("OTLP gRPC endpoint must be a TCP address."));
+        };
+
+        let grpc_server = GrpcServer::new(grpc_socket_addr).add_service(ChecksServer::new(ChecksService {
             events_tx,
             default_hostname,
         }));
 
-        let grpc_socket_addr = match grpc_endpoint {
-            ListenAddress::Tcp(addr) => addr,
-            _ => return Err(generic_error!("OTLP gRPC endpoint must be a TCP address.")),
-        };
         context
-            .topology_context()
-            .global_thread_pool()
-            .spawn_traced_named("checks-ipc-grpc-server", grpc_server.serve(grpc_socket_addr));
+            .spawner()
+            .supervisable(grpc_server)
+            .spawn()
+            .await
+            .error_context("Failed to spawn Checks IPC gRPC server.")?;
 
         health.mark_ready();
         debug!("Checks IPC source started.");

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

@@ -13,7 +13,7 @@ pub mod sources;
 pub mod transforms;
 
 mod spawner;
-pub use self::spawner::{ChildBuilder, ComponentSpawner};
+pub use self::spawner::{BuilderState, ChildBuilder, ComponentSpawner, OneShot, Restartable};
 
 #[cfg(any(test, feature = "test-util"))]
 pub mod test_util;

lib/saluki-core/src/components/sources/context.rs

@@ -52,6 +52,15 @@ impl SourceContext {
         self.shutdown_handle = Some(shutdown_handle);
     }
 
+    /// Installs the shutdown handle for this source context, for tests.
+    ///
+    /// The topology runtime does this itself before a source runs, using the shutdown signal of the component's
+    /// dedicated supervisor. A test that drives a source through a real shutdown has to stand in for it.
+    #[cfg(any(test, feature = "test-util"))]
+    pub fn set_shutdown_handle_for_test(&mut self, shutdown_handle: ShutdownHandle) {
+        self.set_shutdown_handle(shutdown_handle);
+    }
+
     /// Consumes the shutdown handle of this source context.
     ///
     /// # Panics