@@ -7,6 +7,19 @@ use tower_http::cors::CorsLayer;
7
7
use tracing:: info;
8
8
use tracing_subscriber:: { fmt, layer:: SubscriberExt , util:: SubscriberInitExt , EnvFilter } ;
9
9
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
+
10
23
use daily_task:: run_daily_task_at_midnight;
11
24
use graphql:: { Mutation , Query } ;
12
25
use routes:: setup_router;
@@ -37,10 +50,27 @@ impl Config {
37
50
}
38
51
}
39
52
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
+
40
69
#[ tokio:: main]
70
+ #[ tracing:: instrument]
41
71
async fn main ( ) {
42
72
let config = Config :: from_env ( ) ;
43
- setup_tracing ( & config. env ) ;
73
+ let guard = setup_tracing ( & config. env ) ;
44
74
45
75
let pool = setup_database ( & config. database_url ) . await ;
46
76
let schema = build_graphql_schema ( pool. clone ( ) , config. secret_key ) ;
@@ -56,10 +86,84 @@ async fn main() {
56
86
let listener = tokio:: net:: TcpListener :: bind ( format ! ( "0.0.0.0:{}" , config. port) )
57
87
. await
58
88
. 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 ( )
60
160
}
61
161
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
+
63
167
let kolkata_offset = UtcOffset :: from_hms ( 5 , 30 , 0 ) . expect ( "Hardcoded offset must be correct" ) ;
64
168
let timer = fmt:: time:: OffsetTime :: new (
65
169
kolkata_offset,
@@ -75,6 +179,8 @@ fn setup_tracing(env: &str) {
75
179
. with_ansi ( false ) // ANSI encodings are unreadable in the raw file.
76
180
. with_writer ( std:: fs:: File :: create ( "root.log" ) . unwrap ( ) ) ,
77
181
)
182
+ . with ( MetricsLayer :: new ( meter_provider. clone ( ) ) )
183
+ . with ( OpenTelemetryLayer :: new ( tracer) )
78
184
. with ( EnvFilter :: new ( "info" ) )
79
185
. init ( ) ;
80
186
info ! ( "Running in production mode." )
@@ -93,10 +199,17 @@ fn setup_tracing(env: &str) {
93
199
. with_ansi ( false )
94
200
. with_writer ( std:: fs:: File :: create ( "root.log" ) . unwrap ( ) ) ,
95
201
)
202
+ . with ( MetricsLayer :: new ( meter_provider. clone ( ) ) )
203
+ . with ( OpenTelemetryLayer :: new ( tracer) )
96
204
. with ( EnvFilter :: new ( "trace" ) )
97
205
. init ( ) ;
98
206
info ! ( "Running in development mode." ) ;
99
207
}
208
+
209
+ OtelGuard {
210
+ tracer_provider,
211
+ meter_provider,
212
+ }
100
213
}
101
214
102
215
async fn setup_database ( database_url : & str ) -> Arc < PgPool > {
0 commit comments