@@ -3,19 +3,20 @@ pub mod config;
3
3
pub mod workers;
4
4
5
5
use std:: collections:: BTreeMap ;
6
- use std:: thread:: Builder ;
6
+ use std:: fmt:: Display ;
7
+ use std:: thread:: { sleep, Builder , JoinHandle } ;
7
8
use std:: time:: Duration ;
8
9
9
10
use anyhow:: Context ;
10
11
use crossbeam_channel:: { bounded, unbounded} ;
11
- use signal_hook:: consts:: { SIGTERM , SIGUSR1 } ;
12
+ use signal_hook:: consts:: SIGUSR1 ;
12
13
use signal_hook:: iterator:: Signals ;
13
14
14
15
use aquatic_common:: access_list:: update_access_list;
15
16
#[ cfg( feature = "cpu-pinning" ) ]
16
17
use aquatic_common:: cpu_pinning:: { pin_current_if_configured_to, WorkerIndex } ;
17
18
use aquatic_common:: privileges:: PrivilegeDropper ;
18
- use aquatic_common:: { PanicSentinelWatcher , ServerStartInstant } ;
19
+ use aquatic_common:: ServerStartInstant ;
19
20
20
21
use common:: {
21
22
ConnectedRequestSender , ConnectedResponseSender , SocketWorkerIndex , State , SwarmWorkerIndex ,
@@ -28,12 +29,12 @@ pub const APP_NAME: &str = "aquatic_udp: UDP BitTorrent tracker";
28
29
pub const APP_VERSION : & str = env ! ( "CARGO_PKG_VERSION" ) ;
29
30
30
31
pub fn run ( config : Config ) -> :: anyhow:: Result < ( ) > {
31
- let mut signals = Signals :: new ( [ SIGUSR1 , SIGTERM ] ) ?;
32
+ let mut signals = Signals :: new ( [ SIGUSR1 ] ) ?;
32
33
33
34
let state = State :: new ( config. swarm_workers ) ;
34
35
let connection_validator = ConnectionValidator :: new ( & config) ?;
35
- let ( sentinel_watcher, sentinel) = PanicSentinelWatcher :: create_with_sentinel ( ) ;
36
36
let priv_dropper = PrivilegeDropper :: new ( config. privileges . clone ( ) , config. socket_workers ) ;
37
+ let mut join_handles = Vec :: new ( ) ;
37
38
38
39
update_access_list ( & config. access_list , & state. access_list ) ?;
39
40
@@ -62,14 +63,13 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
62
63
}
63
64
64
65
for i in 0 ..config. swarm_workers {
65
- let sentinel = sentinel. clone ( ) ;
66
66
let config = config. clone ( ) ;
67
67
let state = state. clone ( ) ;
68
68
let request_receiver = request_receivers. remove ( & i) . unwrap ( ) . clone ( ) ;
69
69
let response_sender = ConnectedResponseSender :: new ( response_senders. clone ( ) ) ;
70
70
let statistics_sender = statistics_sender. clone ( ) ;
71
71
72
- Builder :: new ( )
72
+ let handle = Builder :: new ( )
73
73
. name ( format ! ( "swarm-{:02}" , i + 1 ) )
74
74
. spawn ( move || {
75
75
#[ cfg( feature = "cpu-pinning" ) ]
@@ -81,7 +81,6 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
81
81
) ;
82
82
83
83
let mut worker = SwarmWorker {
84
- _sentinel : sentinel,
85
84
config,
86
85
state,
87
86
server_start_instant,
@@ -91,13 +90,14 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
91
90
worker_index : SwarmWorkerIndex ( i) ,
92
91
} ;
93
92
94
- worker. run ( ) ;
93
+ worker. run ( )
95
94
} )
96
95
. with_context ( || "spawn swarm worker" ) ?;
96
+
97
+ join_handles. push ( ( WorkerType :: Swarm ( i) , handle) ) ;
97
98
}
98
99
99
100
for i in 0 ..config. socket_workers {
100
- let sentinel = sentinel. clone ( ) ;
101
101
let state = state. clone ( ) ;
102
102
let config = config. clone ( ) ;
103
103
let connection_validator = connection_validator. clone ( ) ;
@@ -106,7 +106,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
106
106
let response_receiver = response_receivers. remove ( & i) . unwrap ( ) ;
107
107
let priv_dropper = priv_dropper. clone ( ) ;
108
108
109
- Builder :: new ( )
109
+ let handle = Builder :: new ( )
110
110
. name ( format ! ( "socket-{:02}" , i + 1 ) )
111
111
. spawn ( move || {
112
112
#[ cfg( feature = "cpu-pinning" ) ]
@@ -118,46 +118,48 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
118
118
) ;
119
119
120
120
workers:: socket:: run_socket_worker (
121
- sentinel,
122
121
state,
123
122
config,
124
123
connection_validator,
125
124
server_start_instant,
126
125
request_sender,
127
126
response_receiver,
128
127
priv_dropper,
129
- ) ;
128
+ )
130
129
} )
131
130
. with_context ( || "spawn socket worker" ) ?;
131
+
132
+ join_handles. push ( ( WorkerType :: Socket ( i) , handle) ) ;
132
133
}
133
134
134
135
if config. statistics . active ( ) {
135
- let sentinel = sentinel. clone ( ) ;
136
136
let state = state. clone ( ) ;
137
137
let config = config. clone ( ) ;
138
138
139
- #[ cfg( feature = "prometheus" ) ]
140
- if config. statistics . run_prometheus_endpoint {
141
- use metrics_exporter_prometheus:: PrometheusBuilder ;
142
- use metrics_util:: MetricKindMask ;
139
+ let handle = Builder :: new ( )
140
+ . name ( "statistics" . into ( ) )
141
+ . spawn ( move || {
142
+ #[ cfg( feature = "cpu-pinning" ) ]
143
+ pin_current_if_configured_to (
144
+ & config. cpu_pinning ,
145
+ config. socket_workers ,
146
+ config. swarm_workers ,
147
+ WorkerIndex :: Util ,
148
+ ) ;
143
149
144
- PrometheusBuilder :: new ( )
145
- . idle_timeout (
146
- MetricKindMask :: ALL ,
147
- Some ( Duration :: from_secs ( config. statistics . interval * 2 ) ) ,
148
- )
149
- . with_http_listener ( config. statistics . prometheus_endpoint_address )
150
- . install ( )
151
- . with_context ( || {
152
- format ! (
153
- "Install prometheus endpoint on {}" ,
154
- config. statistics. prometheus_endpoint_address
155
- )
156
- } ) ?;
157
- }
150
+ workers:: statistics:: run_statistics_worker ( config, state, statistics_receiver)
151
+ } )
152
+ . with_context ( || "spawn statistics worker" ) ?;
158
153
159
- Builder :: new ( )
160
- . name ( "statistics" . into ( ) )
154
+ join_handles. push ( ( WorkerType :: Statistics , handle) ) ;
155
+ }
156
+
157
+ #[ cfg( feature = "prometheus" ) ]
158
+ if config. statistics . active ( ) && config. statistics . run_prometheus_endpoint {
159
+ let config = config. clone ( ) ;
160
+
161
+ let handle = Builder :: new ( )
162
+ . name ( "prometheus" . into ( ) )
161
163
. spawn ( move || {
162
164
#[ cfg( feature = "cpu-pinning" ) ]
163
165
pin_current_if_configured_to (
@@ -167,14 +169,73 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
167
169
WorkerIndex :: Util ,
168
170
) ;
169
171
170
- workers:: statistics:: run_statistics_worker (
171
- sentinel,
172
- config,
173
- state,
174
- statistics_receiver,
172
+ use metrics_exporter_prometheus:: PrometheusBuilder ;
173
+ use metrics_util:: MetricKindMask ;
174
+
175
+ let rt = :: tokio:: runtime:: Builder :: new_current_thread ( )
176
+ . enable_all ( )
177
+ . build ( )
178
+ . context ( "build prometheus tokio runtime" ) ?;
179
+
180
+ rt. block_on ( async {
181
+ let ( recorder, exporter) = PrometheusBuilder :: new ( )
182
+ . idle_timeout (
183
+ MetricKindMask :: ALL ,
184
+ Some ( Duration :: from_secs ( config. statistics . interval * 2 ) ) ,
185
+ )
186
+ . with_http_listener ( config. statistics . prometheus_endpoint_address )
187
+ . build ( )
188
+ . context ( "build prometheus recorder and exporter" ) ?;
189
+
190
+ :: tokio:: spawn ( async move {
191
+ let mut interval = :: tokio:: time:: interval ( Duration :: from_secs ( 5 ) ) ;
192
+
193
+ loop {
194
+ interval. tick ( ) . await ;
195
+
196
+ // Periodically render metrics to make sure
197
+ // idles are cleaned up
198
+ recorder. handle ( ) . render ( ) ;
199
+ }
200
+ } ) ;
201
+
202
+ exporter. await . context ( "run prometheus exporter" )
203
+ } )
204
+ } )
205
+ . with_context ( || "spawn prometheus exporter worker" ) ?;
206
+
207
+ join_handles. push ( ( WorkerType :: Prometheus , handle) ) ;
208
+ }
209
+
210
+ // Spawn signal handler thread
211
+ {
212
+ let config = config. clone ( ) ;
213
+
214
+ let handle: JoinHandle < anyhow:: Result < ( ) > > = Builder :: new ( )
215
+ . name ( "signals" . into ( ) )
216
+ . spawn ( move || {
217
+ #[ cfg( feature = "cpu-pinning" ) ]
218
+ pin_current_if_configured_to (
219
+ & config. cpu_pinning ,
220
+ config. socket_workers ,
221
+ config. swarm_workers ,
222
+ WorkerIndex :: Util ,
175
223
) ;
224
+
225
+ for signal in & mut signals {
226
+ match signal {
227
+ SIGUSR1 => {
228
+ let _ = update_access_list ( & config. access_list , & state. access_list ) ;
229
+ }
230
+ _ => unreachable ! ( ) ,
231
+ }
232
+ }
233
+
234
+ Ok ( ( ) )
176
235
} )
177
- . with_context ( || "spawn statistics worker" ) ?;
236
+ . context ( "spawn signal worker" ) ?;
237
+
238
+ join_handles. push ( ( WorkerType :: Signals , handle) ) ;
178
239
}
179
240
180
241
#[ cfg( feature = "cpu-pinning" ) ]
@@ -185,21 +246,47 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
185
246
WorkerIndex :: Util ,
186
247
) ;
187
248
188
- for signal in & mut signals {
189
- match signal {
190
- SIGUSR1 => {
191
- let _ = update_access_list ( & config. access_list , & state. access_list ) ;
192
- }
193
- SIGTERM => {
194
- if sentinel_watcher. panic_was_triggered ( ) {
195
- return Err ( anyhow:: anyhow!( "worker thread panicked" ) ) ;
249
+ loop {
250
+ for ( i, ( _, handle) ) in join_handles. iter ( ) . enumerate ( ) {
251
+ if handle. is_finished ( ) {
252
+ let ( worker_type, handle) = join_handles. remove ( i) ;
253
+
254
+ match handle. join ( ) {
255
+ Ok ( Ok ( ( ) ) ) => {
256
+ return Err ( anyhow:: anyhow!( "{} stopped" , worker_type) ) ;
257
+ }
258
+ Ok ( Err ( err) ) => {
259
+ return Err ( err. context ( format ! ( "{} stopped" , worker_type) ) ) ;
260
+ }
261
+ Err ( _) => {
262
+ return Err ( anyhow:: anyhow!( "{} panicked" , worker_type) ) ;
263
+ }
196
264
}
197
-
198
- break ;
199
265
}
200
- _ => unreachable ! ( ) ,
201
266
}
267
+
268
+ sleep ( Duration :: from_secs ( 5 ) ) ;
202
269
}
270
+ }
271
+
272
+ enum WorkerType {
273
+ Swarm ( usize ) ,
274
+ Socket ( usize ) ,
275
+ Statistics ,
276
+ Signals ,
277
+ #[ cfg( feature = "prometheus" ) ]
278
+ Prometheus ,
279
+ }
203
280
204
- Ok ( ( ) )
281
+ impl Display for WorkerType {
282
+ fn fmt ( & self , f : & mut std:: fmt:: Formatter < ' _ > ) -> std:: fmt:: Result {
283
+ match self {
284
+ Self :: Swarm ( index) => f. write_fmt ( format_args ! ( "Swarm worker {}" , index + 1 ) ) ,
285
+ Self :: Socket ( index) => f. write_fmt ( format_args ! ( "Socket worker {}" , index + 1 ) ) ,
286
+ Self :: Statistics => f. write_str ( "Statistics worker" ) ,
287
+ Self :: Signals => f. write_str ( "Signals worker" ) ,
288
+ #[ cfg( feature = "prometheus" ) ]
289
+ Self :: Prometheus => f. write_str ( "Prometheus worker" ) ,
290
+ }
291
+ }
205
292
}
0 commit comments