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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

//! An actix-web service to implement [Dockerflow](https://github.com/mozilla-services/Dockerflow).

use std::collections::HashMap;

use actix_web::{
    get,
    web::{self, Data},
    HttpRequest, HttpResponse,
};
use merino_settings::Settings;
use serde::{Deserialize, Serialize};
use tracing::Level;

use crate::errors::HandlerError;

/// Handles required Dockerflow Endpoints.
pub fn configure(config: &mut web::ServiceConfig) {
    config
        .service(lbheartbeat)
        .service(heartbeat)
        .service(version)
        .service(test_error);
}

/// Used by the load balancer to indicate that the server can respond to
/// requests. Should just return OK.
#[get("__lbheartbeat__")]
fn lbheartbeat(_: HttpRequest) -> HttpResponse {
    HttpResponse::Ok().body("")
}

/// Return the contents of the `version.json` file created by CircleCI and stored
/// in the Docker root (or the TBD version stored in the Git repo).
#[get("__version__")]
fn version(_: HttpRequest) -> HttpResponse {
    HttpResponse::Ok()
        .content_type("application/json")
        .body(include_str!("../../../version.json"))
}

/// The status of an individual check, or the whole system, as reported by /__heartbeat__.
#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "snake_case")]
enum CheckStatus {
    /// Everything is OK
    Ok,
    /// The check could not determine the status.
    Unknown,
    /// Something is wrong, but it is not interrupting the system.
    #[allow(dead_code)]
    Warn,
    /// Something is wrong, and it is interrupting the system.
    #[allow(dead_code)]
    Error,
}

impl Default for CheckStatus {
    fn default() -> Self {
        Self::Unknown
    }
}

/// A response to the `/__heartbeat__` endpoint.
#[derive(Debug, Default)]
struct HeartbeatResponse {
    /// Any checks that are relevant to the state of the system.
    checks: HashMap<String, CheckStatus>,
}

impl HeartbeatResponse {
    /// The overall status of all checks.
    ///
    /// Takes the worst state of any check contained, or `CheckStatus::Unknown`
    /// if there are no contained checks.
    fn status(&self) -> CheckStatus {
        self.checks
            .values()
            .copied()
            .max()
            .unwrap_or(CheckStatus::Unknown)
    }

    /// Add the results of a check.
    fn add_check<S: Into<String>>(&mut self, name: S, check: CheckStatus) {
        self.checks.insert(name.into(), check);
    }
}

// Serde doesn't have a concept of "derived" fields for serialization. So
// instead define a concrete type with the calculated field, and delegate
// serialization to that.
impl Serialize for HeartbeatResponse {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        #[derive(Serialize)]
        #[allow(clippy::missing_docs_in_private_items)]
        struct Extended<'a> {
            status: CheckStatus,
            checks: &'a HashMap<String, CheckStatus>,
        }

        let ext = Extended {
            status: self.status(),
            checks: &self.checks,
        };

        ext.serialize(serializer)
    }
}

/// Returns a status message indicating the current state of the server.
#[get("__heartbeat__")]
fn heartbeat(_: HttpRequest) -> HttpResponse {
    let mut checklist = HeartbeatResponse::default();
    checklist.add_check("heartbeat", CheckStatus::Ok);
    HttpResponse::Ok().json(checklist)
}

/// Arguments to the __error__ handler.
#[derive(Debug, Deserialize, Default)]
#[serde(default)]
struct ErrorArgs {
    /// If true, and the server has settings.debug == true, the error handler will panic.
    panic: bool,
}

/// Returning an API error to test error handling.
#[get("__error__")]
async fn test_error(
    params: web::Query<ErrorArgs>,
    settings: Data<Settings>,
) -> Result<HttpResponse, HandlerError> {
    match (params.panic, settings.debug) {
        (true, true) => {
            // allowed panic
            tracing::event!(
                Level::ERROR,
                r#type = "dockerflow.panic_endpoint",
                "The __panic__ endpoint was called"
            );
            panic!("Test panic for debugging");
        }
        (true, false) => Ok(HttpResponse::Forbidden().body("Not permitted in production mode")),
        (false, _) => {
            tracing::event!(
                Level::ERROR,
                r#type = "dockerflow.error_endpoint",
                "The __error__ endpoint was called"
            );
            Err(HandlerError::internal())
        }
    }
}