feat(memory): Introduce a pluggable mem-buffer mechanism and an optimized staging buffer impl (#564)

### Description

Resolves [issues-546](#546) .

This PR introduces a pluggable mem-buffer mechanism, and based on this, this also introduces a optimized staging lookup implementaion of mem-buffer. 

### Changelogs of `optimized staging lookup impl`

1. Converts staging buffer from previous structure to `Vec<Block>` for contiguous memory access
2. Adds `batch_boundaries` vector to track batch start positions
3. Introduces `block_position_index` for O(1) block ID to vector index mapping

###  Performance Impact

- Improves search performance in staging buffers from O(number of staging blocks) to O(1 + number of selected staging blocks) by eliminating linear scans during get operations
- Easily extendable with new memory buffer implementations for benchmarking, testing, etc., and switch to the most efficient approach.

---------

Co-authored-by: thuong <ledacthuong2210>
Co-authored-by: Junfan Zhang <[email protected]>

Author: thuongle2210

riffle-server/src/app_manager.rs

@@ -562,6 +562,7 @@ pub(crate) mod test {
                 capacity: "20B".to_string(),
                 buffer_ticket_timeout_sec: 1,
                 buffer_ticket_check_interval_sec: 1,
+                buffer_type: Default::default(),
             }),
         );
         let _ = std::mem::replace(

riffle-server/src/config.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::block_id_manager::BlockIdManagerType;
+use crate::store::mem::buffer::BufferType;
 use crate::store::ResponseDataIndex::Local;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
@@ -31,6 +32,13 @@ pub struct MemoryStoreConfig {
     pub buffer_ticket_timeout_sec: i64,
     #[serde(default = "as_default_buffer_ticket_timeout_check_interval_sec")]
     pub buffer_ticket_check_interval_sec: i64,
+
+    #[serde(default = "as_default_buffer_type")]
+    pub buffer_type: BufferType,
+}
+
+fn as_default_buffer_type() -> BufferType {
+    BufferType::DEFAULT
 }
 
 fn as_default_buffer_ticket_timeout_check_interval_sec() -> i64 {
@@ -47,6 +55,7 @@ impl MemoryStoreConfig {
             capacity,
             buffer_ticket_timeout_sec: as_default_buffer_ticket_timeout_sec(),
             buffer_ticket_check_interval_sec: as_default_buffer_ticket_timeout_check_interval_sec(),
+            buffer_type: BufferType::DEFAULT,
         }
     }
 
@@ -55,6 +64,7 @@ impl MemoryStoreConfig {
             capacity,
             buffer_ticket_timeout_sec,
             buffer_ticket_check_interval_sec: as_default_buffer_ticket_timeout_check_interval_sec(),
+            buffer_type: BufferType::DEFAULT,
         }
     }
 }

riffle-server/src/store/hybrid.rs

@@ -61,12 +61,15 @@ use crate::app_manager::AppManagerRef;
 use crate::config_reconfigure::ReconfigurableConfManager;
 use crate::runtime::manager::RuntimeManager;
 use crate::store::local::LocalfileStoreStat;
-use crate::store::mem::buffer::MemoryBuffer;
+use crate::store::mem::buffer::default_buffer::DefaultMemoryBuffer;
+use crate::store::mem::buffer::unified_buffer::UnifiedBuffer;
+use crate::store::mem::buffer::{BufferType, MemoryBuffer};
 use crate::store::mem::capacity::CapacitySnapshot;
 use crate::store::spill::hierarchy_event_bus::HierarchyEventBus;
 use crate::store::spill::storage_flush_handler::StorageFlushHandler;
 use crate::store::spill::storage_select_handler::StorageSelectHandler;
 use crate::store::spill::{SpillMessage, SpillWritingViewContext};
+use crate::store::ResponseData::Mem;
 use tokio::time::Instant;
 
 pub trait PersistentStore: Store + Persistent + Send + Sync + Any {
@@ -89,7 +92,7 @@ const DEFAULT_MEMORY_SPILL_MAX_CONCURRENCY: i32 = 20;
 
 pub struct HybridStore {
     // Box<dyn Store> will build fail
-    pub(crate) hot_store: Arc<MemoryStore>,
+    pub(crate) hot_store: Arc<MemoryStore<UnifiedBuffer>>,
 
     pub(crate) warm_store: Option<Box<dyn PersistentStore>>,
     pub(crate) cold_store: Option<Box<dyn PersistentStore>>,
@@ -172,11 +175,12 @@ impl HybridStore {
         }
         let async_watermark_spill_enable = hybrid_conf.async_watermark_spill_trigger_enable;
 
+        // use the unified buffer to delegate the underlying concrete buffer
+        let mem_store: MemoryStore<UnifiedBuffer> =
+            MemoryStore::from(config.memory_store.unwrap(), runtime_manager.clone());
+
         let store = HybridStore {
-            hot_store: Arc::new(MemoryStore::from(
-                config.memory_store.unwrap(),
-                runtime_manager.clone(),
-            )),
+            hot_store: Arc::new(mem_store),
             warm_store: persistent_stores.pop_front(),
             cold_store: persistent_stores.pop_front(),
             config: hybrid_conf,
@@ -426,7 +430,7 @@ impl HybridStore {
         Ok(Default::default())
     }
 
-    pub async fn get_memory_buffer(&self, uid: &PartitionUId) -> Result<Arc<MemoryBuffer>> {
+    pub async fn get_memory_buffer(&self, uid: &PartitionUId) -> Result<Arc<UnifiedBuffer>> {
         self.hot_store.get_buffer(uid)
     }
 
@@ -469,7 +473,7 @@ impl HybridStore {
     async fn buffer_spill_impl(
         &self,
         uid: &PartitionUId,
-        buffer: Arc<MemoryBuffer>,
+        buffer: Arc<UnifiedBuffer>,
     ) -> Result<u64> {
         let spill_result = buffer.spill()?;
         if spill_result.is_none() {