diff --git a/lib/codecs/src/decoding/framing/chunked_gelf.rs b/lib/codecs/src/decoding/framing/chunked_gelf.rs index 1f3046df064d7..6968f5a214d3f 100644 --- a/lib/codecs/src/decoding/framing/chunked_gelf.rs +++ b/lib/codecs/src/decoding/framing/chunked_gelf.rs @@ -140,28 +140,29 @@ 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; - } - - 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(); + fn finish(mut self: Box, sequence_number: u8, final_chunk: Bytes) -> Bytes { + 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() } } @@ -392,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(), @@ -423,6 +428,33 @@ 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"); + // 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))); + } + message_state.add_chunk(sequence_number, chunk); if let Some(max_length) = self.max_length { @@ -442,12 +474,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 +1052,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]