fix(sidecar): Handle backpressure more gracefully (#1682)

Only drop some operations on the floor, but never those setting the state of the connection. Otherwise, under load the state may end up in unexpected state.

Also do not leak file descriptors on partially received messages.

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

Author: bwoebi

datadog-ipc/Cargo.toml

@@ -25,6 +25,7 @@ tarpc = { path = "./tarpc", default-features = false, features = ["serde-transpo
 
 libdd-common = { path = "../libdd-common" }
 datadog-ipc-macros = { path = "../datadog-ipc-macros" }
+tracing = { version = "0.1", default-features = false }
 
 [dev-dependencies]
 criterion = "0.5"

datadog-ipc/src/platform/unix/channel/metadata.rs

@@ -4,6 +4,7 @@
 use std::{
     collections::VecDeque,
     io,
+    os::fd::IntoRawFd,
     os::unix::prelude::{AsRawFd, FromRawFd, RawFd},
 };
 
@@ -17,8 +18,8 @@ use crate::{
 #[derive(Debug)]
 pub struct ChannelMetadata {
     fds_to_send: Vec<PlatformHandle<OwnedFd>>,
-    fds_received: VecDeque<RawFd>,
-    pid: libc::pid_t, // must always be set to current Process ID
+    fds_received: VecDeque<OwnedFd>, // Store as OwnedFd to prevent leaking them
+    pid: libc::pid_t,                // must always be set to current Process ID
 }
 
 impl Default for ChannelMetadata {
@@ -80,7 +81,13 @@ impl ChannelMetadata {
     }
 
     pub(crate) fn receive_fds(&mut self, fds: &[RawFd]) {
-        self.fds_received.append(&mut fds.to_vec().into());
+        self.fds_received.append(
+            &mut fds
+                .iter()
+                .map(|fd| unsafe { OwnedFd::from_raw_fd(*fd) })
+                .collect::<Vec<_>>()
+                .into(),
+        );
     }
 
     pub(crate) fn find_handle<T>(&mut self, hint: &PlatformHandle<T>) -> Option<PlatformHandle<T>> {
@@ -90,6 +97,6 @@ impl ChannelMetadata {
 
         let fd = self.fds_received.pop_front();
 
-        fd.map(|fd| unsafe { PlatformHandle::from_raw_fd(fd) })
+        fd.map(|fd| unsafe { PlatformHandle::from_raw_fd(fd.into_raw_fd()) })
     }
 }

datadog-ipc/src/sequential.rs

@@ -1,21 +1,28 @@
 // Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
-use futures::{ready, Future, Stream};
+use futures::{Future, Stream};
 use std::fmt::Debug;
 use std::{
     pin::Pin,
     task::{Context, Poll},
 };
 use tarpc::server::{Channel, InFlightRequest, Requests, Serve};
+use tokio::sync::mpsc::error::SendError;
+use tokio::sync::mpsc::OwnedPermit;
 
 #[allow(type_alias_bounds)]
 type Request<S, C: Channel> = (S, InFlightRequest<C::Req, C::Resp>);
 
+type PendingPermit<S, C> = Pin<
+    Box<dyn Future<Output = Result<OwnedPermit<Request<S, C>>, SendError<()>>> + Send + 'static>,
+>;
+
 /// Replaces tarpc::server::Channel::execute which spawns one task per message with an executor
 /// that spawns a single worker and queues requests for this task.
 ///
-/// If the queue is full, the requests is dropped and will be cancelled by tarpc.
+/// If the queue is full, the request is dropped and will be cancelled by tarpc unless
+/// `with_backpressure` is configured for that request type.
 pub fn execute_sequential<C, S>(
     reqs: Requests<C>,
     serve: S,
@@ -43,6 +50,8 @@ where
         inner: reqs,
         serve,
         tx,
+        backpressure: |_| false,
+        pending: None,
     }
 }
 
@@ -55,6 +64,11 @@ where
     inner: Requests<C>,
     serve: S,
     tx: tokio::sync::mpsc::Sender<Request<S, C>>,
+    /// Returns true for requests that must not be dropped when the queue is full.
+    /// The executor will pause reading new requests and wait for channel space instead.
+    backpressure: fn(&C::Req) -> bool,
+    /// Pending channel-space reservation for a backpressure request.
+    pending: Option<(PendingPermit<S, C>, Request<S, C>)>,
 }
 
 impl<C, S> Future for SequentialExecutor<C, S>
@@ -68,21 +82,50 @@ where
     type Output = anyhow::Result<()>;
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        while let Some(response_handler) = ready!(self.as_mut().project().inner.poll_next(cx)) {
-            match response_handler {
-                Ok(resp) => {
-                    let server = self.serve.clone();
-                    if let Err(_err) = self.as_ref().tx.try_send((server, resp)) {
-                        // TODO: should we log something in case we drop the request on the floor?
+        loop {
+            // First flush any pending backpressure send before reading new requests.
+            {
+                let this = self.as_mut().project();
+                if let Some((fut, _)) = this.pending.as_mut() {
+                    match fut.as_mut().poll(cx) {
+                        Poll::Pending => return Poll::Pending,
+                        Poll::Ready(Err(_)) => return Poll::Ready(Ok(())), // worker dropped
+                        Poll::Ready(Ok(permit)) => {
+                            #[allow(clippy::unwrap_used)] // we've just checked this
+                            let (_, item) = this.pending.take().unwrap();
+                            permit.send(item);
+                            // fall through to poll next request
+                        }
                     }
                 }
-                Err(e) => {
-                    // TODO: should we log something in case we drop the request on the floor?
-                    return Poll::Ready(Err(e.into()));
+            }
+
+            // Read the next request off the transport.
+            match self.as_mut().project().inner.poll_next(cx) {
+                Poll::Ready(Some(Ok(resp))) => {
+                    let backpressured = (self.backpressure)(&resp.get().message);
+                    match self.as_ref().tx.try_send((self.serve.clone(), resp)) {
+                        Ok(()) => {} // loop to pick up the next request
+                        Err(err) => {
+                            let (_, resp) = err.into_inner();
+                            if backpressured {
+                                let fut = Box::pin(self.as_ref().tx.clone().reserve_owned());
+                                *self.as_mut().project().pending =
+                                    Some((fut, (self.serve.clone(), resp)));
+                            } else {
+                                tracing::warn!(
+                                    "Dropping {:?}: sequential executor queue is full",
+                                    resp.get().message
+                                );
+                            }
+                        }
+                    }
                 }
+                Poll::Ready(Some(Err(e))) => return Poll::Ready(Err(e.into())),
+                Poll::Ready(None) => return Poll::Ready(Ok(())),
+                Poll::Pending => return Poll::Pending,
             }
         }
-        Poll::Ready(Ok(()))
     }
 }
 
@@ -97,4 +140,11 @@ where
         std::mem::swap(&mut self.tx, &mut sender);
         sender
     }
+
+    /// Configures a predicate that identifies requests which must not be dropped when the queue
+    /// is full. For those requests the executor will pause reading and wait for channel space.
+    pub fn with_backpressure(mut self, backpressure: fn(&C::Req) -> bool) -> Self {
+        self.backpressure = backpressure;
+        self
+    }
 }

datadog-sidecar-macros/src/lib.rs

@@ -28,25 +28,46 @@ fn snake_to_camel(ident_str: &str) -> String {
 }
 
 #[proc_macro_attribute]
-pub fn extract_request_id(_attr: TokenStream, mut input: TokenStream) -> TokenStream {
-    let item: ItemTrait = syn::parse(input.clone()).unwrap();
+pub fn extract_request_id(_attr: TokenStream, input: TokenStream) -> TokenStream {
+    let mut item: ItemTrait = syn::parse(input).unwrap();
     let name = &format_ident!("{}Request", item.ident);
     let mut arms: Vec<Arm> = vec![];
-    for inner in item.items {
+    let mut backpressure_variants: Vec<Ident> = vec![];
+
+    for inner in item.items.iter_mut() {
         if let TraitItem::Fn(func) = inner {
-            for any_arg in func.sig.inputs {
+            // Strip #[force_backpressure] and record which methods carry it.
+            let had_force_backpressure = func.attrs.iter().any(|attr| {
+                attr.meta
+                    .path()
+                    .get_ident()
+                    .is_some_and(|i| i == "force_backpressure")
+            });
+            func.attrs.retain(|attr| {
+                attr.meta
+                    .path()
+                    .get_ident()
+                    .is_none_or(|i| i != "force_backpressure")
+            });
+
+            let method = Ident::new(
+                &snake_to_camel(&func.sig.ident.to_string()),
+                Span::mixed_site(),
+            );
+
+            if had_force_backpressure {
+                backpressure_variants.push(method.clone());
+            }
+
+            for any_arg in &func.sig.inputs {
                 if let Typed(arg) = any_arg {
-                    if let Pat::Ident(ident) = *arg.pat {
+                    if let Pat::Ident(ident) = &*arg.pat {
                         let matched_enum_type = match ident.ident.to_string().as_str() {
                             "session_id" => Some(format_ident!("SessionId")),
                             "instance_id" => Some(format_ident!("InstanceId")),
                             _ => None,
                         };
                         if let Some(enum_type) = matched_enum_type {
-                            let method = Ident::new(
-                                &snake_to_camel(&func.sig.ident.to_string()),
-                                Span::mixed_site(),
-                            );
                             arms.push(parse_quote! {
                                 #name::#method { #ident, .. } => RequestIdentifier::#enum_type(#ident.clone())
                             });
@@ -56,7 +77,16 @@ pub fn extract_request_id(_attr: TokenStream, mut input: TokenStream) -> TokenSt
             }
         }
     }
-    input.extend(TokenStream::from(quote! {
+
+    let backpressure_body = if backpressure_variants.is_empty() {
+        quote! { false }
+    } else {
+        quote! { matches!(self, #(#name::#backpressure_variants { .. })|*) }
+    };
+
+    TokenStream::from(quote! {
+        #item
+
         impl RequestIdentification for tarpc::Request<#name> {
             fn extract_identifier(&self) -> RequestIdentifier {
                 match &self.message {
@@ -67,8 +97,14 @@ pub fn extract_request_id(_attr: TokenStream, mut input: TokenStream) -> TokenSt
                 }
             }
         }
-    }));
-    input
+
+        impl #name {
+            /// Returns true if this request variant was annotated with `#[force_backpressure]`.
+            pub fn requires_backpressure(&self) -> bool {
+                #backpressure_body
+            }
+        }
+    })
 }
 
 struct EnvOrDefault {

datadog-sidecar/src/service/sidecar_interface.rs

@@ -60,6 +60,7 @@ pub trait SidecarInterface {
     /// * `session_id` - The ID of the session.
     /// * `pid` - The pid of the sidecar client.
     /// * `config` - The configuration to be set.
+    #[force_backpressure]
     async fn set_session_config(
         session_id: String,
         remote_config_notify_target: RemoteConfigNotifyTarget,
@@ -73,20 +74,23 @@ pub trait SidecarInterface {
     ///
     /// * `session_id` - The ID of the session.
     /// * `process_tags` - The process tags string.
+    #[force_backpressure]
     async fn set_session_process_tags(session_id: String, process_tags: String);
 
     /// Shuts down a runtime.
     ///
     /// # Arguments
     ///
     /// * `instance_id` - The ID of the instance.
+    #[force_backpressure]
     async fn shutdown_runtime(instance_id: InstanceId);
 
     /// Shuts down a session.
     ///
     /// # Arguments
     ///
     /// * `session_id` - The ID of the session.
+    #[force_backpressure]
     async fn shutdown_session(session_id: String);
 
     /// Sends a trace via shared memory.
@@ -165,6 +169,7 @@ pub trait SidecarInterface {
     /// * `global_tags` - Global tags which need to be propagated.
     /// * `dynamic_instrumentation_state` - Whether dynamic instrumentation is enabled, disabled or
     ///   not set.
+    #[force_backpressure]
     async fn set_universal_service_tags(
         instance_id: InstanceId,
         queue_id: QueueId,
@@ -182,6 +187,7 @@ pub trait SidecarInterface {
     /// * `queue_id` - The unique identifier for the trace context.
     /// * `dynamic_instrumentation_state` - Whether dynamic instrumentation is enabled, disabled or
     ///   not set.
+    #[force_backpressure]
     async fn set_request_config(
         instance_id: InstanceId,
         queue_id: QueueId,
@@ -197,6 +203,7 @@ pub trait SidecarInterface {
     async fn send_dogstatsd_actions(instance_id: InstanceId, actions: Vec<DogStatsDActionOwned>);
 
     /// Flushes any outstanding traces queued for sending.
+    #[force_backpressure]
     async fn flush_traces();
 
     /// Sets x-datadog-test-session-token on all requests for the given session.
@@ -208,19 +215,22 @@ pub trait SidecarInterface {
     async fn set_test_session_token(session_id: String, token: String);
 
     /// Sends a ping to the service.
+    #[force_backpressure]
     async fn ping();
 
     /// Dumps the current state of the service.
     ///
     /// # Returns
     ///
     /// A string representation of the current state of the service.
+    #[force_backpressure]
     async fn dump() -> String;
 
     /// Retrieves the current statistics of the service.
     ///
     /// # Returns
     ///
     /// A string representation of the current statistics of the service.
+    #[force_backpressure]
     async fn stats() -> String;
 }

datadog-sidecar/src/service/sidecar_server.rs

@@ -152,7 +152,8 @@ impl SidecarServer {
             server.requests(),
             self.clone().serve(),
             500,
-        );
+        )
+        .with_backpressure(SidecarInterfaceRequest::requires_backpressure);
         let (tx, rx) = tokio::sync::mpsc::channel::<_>(100);
         let tx = executor.swap_sender(tx);