google-cloud-storage 1.21.0

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Retains unacknowledged chunks, trims acknowledged data, and provides chunks
//! for resending upon stream reconnect.

use super::MAX_WRITE_CHUNK_SIZE;
use crate::google::storage::v2::{
    BidiWriteObjectRequest, ChecksummedData, bidi_write_object_request::Data,
};
use bytes::Bytes;
use std::collections::VecDeque;

/// Defines the default capacity of the [`ReplayBuffer`] in bytes (32 MiB).
// TODO(#5716): Remove once ReplayBuffer capacity is configured via CommonOptions.
pub const DEFAULT_REPLAY_BUFFER_SIZE: usize = 32 * 1024 * 1024;

/// Defines the smallest capacity a [`ReplayBuffer`] may have.
///
/// Must exceed `2 * MAX_WRITE_CHUNK_SIZE` so the worker's high watermark
/// (`capacity - 2 * MAX_WRITE_CHUNK_SIZE`) is non-zero and leaves room for at least one chunk.
// TODO(#5716): Enforce this when ReplayBuffer capacity becomes configurable.
pub const MIN_REPLAY_BUFFER_SIZE: usize = 2 * MAX_WRITE_CHUNK_SIZE + 1;

/// Represents an unacknowledged data chunk retained in the [`ReplayBuffer`].
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ReplayChunk {
    /// Holds the logical starting write offset of this chunk.
    pub write_offset: i64,
    /// Contains the raw payload bytes.
    pub data: Bytes,
    /// Stores the precomputed CRC32C checksum of [`data`][Self::data].
    pub crc32c: u32,
}

impl ReplayChunk {
    /// Creates a new [`ReplayChunk`].
    pub fn new(write_offset: i64, data: Bytes, crc32c: u32) -> Self {
        Self {
            write_offset,
            data,
            crc32c,
        }
    }

    /// Returns the ending byte offset (exclusive) of this chunk.
    pub fn end_offset(&self) -> i64 {
        self.write_offset + self.data.len() as i64
    }

    /// Converts this [`ReplayChunk`] into a [`BidiWriteObjectRequest`] for transmission.
    pub fn to_request(&self) -> BidiWriteObjectRequest {
        BidiWriteObjectRequest {
            write_offset: self.write_offset,
            data: Some(Data::ChecksummedData(ChecksummedData {
                content: self.data.clone(),
                crc32c: Some(self.crc32c),
            })),
            ..BidiWriteObjectRequest::default()
        }
    }
}

/// Manages an in-memory FIFO queue of unacknowledged chunks up to a configurable capacity.
#[derive(Debug)]
pub struct ReplayBuffer {
    queue: VecDeque<ReplayChunk>,
    unpersisted_bytes: usize,
    capacity: usize,
}

impl Default for ReplayBuffer {
    fn default() -> Self {
        Self::new()
    }
}

impl ReplayBuffer {
    /// Creates a new, empty [`ReplayBuffer`] with the default capacity ([`DEFAULT_REPLAY_BUFFER_SIZE`]).
    pub fn new() -> Self {
        Self::with_capacity(DEFAULT_REPLAY_BUFFER_SIZE)
    }

