common_procedure/
test_util.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::collections::hash_map::Entry;
16use std::collections::HashMap;
17use std::sync::{Arc, RwLock};
18
19use snafu::ensure;
20
21use super::*;
22use crate::error;
23use crate::store::poison_store::PoisonStore;
24
25/// A poison store that uses an in-memory map to store the poison state.
26#[derive(Debug, Default)]
27pub struct InMemoryPoisonStore {
28    map: Arc<RwLock<HashMap<String, String>>>,
29}
30
31impl InMemoryPoisonStore {
32    /// Create a new in-memory poison manager.
33    pub fn new() -> Self {
34        Self::default()
35    }
36}
37
38#[async_trait::async_trait]
39impl PoisonStore for InMemoryPoisonStore {
40    async fn try_put_poison(&self, key: String, token: String) -> Result<()> {
41        let mut map = self.map.write().unwrap();
42        match map.entry(key) {
43            Entry::Vacant(v) => {
44                v.insert(token.to_string());
45            }
46            Entry::Occupied(o) => {
47                let value = o.get();
48                ensure!(
49                    value == &token,
50                    error::UnexpectedSnafu {
51                        err_msg: format!("The poison is already set by other token {}", value)
52                    }
53                );
54            }
55        }
56        Ok(())
57    }
58
59    async fn delete_poison(&self, key: String, token: String) -> Result<()> {
60        let mut map = self.map.write().unwrap();
61        match map.entry(key) {
62            Entry::Vacant(_) => {
63                // do nothing
64            }
65            Entry::Occupied(o) => {
66                let value = o.get();
67                ensure!(
68                    value == &token,
69                    error::UnexpectedSnafu {
70                        err_msg: format!("The poison is not set by the token {}", value)
71                    }
72                );
73
74                o.remove();
75            }
76        }
77        Ok(())
78    }
79
80    async fn get_poison(&self, key: &str) -> Result<Option<String>> {
81        let map = self.map.read().unwrap();
82        let key = key.to_string();
83        Ok(map.get(&key).cloned())
84    }
85}
common_procedure/test_util.rs

common_procedure/
test_util.rs
