1- use crate :: telemetry:: events:: TelemetryEvent ;
2-
3- use ddcommon:: hyper_migration;
1+ use crate :: {
2+ http:: { extract_request_body, handler_not_found} ,
3+ telemetry:: events:: TelemetryEvent ,
4+ } ;
5+
6+ use axum:: {
7+ extract:: { Request , State } ,
8+ http:: StatusCode ,
9+ response:: { IntoResponse , Response } ,
10+ routing:: post,
11+ Router ,
12+ } ;
413use std:: net:: SocketAddr ;
5- use tokio:: sync:: mpsc:: Sender ;
6-
7- use http_body_util:: BodyExt ;
8- use hyper:: service:: service_fn;
9- use hyper:: Response ;
10- use std:: io;
11- use tracing:: { debug, error} ;
14+ use tokio:: { net:: TcpListener , sync:: mpsc:: Sender } ;
15+ use tokio_util:: sync:: CancellationToken ;
16+ use tracing:: debug;
1217
1318#[ allow( clippy:: module_name_repetitions) ]
14- #[ derive( Debug , Clone , Copy ) ]
15- pub struct TelemetryListener { }
16-
17- pub struct TelemetryListenerConfig {
18- pub host : String ,
19- pub port : u16 ,
19+ #[ derive( Debug , Clone ) ]
20+ pub struct TelemetryListener {
21+ host : [ u8 ; 4 ] ,
22+ port : u16 ,
23+ cancel_token : CancellationToken ,
24+ event_bus : Sender < TelemetryEvent > ,
2025}
2126
2227impl TelemetryListener {
23- pub async fn spin (
24- config : & TelemetryListenerConfig ,
25- event_bus : Sender < TelemetryEvent > ,
26- _cancel_token : tokio_util:: sync:: CancellationToken , // todo cancel token
27- ) -> Result < ( ) , Box < dyn std:: error:: Error > > {
28- let addr = SocketAddr :: from ( ( [ 0 , 0 , 0 , 0 ] , config. port ) ) ;
29-
30- let service = service_fn ( move |req| {
31- let event_bus = event_bus. clone ( ) ;
32- Self :: handle ( req. map ( hyper_migration:: Body :: incoming) , event_bus. clone ( ) )
28+ #[ must_use]
29+ pub fn new ( host : [ u8 ; 4 ] , port : u16 , event_bus : Sender < TelemetryEvent > ) -> Self {
30+ let cancel_token = CancellationToken :: new ( ) ;
31+ Self {
32+ host,
33+ port,
34+ cancel_token,
35+ event_bus,
36+ }
37+ }
38+
39+ #[ must_use]
40+ pub fn cancel_token ( & self ) -> CancellationToken {
41+ self . cancel_token . clone ( )
42+ }
43+
44+ pub fn start ( & self ) -> Result < ( ) , Box < dyn std:: error:: Error > > {
45+ let socket = SocketAddr :: from ( ( self . host , self . port ) ) ;
46+ let router = self . make_router ( ) ;
47+
48+ let cancel_token_clone = self . cancel_token ( ) ;
49+ tokio:: spawn ( async move {
50+ let listener = TcpListener :: bind ( & socket)
51+ . await
52+ . expect ( "Failed to bind socket" ) ;
53+ debug ! ( "Telemetry API | Starting listener on {}" , socket) ;
54+ axum:: serve ( listener, router)
55+ . with_graceful_shutdown ( Self :: graceful_shutdown ( cancel_token_clone) )
56+ . await
57+ . expect ( "Failed to start telemetry listener" ) ;
3358 } ) ;
3459
35- let listener = tokio:: net:: TcpListener :: bind ( & addr) . await ?;
36-
37- let server = hyper:: server:: conn:: http1:: Builder :: new ( ) ;
38- let mut joinset = tokio:: task:: JoinSet :: new ( ) ;
39- loop {
40- let conn = tokio:: select! {
41- con_res = listener. accept( ) => match con_res {
42- Err ( e)
43- if matches!(
44- e. kind( ) ,
45- io:: ErrorKind :: ConnectionAborted
46- | io:: ErrorKind :: ConnectionReset
47- | io:: ErrorKind :: ConnectionRefused
48- ) =>
49- {
50- continue ;
51- }
52- Err ( e) => {
53- error!( "Server error: {e}" ) ;
54- return Err ( e. into( ) ) ;
55- }
56- Ok ( ( conn, _) ) => conn,
57- } ,
58- finished = async {
59- match joinset. join_next( ) . await {
60- Some ( finished) => finished,
61- None => std:: future:: pending( ) . await ,
62- }
63- } => match finished {
64- Err ( e) if e. is_panic( ) => {
65- std:: panic:: resume_unwind( e. into_panic( ) ) ;
66- } ,
67- Ok ( ( ) ) | Err ( _) => continue ,
68- } ,
69- } ;
70- let conn = hyper_util:: rt:: TokioIo :: new ( conn) ;
71- let server = server. clone ( ) ;
72- let service = service. clone ( ) ;
73- joinset. spawn ( async move {
74- if let Err ( e) = server. serve_connection ( conn, service) . await {
75- debug ! ( "Telemetry Connection error: {e}" ) ;
76- }
77- } ) ;
78- }
60+ Ok ( ( ) )
7961 }
8062
81- pub async fn handle (
82- req : hyper_migration:: HttpRequest ,
83- event_bus : Sender < TelemetryEvent > ,
84- ) -> Result < hyper_migration:: HttpResponse , hyper:: Error > {
85- let body = match req. collect ( ) . await {
86- Ok ( body_bytes_collected) => body_bytes_collected. to_bytes ( ) . to_vec ( ) ,
63+ fn make_router ( & self ) -> Router {
64+ let event_bus = self . event_bus . clone ( ) ;
65+
66+ Router :: new ( )
67+ . route ( "/" , post ( Self :: handle) )
68+ . fallback ( handler_not_found)
69+ . with_state ( event_bus)
70+ }
71+
72+ async fn graceful_shutdown ( cancel_token : CancellationToken ) {
73+ cancel_token. cancelled ( ) . await ;
74+ debug ! ( "Telemetry API | Shutdown signal received, shutting down" ) ;
75+ }
76+
77+ async fn handle ( State ( event_bus) : State < Sender < TelemetryEvent > > , request : Request ) -> Response {
78+ let ( _, body) = match extract_request_body ( request) . await {
79+ Ok ( r) => r,
8780 Err ( e) => {
88- error ! ( "Failed to collect body: {:?}" , e ) ;
89- return Ok ( Response :: builder ( )
90- . status ( hyper :: StatusCode :: INTERNAL_SERVER_ERROR )
91- . body ( hyper_migration :: Body :: from ( "Failed to collect body" ) )
92- . expect ( "infallible" ) ) ;
81+ return (
82+ StatusCode :: INTERNAL_SERVER_ERROR ,
83+ format ! ( "Failed to extract request body: {e}" ) ,
84+ )
85+ . into_response ( ) ;
9386 }
9487 } ;
88+
9589 let body = std:: str:: from_utf8 ( & body) . expect ( "infallible" ) ;
9690
9791 let mut telemetry_events: Vec < TelemetryEvent > = match serde_json:: from_str ( body) {
@@ -103,43 +97,49 @@ impl TelemetryListener {
10397 // This will result in a dropped payload, but may be from
10498 // events we haven't added support for yet
10599 debug ! ( "Failed to parse telemetry events: {:?}" , e) ;
106- return Ok ( Response :: builder ( )
107- . status ( hyper:: StatusCode :: OK )
108- . body ( hyper_migration:: Body :: from (
109- "Failed to parse telemetry events" ,
110- ) )
111- . expect ( "infallible" ) ) ;
100+ return ( StatusCode :: OK , "Failed to parse telemetry events" ) . into_response ( ) ;
112101 }
113102 } ;
103+
114104 for event in telemetry_events. drain ( ..) {
115105 event_bus. send ( event) . await . expect ( "infallible" ) ;
116106 }
117107
118- Ok ( Response :: new ( hyper_migration :: Body :: from ( "OK" ) ) )
108+ ( StatusCode :: OK , "OK" ) . into_response ( )
119109 }
120110}
121111
122112#[ cfg( test) ]
123113mod tests {
114+ use super :: * ;
115+ use axum:: body:: Body ;
116+ use axum:: http:: Request ;
124117 use chrono:: DateTime ;
125- use ddcommon:: hyper_migration;
126118
127119 use crate :: telemetry:: events:: { InitPhase , InitType , TelemetryRecord } ;
128120
129121 #[ tokio:: test]
130122 #[ allow( clippy:: unwrap_used) ]
131123 async fn test_handle ( ) {
132- let event_body = hyper_migration :: Body :: from (
124+ let event_body = Body :: from (
133125 r#"[{"time":"2024-04-25T17:35:59.944Z","type":"platform.initStart","record":{"initializationType":"on-demand","phase":"init","runtimeVersion":"nodejs:20.v22","runtimeVersionArn":"arn:aws:lambda:us-east-1::runtime:da57c20c4b965d5b75540f6865a35fc8030358e33ec44ecfed33e90901a27a72","functionName":"hello-world","functionVersion":"$LATEST"}}]"# ,
134126 ) ;
135- let req = hyper :: Request :: builder ( )
127+ let req = Request :: builder ( )
136128 . method ( "POST" )
137129 . uri ( "http://localhost:8080" )
138130 . body ( event_body)
139131 . unwrap ( ) ;
140132
141133 let ( tx, mut rx) = tokio:: sync:: mpsc:: channel ( 1 ) ;
142- let _ = super :: TelemetryListener :: handle ( req, tx) . await ;
134+
135+ // Create a new request with the body for testing
136+ let ( parts, body) = req. into_parts ( ) ;
137+ let req = Request :: from_parts ( parts, body) ;
138+
139+ let response = TelemetryListener :: handle ( axum:: extract:: State ( tx) , req) . await ;
140+
141+ // Check that the response is OK
142+ assert_eq ! ( response. status( ) , axum:: http:: StatusCode :: OK ) ;
143143
144144 let telemetry_event = rx. recv ( ) . await . unwrap ( ) ;
145145 let expected_time =
0 commit comments