witchcraft_server/
readiness.rs

1// Copyright 2022 Palantir Technologies, Inc.
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//! Readiness checks.
15use once_cell::sync::Lazy;
16use parking_lot::Mutex;
17use regex::Regex;
18use serde::Serialize;
19use staged_builder::staged_builder;
20use std::collections::btree_map;
21use std::collections::hash_map;
22use std::collections::{BTreeMap, HashMap};
23use std::sync::Arc;
24
25static TYPE_PATTERN: Lazy<Regex> = Lazy::new(|| Regex::new("^[A-Z_]+$").unwrap());
26
27/// A readiness check.
28pub trait ReadinessCheck: 'static + Sync + Send {
29    /// Returns the check's type.
30    ///
31    /// The type must be `SCREAMING_SNAKE_CASE`.
32    fn type_(&self) -> &str;
33
34    /// Performs the check, returning its result.
35    fn result(&self) -> ReadinessCheckResult;
36}
37
38/// The result of a readiness check.
39#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
40#[staged_builder]
41pub struct ReadinessCheckResult {
42    successful: bool,
43}
44
45impl ReadinessCheckResult {
46    /// The success of the check.
47    #[inline]
48    pub fn successful(&self) -> bool {
49        self.successful
50    }
51}
52
53/// A registry of readiness checks for the server.
54pub struct ReadinessCheckRegistry {
55    checks: Mutex<HashMap<String, Arc<dyn ReadinessCheck>>>,
56}
57
58impl ReadinessCheckRegistry {
59    pub(crate) fn new() -> Self {
60        ReadinessCheckRegistry {
61            checks: Mutex::new(HashMap::new()),
62        }
63    }
64
65    /// Registers a new readiness check.
66    ///
67    /// # Panics
68    ///
69    /// Panics if the check's type is not `SCREAMING_SNAKE_CASE` or if a check with the same type is already registered.
70    pub fn register<T>(&self, check: T)
71    where
72        T: ReadinessCheck,
73    {
74        self.register_inner(Arc::new(check))
75    }
76
77    fn register_inner(&self, check: Arc<dyn ReadinessCheck>) {
78        let type_ = check.type_();
79
80        assert!(
81            TYPE_PATTERN.is_match(type_),
82            "{type_} must `SCREAMING_SNAKE_CASE",
83        );
84
85        match self.checks.lock().entry(type_.to_string()) {
86            hash_map::Entry::Occupied(_) => {
87                panic!("a check has already been registered for type {type_}")
88            }
89            hash_map::Entry::Vacant(e) => {
90                e.insert(check);
91            }
92        }
93    }
94
95    pub(crate) fn run_checks(&self) -> BTreeMap<String, ReadinessCheckMetadata> {
96        // A bit of extra complexity to allow registration while we're running checks.
97        let mut results = BTreeMap::new();
98
99        let mut progress = true;
100        while progress {
101            progress = false;
102            let checks = self.checks.lock().clone();
103
104            for (type_, check) in checks {
105                if let btree_map::Entry::Vacant(e) = results.entry(type_.clone()) {
106                    let result = check.result();
107                    e.insert(ReadinessCheckMetadata {
108                        r#type: type_,
109                        successful: result.successful,
110                    });
111                    progress = true;
112                }
113            }
114        }
115
116        results
117    }
118}
119
120#[derive(Serialize)]
121pub(crate) struct ReadinessCheckMetadata {
122    r#type: String,
123    pub(crate) successful: bool,
124}
witchcraft_server/readiness.rs

witchcraft_server/
readiness.rs
