fix: revert memtable pk rb cache to rwlock

src/mito2/src/memtable/time_series.rs

@@ -19,7 +19,6 @@ use std::sync::atomic::{AtomicI64, AtomicU32, Ordering};
 use std::sync::{Arc, RwLock};
 
 use api::v1::OpType;
-use arc_swap::ArcSwapOption;
 use common_telemetry::debug;
 use datatypes::arrow;
 use datatypes::arrow::array::ArrayRef;
@@ -260,7 +259,7 @@ impl Memtable for TimeSeriesMemtable {
 
 struct PrimaryKey {
     bytes: Vec<u8>,
-    record_batch: ArcSwapOption<RecordBatch>,
+    record_batch: RwLock<Option<Arc<RecordBatch>>>,
 }
 
 impl Clone for PrimaryKey {
@@ -276,21 +275,25 @@ impl PrimaryKey {
     fn new(bytes: Vec<u8>) -> Self {
         Self {
             bytes,
-            record_batch: ArcSwapOption::empty(),
+            record_batch: RwLock::new(None),
         }
     }
 
     fn get_or_update_record_batch_with<F: FnMut() -> Result<RecordBatch>>(
         &self,
         mut f: F,
     ) -> Result<Arc<RecordBatch>> {
-        if let Some(rb) = self.record_batch.load_full() {
-            return Ok(rb);
+        if let Some(rb) = self.record_batch.read().unwrap().as_ref() {
+            return Ok(rb.clone());
         }
 
         let batch = Arc::new(f()?);
-        self.record_batch.store(Some(batch.clone()));
-        Ok(batch)
+        Ok(self
+            .record_batch
+            .write()
+            .unwrap()
+            .get_or_insert(batch)
+            .clone())
     }
 }