    /// Creates a new, empty [`ReplayBuffer`] with a specified capacity in bytes.
    ///
    /// Clamps `capacity` to at least [`MIN_REPLAY_BUFFER_SIZE`].
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            queue: VecDeque::new(),
            unpersisted_bytes: 0,
            capacity: std::cmp::max(capacity, MIN_REPLAY_BUFFER_SIZE),
        }
    }

    /// Returns the configured capacity of the [`ReplayBuffer`] in bytes.
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// Enqueues an unacknowledged [`ReplayChunk`] to the [`ReplayBuffer`].
    pub fn push(&mut self, chunk: ReplayChunk) {
        self.unpersisted_bytes += chunk.data.len();
        self.queue.push_back(chunk);
    }

    /// Trims acknowledged chunks up to `persisted_size`.
    ///
    /// If `persisted_size` lands inside a chunk, that chunk is sliced in-place
    /// and its CRC32C is recomputed for the unpersisted sub-slice only.
    pub fn ack(&mut self, persisted_size: i64) {
        while let Some(front) = self.queue.front()
            && front.end_offset() <= persisted_size
        {
            self.unpersisted_bytes -= front.data.len();
            self.queue.pop_front();
        }

        if let Some(front) = self.queue.front_mut()
            && front.write_offset < persisted_size
        {
            let trimmed_bytes = (persisted_size - front.write_offset) as usize;
            front.data = front.data.slice(trimmed_bytes..);
            front.write_offset = persisted_size;
            front.crc32c = crc32c::crc32c(&front.data);
            self.unpersisted_bytes -= trimmed_bytes;
        }
    }

    /// Returns `true` if the buffered byte count has reached or exceeded [`capacity`][Self::capacity].
    pub fn is_full(&self) -> bool {
        self.unpersisted_bytes >= self.capacity
    }

    /// Returns `true` if the [`ReplayBuffer`] contains no chunks.
    pub fn is_empty(&self) -> bool {
        self.queue.is_empty()
    }

    /// Returns the number of chunks currently held in the [`ReplayBuffer`].
    #[cfg(test)]
    pub fn num_chunks(&self) -> usize {
        self.queue.len()
    }

    /// Returns the total unpersisted byte count currently retained in the [`ReplayBuffer`].
    pub fn unpersisted_bytes(&self) -> usize {
        self.unpersisted_bytes
    }

    /// Returns an iterator over the unpersisted [`ReplayChunk`]s in FIFO order for replay.
    pub fn chunks_to_replay(&self) -> impl Iterator<Item = &ReplayChunk> {
        self.queue.iter()
    }

    /// Returns the end offset (exclusive) of the newest unacknowledged chunk, if any.
    pub fn end_offset(&self) -> Option<i64> {
        self.queue.back().map(|chunk| chunk.end_offset())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn new_buffer_state() {
        // Arrange.
        let mut buf = ReplayBuffer::new();

        // Assert.
        assert!(buf.is_empty());
        assert_eq!(buf.num_chunks(), 0);
        assert_eq!(buf.unpersisted_bytes(), 0);
        assert!(!buf.is_full());

        // Act.
        buf.ack(100);

        // Assert.
        assert!(buf.is_empty());
    }

    #[test]
    fn push_and_ack_full_chunks() {
        // Arrange.
        let mut buf = ReplayBuffer::new();
        let chunk1 = Bytes::from_static(b"hello ");
        let chunk2 = Bytes::from_static(b"world!");

        // Act.
        buf.push(ReplayChunk::new(0, chunk1.clone(), crc32c::crc32c(&chunk1)));
        buf.push(ReplayChunk::new(6, chunk2.clone(), crc32c::crc32c(&chunk2)));

        // Assert.
        assert_eq!(buf.num_chunks(), 2);
        assert_eq!(buf.unpersisted_bytes(), 12);

        // Act.
        // Acknowledge partially up to 4 bytes (within chunk1)
        buf.ack(4);

        // Assert.
        assert_eq!(buf.num_chunks(), 2);
        assert_eq!(buf.unpersisted_bytes(), 8);

        let chunks: Vec<_> = buf.chunks_to_replay().cloned().collect();
        assert_eq!(chunks[0].write_offset, 4);
        assert_eq!(chunks[0].data.as_ref(), b"o ");
        assert_eq!(chunks[0].crc32c, crc32c::crc32c(b"o "));
        assert_eq!(chunks[1].write_offset, 6);
        assert_eq!(chunks[1].data.as_ref(), b"world!");
        assert_eq!(chunks[1].crc32c, crc32c::crc32c(b"world!"));

        // Act.
        // Acknowledge fully past chunk1 up to 10 (within chunk2)
        buf.ack(10);

        // Assert.
        assert_eq!(buf.num_chunks(), 1);
        assert_eq!(buf.unpersisted_bytes(), 2);

        let chunks: Vec<_> = buf.chunks_to_replay().cloned().collect();
        assert_eq!(chunks[0].write_offset, 10);
        assert_eq!(chunks[0].data.as_ref(), b"d!");
        assert_eq!(chunks[0].crc32c, crc32c::crc32c(b"d!"));

        // Act.
        // Acknowledge all remaining bytes.
        buf.ack(12);

        // Assert.
        assert!(buf.is_empty());
        assert_eq!(buf.unpersisted_bytes(), 0);
    }

    #[test]
    fn ack_duplicate_or_earlier_offset() {
        // Arrange.
        let mut buf = ReplayBuffer::new();
        let chunk = Bytes::from_static(b"abcdef");
        buf.push(ReplayChunk::new(10, chunk.clone(), crc32c::crc32c(&chunk)));

        // Act.
        // Acknowledge offset earlier than front write_offset.
        buf.ack(5);

        // Assert.
        assert_eq!(buf.num_chunks(), 1);
        assert_eq!(buf.unpersisted_bytes(), 6);

        // Act.
        // Acknowledge exact write_offset of the front chunk.
        buf.ack(10);

        // Assert.
        assert_eq!(buf.num_chunks(), 1);
        assert_eq!(buf.unpersisted_bytes(), 6);
    }

    #[test]
    fn is_full_threshold() {
        // Arrange.
        let mut buf = ReplayBuffer::new();
        let huge_chunk = Bytes::from(vec![0u8; DEFAULT_REPLAY_BUFFER_SIZE]);

        // Act.
        buf.push(ReplayChunk::new(0, huge_chunk, 0));

        // Assert.
        assert!(buf.is_full());

        // Act.
        buf.ack(1);

        // Assert.
        assert!(!buf.is_full());
    }

    #[test]
    fn chunk_to_request_conversion() {
        // Arrange.
        let data = Bytes::from_static(b"replay data");
        let crc = crc32c::crc32c(&data);
        let chunk = ReplayChunk::new(42, data.clone(), crc);

        // Assert end_offset calculation.
        assert_eq!(chunk.end_offset(), 42 + data.len() as i64);

        // Act.
        let req = chunk.to_request();

        // Assert.
        assert_eq!(req.write_offset, 42);
        if let Some(Data::ChecksummedData(cd)) = req.data {
            assert_eq!(cd.content, data);
            assert_eq!(cd.crc32c, Some(crc));
        } else {
            panic!("expected ChecksummedData");
        }
    }

    #[test]
    fn with_capacity_default() {
        // Arrange & Act.
        let buf = ReplayBuffer::new();

        // Assert.
        assert_eq!(buf.capacity(), DEFAULT_REPLAY_BUFFER_SIZE);
    }

    #[test]
    fn with_capacity_custom() {
        // Arrange.
        const CUSTOM_CAPACITY: usize = 64 * 1024 * 1024; // 64 MiB

        // Act.
        let buf = ReplayBuffer::with_capacity(CUSTOM_CAPACITY);

        // Assert.
        assert_eq!(buf.capacity(), CUSTOM_CAPACITY);
    }

    #[test]
    fn is_full_with_custom_capacity() {
        // Arrange.
        // The smallest capacity the buffer accepts, so the test stays deterministic.
        let mut buf = ReplayBuffer::with_capacity(MIN_REPLAY_BUFFER_SIZE);
        let chunk = Bytes::from(vec![0u8; MIN_REPLAY_BUFFER_SIZE]);

        // Act.
        buf.push(ReplayChunk::new(0, chunk, 0));

        // Assert.
        assert!(buf.is_full());

        // Act.
        buf.ack(1);

        // Assert.
        assert!(!buf.is_full());
    }

    #[test]
    fn with_capacity_enforces_minimum() {
        // Arrange & Act. A capacity at or below the worker's two-chunk headroom would saturate its
        // watermark to zero, and a zero capacity would make the buffer permanently full.
        let buf = ReplayBuffer::with_capacity(0);
        let two_chunks = ReplayBuffer::with_capacity(2 * MAX_WRITE_CHUNK_SIZE);

        // Assert.
        assert_eq!(buf.capacity(), MIN_REPLAY_BUFFER_SIZE);
        assert!(!buf.is_full());
        assert_eq!(two_chunks.capacity(), MIN_REPLAY_BUFFER_SIZE);
    }
}