meta_srv/service/store/
cached_kv.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::any::Any;
16use std::collections::HashSet;
17use std::sync::atomic::{AtomicUsize, Ordering};
18use std::sync::{Arc, RwLock};
19
20use common_meta::error::{Error, Result};
21use common_meta::key::CACHE_KEY_PREFIXES;
22use common_meta::kv_backend::memory::MemoryKvBackend;
23use common_meta::kv_backend::txn::{Txn, TxnOp, TxnRequest, TxnResponse};
24use common_meta::kv_backend::{
25    KvBackend, KvBackendRef, ResettableKvBackend, ResettableKvBackendRef, TxnService,
26};
27use common_meta::range_stream::{DEFAULT_PAGE_SIZE, PaginationStream};
28use common_meta::rpc::KeyValue;
29use common_meta::rpc::store::{
30    BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest,
31    BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, DeleteRangeRequest,
32    DeleteRangeResponse, PutRequest, PutResponse, RangeRequest, RangeResponse,
33};
34use futures::TryStreamExt;
35
36use crate::metrics;
37use crate::state::State;
38
39pub type CheckLeaderRef = Arc<dyn CheckLeader>;
40
41pub trait CheckLeader: Sync + Send {
42    fn check(&self) -> bool;
43}
44
45struct AlwaysLeader;
46
47impl CheckLeader for AlwaysLeader {
48    fn check(&self) -> bool {
49        true
50    }
51}
52
53impl CheckLeader for RwLock<State> {
54    fn check(&self) -> bool {
55        self.read().unwrap().enable_leader_cache()
56    }
57}
58
59/// A cache dedicated to a Leader node, in order to cache some metadata.
60///
61/// To use this cache, the following constraints must be followed:
62///   1. The leader node can create this metadata.
63///   2. The follower node can create this metadata. The leader node can lazily retrieve
64///      the corresponding data through the caching loading mechanism.
65///   3. Only the leader node can update this metadata, as the cache cannot detect
66///      modifications made to the data on the follower node.
67///   4. Only the leader node can delete this metadata for the same reason mentioned above.
68pub struct LeaderCachedKvBackend {
69    check_leader: CheckLeaderRef,
70    store: KvBackendRef,
71    cache: ResettableKvBackendRef,
72    /// Incremented before and after each store mutation, also while the cache is
73    /// disabled. The first increment detects overlapping writers; the second
74    /// invalidates cache fills and loads that read the old value while the
75    /// store mutation was in flight.
76    version: AtomicUsize,
77    name: String,
78}
79
80impl LeaderCachedKvBackend {
81    pub fn new(check_leader: CheckLeaderRef, store: KvBackendRef) -> Self {
82        let name = format!("LeaderCached({})", store.name());
83        Self {
84            check_leader,
85            store,
86            cache: Arc::new(MemoryKvBackend::new()),
87            version: AtomicUsize::new(0),
88            name,
89        }
90    }
91
92    /// With a leader checker which always returns true when checking,
93    /// mainly used in test scenarios.
94    pub fn with_always_leader(store: KvBackendRef) -> Self {
95        Self::new(Arc::new(AlwaysLeader), store)
96    }
97
98    /// Loads the cached prefixes. Writes may run meanwhile: a prefix whose load overlapped a
99    /// write is dropped from the cache again, and read through later.
100    pub async fn load(&self) -> Result<()> {
101        for prefix in &CACHE_KEY_PREFIXES[..] {
102            let _timer =
103                metrics::METRIC_META_LEADER_CACHED_KV_LOAD_ELAPSED.with_label_values(&[prefix]);
104            let ver = self.get_version();
105
106            // TODO(weny): Refactors PaginationStream's output to unary output.
107            let stream = PaginationStream::new(
108                self.store.clone(),
109                RangeRequest::new().with_prefix(prefix.as_bytes()),
110                DEFAULT_PAGE_SIZE,
111                Ok,
112            )
113            .into_stream();
114
115            let kvs = stream.try_collect::<Vec<_>>().await?;
116            let keys = kvs.iter().map(|kv| kv.key.clone()).collect::<Vec<_>>();
117
118            self.cache
119                .batch_put(BatchPutRequest {
120                    kvs,
121                    prev_kv: false,
122                })
123                .await?;
124            // A write while loading may have changed what was read.
125            if !self.validate_version(ver) {
126                self.invalid_keys(keys).await?;
127            }
128        }
129
130        Ok(())
131    }
132
133    #[inline]
134    fn is_leader(&self) -> bool {
135        self.check_leader.check()
136    }
137
138    #[inline]
139    async fn invalid_key(&self, key: Vec<u8>) -> Result<()> {
140        let _ = self.cache.delete(&key, false).await?;
141        Ok(())
142    }
143
144    #[inline]
145    async fn invalid_keys(&self, keys: Vec<Vec<u8>>) -> Result<()> {
146        let txn = Txn::new().and_then(keys.into_iter().map(TxnOp::Delete).collect::<Vec<_>>());
147        let _ = self.cache.txn(txn).await?;
148        Ok(())
149    }
150
151    #[inline]
152    fn get_version(&self) -> usize {
153        self.version.load(Ordering::Relaxed)
154    }
155
156    #[inline]
157    fn create_new_version(&self) -> usize {
158        self.version.fetch_add(1, Ordering::Relaxed) + 1
159    }
160
161    #[inline]
162    fn validate_version(&self, version: usize) -> bool {
163        version == self.version.load(Ordering::Relaxed)
164    }
165}
166
167#[async_trait::async_trait]
168impl KvBackend for LeaderCachedKvBackend {
169    fn name(&self) -> &str {
170        &self.name
171    }
172
173    fn as_any(&self) -> &dyn Any {
174        self
175    }
176
177    async fn range(&self, req: RangeRequest) -> Result<RangeResponse> {
178        if !self.is_leader() {
179            return self.store.range(req).await;
180        }
181
182        // We can only cache for exact key queries (i.e. get requests)
183        // because we cannot confirm if a range response is complete.
184        if !req.range_end.is_empty() {
185            return self.store.range(req).await;
186        }
187
188        let res = self.cache.range(req.clone()).await?;
189        if !res.kvs.is_empty() {
190            return Ok(res);
191        }
192
193        let ver = self.get_version();
194
195        let res = self
196            .store
197            .range(RangeRequest {
198                // ignores `keys_only`
199                keys_only: false,
200                ..req.clone()
201            })
202            .await?;
203        if !res.kvs.is_empty() {
204            let KeyValue { key, value } = res.kvs[0].clone();
205            let put_req = PutRequest {
206                key: key.clone(),
207                value,
208                ..Default::default()
209            };
210            let _ = self.cache.put(put_req).await?;
211
212            if !self.validate_version(ver) {
213                self.invalid_key(key).await?;
214            }
215        }
216
217        return Ok(res);
218    }
219
220    async fn put(&self, req: PutRequest) -> Result<PutResponse> {
221        let ver = self.create_new_version();
222
223        let res = self.store.put(req.clone()).await?;
224        self.create_new_version();
225        let _ = self.cache.put(req.clone()).await?;
226
227        // Only our own completion increment may follow the starting version.
228        if !self.validate_version(ver + 1) {
229            self.invalid_key(req.key).await?;
230        }
231
232        Ok(res)
233    }
234
235    async fn batch_put(&self, req: BatchPutRequest) -> Result<BatchPutResponse> {
236        let ver = self.create_new_version();
237
238        let res = self.store.batch_put(req.clone()).await?;
239        self.create_new_version();
240        let _ = self.cache.batch_put(req.clone()).await?;
241
242        // Only our own completion increment may follow the starting version.
243        if !self.validate_version(ver + 1) {
244            let keys = req.kvs.into_iter().map(|kv| kv.key).collect::<Vec<_>>();
245            self.invalid_keys(keys).await?;
246        }
247
248        Ok(res)
249    }
250
251    async fn batch_get(&self, req: BatchGetRequest) -> Result<BatchGetResponse> {
252        if !self.is_leader() {
253            return self.store.batch_get(req).await;
254        }
255
256        let cached_res = self.cache.batch_get(req.clone()).await?;
257        let hit_keys = cached_res
258            .kvs
259            .iter()
260            .map(|kv| kv.key.clone())
261            .collect::<HashSet<_>>();
262
263        metrics::METRIC_META_KV_CACHE_HIT
264            .with_label_values(&["batch_get"])
265            .inc_by(hit_keys.len() as u64);
266
267        // The cache hit all keys
268        if cached_res.kvs.len() == req.keys.len() {
269            return Ok(cached_res);
270        }
271
272        let missed_keys = req
273            .keys
274            .iter()
275            .filter(|key| !hit_keys.contains(*key))
276            .cloned()
277            .collect::<Vec<_>>();
278        metrics::METRIC_META_KV_CACHE_MISS
279            .with_label_values(&["batch_get"])
280            .inc_by(missed_keys.len() as u64);
281
282        let remote_req = BatchGetRequest { keys: missed_keys };
283
284        let ver = self.get_version();
285
286        let remote_res = self.store.batch_get(remote_req).await?;
287        let put_req = BatchPutRequest {
288            kvs: remote_res.kvs.clone().into_iter().collect(),
289            ..Default::default()
290        };
291        let _ = self.cache.batch_put(put_req).await?;
292
293        if !self.validate_version(ver) {
294            let keys = remote_res
295                .kvs
296                .iter()
297                .map(|kv| kv.key().to_vec())
298                .collect::<Vec<_>>();
299            self.invalid_keys(keys).await?;
300        }
301
302        let mut merged_res = cached_res;
303        merged_res.kvs.extend(remote_res.kvs);
304        Ok(merged_res)
305    }
306
307    async fn compare_and_put(&self, req: CompareAndPutRequest) -> Result<CompareAndPutResponse> {
308        let _ = self.create_new_version();
309
310        let key = req.key.clone();
311        let res = self.store.compare_and_put(req).await?;
312        self.create_new_version();
313        // Delete key in the cache.
314        //
315        // Cache can not deal with the CAS operation, because it does
316        // not contain full data, so we need to delete the key.
317        self.invalid_key(key).await?;
318        Ok(res)
319    }
320
321    async fn delete_range(&self, req: DeleteRangeRequest) -> Result<DeleteRangeResponse> {
322        let _ = self.create_new_version();
323
324        let res = self.store.delete_range(req.clone()).await?;
325        self.create_new_version();
326        let _ = self.cache.delete_range(req).await?;
327        Ok(res)
328    }
329
330    async fn batch_delete(&self, req: BatchDeleteRequest) -> Result<BatchDeleteResponse> {
331        let _ = self.create_new_version();
332
333        let res = self.store.batch_delete(req.clone()).await?;
334        self.create_new_version();
335        let _ = self.cache.batch_delete(req).await?;
336        Ok(res)
337    }
338}
339
340#[async_trait::async_trait]
341impl TxnService for LeaderCachedKvBackend {
342    type Error = Error;
343
344    async fn txn(&self, txn: Txn) -> Result<TxnResponse> {
345        let _ = self.create_new_version();
346
347        let res = self.store.txn(txn.clone()).await?;
348        self.create_new_version();
349        let TxnRequest {
350            success, failure, ..
351        } = txn.into();
352        let mut all = success;
353        all.extend(failure);
354        // Delete all keys in the cache.
355        //
356        // Cache can not deal with the txn operation, because it does
357        // not contain full data, so we need to delete both keys.
358        let mut keys = Vec::with_capacity(all.len());
359        for txn_op in all {
360            match txn_op {
361                TxnOp::Put(key, _) => {
362                    keys.push(key);
363                }
364                TxnOp::Delete(key) => {
365                    keys.push(key);
366                }
367                TxnOp::Get(_) => {}
368            }
369        }
370        self.invalid_keys(keys).await?;
371
372        Ok(res)
373    }
374
375    fn max_txn_ops(&self) -> usize {
376        self.store.max_txn_ops()
377    }
378}
379
380impl ResettableKvBackend for LeaderCachedKvBackend {
381    fn reset(&self) {
382        self.create_new_version();
383        self.cache.reset()
384    }
385
386    fn as_kv_backend_ref(self: Arc<Self>) -> KvBackendRef<Self::Error> {
387        self
388    }
389}
390
391#[cfg(test)]
392mod tests {
393    use std::sync::Mutex;
394
395    use common_meta::kv_backend::txn::{Compare, CompareOp};
396    use common_meta::rpc::KeyValue;
397    use futures::poll;
398    use tokio::sync::oneshot;
399
400    use super::*;
401
402    /// Pauses one backend call at a known point, without relying on task scheduling.
403    #[derive(Default)]
404    struct Pause {
405        receiver: Mutex<Option<oneshot::Receiver<()>>>,
406    }
407
408    impl Pause {
409        fn arm(&self) -> oneshot::Sender<()> {
410            let (sender, receiver) = oneshot::channel();
411            assert!(self.receiver.lock().unwrap().replace(receiver).is_none());
412            sender
413        }
414
415        async fn wait(&self) {
416            let receiver = self.receiver.lock().unwrap().take();
417            if let Some(receiver) = receiver {
418                receiver.await.unwrap();
419            }
420        }
421    }
422
423    #[derive(Default)]
424    struct ControlledBackend {
425        inner: MemoryKvBackend<Error>,
426        before_write: Pause,
427        after_write: Pause,
428        after_read: Pause,
429    }
430
431    #[async_trait::async_trait]
432    impl TxnService for ControlledBackend {
433        type Error = Error;
434
435        async fn txn(&self, txn: Txn) -> Result<TxnResponse> {
436            self.before_write.wait().await;
437            let res = self.inner.txn(txn).await;
438            self.after_write.wait().await;
439            res
440        }
441
442        fn max_txn_ops(&self) -> usize {
443            self.inner.max_txn_ops()
444        }
445    }
446
447    #[async_trait::async_trait]
448    impl KvBackend for ControlledBackend {
449        fn name(&self) -> &str {
450            "ControlledBackend"
451        }
452
453        fn as_any(&self) -> &dyn Any {
454            self
455        }
456
457        async fn range(&self, req: RangeRequest) -> Result<RangeResponse> {
458            let res = self.inner.range(req).await;
459            self.after_read.wait().await;
460            res
461        }
462
463        async fn batch_get(&self, req: BatchGetRequest) -> Result<BatchGetResponse> {
464            let res = self.inner.batch_get(req).await;
465            self.after_read.wait().await;
466            res
467        }
468
469        async fn put(&self, req: PutRequest) -> Result<PutResponse> {
470            self.before_write.wait().await;
471            let res = self.inner.put(req).await;
472            self.after_write.wait().await;
473            res
474        }
475
476        async fn batch_put(&self, req: BatchPutRequest) -> Result<BatchPutResponse> {
477            self.before_write.wait().await;
478            let res = self.inner.batch_put(req).await;
479            self.after_write.wait().await;
480            res
481        }
482
483        async fn delete_range(&self, req: DeleteRangeRequest) -> Result<DeleteRangeResponse> {
484            self.before_write.wait().await;
485            let res = self.inner.delete_range(req).await;
486            self.after_write.wait().await;
487            res
488        }
489
490        async fn batch_delete(&self, req: BatchDeleteRequest) -> Result<BatchDeleteResponse> {
491            self.before_write.wait().await;
492            let res = self.inner.batch_delete(req).await;
493            self.after_write.wait().await;
494            res
495        }
496    }
497
498    #[derive(Clone, Copy, Debug)]
499    enum Write {
500        Txn,
501        TxnFailure,
502        Put,
503        BatchPut,
504        CompareAndPut,
505        DeleteRange,
506        BatchDelete,
507    }
508
509    const WRITES: [Write; 7] = [
510        Write::Txn,
511        Write::TxnFailure,
512        Write::Put,
513        Write::BatchPut,
514        Write::CompareAndPut,
515        Write::DeleteRange,
516        Write::BatchDelete,
517    ];
518    /// Under the first prefix the leader cache loads.
519    const KEY: &[u8] = b"__table_name/1";
520
521    impl Write {
522        async fn execute(self, store: &LeaderCachedKvBackend, value: &[u8]) {
523            match self {
524                Self::Txn => {
525                    store
526                        .txn(Txn::new().and_then(vec![TxnOp::Put(KEY.to_vec(), value.to_vec())]))
527                        .await
528                        .unwrap();
529                }
530                Self::TxnFailure => {
531                    let res = store
532                        .txn(
533                            Txn::new()
534                                .when(vec![Compare::with_value_not_exists(
535                                    KEY.to_vec(),
536                                    CompareOp::Equal,
537                                )])
538                                .or_else(vec![TxnOp::Put(KEY.to_vec(), value.to_vec())]),
539                        )
540                        .await
541                        .unwrap();
542                    assert!(!res.succeeded);
543                }
544                Self::Put => {
545                    store
546                        .put(PutRequest::new().with_key(KEY).with_value(value))
547                        .await
548                        .unwrap();
549                }
550                Self::BatchPut => {
551                    store
552                        .batch_put(BatchPutRequest {
553                            kvs: vec![KeyValue {
554                                key: KEY.to_vec(),
555                                value: value.to_vec(),
556                            }],
557                            ..Default::default()
558                        })
559                        .await
560                        .unwrap();
561                }
562                Self::CompareAndPut => {
563                    let res = store
564                        .compare_and_put(
565                            CompareAndPutRequest::new()
566                                .with_key(KEY)
567                                .with_expect(b"old")
568                                .with_value(value),
569                        )
570                        .await
571                        .unwrap();
572                    assert!(res.success);
573                }
574                Self::DeleteRange => {
575                    store
576                        .delete_range(DeleteRangeRequest::new().with_key(KEY))
577                        .await
578                        .unwrap();
579                }
580                Self::BatchDelete => {
581                    store
582                        .batch_delete(BatchDeleteRequest {
583                            keys: vec![KEY.to_vec()],
584                            prev_kv: false,
585                        })
586                        .await
587                        .unwrap();
588                }
589            }
590        }
591    }
592
593    async fn read(store: &LeaderCachedKvBackend, batch: bool) -> Option<KeyValue> {
594        if batch {
595            store
596                .batch_get(BatchGetRequest {
597                    keys: vec![KEY.to_vec()],
598                })
599                .await
600                .unwrap()
601                .kvs
602                .pop()
603        } else {
604            store.get(KEY).await.unwrap()
605        }
606    }
607
608    /// A cache miss reads the old value after a writer starts, then returns after
609    /// the writer has committed and updated/invalidated the cache (#9548).
610    #[tokio::test]
611    async fn test_cache_miss_overlaps_write() {
612        for write in WRITES {
613            for batch in [true, false] {
614                let backing = Arc::new(ControlledBackend::default());
615                backing
616                    .inner
617                    .put(PutRequest::new().with_key(KEY).with_value(b"old"))
618                    .await
619                    .unwrap();
620                let store = LeaderCachedKvBackend::with_always_leader(backing.clone());
621
622                let commit = backing.before_write.arm();
623                let writer = write.execute(&store, b"new");
624                tokio::pin!(writer);
625                assert!(poll!(writer.as_mut()).is_pending());
626
627                let return_read = backing.after_read.arm();
628                let reader = read(&store, batch);
629                tokio::pin!(reader);
630                assert!(poll!(reader.as_mut()).is_pending());
631
632                commit.send(()).unwrap();
633                writer.await;
634                return_read.send(()).unwrap();
635                assert_eq!(reader.await.unwrap().value(), b"old");
636
637                let durable = backing.inner.get(KEY).await.unwrap();
638                assert_eq!(
639                    read(&store, batch).await,
640                    durable,
641                    "{write:?}, batch={batch}"
642                );
643            }
644        }
645    }
646
647    /// Start order, commit order and response order can all differ. An older
648    /// put response must not overwrite the cache after a newer mutation.
649    #[tokio::test]
650    async fn test_overlapping_writes() {
651        for first in [Write::Put, Write::BatchPut] {
652            for second in WRITES {
653                for reverse_commit in [false, true] {
654                    let backing = Arc::new(ControlledBackend::default());
655                    backing
656                        .inner
657                        .put(PutRequest::new().with_key(KEY).with_value(b"old"))
658                        .await
659                        .unwrap();
660                    let store = LeaderCachedKvBackend::with_always_leader(backing.clone());
661
662                    let resume = if reverse_commit {
663                        backing.before_write.arm()
664                    } else {
665                        backing.after_write.arm()
666                    };
667                    // Keep the initial value so the second CAS succeeds.
668                    let writer = first.execute(&store, b"old");
669                    tokio::pin!(writer);
670                    assert!(poll!(writer.as_mut()).is_pending());
671                    second.execute(&store, b"new").await;
672                    resume.send(()).unwrap();
673                    writer.await;
674
675                    let durable = backing.inner.get(KEY).await.unwrap();
676                    assert_eq!(
677                        read(&store, true).await,
678                        durable,
679                        "{first:?}, {second:?}, reverse_commit={reverse_commit}"
680                    );
681                }
682            }
683        }
684    }
685
686    /// The later-started put commits first but returns last, after its stored
687    /// value has already been replaced by the first writer.
688    #[tokio::test]
689    async fn test_write_response_reordering() {
690        for first in WRITES {
691            for second in [Write::Put, Write::BatchPut] {
692                let backing = Arc::new(ControlledBackend::default());
693                backing
694                    .inner
695                    .put(PutRequest::new().with_key(KEY).with_value(b"old"))
696                    .await
697                    .unwrap();
698                let store = LeaderCachedKvBackend::with_always_leader(backing.clone());
699
700                let commit_first = backing.before_write.arm();
701                let first_writer = first.execute(&store, b"new");
702                tokio::pin!(first_writer);
703                assert!(poll!(first_writer.as_mut()).is_pending());
704
705                let return_second = backing.after_write.arm();
706                let second_writer = second.execute(&store, b"old");
707                tokio::pin!(second_writer);
708                assert!(poll!(second_writer.as_mut()).is_pending());
709
710                commit_first.send(()).unwrap();
711                first_writer.await;
712                return_second.send(()).unwrap();
713                second_writer.await;
714
715                let durable = backing.inner.get(KEY).await.unwrap();
716                assert_eq!(read(&store, true).await, durable, "{first:?}, {second:?}");
717            }
718        }
719    }
720
721    struct EnableCache(std::sync::atomic::AtomicBool);
722
723    impl CheckLeader for EnableCache {
724        fn check(&self) -> bool {
725            self.0.load(Ordering::Relaxed)
726        }
727    }
728
729    /// A write made while the cache is disabled overlaps the load: the load reads the old value
730    /// either before the write commits and installs it afterwards, or while the write is in
731    /// flight and completes, with the cache enabled, before the write commits.
732    #[tokio::test]
733    async fn test_load_overlaps_write() {
734        for write in WRITES {
735            for write_first in [false, true] {
736                let backing = Arc::new(ControlledBackend::default());
737                backing
738                    .inner
739                    .put(PutRequest::new().with_key(KEY).with_value(b"old"))
740                    .await
741                    .unwrap();
742                let enabled = Arc::new(EnableCache(Default::default()));
743                let store = LeaderCachedKvBackend::new(enabled.clone(), backing.clone());
744
745                if write_first {
746                    let commit = backing.before_write.arm();
747                    let writer = write.execute(&store, b"new");
748                    tokio::pin!(writer);
749                    assert!(poll!(writer.as_mut()).is_pending());
750                    store.load().await.unwrap();
751                    enabled.0.store(true, Ordering::Relaxed);
752                    commit.send(()).unwrap();
753                    writer.await;
754                } else {
755                    let return_read = backing.after_read.arm();
756                    let loader = store.load();
757                    tokio::pin!(loader);
758                    assert!(poll!(loader.as_mut()).is_pending());
759                    write.execute(&store, b"new").await;
760                    return_read.send(()).unwrap();
761                    loader.await.unwrap();
762                    enabled.0.store(true, Ordering::Relaxed);
763                }
764
765                let durable = backing.inner.get(KEY).await.unwrap();
766                assert_eq!(
767                    read(&store, false).await,
768                    durable,
769                    "{write:?}, write_first={write_first}"
770                );
771            }
772        }
773    }
774
775    #[tokio::test]
776    async fn test_cache_miss_overlaps_reset() {
777        for batch in [true, false] {
778            let backing = Arc::new(ControlledBackend::default());
779            backing
780                .inner
781                .put(PutRequest::new().with_key(KEY).with_value(b"old"))
782                .await
783                .unwrap();
784            let store = LeaderCachedKvBackend::with_always_leader(backing.clone());
785
786            let return_read = backing.after_read.arm();
787            let reader = read(&store, batch);
788            tokio::pin!(reader);
789            assert!(poll!(reader.as_mut()).is_pending());
790            store.reset();
791            return_read.send(()).unwrap();
792            reader.await;
793            assert!(store.cache.get(KEY).await.unwrap().is_none());
794        }
795    }
796
797    fn create_leader_cached_kv_backend() -> LeaderCachedKvBackend {
798        let store = Arc::new(MemoryKvBackend::new());
799        LeaderCachedKvBackend::with_always_leader(store)
800    }
801
802    #[tokio::test]
803    async fn test_get_put_delete() {
804        let cached_store = create_leader_cached_kv_backend();
805        let inner_store = cached_store.store.clone();
806        let inner_cache = cached_store.cache.clone();
807
808        let key = "test_key".to_owned().into_bytes();
809        let value = "value".to_owned().into_bytes();
810
811        let put_req = PutRequest {
812            key: key.clone(),
813            value: value.clone(),
814            ..Default::default()
815        };
816        let _ = inner_store.put(put_req).await.unwrap();
817
818        let cached_value = inner_cache.get(&key).await.unwrap();
819        assert!(cached_value.is_none());
820
821        let cached_value = cached_store.get(&key).await.unwrap().unwrap();
822        assert_eq!(cached_value.value(), value);
823
824        let cached_value = inner_cache.get(&key).await.unwrap().unwrap();
825        assert_eq!(cached_value.value(), value);
826
827        let res = cached_store.delete(&key, true).await.unwrap().unwrap();
828        assert_eq!(res.value(), value);
829
830        let cached_value = inner_cache.get(&key).await.unwrap();
831        assert!(cached_value.is_none());
832    }
833
834    #[tokio::test]
835    async fn test_batch_get_put_delete() {
836        let cached_store = create_leader_cached_kv_backend();
837        let inner_store = cached_store.store.clone();
838        let inner_cache = cached_store.cache.clone();
839
840        let kvs = (1..3)
841            .map(|i| {
842                let key = format!("test_key_{}", i).into_bytes();
843                let value = format!("value_{}", i).into_bytes();
844                KeyValue { key, value }
845            })
846            .collect::<Vec<_>>();
847
848        let batch_put_req = BatchPutRequest {
849            kvs: kvs.clone(),
850            ..Default::default()
851        };
852
853        let _ = inner_store.batch_put(batch_put_req).await.unwrap();
854
855        let keys = (1..5)
856            .map(|i| format!("test_key_{}", i).into_bytes())
857            .collect::<Vec<_>>();
858
859        let batch_get_req = BatchGetRequest { keys };
860
861        let cached_values = inner_cache.batch_get(batch_get_req.clone()).await.unwrap();
862        assert!(cached_values.kvs.is_empty());
863
864        let cached_values = cached_store.batch_get(batch_get_req.clone()).await.unwrap();
865        assert_eq!(cached_values.kvs.len(), 2);
866
867        let cached_values = inner_cache.batch_get(batch_get_req.clone()).await.unwrap();
868        assert_eq!(cached_values.kvs.len(), 2);
869
870        cached_store.reset();
871
872        let cached_values = inner_cache.batch_get(batch_get_req).await.unwrap();
873        assert!(cached_values.kvs.is_empty());
874    }
875
876    #[tokio::test]
877    async fn test_txn() {
878        let cached_store = create_leader_cached_kv_backend();
879        let inner_cache = cached_store.cache.clone();
880
881        let kvs = (1..5)
882            .map(|i| {
883                let key = format!("test_key_{}", i).into_bytes();
884                let value = format!("value_{}", i).into_bytes();
885                KeyValue { key, value }
886            })
887            .collect::<Vec<_>>();
888
889        let batch_put_req = BatchPutRequest {
890            kvs: kvs.clone(),
891            ..Default::default()
892        };
893        let _ = cached_store.batch_put(batch_put_req).await.unwrap();
894
895        let keys = (1..5)
896            .map(|i| format!("test_key_{}", i).into_bytes())
897            .collect::<Vec<_>>();
898        let batch_get_req = BatchGetRequest { keys };
899        let cached_values = inner_cache.batch_get(batch_get_req.clone()).await.unwrap();
900        assert_eq!(cached_values.kvs.len(), 4);
901
902        let put_ops = (1..5)
903            .map(|i| {
904                let key = format!("test_key_{}", i).into_bytes();
905                let value = format!("value_{}", i).into_bytes();
906                TxnOp::Put(key, value)
907            })
908            .collect::<Vec<_>>();
909        let txn = Txn::new().and_then(put_ops);
910        let _ = cached_store.txn(txn).await.unwrap();
911
912        let cached_values = inner_cache.batch_get(batch_get_req).await.unwrap();
913        assert!(cached_values.kvs.is_empty());
914    }
915}
meta_srv/service/store/cached_kv.rs

meta_srv/service/store/
cached_kv.rs
