Grcov report - health_check

1

use std::{

2

    collections::HashMap,

3

    sync::{Arc, Mutex},

4

    time::Duration,

5

};

6

7

use futures_util::future::join_all;

8

use tokio::sync::{mpsc, oneshot};

9

10

pub type HealthCheckRequest = oneshot::Sender<Result<(), String>>;

11

12

#[derive(Default, Clone)]

13

pub struct HealthCheckRegistry {

14

    checks: Arc<Mutex<HashMap<&'static str, mpsc::Sender<HealthCheckRequest>>>>,

15

16

17

impl HealthCheckRegistry {

18

3

    pub fn new() -> Self {

19

3

        Self::default()

20

3

21

22

    /// Registers a new health check with the given name and returns a receiver for health check requests.

23

3

    pub fn add_petter(&self, name: &'static str) -> mpsc::Receiver<HealthCheckRequest> {

24

3

        let (tx, rx) = mpsc::channel(1);

25

3

        self.checks.lock().unwrap().insert(name, tx);

26

3

rx

27

3

28

29

    /// Checks that all registered health checks respond within the given timeout.

30

3

    pub async fn request_pets(&self, timeout: Duration) -> Result<(), &'static str> {

31

3

        let checks = self.checks.lock().unwrap().clone();

32

33

3

        let results = join_all(checks.iter().map(|(name, tx)| async move {

34

3

            tokio::time::timeout(timeout, async {

35

3

                let (reply_tx, reply_rx) = oneshot::channel();

36

3

                tx.send(reply_tx)

37

3

                    .await

38

3

                    .map_err(|_| "check request channel closed".to_string())?;

39

3

                reply_rx

40

3

                    .await

41

2

                    .map_err(|_| "check response channel closed".to_string())?

42

2

})

43

3

            .await

44

3

            .map_err(|_| format!("timed out after {} milliseconds", timeout.as_millis()))

45

3

            .and_then(|result| result)

46

3

            .map_err(|reason| {

47

2

                tracing::warn!("Health check {:?} failed: {}", name, reason);

48

2

                *name

49

2

})

50

6

}))

51

3

        .await;

52

53

3

        results.into_iter().collect()

54

3

55

56

57

#[cfg(test)]

58

mod tests {

59

    use super::*;

60

61

2

    fn respond_once(

62

2

        mut requests: mpsc::Receiver<HealthCheckRequest>,

63

2

        response: Result<(), String>,

64

2

) {

65

2

        tokio::spawn(async move {

66

2

            if let Some(reply_tx) = requests.recv().await {

67

2

                let _ = reply_tx.send(response);

68

2

69

2

});

70

2

71

72

    #[tokio::test]

73

1

    async fn passing_check_is_healthy() {

74

1

        let registry = HealthCheckRegistry::new();

75

1

        let requests = registry.add_petter("ok-check");

76

1

        respond_once(requests, Ok(()));

77

78

1

        assert_eq!(

79

1

            registry.request_pets(Duration::from_millis(50)).await,

80

1

            Ok(())

81

1

);

82

1

83

84

    #[tokio::test]

85

1

    async fn failing_check_is_unhealthy() {

86

1

        let registry = HealthCheckRegistry::new();

87

1

        let requests = registry.add_petter("bad-check");

88

1

        respond_once(requests, Err("something broke".to_string()));

89

90

1

        assert_eq!(

91

1

            registry.request_pets(Duration::from_millis(50)).await,

92

1

            Err("bad-check")

93

1

);

94

1

95

96

    #[tokio::test]

97

1

    async fn unresponsive_check_times_out_as_unhealthy() {

98

1

        let registry = HealthCheckRegistry::new();

99

1

        let _requests = registry.add_petter("slow-check");

100

        // Nothing ever answers `_requests`, so the check should time out.

101

102

1

        assert_eq!(

103

1

            registry.request_pets(Duration::from_millis(20)).await,

104

1

            Err("slow-check")

105

1

);

106

1

107