Make Batch a customizable type in RaftTypeConfig

Author: Lymah123

openraft/src/base/batch/display.rs

@@ -3,40 +3,46 @@
 use std::fmt;
 use std::fmt::Formatter;
 
-use super::Batch;
+use super::RaftBatch;
+use crate::OptionalSend;
 
-impl<T: fmt::Display> Batch<T> {
-    /// Returns a display helper that shows all elements.
-    pub fn display(&self) -> impl fmt::Display + '_ {
-        BatchDisplay {
-            elements: self,
-            max: None,
-        }
-    }
+/// Display helper for types implementing `RaftBatch`.
+pub struct DisplayBatch<'a, T, B>
+where
+    T: fmt::Display + OptionalSend + 'static + fmt::Debug,
+    B: RaftBatch<T>,
+{
+    pub(super) elements: &'a B,
+    pub(super) max: Option<usize>,
+    pub(super) _phantom: std::marker::PhantomData<T>,
+}
 
-    /// Returns a display helper that shows at most `max` elements.
-    pub fn display_n(&self, max: usize) -> impl fmt::Display + '_ {
-        BatchDisplay {
-            elements: self,
-            max: Some(max),
+impl<'a, T, B> DisplayBatch<'a, T, B>
+where
+    T: fmt::Display + OptionalSend + 'static + fmt::Debug,
+    B: RaftBatch<T>,
+{
+    pub(super) fn new(elements: &'a B, max: Option<usize>) -> Self {
+        Self {
+            elements,
+            max,
+            _phantom: std::marker::PhantomData,
         }
     }
 }
 
-struct BatchDisplay<'a, T> {
-    elements: &'a Batch<T>,
-    max: Option<usize>,
-}
-
-impl<'a, T: fmt::Display> fmt::Display for BatchDisplay<'a, T> {
+impl<'a, T, B> fmt::Display for DisplayBatch<'a, T, B>
+where
+    T: fmt::Display + OptionalSend + 'static + fmt::Debug,
+    B: RaftBatch<T>,
+{
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        let slice = self.elements.as_slice();
-        let max = self.max.unwrap_or(slice.len());
-        let len = slice.len();
+        let len = self.elements.len();
+        let max = self.max.unwrap_or(len);
         let shown = max.min(len);
 
         write!(f, "[")?;
-        for (i, e) in slice.iter().take(max).enumerate() {
+        for (i, e) in self.elements.iter().take(max).enumerate() {
             if i > 0 {
                 write!(f, ", ")?;
             }
@@ -55,25 +61,53 @@ impl<'a, T: fmt::Display> fmt::Display for BatchDisplay<'a, T> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::impls::Batch;
 
     #[test]
     fn test_display() {
-        assert_eq!(format!("{}", Batch::Single(42).display()), "[42]");
-        assert_eq!(format!("{}", Batch::Vec(vec![1, 2]).display()), "[1, 2]");
-        assert_eq!(format!("{}", Batch::<i32>::Vec(vec![]).display()), "[]");
+        assert_eq!(
+            format!("{}", <Batch<i32> as RaftBatch<i32>>::display(&Batch::Single(42))),
+            "[42]"
+        );
+        assert_eq!(
+            format!("{}", <Batch<i32> as RaftBatch<i32>>::display(&Batch::Vec(vec![1, 2]))),
+            "[1, 2]"
+        );
+        assert_eq!(
+            format!(
+                "{}",
+                <Batch<i32> as RaftBatch<i32>>::display(&Batch::<i32>::Vec(vec![]))
+            ),
+            "[]"
+        );
     }
 
     #[test]
     fn test_display_n() {
         let v: Batch<i32> = [1, 2, 3, 4, 5].into();
 
-        assert_eq!(format!("{}", v.display_n(3)), "[1, 2, 3, ... 2 more]");
-        assert_eq!(format!("{}", v.display_n(5)), "[1, 2, 3, 4, 5]");
-        assert_eq!(format!("{}", v.display_n(10)), "[1, 2, 3, 4, 5]");
-        assert_eq!(format!("{}", v.display_n(0)), "[... 5 more]");
+        assert_eq!(
+            format!("{}", <Batch<i32> as RaftBatch<i32>>::display_n(&v, 3)),
+            "[1, 2, 3, ... 2 more]"
+        );
+        assert_eq!(
+            format!("{}", <Batch<i32> as RaftBatch<i32>>::display_n(&v, 5)),
+            "[1, 2, 3, 4, 5]"
+        );
+        assert_eq!(
+            format!("{}", <Batch<i32> as RaftBatch<i32>>::display_n(&v, 10)),
+            "[1, 2, 3, 4, 5]"
+        );
+        assert_eq!(
+            format!("{}", <Batch<i32> as RaftBatch<i32>>::display_n(&v, 0)),
+            "[... 5 more]"
+        );
 
         let v2: Batch<i32> = 42.into();
-        assert_eq!(format!("{}", v2.display_n(0)), "[... 1 more]");
-        assert_eq!(format!("{}", v2.display_n(1)), "[42]");
+        assert_eq!(
+            format!("{}", <Batch<i32> as RaftBatch<i32>>::display_n(&v2, 0)),
+            "[... 1 more]"
+        );
+        assert_eq!(format!("{}", <Batch<i32> as RaftBatch<i32>>::display_n(&v2, 1)), "[42]");
     }
 }

openraft/src/base/batch/iter.rs

@@ -44,6 +44,9 @@ impl<T> Iterator for BatchIter<T> {
 
 impl<T> ExactSizeIterator for BatchIter<T> {}
 
+// Safety: BatchIter<T> is Send when T is Send because both Option<T> and vec::IntoIter<T> are Send
+unsafe impl<T: Send> Send for BatchIter<T> {}
+
 #[cfg(test)]
 mod tests {
     use super::*;

openraft/src/base/batch/mod.rs

@@ -2,17 +2,20 @@
 
 mod display;
 mod iter;
+mod raft_batch;
 
 use std::ops::Index;
 use std::slice;
 
+pub use raft_batch::RaftBatch;
+
 /// A container that stores elements efficiently by avoiding heap allocation for single elements.
 ///
 /// This type uses an enum with two variants:
 /// - `Single`: stores exactly one element inline (no heap allocation)
 /// - `Vec`: stores zero or more elements using a `Vec`
 #[derive(Debug, Clone, Eq)]
-pub(crate) enum Batch<T> {
+pub enum Batch<T> {
     /// A single element stored inline without heap allocation.
     Single(T),
     /// Multiple elements stored in a Vec.
@@ -135,6 +138,66 @@ impl<T> Batch<T> {
     }
 }
 
+// Implement RaftBatch trait for Batch
+impl<T> RaftBatch<T> for Batch<T>
+where T: crate::OptionalSend + 'static + std::fmt::Debug
+{
+    type Iter<'a>
+        = slice::Iter<'a, T>
+    where T: 'a;
+    type IterMut<'a>
+        = slice::IterMut<'a, T>
+    where T: 'a;
+    type IntoIter = iter::BatchIter<T>;
+
+    fn from_item(item: T) -> Self {
+        Batch::Single(item)
+    }
+
+    fn from_vec(vec: Vec<T>) -> Self {
+        Batch::from(vec)
+    }
+
+    fn from_exact_iter<I>(iter: I) -> Self
+    where I: ExactSizeIterator<Item = T> {
+        match iter.len() {
+            0 => Batch::Vec(Vec::new()),
+            1 => Batch::Single(iter.into_iter().next().unwrap()),
+            _ => Batch::Vec(iter.collect()),
+        }
+    }
+
+    fn len(&self) -> usize {
+        Batch::len(self)
+    }
+
+    fn first(&self) -> Option<&T> {
+        Batch::first(self)
+    }
+
+    fn last(&self) -> Option<&T> {
+        Batch::last(self)
+    }
+
+    fn iter(&self) -> Self::Iter<'_> {
+        self.as_slice().iter()
+    }
+
+    fn iter_mut(&mut self) -> Self::IterMut<'_> {
+        self.as_mut_slice().iter_mut()
+    }
+
+    fn into_iter(self) -> Self::IntoIter {
+        // Use the existing IntoIterator impl
+        IntoIterator::into_iter(self)
+    }
+
+    fn extend(&mut self, other: Self) {
+        Batch::extend(self, other)
+    }
+}
+
+// Index
 impl<T> Index<usize> for Batch<T> {
     type Output = T;
 

openraft/src/base/batch/raft_batch.rs

@@ -0,0 +1,105 @@
+//! Trait definition for customizable batch containers in Raft.
+
+use std::fmt::Debug;
+use std::fmt::Display;
+
+use super::display::DisplayBatch;
+use crate::OptionalSend;
+
+/// Trait for batch containers used throughout Raft for efficient element grouping.
+///
+/// This trait abstracts over different batch container implementations while
+/// preserving performance characteristics required by Raft.
+///
+/// Implementations may optimize for:
+/// - Avoiding heap allocation for small batches
+/// - Arena or pool-backed allocation
+/// - Cache locality
+///
+/// # Design Notes
+///
+/// - The trait is `Sized` to enable static dispatch and full monomorphization.
+/// - Iteration is explicit via `iter()` / `iter_mut()` to preserve lifetime and `ExactSizeIterator`
+///   guarantees.
+/// - Consuming iteration via `into_iter()` method with explicit `IntoIter` associated type.
+pub trait RaftBatch<T>: OptionalSend + Sized + Default + Debug + 'static
+where T: OptionalSend + Debug + 'static
+{
+    /// Iterator type for immutable element references.
+    ///
+    /// Must implement `ExactSizeIterator` to allow efficient size queries.
+    type Iter<'a>: Iterator<Item = &'a T> + ExactSizeIterator
+    where
+        T: 'a,
+        Self: 'a;
+
+    /// Iterator type for mutable element references.
+    ///
+    /// Must implement `ExactSizeIterator` to allow efficient size queries.
+    type IterMut<'a>: Iterator<Item = &'a mut T> + ExactSizeIterator
+    where
+        T: 'a,
+        Self: 'a;
+
+    /// Iterator type for consuming iteration.
+    ///
+    /// Must implement `ExactSizeIterator` to allow efficient size queries.
+    type IntoIter: Iterator<Item = T> + ExactSizeIterator + OptionalSend;
+
+    /// Creates a batch containing a single item.
+    ///
+    /// Implementations may optimize this case to avoid heap allocation.
+    fn from_item(item: T) -> Self;
+
+    /// Creates a batch from a `Vec`.
+    ///
+    /// Implementations may optimize for the single-element case.
+    fn from_vec(vec: Vec<T>) -> Self;
+
+    /// Creates a batch from an exact-size iterator.
+    ///
+    /// The exact size allows implementations to pre-allocate or select
+    /// efficient internal representations.
+    fn from_exact_iter<I>(iter: I) -> Self
+    where I: ExactSizeIterator<Item = T>;
+
+    /// Returns the number of elements in the batch.
+    fn len(&self) -> usize;
+
+    /// Returns `true` if the batch contains no elements.
+    fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    /// Returns a reference to the first element, or `None` if empty.
+    fn first(&self) -> Option<&T>;
+
+    /// Returns a reference to the last element, or `None` if empty.
+    fn last(&self) -> Option<&T>;
+
+    /// Returns an iterator over immutable element references.
+    fn iter(&self) -> Self::Iter<'_>;
+
+    /// Returns an iterator over mutable element references.
+    fn iter_mut(&mut self) -> Self::IterMut<'_>;
+
+    /// Consumes the batch and returns an iterator over the elements.
+    fn into_iter(self) -> Self::IntoIter;
+
+    /// Appends all elements from another batch to this batch.
+    ///
+    /// This method consumes `other` to allow efficient transfer of ownership.
+    fn extend(&mut self, other: Self);
+
+    /// Returns a display helper that formats all elements.
+    fn display(&self) -> DisplayBatch<'_, T, Self>
+    where T: Display {
+        DisplayBatch::new(self, None)
+    }
+
+    /// Returns a display helper that formats at most `max` elements.
+    fn display_n(&self, max: usize) -> DisplayBatch<'_, T, Self>
+    where T: Display {
+        DisplayBatch::new(self, Some(max))
+    }
+}

openraft/src/base/mod.rs

@@ -32,7 +32,7 @@ pub(crate) mod finalized;
 pub(crate) mod histogram;
 pub(crate) mod shared_id_generator;
 
-pub(crate) use batch::Batch;
+pub use batch::RaftBatch;
 pub use openraft_rt::BoxAny;
 pub use openraft_rt::BoxAsyncOnceMut;
 pub use openraft_rt::BoxFuture;

openraft/src/core/merged_raft_msg_receiver.rs

@@ -8,6 +8,7 @@
 use crate::RaftTypeConfig;
 use crate::async_runtime::MpscReceiver;
 use crate::async_runtime::TryRecvError;
+use crate::base::RaftBatch;
 use crate::core::raft_msg::RaftMsg;
 use crate::error::Fatal;
 use crate::type_config::alias::MpscReceiverOf;
@@ -196,7 +197,6 @@ where C: RaftTypeConfig
 mod tests {
     use super::*;
     use crate::async_runtime::MpscSender;
-    use crate::base::Batch;
     use crate::engine::testing::UTConfig;
     use crate::engine::testing::log_id;
     use crate::entry::EntryPayload;
@@ -211,15 +211,14 @@ mod tests {
 
     fn client_write(data: u64, leader: Option<CommittedLeaderIdOf<C>>) -> RaftMsg<C> {
         RaftMsg::ClientWrite {
-            payloads: Batch::Single(EntryPayload::Normal(data)),
-            responders: Batch::Single(None),
+            payloads: C::Batch::from_item(EntryPayload::Normal(data)),
+            responders: C::Batch::from_item(None),
             expected_leader: leader,
         }
     }
 
-    fn extract_payload_data(payloads: &Batch<EntryPayload<C>>) -> Vec<u64> {
+    fn extract_payload_data(payloads: &C::Batch<EntryPayload<C>>) -> Vec<u64> {
         payloads
-            .as_slice()
             .iter()
             .map(|p| match p {
                 EntryPayload::Normal(d) => *d,