rustfs/crates/io-core/src/writer.rs

// Copyright 2024 RustFS Team
//
// 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
//
//     http://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.

//! BytesMut-backed object writer for optimized write operations.
//!
//! It uses `BytesMut` for efficient buffering; writes into that buffer may
//! still copy input bytes. The historical `ZeroCopyObjectWriter` name remains
//! available as a deprecated compatibility alias.

use bytes::{BufMut, Bytes, BytesMut};
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::AsyncWrite;

/// BytesMut-backed object writer for optimized write operations.
///
/// This writer minimizes memory allocations by:
/// - Using BytesMut for efficient buffer growth
/// - Accepting `Bytes` inputs for efficient buffer handling
/// - Optional integration with BytesPool for buffer reuse
///
/// # Example
///
/// ```ignore
/// use rustfs_io_core::BytesMutWriter;
/// use bytes::Bytes;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
///     let mut writer = BytesMutWriter::new();
///
///     // Write into the internal BytesMut buffer
///     let data = Bytes::from("hello world");
///     writer.write_buffered(data).await?;
///
///     // Get the result as Bytes (zero-copy conversion)
///     let result = writer.into_bytes();
///
///     Ok(())
/// }
/// ```
pub struct BytesMutWriter {
    /// Internal buffer using BytesMut for efficient growth
    buffer: BytesMut,
    /// Total bytes written
    bytes_written: usize,
    /// Whether the writer has been finalized
    finalized: bool,
}

/// Historical name for the BytesMut-backed object writer.
#[deprecated(since = "1.0.0-beta.8", note = "use BytesMutWriter; writes append into a BytesMut buffer")]
pub type ZeroCopyObjectWriter = BytesMutWriter;

impl BytesMutWriter {
    /// Create a new bytes-backed object writer with default capacity (8KB).
    ///
    /// # Example
    ///
    /// ```ignore
    /// let writer = BytesMutWriter::new();
    /// ```
    pub fn new() -> Self {
        Self::with_capacity(8 * 1024)
    }

