3
3
//! This API is intended to be used by the container infrastructure to check if
4
4
//! the whole application is healthy.
5
5
use std:: net:: SocketAddr ;
6
+ use std:: time:: Duration ;
6
7
8
+ use axum:: http:: { HeaderName , HeaderValue } ;
9
+ use axum:: response:: Response ;
7
10
use axum:: routing:: get;
8
11
use axum:: { Json , Router } ;
9
12
use axum_server:: Handle ;
10
13
use futures:: Future ;
14
+ use hyper:: Request ;
11
15
use log:: debug;
12
16
use serde_json:: json;
13
17
use tokio:: sync:: oneshot:: { Receiver , Sender } ;
18
+ use tower_http:: compression:: CompressionLayer ;
19
+ use tower_http:: propagate_header:: PropagateHeaderLayer ;
20
+ use tower_http:: request_id:: { MakeRequestId , RequestId , SetRequestIdLayer } ;
21
+ use tower_http:: trace:: { DefaultMakeSpan , TraceLayer } ;
22
+ use tracing:: { Level , Span } ;
23
+ use uuid:: Uuid ;
14
24
15
25
use crate :: bootstrap:: jobs:: Started ;
16
26
use crate :: servers:: health_check_api:: handlers:: health_check_handler;
@@ -31,14 +41,48 @@ pub fn start(
31
41
let router = Router :: new ( )
32
42
. route ( "/" , get ( || async { Json ( json ! ( { } ) ) } ) )
33
43
. route ( "/health_check" , get ( health_check_handler) )
34
- . with_state ( register) ;
44
+ . with_state ( register)
45
+ . layer ( CompressionLayer :: new ( ) )
46
+ . layer ( SetRequestIdLayer :: x_request_id ( RequestIdGenerator ) )
47
+ . layer ( PropagateHeaderLayer :: new ( HeaderName :: from_static ( "x-request-id" ) ) )
48
+ . layer (
49
+ TraceLayer :: new_for_http ( )
50
+ . make_span_with ( DefaultMakeSpan :: new ( ) . level ( Level :: INFO ) )
51
+ . on_request ( |request : & Request < axum:: body:: Body > , _span : & Span | {
52
+ let method = request. method ( ) . to_string ( ) ;
53
+ let uri = request. uri ( ) . to_string ( ) ;
54
+ let request_id = request
55
+ . headers ( )
56
+ . get ( "x-request-id" )
57
+ . map ( |v| v. to_str ( ) . unwrap_or_default ( ) )
58
+ . unwrap_or_default ( ) ;
59
+
60
+ tracing:: span!(
61
+ target: "HEALTH CHECK API" ,
62
+ tracing:: Level :: INFO , "request" , method = %method, uri = %uri, request_id = %request_id) ;
63
+ } )
64
+ . on_response ( |response : & Response , latency : Duration , _span : & Span | {
65
+ let status_code = response. status ( ) ;
66
+ let request_id = response
67
+ . headers ( )
68
+ . get ( "x-request-id" )
69
+ . map ( |v| v. to_str ( ) . unwrap_or_default ( ) )
70
+ . unwrap_or_default ( ) ;
71
+ let latency_ms = latency. as_millis ( ) ;
72
+
73
+ tracing:: span!(
74
+ target: "HEALTH CHECK API" ,
75
+ tracing:: Level :: INFO , "response" , latency = %latency_ms, status = %status_code, request_id = %request_id) ;
76
+ } ) ,
77
+ )
78
+ . layer ( SetRequestIdLayer :: x_request_id ( RequestIdGenerator ) ) ;
35
79
36
80
let socket = std:: net:: TcpListener :: bind ( bind_to) . expect ( "Could not bind tcp_listener to address." ) ;
37
81
let address = socket. local_addr ( ) . expect ( "Could not get local_addr from tcp_listener." ) ;
38
82
39
83
let handle = Handle :: new ( ) ;
40
84
41
- debug ! ( target: "Health Check API" , "Starting service with graceful shutdown in a spawned task ..." ) ;
85
+ debug ! ( target: "HEALTH CHECK API" , "Starting service with graceful shutdown in a spawned task ..." ) ;
42
86
43
87
tokio:: task:: spawn ( graceful_shutdown (
44
88
handle. clone ( ) ,
@@ -55,3 +99,13 @@ pub fn start(
55
99
56
100
running
57
101
}
102
+
103
+ #[ derive( Clone , Default ) ]
104
+ struct RequestIdGenerator ;
105
+
106
+ impl MakeRequestId for RequestIdGenerator {
107
+ fn make_request_id < B > ( & mut self , _request : & Request < B > ) -> Option < RequestId > {
108
+ let id = HeaderValue :: from_str ( & Uuid :: new_v4 ( ) . to_string ( ) ) . expect ( "UUID is a valid HTTP header value" ) ;
109
+ Some ( RequestId :: new ( id) )
110
+ }
111
+ }
0 commit comments