@@ -7,6 +7,19 @@ use tower_http::cors::CorsLayer;
77use tracing:: info;
88use tracing_subscriber:: { fmt, layer:: SubscriberExt , util:: SubscriberInitExt , EnvFilter } ;
99
10+ use opentelemetry:: { global, trace:: TracerProvider as _, KeyValue } ;
11+ use opentelemetry_sdk:: {
12+ metrics:: { MeterProviderBuilder , PeriodicReader , SdkMeterProvider } ,
13+ trace:: { RandomIdGenerator , Sampler , SdkTracerProvider } ,
14+ Resource ,
15+ } ;
16+ use opentelemetry_semantic_conventions:: {
17+ attribute:: { DEPLOYMENT_ENVIRONMENT_NAME , SERVICE_NAME , SERVICE_VERSION } ,
18+ SCHEMA_URL ,
19+ } ;
20+ use tracing_opentelemetry:: { MetricsLayer , OpenTelemetryLayer } ;
21+
22+
1023use daily_task:: run_daily_task_at_midnight;
1124use graphql:: { Mutation , Query } ;
1225use routes:: setup_router;
@@ -37,10 +50,27 @@ impl Config {
3750 }
3851}
3952
53+ struct OtelGuard {
54+ tracer_provider : SdkTracerProvider ,
55+ meter_provider : SdkMeterProvider ,
56+ }
57+
58+ impl Drop for OtelGuard {
59+ fn drop ( & mut self ) {
60+ if let Err ( err) = self . tracer_provider . shutdown ( ) {
61+ eprintln ! ( "{err:?}" ) ;
62+ }
63+ if let Err ( err) = self . meter_provider . shutdown ( ) {
64+ eprintln ! ( "{err:?}" ) ;
65+ }
66+ }
67+ }
68+
4069#[ tokio:: main]
70+ #[ tracing:: instrument]
4171async fn main ( ) {
4272 let config = Config :: from_env ( ) ;
43- setup_tracing ( & config. env ) ;
73+ let guard = setup_tracing ( & config. env ) ;
4474
4575 let pool = setup_database ( & config. database_url ) . await ;
4676 let schema = build_graphql_schema ( pool. clone ( ) , config. secret_key ) ;
@@ -56,10 +86,84 @@ async fn main() {
5686 let listener = tokio:: net:: TcpListener :: bind ( format ! ( "0.0.0.0:{}" , config. port) )
5787 . await
5888 . unwrap ( ) ;
59- axum:: serve ( listener, router) . await . unwrap ( ) ;
89+
90+ axum:: serve ( listener, router)
91+ . with_graceful_shutdown ( shutdown_signal ( ) )
92+ . await
93+ . unwrap ( ) ;
94+
95+ drop ( guard) ;
96+
97+ }
98+
99+ #[ tracing:: instrument]
100+ async fn shutdown_signal ( ) {
101+ // Wait for Ctrl-C
102+ tokio:: signal:: ctrl_c ( )
103+ . await
104+ . expect ( "failed to install Ctrl+C handler" ) ;
105+
106+ tracing:: info!( "Shutdown signal received. Flushing telemetry..." ) ;
107+
108+ // Flush traces and metrics
109+ // guard.tracer_provider.shutdown().unwrap();
110+ }
111+
112+ fn resource ( ) -> Resource {
113+ Resource :: builder ( )
114+ . with_attributes ( vec ! [
115+ KeyValue :: new( SERVICE_NAME , env!( "CARGO_PKG_NAME" ) ) ,
116+ KeyValue :: new( SERVICE_VERSION , env!( "CARGO_PKG_VERSION" ) ) ,
117+ KeyValue :: new( DEPLOYMENT_ENVIRONMENT_NAME , "develop" ) ,
118+ ] )
119+ . with_schema_url ( Vec :: new ( ) , SCHEMA_URL )
120+ . build ( )
121+ }
122+
123+ fn init_meter_provider ( ) -> SdkMeterProvider {
124+ let exporter = opentelemetry_otlp:: MetricExporter :: builder ( )
125+ . with_tonic ( )
126+ . with_temporality ( opentelemetry_sdk:: metrics:: Temporality :: default ( ) )
127+ . build ( )
128+ . unwrap ( ) ;
129+
130+ let reader = PeriodicReader :: builder ( exporter)
131+ . with_interval ( std:: time:: Duration :: from_secs ( 30 ) )
132+ . build ( ) ;
133+
134+ let stdout_reader =
135+ PeriodicReader :: builder ( opentelemetry_stdout:: MetricExporter :: default ( ) ) . build ( ) ;
136+
137+ let meter_provider = MeterProviderBuilder :: default ( )
138+ . with_resource ( resource ( ) )
139+ . with_reader ( reader)
140+ . with_reader ( stdout_reader)
141+ . build ( ) ;
142+
143+ global:: set_meter_provider ( meter_provider. clone ( ) ) ;
144+
145+ meter_provider
146+ }
147+
148+ fn init_tracer_provider ( ) -> SdkTracerProvider {
149+ let exporter = opentelemetry_otlp:: SpanExporter :: builder ( )
150+ . with_tonic ( )
151+ . build ( )
152+ . unwrap ( ) ;
153+
154+ SdkTracerProvider :: builder ( )
155+ . with_sampler ( Sampler :: ParentBased ( Box :: new ( Sampler :: TraceIdRatioBased ( 1.0 ) ) ) )
156+ . with_id_generator ( RandomIdGenerator :: default ( ) )
157+ . with_resource ( resource ( ) )
158+ . with_batch_exporter ( exporter)
159+ . build ( )
60160}
61161
62- fn setup_tracing ( env : & str ) {
162+ fn setup_tracing ( env : & str ) -> OtelGuard {
163+ let tracer_provider = init_tracer_provider ( ) ;
164+ let meter_provider = init_meter_provider ( ) ;
165+ let tracer = tracer_provider. tracer ( "tracing-otel-subscriber" ) ;
166+
63167 let kolkata_offset = UtcOffset :: from_hms ( 5 , 30 , 0 ) . expect ( "Hardcoded offset must be correct" ) ;
64168 let timer = fmt:: time:: OffsetTime :: new (
65169 kolkata_offset,
@@ -75,6 +179,8 @@ fn setup_tracing(env: &str) {
75179 . with_ansi ( false ) // ANSI encodings are unreadable in the raw file.
76180 . with_writer ( std:: fs:: File :: create ( "root.log" ) . unwrap ( ) ) ,
77181 )
182+ . with ( MetricsLayer :: new ( meter_provider. clone ( ) ) )
183+ . with ( OpenTelemetryLayer :: new ( tracer) )
78184 . with ( EnvFilter :: new ( "info" ) )
79185 . init ( ) ;
80186 info ! ( "Running in production mode." )
@@ -93,10 +199,17 @@ fn setup_tracing(env: &str) {
93199 . with_ansi ( false )
94200 . with_writer ( std:: fs:: File :: create ( "root.log" ) . unwrap ( ) ) ,
95201 )
202+ . with ( MetricsLayer :: new ( meter_provider. clone ( ) ) )
203+ . with ( OpenTelemetryLayer :: new ( tracer) )
96204 . with ( EnvFilter :: new ( "trace" ) )
97205 . init ( ) ;
98206 info ! ( "Running in development mode." ) ;
99207 }
208+
209+ OtelGuard {
210+ tracer_provider,
211+ meter_provider,
212+ }
100213}
101214
102215async fn setup_database ( database_url : & str ) -> Arc < PgPool > {
0 commit comments