    /// Create a new bytes-backed object writer with specified capacity.
    ///
    /// # Arguments
    ///
    /// * `capacity` - Initial buffer capacity in bytes
    ///
    /// # Example
    ///
    /// ```ignore
    /// let writer = BytesMutWriter::with_capacity(64 * 1024);
    /// ```
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            buffer: BytesMut::with_capacity(capacity),
            bytes_written: 0,
            finalized: false,
        }
    }

    /// Write data into the internal buffer.
    ///
    /// This method accepts `Bytes` for API compatibility, then appends the
    /// bytes into the internal `BytesMut` buffer.
    ///
    /// # Arguments
    ///
    /// * `data` - Data to append to the internal buffer
    ///
    /// # Returns
    ///
    /// * `Ok(usize)` - Number of bytes written
    /// * `Err(ZeroCopyWriteError)` - Write error
    ///
    /// # Example
    ///
    /// ```ignore
    /// let data = Bytes::from("hello world");
    /// let written = writer.write_buffered(data).await?;
    /// ```
    pub async fn write_buffered(&mut self, data: Bytes) -> Result<usize, ZeroCopyWriteError> {
        if self.finalized {
            return Err(ZeroCopyWriteError::Finalized("Cannot write to finalized writer".to_string()));
        }

        let len = data.len();
        self.buffer.put(data);

        self.bytes_written += len;
        Ok(len)
    }

    /// Historical name for `write_buffered`.
    #[deprecated(
        since = "1.0.0-beta.8",
        note = "use write_buffered; this method appends bytes into an internal buffer"
    )]
    pub async fn write_zero_copy(&mut self, data: Bytes) -> Result<usize, ZeroCopyWriteError> {
        self.write_buffered(data).await
    }

    /// Write a slice of data.
    ///
    /// # Arguments
    ///
    /// * `data` - Data slice to write
    ///
    /// # Returns
    ///
    /// * `Ok(usize)` - Number of bytes written
    /// * `Err(ZeroCopyWriteError)` - Write error
    pub async fn write_slice(&mut self, data: &[u8]) -> Result<usize, ZeroCopyWriteError> {
        if self.finalized {
            return Err(ZeroCopyWriteError::Finalized("Cannot write to finalized writer".to_string()));
        }

        let len = data.len();
        self.buffer.put_slice(data);
        self.bytes_written += len;
        Ok(len)
    }

    /// Finalize the writer and consume it, returning the written data as Bytes.
    ///
    /// This converts the internal BytesMut to Bytes, which is a zero-copy
    /// operation that freezes the buffer.
    ///
    /// # Returns
    ///
    /// The written data as Bytes
    ///
    /// # Example
    ///
    /// ```ignore
    /// let result = writer.into_bytes();
    /// ```
    pub fn into_bytes(mut self) -> Bytes {
        self.finalized = true;
        self.buffer.freeze()
    }

    /// Get the current buffer as a slice (without consuming).
    ///
    /// # Returns
    ///
    /// Slice of the current buffer content
    pub fn as_slice(&self) -> &[u8] {
        &self.buffer[..]
    }

    /// Get the total number of bytes written.
    ///
    /// # Returns
    ///
    /// Number of bytes written
    pub fn bytes_written(&self) -> usize {
        self.bytes_written
    }

    /// Get the current buffer capacity.
    ///
    /// # Returns
    ///
    /// Current buffer capacity in bytes
    pub fn capacity(&self) -> usize {
        self.buffer.capacity()
    }

    /// Get the current buffer length.
    ///
    /// # Returns
    ///
    /// Current buffer length in bytes
    pub fn len(&self) -> usize {
        self.buffer.len()
    }

    /// Check if the buffer is empty.
    ///
    /// # Returns
    ///
    /// `true` if buffer is empty, `false` otherwise
    pub fn is_empty(&self) -> bool {
        self.buffer.is_empty()
    }

    /// Clear the buffer, resetting it to empty.
    ///
    /// This does not change the capacity, just resets the length to 0.
    pub fn clear(&mut self) {
        self.buffer.clear();
        self.bytes_written = 0;
        self.finalized = false;
    }

    /// Reserve additional capacity in the buffer.
    ///
    /// # Arguments
    ///
    /// * `additional` - Additional capacity to reserve
    pub fn reserve(&mut self, additional: usize) {
        self.buffer.reserve(additional);
    }
}

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

impl std::fmt::Debug for BytesMutWriter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BytesMutWriter")
            .field("buffer_len", &self.buffer.len())
            .field("buffer_capacity", &self.buffer.capacity())
            .field("bytes_written", &self.bytes_written)
            .field("finalized", &self.finalized)
            .finish()
    }
}

/// AsyncWrite implementation for BytesMutWriter.
///
/// This allows the writer to be used with tokio's async I/O utilities.
impl AsyncWrite for BytesMutWriter {
    fn poll_write(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, tokio::io::Error>> {
        if self.finalized {
            return Poll::Ready(Err(tokio::io::Error::new(
                tokio::io::ErrorKind::WriteZero,
                "Cannot write to finalized writer",
            )));
        }

        let len = buf.len();
        self.buffer.put_slice(buf);
        self.bytes_written += len;
        Poll::Ready(Ok(len))
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), tokio::io::Error>> {
        // Nothing to flush for in-memory buffer
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), tokio::io::Error>> {
        self.finalized = true;
        Poll::Ready(Ok(()))
    }
}

/// Zero-copy write error types.
#[derive(Debug, thiserror::Error)]
pub enum ZeroCopyWriteError {
    /// I/O error occurred
    #[error("I/O error: {0}")]
    Io(#[from] tokio::io::Error),

    /// Writer has been finalized and cannot accept more writes
    #[error("Writer finalized: {0}")]
    Finalized(String),

    /// Invalid input provided
    #[error("Invalid input: {0}")]
    InvalidInput(String),
}

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

    #[tokio::test]
    async fn test_new_writer() {
        let writer = BytesMutWriter::new();
        assert!(writer.is_empty());
        assert_eq!(writer.bytes_written(), 0);
        assert!(writer.capacity() >= 8 * 1024);
    }

