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
use std::ops::BitAnd; use std::collections::HashMap; use ::iron::prelude::*; use ::iron::status; use ::serde::{Serialize, Serializer}; use ::json; #[derive(Debug, Clone)] pub enum HealthCheckStatus { Healthy, Unhealthy, } impl PartialEq for HealthCheckStatus { fn eq(&self, other: &HealthCheckStatus) -> bool { match (self, other) { (&HealthCheckStatus::Healthy, &HealthCheckStatus::Healthy) => true, (&HealthCheckStatus::Unhealthy, &HealthCheckStatus::Unhealthy) => true, _ => false, } } } impl Into<status::Status> for HealthCheckStatus { fn into(self) -> status::Status { match self { HealthCheckStatus::Healthy => status::Ok, HealthCheckStatus::Unhealthy => status::InternalServerError, } } } impl Serialize for HealthCheckStatus { fn serialize<S>(&self, serializer: S) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error> where S: Serializer { match *self { HealthCheckStatus::Healthy => serializer.serialize_str("Ok"), HealthCheckStatus::Unhealthy => serializer.serialize_str("Failed"), } } } impl BitAnd for HealthCheckStatus { type Output = HealthCheckStatus; fn bitand(self, other: HealthCheckStatus) -> HealthCheckStatus { match (self, other) { (HealthCheckStatus::Healthy, HealthCheckStatus::Healthy) => HealthCheckStatus::Healthy, _ => HealthCheckStatus::Unhealthy, } } } pub trait HealthCheck: Send { fn name(&self) -> String; fn check_health(&mut self) -> HealthCheckStatus; } #[derive(Default)] pub struct HealthCheckService { checks: Vec<Box<HealthCheck + 'static>>, } impl HealthCheckService { pub fn register_check<H: HealthCheck + 'static>(&mut self, check: H) { self.checks.push(Box::new(check)); } pub fn check_service_health(&mut self, _: &mut Request) -> IronResult<Response> { let (global, health) = self.execute(); let payload = json::to_string(&health).unwrap(); let status: status::Status = global.into(); Ok(Response::with((status, payload))) } pub fn execute(&mut self) -> (HealthCheckStatus, HashMap<String, HealthCheckStatus>) { let mut map = HashMap::new(); for check in &mut self.checks { let res = check.check_health(); map.insert(check.name(), res); } let global_health = map.values() .fold(HealthCheckStatus::Healthy, |check, val| check & val.clone()); (global_health, map) } }