diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index 9ea4e04ec36f..02955aa5c38d 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/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, - record_batch: ArcSwapOption, + record_batch: RwLock>>, } impl Clone for PrimaryKey { @@ -276,7 +275,7 @@ impl PrimaryKey { fn new(bytes: Vec) -> Self { Self { bytes, - record_batch: ArcSwapOption::empty(), + record_batch: RwLock::new(None), } } @@ -284,13 +283,17 @@ impl PrimaryKey { &self, mut f: F, ) -> Result> { - 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()) } }