common_meta/key/
maintenance.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::sync::Arc;
16
17use crate::error::Result;
18use crate::key::MAINTENANCE_KEY;
19use crate::kv_backend::KvBackendRef;
20use crate::rpc::store::PutRequest;
21
22pub type MaintenanceModeManagerRef = Arc<MaintenanceModeManager>;
23
24/// The maintenance mode manager.
25///
26/// Used to enable or disable maintenance mode.
27#[derive(Clone)]
28pub struct MaintenanceModeManager {
29    kv_backend: KvBackendRef,
30}
31
32impl MaintenanceModeManager {
33    pub fn new(kv_backend: KvBackendRef) -> Self {
34        Self { kv_backend }
35    }
36
37    /// Enables maintenance mode.
38    pub async fn set_maintenance_mode(&self) -> Result<()> {
39        let req = PutRequest {
40            key: Vec::from(MAINTENANCE_KEY),
41            value: vec![],
42            prev_kv: false,
43        };
44        self.kv_backend.put(req).await?;
45        Ok(())
46    }
47
48    /// Unsets maintenance mode.
49    pub async fn unset_maintenance_mode(&self) -> Result<()> {
50        self.kv_backend
51            .delete(MAINTENANCE_KEY.as_bytes(), false)
52            .await?;
53        Ok(())
54    }
55
56    /// Returns true if maintenance mode is enabled.
57    pub async fn maintenance_mode(&self) -> Result<bool> {
58        self.kv_backend.exists(MAINTENANCE_KEY.as_bytes()).await
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use std::sync::Arc;
65
66    use crate::key::maintenance::MaintenanceModeManager;
67    use crate::kv_backend::memory::MemoryKvBackend;
68
69    #[tokio::test]
70    async fn test_maintenance_mode_manager() {
71        let maintenance_mode_manager = Arc::new(MaintenanceModeManager::new(Arc::new(
72            MemoryKvBackend::new(),
73        )));
74        assert!(!maintenance_mode_manager.maintenance_mode().await.unwrap());
75        maintenance_mode_manager
76            .set_maintenance_mode()
77            .await
78            .unwrap();
79        assert!(maintenance_mode_manager.maintenance_mode().await.unwrap());
80        maintenance_mode_manager
81            .unset_maintenance_mode()
82            .await
83            .unwrap();
84        assert!(!maintenance_mode_manager.maintenance_mode().await.unwrap());
85    }
86}
common_meta/key/maintenance.rs

common_meta/key/
maintenance.rs
