WIP Get UniPS write batch size

proxy_components/engine_tiflash/src/proxy_utils/engine_ext.rs

@@ -45,6 +45,10 @@ impl PageStorageExt {
         self.helper().get_wb_size(wb) as usize
     }
 
+    pub fn write_batch_count(&self, wb: RawVoidPtr) -> usize {
+        self.helper().get_wb_count(wb) != 0
+    }
+
     pub fn write_batch_is_empty(&self, wb: RawVoidPtr) -> bool {
         self.helper().is_wb_empty(wb) != 0
     }

proxy_components/engine_tiflash/src/ps_engine/ps_write_batch.rs

@@ -54,9 +54,7 @@ impl ElementaryWriteBatch for PSElementWriteBatch {
     }
 
     fn count(&self) -> usize {
-        // FIXME
-        // TODO
-        0
+        self.ps_ext.write_batch_count(self.ps_wb.ptr)
     }
 
     fn is_empty(&self) -> bool {

proxy_components/mock-engine-store/src/mock_store/mock_engine_store_server.rs

@@ -400,6 +400,7 @@ pub fn gen_engine_store_server_helper(
             fn_wb_put_page: Some(ffi_mockps_wb_put_page),
             fn_wb_del_page: Some(ffi_mockps_wb_del_page),
             fn_get_wb_size: Some(ffi_mockps_get_wb_size),
+            fn_get_wb_count: Some(ffi_mockps_get_wb_count),
             fn_is_wb_empty: Some(ffi_mockps_is_wb_empty),
             fn_handle_merge_wb: Some(ffi_mockps_handle_merge_wb),
             fn_handle_clear_wb: Some(ffi_mockps_handle_clear_wb),

proxy_components/mock-engine-store/src/mock_store/mock_page_storage.rs

@@ -115,6 +115,11 @@ pub unsafe extern "C" fn ffi_mockps_get_wb_size(wb: RawVoidPtr) -> u64 {
     wb.data.len() as u64
 }
 
+pub unsafe extern "C" fn ffi_mockps_get_wb_count(wb: RawVoidPtr) -> u64 {
+    let wb = <&mut MockPSWriteBatch as From<RawVoidPtr>>::from(wb);
+    wb.data.len() as u64
+}
+
 pub unsafe extern "C" fn ffi_mockps_is_wb_empty(wb: RawVoidPtr) -> u8 {
     let wb = <&mut MockPSWriteBatch as From<RawVoidPtr>>::from(wb);
     u8::from(wb.data.is_empty())

proxy_components/proxy_ffi/src/engine_store_helper_impls.rs

@@ -531,6 +531,11 @@ impl EngineStoreServerHelper {
         unsafe { (self.ps.fn_get_wb_size.into_inner())(wb) }
     }
 
+    pub fn get_wb_count(&self, wb: RawVoidPtr) -> usize {
+        debug_assert!(self.ps.fn_get_wb_count.is_some());
+        unsafe { (self.ps.fn_get_wb_count.into_inner())(wb) }
+    }
+
     pub fn is_wb_empty(&self, wb: RawVoidPtr) -> u8 {
         debug_assert!(self.ps.fn_is_wb_empty.is_some());
         unsafe { (self.ps.fn_is_wb_empty.into_inner())(wb) }

proxy_components/proxy_ffi/src/interfaces.rs

@@ -519,6 +519,8 @@ pub mod root {
             >,
             pub fn_get_wb_size:
                 ::std::option::Option<unsafe extern "C" fn(arg1: root::DB::RawVoidPtr) -> u64>,
+            pub fn_get_wb_count:
+                ::std::option::Option<unsafe extern "C" fn(arg1: root::DB::RawVoidPtr) -> u64>,
             pub fn_is_wb_empty:
                 ::std::option::Option<unsafe extern "C" fn(arg1: root::DB::RawVoidPtr) -> u8>,
             pub fn_handle_merge_wb: ::std::option::Option<
@@ -789,7 +791,7 @@ pub mod root {
                 arg3: root::DB::RawVoidPtr,
             ) -> u32;
         }
-        pub const RAFT_STORE_PROXY_VERSION: u64 = 9679186549381427051;
+        pub const RAFT_STORE_PROXY_VERSION: u64 = 4376098841976428098;
         pub const RAFT_STORE_PROXY_MAGIC_NUMBER: u32 = 324508639;
     }
 }

raftstore-proxy/ffi/src/RaftStoreProxyFFI/@version

@@ -1,3 +1,3 @@
 #pragma once
 #include <cstdint>
-namespace DB { constexpr uint64_t RAFT_STORE_PROXY_VERSION = 9679186549381427051ull; }
+namespace DB { constexpr uint64_t RAFT_STORE_PROXY_VERSION = 4376098841976428098ull; }

raftstore-proxy/ffi/src/RaftStoreProxyFFI/ProxyFFI.h

@@ -297,6 +297,7 @@ struct PageStorageInterfaces {
   void (*fn_wb_put_page)(RawVoidPtr, BaseBuffView, BaseBuffView);
   void (*fn_wb_del_page)(RawVoidPtr, BaseBuffView);
   uint64_t (*fn_get_wb_size)(RawVoidPtr);
+  uint64_t (*fn_get_wb_count)(RawVoidPtr);
   uint8_t (*fn_is_wb_empty)(RawVoidPtr);
   void (*fn_handle_merge_wb)(RawVoidPtr, RawVoidPtr);
   void (*fn_handle_clear_wb)(RawVoidPtr);