From 1961f0ae35eba402f9cea913bd9f5f7c20211be2 Mon Sep 17 00:00:00 2001 From: Pavlos Rontidis Date: Thu, 3 Sep 2026 12:41:42 -0400 Subject: [PATCH 1/3] refactor(codecs): assemble final GELF chunk directly --- .../src/decoding/framing/chunked_gelf.rs | 78 +++++++++++-------- 1 file changed, 45 insertions(+), 33 deletions(-) diff --git a/lib/codecs/src/decoding/framing/chunked_gelf.rs b/lib/codecs/src/decoding/framing/chunked_gelf.rs index 1f3046df064d7..bf2981522e056 100644 --- a/lib/codecs/src/decoding/framing/chunked_gelf.rs +++ b/lib/codecs/src/decoding/framing/chunked_gelf.rs @@ -140,28 +140,30 @@ impl MessageState { self.chunks[sequence_number as usize] = Bytes::copy_from_slice(&chunk); } - fn is_complete(&self) -> bool { - self.chunks_bitmap.count_ones() == self.total_chunks as u32 + /// Callers must have ruled out a duplicate, which would not raise the count. + fn is_final_missing_chunk(&self) -> bool { + self.chunks_bitmap.count_ones() + 1 == self.total_chunks as u32 } fn current_length(&self) -> usize { self.current_length } - /// Peak is ~2x the message: a contiguous destination coexists with the chunks it copies - /// from. Reserving exactly keeps a growing buffer from adding slack on top of that. - fn retrieve_message(&mut self) -> Option { - if !self.is_complete() { - return None; - } - + fn finish(mut self: Box, sequence_number: u8, final_chunk: Bytes) -> Bytes { self.timeout_task.abort(); - let mut message = BytesMut::with_capacity(self.current_length); - for chunk in &mut self.chunks[0..self.total_chunks as usize] { - message.extend_from_slice(chunk); - *chunk = Bytes::new(); + let mut message = BytesMut::with_capacity(self.current_length + final_chunk.len()); + for (index, chunk) in self.chunks[0..self.total_chunks as usize] + .iter_mut() + .enumerate() + { + if index == sequence_number as usize { + message.extend_from_slice(&final_chunk); + } else { + message.extend_from_slice(chunk); + *chunk = Bytes::new(); + } } - Some(message.freeze()) + message.freeze() } } @@ -423,6 +425,29 @@ impl ChunkedGelfDecoder { return Ok(None); } + // Remove a complete message before assembling it. The final chunk goes straight into + // the output, so it needs neither an intermediate copy nor shared state during assembly. + if message_state.is_final_missing_chunk() { + let length = message_state.current_length().saturating_add(chunk.len()); + if let Some(max_length) = self.max_length + && length > max_length + { + if let Some(dropped) = state_lock.remove(&message_id) { + dropped.timeout_task.abort(); + } + return Err(ChunkedGelfDecoderError::MaxLengthExceed { + message_id, + sequence_number, + length, + max_length, + }); + } + + let message_state = state_lock.remove(&message_id).expect("entry must exist"); + drop(state_lock); + return Ok(Some(message_state.finish(sequence_number, chunk))); + } + message_state.add_chunk(sequence_number, chunk); if let Some(max_length) = self.max_length { @@ -442,12 +467,7 @@ impl ChunkedGelfDecoder { } } - if let Some(message) = message_state.retrieve_message() { - state_lock.remove(&message_id); - Ok(Some(message)) - } else { - Ok(None) - } + Ok(None) } /// Decode a GELF message that may be chunked or not. The source bytes are expected to be @@ -1025,22 +1045,14 @@ mod tests { } #[tokio::test] - async fn retrieve_message_releases_chunks_while_assembling() { - // Does not lower the 2x peak, just shortens how long the source side is held. - let mut state = MessageState::new(2, tokio::spawn(async {})); + async fn finish_assembles_the_final_chunk_in_sequence() { + let mut state = Box::new(MessageState::new(3, tokio::spawn(async {}))); state.add_chunk(0, Bytes::from_static(b"foo")); - state.add_chunk(1, Bytes::from_static(b"bar")); + state.add_chunk(2, Bytes::from_static(b"baz")); - let message = state - .retrieve_message() - .expect("message should be complete"); + let message = state.finish(1, Bytes::from_static(b"bar")); - assert_eq!(message, Bytes::from_static(b"foobar")); - assert!( - state.chunks[..2].iter().all(Bytes::is_empty), - "each chunk must be released as it is copied" - ); - assert_eq!(state.current_length(), 6); + assert_eq!(message, Bytes::from_static(b"foobarbaz")); } #[rstest] From c4d333a1383b7aa57f3641b008ae632c37ae5ca3 Mon Sep 17 00:00:00 2001 From: Pavlos Rontidis Date: Thu, 3 Sep 2026 13:34:13 -0400 Subject: [PATCH 2/3] fix(codecs): cancel GELF timeout before releasing ID --- lib/codecs/src/decoding/framing/chunked_gelf.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/lib/codecs/src/decoding/framing/chunked_gelf.rs b/lib/codecs/src/decoding/framing/chunked_gelf.rs index bf2981522e056..8c79e2864e6ad 100644 --- a/lib/codecs/src/decoding/framing/chunked_gelf.rs +++ b/lib/codecs/src/decoding/framing/chunked_gelf.rs @@ -150,7 +150,6 @@ impl MessageState { } fn finish(mut self: Box, sequence_number: u8, final_chunk: Bytes) -> Bytes { - self.timeout_task.abort(); let mut message = BytesMut::with_capacity(self.current_length + final_chunk.len()); for (index, chunk) in self.chunks[0..self.total_chunks as usize] .iter_mut() @@ -444,6 +443,10 @@ impl ChunkedGelfDecoder { } let message_state = state_lock.remove(&message_id).expect("entry must exist"); + // Abort while the message ID is still protected. Otherwise another decoder clone + // can reuse the ID before the old timeout is canceled, and that callback can remove + // the new message. + message_state.timeout_task.abort(); drop(state_lock); return Ok(Some(message_state.finish(sequence_number, chunk))); } From 415596a852108335ddd593990d3fc828ea199577 Mon Sep 17 00:00:00 2001 From: Pavlos Rontidis Date: Thu, 3 Sep 2026 14:50:46 -0400 Subject: [PATCH 3/3] fix(codecs): ignore stale GELF timeout callbacks --- lib/codecs/src/decoding/framing/chunked_gelf.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/lib/codecs/src/decoding/framing/chunked_gelf.rs b/lib/codecs/src/decoding/framing/chunked_gelf.rs index 8c79e2864e6ad..6968f5a214d3f 100644 --- a/lib/codecs/src/decoding/framing/chunked_gelf.rs +++ b/lib/codecs/src/decoding/framing/chunked_gelf.rs @@ -393,8 +393,12 @@ impl ChunkedGelfDecoder { let timeout = self.timeout; let timeout_handle = tokio::spawn(async move { tokio::time::sleep(timeout).await; + let timeout_task_id = tokio::task::id(); let mut state_lock = state.lock().expect("poisoned lock"); - if state_lock.remove(&message_id).is_some() { + let owns_message = state_lock + .get(&message_id) + .is_some_and(|message| message.timeout_task.id() == timeout_task_id); + if owns_message && state_lock.remove(&message_id).is_some() { warn!( message_id = message_id, timeout_secs = timeout.as_secs_f64(),