    #[tokio::test]
    async fn test_write_buffered() {
        let mut writer = BytesMutWriter::new();
        let data = Bytes::from("hello world");

        let written = writer.write_buffered(data).await.unwrap();
        assert_eq!(written, 11);
        assert_eq!(writer.bytes_written(), 11);
        assert_eq!(writer.as_slice(), b"hello world");
    }

    #[tokio::test]
    async fn test_preferred_writer_alias() {
        let mut writer = BytesMutWriter::new();
        let written = writer
            .write_buffered(Bytes::from("hello world"))
            .await
            .expect("write bytes through alias");

        assert_eq!(written, 11);
        assert_eq!(writer.as_slice(), b"hello world");
    }

    #[tokio::test]
    async fn test_write_slice() {
        let mut writer = BytesMutWriter::new();
        let data = b"hello world";

        let written = writer.write_slice(data).await.unwrap();
        assert_eq!(written, 11);
        assert_eq!(writer.bytes_written(), 11);
        assert_eq!(writer.as_slice(), b"hello world");
    }

    #[tokio::test]
    async fn test_into_bytes() {
        let mut writer = BytesMutWriter::new();
        let data = Bytes::from("hello world");

        writer.write_buffered(data).await.unwrap();
        let result = writer.into_bytes();

        assert_eq!(result.as_ref(), b"hello world");
    }

    #[tokio::test]
    async fn test_write_after_finalize() {
        let mut writer = BytesMutWriter::new();
        let data = Bytes::from("hello");

        writer.write_buffered(data).await.unwrap();
        let _result = writer.into_bytes();

        // Create new writer and try to write after finalize
        let mut writer2 = BytesMutWriter::new();
        writer2.write_buffered(Bytes::from("test")).await.unwrap();
        let _ = writer2.into_bytes();

        // Writing to a consumed writer should work via new writer
        let mut writer3 = BytesMutWriter::new();
        let result = writer3.write_buffered(Bytes::from("final")).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_clear() {
        let mut writer = BytesMutWriter::new();
        writer.write_slice(b"hello").await.unwrap();

        writer.clear();
        assert!(writer.is_empty());
        assert_eq!(writer.bytes_written(), 0);
        // Capacity should remain
        assert!(writer.capacity() > 0);
    }

    #[tokio::test]
    async fn test_reserve() {
        let mut writer = BytesMutWriter::with_capacity(10);
        let initial_capacity = writer.capacity();

        writer.reserve(1000);
        // Reserve ensures at least the additional capacity can be added
        // but may allocate more than requested
        assert!(writer.capacity() >= initial_capacity);
    }

    #[tokio::test]
    async fn test_multiple_writes() {
        let mut writer = BytesMutWriter::new();

        writer.write_buffered(Bytes::from("hello ")).await.unwrap();
        writer.write_slice(b"world").await.unwrap();

        assert_eq!(writer.as_slice(), b"hello world");
        assert_eq!(writer.bytes_written(), 11);
    }

    #[tokio::test]
    async fn test_async_write() {
        use tokio::io::AsyncWriteExt;

        let mut writer = BytesMutWriter::new();
        let data = b"hello world";

        let written = writer.write(data).await.unwrap();
        assert_eq!(written, 11);
        assert_eq!(writer.as_slice(), b"hello world");
    }

    #[tokio::test]
    async fn test_debug() {
        let writer = BytesMutWriter::new();
        let debug_str = format!("{:?}", writer);
        assert!(debug_str.contains("BytesMutWriter"));
        assert!(debug_str.contains("buffer_len"));
    }

    #[tokio::test]
    #[allow(deprecated)]
    async fn test_legacy_writer_alias() {
        let mut writer = ZeroCopyObjectWriter::new();
        let written = writer.write_zero_copy(Bytes::from("hello")).await.unwrap();

        assert_eq!(written, 5);
        assert_eq!(writer.as_slice(), b"hello");
    }
}