@@ -7,6 +7,18 @@ 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+
1022use daily_task:: run_daily_task_at_midnight;
1123use graphql:: { Mutation , Query } ;
1224use routes:: setup_router;
@@ -37,10 +49,27 @@ impl Config {
3749 }
3850}
3951
52+ struct OtelGuard {
53+ tracer_provider : SdkTracerProvider ,
54+ meter_provider : SdkMeterProvider ,
55+ }
56+
57+ impl Drop for OtelGuard {
58+ fn drop ( & mut self ) {
59+ if let Err ( err) = self . tracer_provider . shutdown ( ) {
60+ eprintln ! ( "{err:?}" ) ;
61+ }
62+ if let Err ( err) = self . meter_provider . shutdown ( ) {
63+ eprintln ! ( "{err:?}" ) ;
64+ }
65+ }
66+ }
67+
4068#[ tokio:: main]
69+ #[ tracing:: instrument]
4170async fn main ( ) {
4271 let config = Config :: from_env ( ) ;
43- setup_tracing ( & config. env ) ;
72+ let guard = setup_tracing ( & config. env ) ;
4473
4574 let pool = setup_database ( & config. database_url ) . await ;
4675 let schema = build_graphql_schema ( pool. clone ( ) , config. secret_key ) ;
@@ -56,10 +85,81 @@ async fn main() {
5685 let listener = tokio:: net:: TcpListener :: bind ( format ! ( "0.0.0.0:{}" , config. port) )
5786 . await
5887 . unwrap ( ) ;
59- axum:: serve ( listener, router) . await . unwrap ( ) ;
88+
89+ axum:: serve ( listener, router)
90+ . with_graceful_shutdown ( shutdown_signal ( ) )
91+ . await
92+ . unwrap ( ) ;
93+
94+ drop ( guard) ;
95+ }
96+
97+ #[ tracing:: instrument]
98+ async fn shutdown_signal ( ) {
99+ tokio:: signal:: ctrl_c ( )
100+ . await
101+ . expect ( "failed to install Ctrl+C handler" ) ;
102+
103+ tracing:: info!( "Shutdown signal received. Flushing telemetry..." ) ;
60104}
61105
62- fn setup_tracing ( env : & str ) {
106+ fn resource ( ) -> Resource {
107+ Resource :: builder ( )
108+ . with_attributes ( vec ! [
109+ KeyValue :: new( SERVICE_NAME , env!( "CARGO_PKG_NAME" ) ) ,
110+ KeyValue :: new( SERVICE_VERSION , env!( "CARGO_PKG_VERSION" ) ) ,
111+ KeyValue :: new( DEPLOYMENT_ENVIRONMENT_NAME , "develop" ) ,
112+ ] )
113+ . with_schema_url ( Vec :: new ( ) , SCHEMA_URL )
114+ . build ( )
115+ }
116+
117+ fn init_meter_provider ( ) -> SdkMeterProvider {
118+ let exporter = opentelemetry_otlp:: MetricExporter :: builder ( )
119+ . with_tonic ( )
120+ . with_temporality ( opentelemetry_sdk:: metrics:: Temporality :: default ( ) )
121+ . build ( )
122+ . unwrap ( ) ;
123+
124+ let reader = PeriodicReader :: builder ( exporter)
125+ . with_interval ( std:: time:: Duration :: from_secs ( 30 ) )
126+ . build ( ) ;
127+
128+ let stdout_reader =
129+ PeriodicReader :: builder ( opentelemetry_stdout:: MetricExporter :: default ( ) ) . build ( ) ;
130+
131+ let meter_provider = MeterProviderBuilder :: default ( )
132+ . with_resource ( resource ( ) )
133+ . with_reader ( reader)
134+ . with_reader ( stdout_reader)
135+ . build ( ) ;
136+
137+ global:: set_meter_provider ( meter_provider. clone ( ) ) ;
138+
139+ meter_provider
140+ }
141+
142+ fn init_tracer_provider ( ) -> SdkTracerProvider {
143+ let exporter = opentelemetry_otlp:: SpanExporter :: builder ( )
144+ . with_tonic ( )
145+ . build ( )
146+ . unwrap ( ) ;
147+
148+ SdkTracerProvider :: builder ( )
149+ . with_sampler ( Sampler :: ParentBased ( Box :: new ( Sampler :: TraceIdRatioBased (
150+ 1.0 ,
151+ ) ) ) )
152+ . with_id_generator ( RandomIdGenerator :: default ( ) )
153+ . with_resource ( resource ( ) )
154+ . with_batch_exporter ( exporter)
155+ . build ( )
156+ }
157+
158+ fn setup_tracing ( env : & str ) -> OtelGuard {
159+ let tracer_provider = init_tracer_provider ( ) ;
160+ let meter_provider = init_meter_provider ( ) ;
161+ let tracer = tracer_provider. tracer ( "tracing-otel-subscriber" ) ;
162+
63163 let kolkata_offset = UtcOffset :: from_hms ( 5 , 30 , 0 ) . expect ( "Hardcoded offset must be correct" ) ;
64164 let timer = fmt:: time:: OffsetTime :: new (
65165 kolkata_offset,
@@ -75,6 +175,8 @@ fn setup_tracing(env: &str) {
75175 . with_ansi ( false ) // ANSI encodings are unreadable in the raw file.
76176 . with_writer ( std:: fs:: File :: create ( "root.log" ) . unwrap ( ) ) ,
77177 )
178+ . with ( MetricsLayer :: new ( meter_provider. clone ( ) ) )
179+ . with ( OpenTelemetryLayer :: new ( tracer) )
78180 . with ( EnvFilter :: new ( "info" ) )
79181 . init ( ) ;
80182 info ! ( "Running in production mode." )
@@ -93,10 +195,17 @@ fn setup_tracing(env: &str) {
93195 . with_ansi ( false )
94196 . with_writer ( std:: fs:: File :: create ( "root.log" ) . unwrap ( ) ) ,
95197 )
198+ . with ( MetricsLayer :: new ( meter_provider. clone ( ) ) )
199+ . with ( OpenTelemetryLayer :: new ( tracer) )
96200 . with ( EnvFilter :: new ( "trace" ) )
97201 . init ( ) ;
98202 info ! ( "Running in development mode." ) ;
99203 }
204+
205+ OtelGuard {
206+ tracer_provider,
207+ meter_provider,
208+ }
100209}
101210
102211async fn setup_database ( database_url : & str ) -> Arc < PgPool > {
0 commit comments