1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78

/*
 * Copyright 2021 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

use std::sync::atomic::AtomicBool;

use hyper::{Body, Response, StatusCode};
use std::panic;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::Arc;

#[derive(Clone)]
pub struct Health {
    healthy: Arc<AtomicBool>,
}

impl Health {
    pub fn new() -> Self {
        let health = Self {
            healthy: Arc::new(AtomicBool::new(true)),
        };

        let healthy = health.healthy.clone();
        let default_hook = panic::take_hook();
        panic::set_hook(Box::new(move |panic_info| {
            tracing::error!(%panic_info, "Panic has occurred. Moving to Unhealthy");
            healthy.swap(false, Relaxed);
            default_hook(panic_info);
        }));

        health
    }

    /// returns a HTTP 200 response if the proxy is healthy.
    pub fn check_healthy(&self) -> Response<Body> {
        if self.healthy.load(Relaxed) {
            return Response::new("ok".into());
        };

        let mut response = Response::new(Body::empty());
        *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
        response
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use hyper::StatusCode;

    #[test]
    fn panic_hook() {
        let health = Health::new();

        let response = health.check_healthy();
        assert_eq!(response.status(), StatusCode::OK);

        let _ = std::panic::catch_unwind(|| {
            panic!("oh no!");
        });

        let response = health.check_healthy();
        assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
    }
}