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