saluki_components/encoders/buffered_incremental/
mod.rs

1use std::time::{Duration, Instant};
2
3use agent_data_plane_config::defaults::DEFAULT_ENCODER_FLUSH_TIMEOUT;
4use async_trait::async_trait;
5use saluki_common::sync::shutdown::ShutdownCoordinator;
6use saluki_core::accounting::{MemoryBounds, MemoryBoundsBuilder};
7use saluki_core::runtime;
8use saluki_core::{
9    components::{encoders::*, BuildContext},
10    data_model::{event::EventType, payload::PayloadType},
11    observability::ComponentMetricsExt,
12    topology::PayloadsDispatcher,
13};
14use saluki_error::GenericError;
15use saluki_metrics::MetricsBuilder;
16use tokio::{pin, select, time::sleep};
17use tracing::{debug, error};
18
19mod telemetry;
20use self::telemetry::ComponentTelemetry;
21
22/// Buffered incremental encoder.
23///
24/// Wraps an `IncrementalEncoder` and drives it with incoming events, allowing buffering by utilizing a configurable
25/// flush timeout. Payloads are encoded on the global thread pool to avoid affecting latency-sensitive tasks.
26pub struct BufferedIncrementalConfiguration<EB> {
27    /// Flush timeout for pending requests.
28    ///
29    /// When the encoder has written events to the in-flight request payload, but it hasn't yet reached the
30    /// payload size limits that would force the payload to be flushed, the encoder will wait for a period of time
31    /// before flushing the in-flight request payload. This allows for the possibility of other events to be processed
32    /// and written into the request payload, thereby maximizing the payload size and reducing the number of requests
33    /// generated and sent overall.
34    ///
35    /// Defaults to 2 seconds.
36    flush_timeout: Duration,
37
38    encoder_builder: EB,
39}
40
41impl<EB> BufferedIncrementalConfiguration<EB>
42where
43    EB: IncrementalEncoderBuilder,
44{
45    /// Creates a new `BufferedIncrementalConfiguration` from the given incremental encoder builder.
46    pub fn from_encoder_builder(encoder_builder: EB) -> Self {
47        Self {
48            flush_timeout: DEFAULT_ENCODER_FLUSH_TIMEOUT,
49            encoder_builder,
50        }
51    }
52}
53
54#[async_trait]
55impl<EB> EncoderBuilder for BufferedIncrementalConfiguration<EB>
56where
57    EB: IncrementalEncoderBuilder + Sync,
58    EB::Output: Send + 'static,
59{
60    fn input_event_type(&self) -> EventType {
61        self.encoder_builder.input_event_type()
62    }
63
64    fn output_payload_type(&self) -> PayloadType {
65        self.encoder_builder.output_payload_type()
66    }
67
68    async fn build(&self, context: BuildContext) -> Result<Box<dyn Encoder + Send>, GenericError> {
69        let metrics_builder = MetricsBuilder::from_component_context(context.component_context());
70        let telemetry = ComponentTelemetry::from_builder(&metrics_builder);
71
72        let encoder = self.encoder_builder.build(context).await?;
73
74        let flush_timeout = match self.flush_timeout {
75            // We always give ourselves a minimum flush timeout of 10ms to allow for some very minimal amount of
76            // batching, while still practically flushing things almost immediately.
77            Duration::ZERO => Duration::from_millis(10),
78            dur => dur,
79        };
80
81        Ok(Box::new(BufferedIncremental {
82            encoder,
83            telemetry,
84            flush_timeout,
85        }))
86    }
87}
88
89impl<EB> MemoryBounds for BufferedIncrementalConfiguration<EB>
90where
91    EB: IncrementalEncoderBuilder,
92{
93    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder) {
94        builder
95            .minimum()
96            .with_single_value::<BufferedIncremental<EB::Output>>("component struct");
97
98        self.encoder_builder.specify_bounds(builder);
99    }
100}
101
102pub struct BufferedIncremental<E> {
103    encoder: E,
104    telemetry: ComponentTelemetry,
105    flush_timeout: Duration,
106}
107
108#[async_trait]
109impl<E> Encoder for BufferedIncremental<E>
110where
111    E: IncrementalEncoder + Send + 'static,
112{
113    async fn run(mut self: Box<Self>, context: EncoderContext) -> Result<(), GenericError> {
114        let Self {
115            encoder,
116            telemetry,
117            flush_timeout,
118        } = *self;
119
120        // Run our encoder task on the worker pool.
121        //
122        // The encoder task owns _everything_ -- health checking, encoding, dispatching payloads, etc -- so we just
123        // end up as a glorified watchdog here, where we block until the encoder task has completed. The encoder task
124        // itself is what waits for the shutdown signal and all of that.
125        //
126        // We're abusing `ShutdownCoordinator` here to basically detect when the encoder task exits: we immediately
127        // call `shutdown_and_wait` after spawning the task, which will wait until the `ShutdownHandle` we've given to
128        // the encoder tasks is dropped... which only happens when the task exits.
129        let worker_pool = context.topology_context().global_thread_pool().clone();
130
131        let mut shutdown_coordinator = ShutdownCoordinator::default();
132        let task_shutdown_handle = shutdown_coordinator.register();
133
134        runtime::worker("incremental_encoder", async move {
135            // TODO: Conceptually, all components downstream of sources/relays should drain their incoming events
136            // consumer until it's empty, so on and so forth until forwarders/destinations have done so and nothing
137            // is left to process.
138            //
139            // However, this isn't feasible if we ever want to support restarting individual components or stitching
140            // in new components to a topology... which we eventually do. We'd have to respond to shutdown (which
141            // we're ignoring here by aliasing the handle as `_shutdown`) to shut down in a timely fashion.
142            //
143            // We should consider a mechanism to expose consumers/dispatchers such that they're borrowed in an owned
144            // fashion, allowing a component to hold on to them when running, but then effectively release them when
145            // they exit. This would allow us to stop a component, and without dropping its interconnect channel and
146            // losing all of the events or payloads within it, reconnect it to the new instance of the component (or
147            // whatever component takes its place, etc) which would allow us to more cleanly respond to shutdown
148            // here as more of an exceptional thing: an immediate stoppage of work, without throwing away _pending_
149            // work.
150            let _task_shutdown_handle = task_shutdown_handle;
151            run_incremental_encoder(context, encoder, telemetry, flush_timeout).await
152        })
153        .on_runtime(worker_pool)
154        .needed()
155        .spawn();
156
157        debug!("Buffered Incremental encoder started.");
158
159        shutdown_coordinator.shutdown_and_wait().await;
160
161        debug!("Buffered Incremental encoder stopped.");
162
163        Ok(())
164    }
165}
166
167async fn run_incremental_encoder<E>(
168    mut context: EncoderContext, mut encoder: E, telemetry: ComponentTelemetry, flush_timeout: Duration,
169) -> Result<(), GenericError>
170where
171    E: IncrementalEncoder,
172{
173    let mut health = context.take_health_handle();
174
175    health.mark_ready();
176
177    let mut pending_flush = false;
178    let mut pending_events = 0;
179    let pending_flush_timeout = sleep(flush_timeout);
180    pin!(pending_flush_timeout);
181
182    loop {
183        select! {
184            _ = health.live() => continue,
185            maybe_event_buffer = context.events().next() => {
186                // Break out of our loop when the events channel is closed.
187                let event_buffer = match maybe_event_buffer {
188                    Some(event_buffer) => event_buffer,
189                    None => break,
190                };
191
192                for event in event_buffer {
193                    // Try to process the event.
194                    //
195                    // If we're informed that we need to flush, we'll hold on to this event before triggering a flush and then
196                    // retry processing it after flushing.
197                    let event_to_retry = match encoder.process_event(event).await? {
198                        ProcessResult::Continue => {
199                            pending_events += 1;
200                            continue;
201                        }
202                        ProcessResult::FlushRequired(event) => event,
203                    };
204
205                    // Flush the encoder, waiting any payloads it has generated.
206                    flush_encoder(&mut encoder, context.dispatcher(), &telemetry, &mut pending_events).await?;
207
208                    // Now try to process the event again.
209                    //
210                    // If this fails, then we drop the event because it's a logical bug to not be able to encode an event after
211                    // flushing, and we don't want to get stuck in an infinite loop.
212                    match encoder.process_event(event_to_retry).await? {
213                        ProcessResult::Continue => pending_events += 1,
214                        ProcessResult::FlushRequired(_) => {
215                            error!("Failed to process event after flushing.");
216                            telemetry.events_dropped_encoder().increment(1);
217                        },
218                    }
219                }
220
221                debug!("Processed event buffer.");
222
223                // If we're not already pending a flush, we'll start the countdown.
224                if !pending_flush {
225                    pending_flush_timeout.as_mut().reset(tokio::time::Instant::now() + flush_timeout);
226                    pending_flush = true;
227                }
228            },
229            _ = &mut pending_flush_timeout, if pending_flush => {
230                debug!("Flushing encoder of any pending payload(s).");
231
232                pending_flush = false;
233
234                flush_encoder(&mut encoder, context.dispatcher(), &telemetry, &mut pending_events).await?;
235
236                debug!("All pending payloads flushed.");
237            }
238        }
239    }
240
241    // Do a final flush since we may have had a pending payloads before breaking out of the loop.
242    flush_encoder(&mut encoder, context.dispatcher(), &telemetry, &mut pending_events).await?;
243
244    Ok(())
245}
246
247/// Flushes the encoder, recording flush telemetry for the events processed since the last flush.
248///
249/// Flushes with no pending events are not recorded, so the last-flush telemetry always describes a flush that sent
250/// something.
251async fn flush_encoder<E>(
252    encoder: &mut E, dispatcher: &PayloadsDispatcher, telemetry: &ComponentTelemetry, pending_events: &mut u64,
253) -> Result<(), GenericError>
254where
255    E: IncrementalEncoder,
256{
257    let flush_start = Instant::now();
258    encoder.flush(dispatcher).await?;
259
260    if *pending_events > 0 {
261        telemetry.record_flush(*pending_events, flush_start.elapsed());
262        *pending_events = 0;
263    }
264
265    Ok(())
266}
saluki_components/encoders/buffered_incremental/mod.rs

saluki_components/encoders/buffered_incremental/
mod.rs
