0.2.7: tracing support, dynamic call window

Author: hackermondev

Cargo.lock

@@ -9,13 +9,14 @@ cargo add batched
 Or add this to your `Cargo.toml`:
 ```toml
 [dependencies]
-batched = "0.2.3"
+batched = "0.2.7"
 ```
 
 ## #[batched]
-- **window**: Minimum amount of time (in milliseconds) the background thread waits before processing a batch.
 - **limit**: Maximum amount of items that can be grouped and processed in a single batch.
 - **concurrent**: Maximum amount of concurrent batched tasks running (default: `Infinity`)
+- **window**: Minimum amount of time (in milliseconds) the background thread waits before processing a batch.
+- **window[x]**: Minimum amount of time (in milliseconds) the background thread waits before processing a batch when latest buffer size is <= x
 
 The target function must have a single argument, a vector of items (`Vec<T>`). 
 
@@ -39,6 +40,12 @@ impl A {
 }
 ```
 
+## Tracing
+### [`tracing_span`]
+This feature automatically adds tracing spans to call functions for batched requests (`x`, `x_multiple`).
+
+## [`tracing_opentelemetry`]
+This feature adds support for linking spans from callers to the inner batched call when using OpenTelemetry. Depending on whether your OpenTelemetry client supports it, you should be able to see the linked span to the batched call. 
 
 ## Examples
 

README.md

@@ -1,14 +1,21 @@
 [package]
 name = "batched"
 description = "rust macro util for batching expensive operations"
-version = "0.2.5"
+version = "0.2.7"
 edition = "2024"
 license = "MIT"
 readme = "../README.md"
 repository = "https://github.com/hackermondev/batched"
 keywords = ["batch", "performance", "efficient"]
 
+[features]
+default = []
+tracing_span = ["batched_derive/tracing_span"]
+tracing_opentelemetry = ["opentelemetry", "tracing-opentelemetry"]
 
 [dependencies]
 anyhow = "1.0.98"
-batched_derive = { version = "0.2.5", path = "../batched_derive" }
+batched_derive = { version = "0.2.7", path = "../batched_derive" }
+opentelemetry = { version = "0.30.0", optional = true }
+tracing = "0.1.41"
+tracing-opentelemetry = { version = "0.31.0", optional = true }

batched/Cargo.toml

@@ -1,5 +1,7 @@
 #![feature(try_trait_v2)]
 #![feature(negative_impls)]
+#![feature(specialization)]
 
 pub mod error;
 pub use batched_derive::batched;
+pub mod tracing;

batched/src/lib.rs

@@ -0,0 +1,22 @@
+pub use tracing::*;
+
+pub trait TracingSpan {
+    fn link_span(&mut self, span: &Span);
+}
+
+impl<N> TracingSpan for N {
+    default fn link_span(&mut self, _span: &Span) {}
+}
+
+#[cfg(feature = "tracing_opentelemetry")]
+impl<T: tracing_opentelemetry::OpenTelemetrySpanExt> TracingSpan for T {
+    fn link_span(&mut self, span: &Span) {
+        self.add_link(
+            opentelemetry::trace::TraceContextExt::span(
+                &tracing_opentelemetry::OpenTelemetrySpanExt::context(span),
+            )
+            .span_context()
+            .clone(),
+        );
+    }
+}

batched/src/tracing.rs

@@ -1,13 +1,17 @@
 [package]
 name = "batched_derive"
 description = "rust macro util for batching expensive operations"
-version = "0.2.5"
+version = "0.2.7"
 edition = "2024"
 license = "MIT"
 readme = "../README.md"
 repository = "https://github.com/hackermondev/batched"
 keywords = ["batch", "performance", "efficient"]
 
+[features]
+default = []
+tracing_span = []
+
 [lib]
 proc-macro = true
 

batched_derive/Cargo.toml

@@ -11,15 +11,15 @@ struct Identifiers {
     public_interface_multiple: Ident,
     inner_batched: Ident,
     executor_producer_channel: Ident,
-    executor_background_fn: Ident
+    executor_background_fn: Ident,
 }
 
 fn build_identifiers(call_function: &Function) -> Identifiers {
     let id = &call_function.identifier;
 
     let public_interface = format_ident!("{id}");
     let public_interface_multiple = format_ident!("{id}_multiple");
-    let inner_batched = format_ident!("__{id}__");
+    let inner_batched = format_ident!("{id}__batched");
 
     let executor_producer_channel = format_ident!("BATCHED_{}", id.to_uppercase());
     let executor_background_fn = format_ident!("spawn_executor_{id}");
@@ -29,7 +29,7 @@ fn build_identifiers(call_function: &Function) -> Identifiers {
         public_interface_multiple,
         inner_batched,
         executor_producer_channel,
-        executor_background_fn
+        executor_background_fn,
     }
 }
 
@@ -104,14 +104,9 @@ pub fn build_code(call_function: Function, options: Attributes) -> TokenStream {
     let cast_result_error = match &call_function.returned.result_type {
         FunctionResultType::Raw(_) => None,
         FunctionResultType::VectorRaw(_) => None,
-        FunctionResultType::Result(_, _, inner_shared_error) => {
-            match inner_shared_error {
-                Some(inner_shared_error) => Some(quote! {
-                    let result = result.map_err(|e: #inner_shared_error| e.into());
-                }),
-                _ => None
-            }
-        }
+        FunctionResultType::Result(_, _, inner_shared_error) => inner_shared_error.as_ref().map(|inner_shared_error| quote! {
+                let result = result.map_err(|e: #inner_shared_error| e.into());
+            }),
     };
 
     let executor = build_executor(&identifiers, &call_function, &options);
@@ -121,6 +116,11 @@ pub fn build_code(call_function: Function, options: Attributes) -> TokenStream {
     let public_interface = identifiers.public_interface;
     let public_interface_multiple = identifiers.public_interface_multiple;
 
+    #[cfg(feature = "tracing_span")]
+    let tracing_span = quote! { #[tracing::instrument(skip_all)] };
+    #[cfg(not(feature = "tracing_span"))]
+    let tracing_span = quote! {};
+
     quote! {
         #executor
 
@@ -132,17 +132,20 @@ pub fn build_code(call_function: Function, options: Attributes) -> TokenStream {
             result
         }
 
+        #tracing_span
         #visibility async fn #public_interface(#arg_name: #arg_type) -> #public_interface_returned_type {
             let mut result = #public_interface_multiple(vec![#arg_name]).await;
             #public_interface_return_result
         }
 
+        #tracing_span
         #visibility async fn #public_interface_multiple(#arg_name: Vec<#arg_type>) -> #public_interface_multiple_returned_type {
             let channel = &#executor_producer_channel;
             let channel = channel.get_or_init(async || { #executor_background_fn().await }).await;
 
             let (response_channel_sender, mut response_channel_recv) = ::tokio::sync::mpsc::channel(1);
-            channel.send((#arg_name, response_channel_sender)).await
+            let span = ::batched::tracing::Span::current();
+            channel.send((#arg_name, span, response_channel_sender)).await
                 .expect("batched function panicked (send)");
 
             let result = response_channel_recv.recv().await
@@ -159,8 +162,17 @@ fn build_executor(
     options: &Attributes,
 ) -> TokenStream {
     let capacity = options.limit;
-    let window = options.window;
     let concurrent_limit = options.concurrent_limit.unwrap_or(SEMAPHORE_MAX_PERMITS);
+    let default_window = options.default_window;
+
+    let windows = options.windows.iter();
+    let windows = windows.map(|(call_size, call_window)| {
+        quote! { windows.insert(#call_size, #call_window); }
+    });
+    let windows = quote! {
+        let mut windows = ::std::collections::BTreeMap::<u64, u64>::new();
+        #(#windows)*
+    };
 
     let arg_type = &call_function.batched_arg_type;
     let returned_type_plural = match &call_function.returned.result_type {
@@ -196,19 +208,22 @@ fn build_executor(
             let result = result.clone();
         }
     };
-    let channel_type =
-        quote! { (Vec<#arg_type>, ::tokio::sync::mpsc::Sender<#returned_type_plural>) };
+    let channel_type = quote! { (Vec<#arg_type>, ::batched::tracing::Span, ::tokio::sync::mpsc::Sender<#returned_type_plural>) };
 
     let inner_batched = &identifiers.inner_batched;
+    let batched_span_name = inner_batched.to_string();
     let executor_producer_channel = &identifiers.executor_producer_channel;
     let executor_background_fn = &identifiers.executor_background_fn;
+
     quote! {
         static #executor_producer_channel:
             ::tokio::sync::OnceCell<::tokio::sync::mpsc::Sender<#channel_type>> = ::tokio::sync::OnceCell::const_new();
 
         async fn #executor_background_fn() -> ::tokio::sync::mpsc::Sender<#channel_type> {
             let capacity = #capacity;
-            let window = ::tokio::time::Duration::from_millis(#window);
+            let default_window = #default_window;
+            #windows
+
 
             let (sender, mut receiver) = tokio::sync::mpsc::channel(capacity);
             tokio::task::spawn(async move {
@@ -217,31 +232,39 @@ fn build_executor(
                 loop {
                     let mut data_buffer = Vec::with_capacity(capacity);
                     let mut return_channels: Vec<(::tokio::sync::mpsc::Sender<#returned_type_plural>, usize)> = vec![];
-                    let mut timer = tokio::time::interval(window);
+                    let mut waiting_spans: Vec<::batched::tracing::Span> = vec![];
+
+                    let window_start = ::std::time::Instant::now();
 
                     loop {
+                        let window = windows.iter()
+                            .find(|(max_calls, _)|  **max_calls >= data_buffer.len() as u64)
+                            .map(|(_, window)| window);
+                        let window = window.unwrap_or(&default_window);
+                        let window = ::std::time::Duration::from_millis(*window as u64);
+
+                        let window_end = window_start + window;
+                        let remaining_duration = window_end.duration_since(std::time::Instant::now());
+
                         tokio::select! {
                             event = receiver.recv() => {
                                 if event.is_none() {
                                     return;
                                 }
 
-                                if data_buffer.is_empty() {
-                                    timer.reset();
-                                }
-
                                 let event: #channel_type = event.unwrap();
-                                let (mut data, channel) = event;
+                                let (mut data, span, channel) = event;
 
                                 return_channels.push((channel, data.len()));
+                                waiting_spans.push(span);
                                 data_buffer.append(&mut data);
 
                                 if data_buffer.len() >= capacity {
                                     break;
                                 }
                             }
 
-                            _ = timer.tick() => {
+                            _ = ::tokio::time::sleep(remaining_duration) => {
                                 break;
                             }
                         }
@@ -252,15 +275,23 @@ fn build_executor(
                     }
 
                     let mut data = vec![];
+                    let mut spans = vec![];
                     let mut channels = vec![];
 
                     std::mem::swap(&mut data, &mut data_buffer);
+                    std::mem::swap(&mut spans, &mut waiting_spans);
                     std::mem::swap(&mut channels, &mut return_channels);
 
                     let permit = semaphore.clone().acquire_owned().await.unwrap();
                     tokio::task::spawn(async move {
                         let _permit = permit;
-                        let mut result = #inner_batched(data).await;
+                        let batched_span = ::batched::tracing::info_span!(#batched_span_name, count = data.len());
+                        for mut span in spans {
+                            ::batched::tracing::TracingSpan::link_span(&mut span, &batched_span);
+                        }
+
+                        let future = #inner_batched(data);
+                        let mut result = ::batched::tracing::Instrument::instrument(future, batched_span).await;
                         for (channel, count) in channels {
                             #handle_result
                             let _ = channel.try_send(result);
@@ -282,7 +313,9 @@ fn function_flags(function: &Function) -> (bool, bool) {
     match &function.returned.result_type {
         FunctionResultType::Result(result, _, _) => {
             is_result = true;
-            if let FunctionResultType::VectorRaw(_) = result.result_type { is_vec = true };
+            if let FunctionResultType::VectorRaw(_) = result.result_type {
+                is_vec = true
+            };
         }
         FunctionResultType::VectorRaw(_) => is_vec = true,
         _ => {}