Skip to main content

arminius_lib_nats/client/
maintenance.rs

1// SPDX-FileCopyrightText: 2025 BWI GmbH
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! Contains logic related to the maintenance mode.
6//!
7//! The maintenance mode enables us to notify all services at once at startup
8//! that necessary database operations are scheduled at the moment.
9use async_nats::jetstream;
10
11use crate::{NatsClient, NatsError};
12
13impl NatsClient {
14	/// Subject for the maintenance mode.
15	///
16	/// Maintenance mode is considered enabled if its value is 1 or unset,
17	/// disabled if zero.
18	const MAINTENANCE_MODE_KEY: &str = "maintenance_mode";
19
20	/// Sets maintenance mode to the boolean value
21	pub async fn set_maintenance_mode(&self, maintenance_mode: bool) -> Result<(), NatsError> {
22		let jetstream = jetstream::new(self.client.clone());
23		let store = Self::connect_to_store(&jetstream).await?;
24		store.put(Self::MAINTENANCE_MODE_KEY, vec![u8::from(maintenance_mode)].into()).await?;
25		Ok(())
26	}
27
28	/// Returns true if the NATS K/V store says we are in maintenance mode
29	pub async fn is_in_maintenance_mode(&self) -> Result<bool, NatsError> {
30		let jetstream = jetstream::new(self.client.clone());
31		let store = Self::connect_to_store(&jetstream).await?;
32		Ok(store
33			.entry(Self::MAINTENANCE_MODE_KEY)
34			.await?
35			.is_none_or(|entry| entry.value.first().unwrap_or(&1) == &1))
36	}
37
38	/// Returns a future that completes when maintenance mode has this status
39	pub async fn wait_for_maintenance_mode(&self, status: bool) -> Result<(), NatsError> {
40		let jetstream = jetstream::new(self.client.clone());
41		let store = Self::connect_to_store(&jetstream).await?;
42		while self.is_in_maintenance_mode().await? != status {
43			store.watch(Self::MAINTENANCE_MODE_KEY).await?;
44		}
45		Ok(())
46	}
47}
48
49#[cfg(test)]
50mod tests;