tonic-rails 0.1.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,146 @@
1
+ //! Listener setup and the global connection cap.
2
+
3
+ use std::net::SocketAddr;
4
+ use std::pin::Pin;
5
+ use std::sync::Arc;
6
+ use std::sync::atomic::{AtomicUsize, Ordering};
7
+ use std::task::{Context, Poll};
8
+ use std::time::Duration;
9
+
10
+ use socket2::{Domain, Protocol, Socket, Type};
11
+ use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
12
+ use tokio::net::{TcpListener, TcpStream};
13
+ use tokio::sync::{OwnedSemaphorePermit, Semaphore};
14
+ use tokio_stream::Stream;
15
+ use tonic::transport::server::{Connected, TcpConnectInfo};
16
+
17
+ /// Bind with SO_REUSEADDR + SO_REUSEPORT so several worker processes can
18
+ /// share a port (the kernel load-balances accepts) and restarts don't
19
+ /// fail on TIME_WAIT. Must be called inside a tokio runtime context.
20
+ pub fn bind(addr: SocketAddr, backlog: i32) -> std::io::Result<TcpListener> {
21
+ let domain = if addr.is_ipv4() {
22
+ Domain::IPV4
23
+ } else {
24
+ Domain::IPV6
25
+ };
26
+ let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))?;
27
+ socket.set_reuse_address(true)?;
28
+ #[cfg(all(
29
+ unix,
30
+ not(any(target_os = "solaris", target_os = "illumos", target_os = "cygwin"))
31
+ ))]
32
+ socket.set_reuse_port(true)?;
33
+ socket.set_nonblocking(true)?;
34
+ socket.bind(&addr.into())?;
35
+ socket.listen(backlog)?;
36
+ TcpListener::from_std(socket.into())
37
+ }
38
+
39
+ /// TCP_NODELAY for latency; TCP keepalive to reap half-open peers that
40
+ /// never answer HTTP/2 PINGs (e.g. behind a NAT that dropped state).
41
+ fn configure_accepted(stream: &TcpStream) {
42
+ let _ = stream.set_nodelay(true);
43
+ let keepalive = socket2::TcpKeepalive::new().with_time(Duration::from_secs(60));
44
+ let _ = socket2::SockRef::from(stream).set_tcp_keepalive(&keepalive);
45
+ }
46
+
47
+ /// A connection holding a slot in the global connection cap.
48
+ pub struct LimitedConn {
49
+ inner: TcpStream,
50
+ _permit: OwnedSemaphorePermit,
51
+ active: Arc<AtomicUsize>,
52
+ }
53
+
54
+ impl Drop for LimitedConn {
55
+ fn drop(&mut self) {
56
+ self.active.fetch_sub(1, Ordering::Relaxed);
57
+ }
58
+ }
59
+
60
+ impl AsyncRead for LimitedConn {
61
+ fn poll_read(
62
+ mut self: Pin<&mut Self>,
63
+ cx: &mut Context<'_>,
64
+ buf: &mut ReadBuf<'_>,
65
+ ) -> Poll<std::io::Result<()>> {
66
+ Pin::new(&mut self.inner).poll_read(cx, buf)
67
+ }
68
+ }
69
+
70
+ impl AsyncWrite for LimitedConn {
71
+ fn poll_write(
72
+ mut self: Pin<&mut Self>,
73
+ cx: &mut Context<'_>,
74
+ buf: &[u8],
75
+ ) -> Poll<std::io::Result<usize>> {
76
+ Pin::new(&mut self.inner).poll_write(cx, buf)
77
+ }
78
+
79
+ fn poll_write_vectored(
80
+ mut self: Pin<&mut Self>,
81
+ cx: &mut Context<'_>,
82
+ bufs: &[std::io::IoSlice<'_>],
83
+ ) -> Poll<std::io::Result<usize>> {
84
+ Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
85
+ }
86
+
87
+ fn is_write_vectored(&self) -> bool {
88
+ self.inner.is_write_vectored()
89
+ }
90
+
91
+ fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
92
+ Pin::new(&mut self.inner).poll_flush(cx)
93
+ }
94
+
95
+ fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
96
+ Pin::new(&mut self.inner).poll_shutdown(cx)
97
+ }
98
+ }
99
+
100
+ impl Connected for LimitedConn {
101
+ type ConnectInfo = TcpConnectInfo;
102
+
103
+ fn connect_info(&self) -> Self::ConnectInfo {
104
+ self.inner.connect_info()
105
+ }
106
+ }
107
+
108
+ /// Accept stream that closes connections beyond the cap immediately.
109
+ pub struct LimitedIncoming<S> {
110
+ pub inner: S,
111
+ pub semaphore: Arc<Semaphore>,
112
+ pub active: Arc<AtomicUsize>,
113
+ }
114
+
115
+ impl<S> Stream for LimitedIncoming<S>
116
+ where
117
+ S: Stream<Item = std::io::Result<TcpStream>> + Unpin,
118
+ {
119
+ type Item = std::io::Result<LimitedConn>;
120
+
121
+ fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
122
+ loop {
123
+ let stream = match Pin::new(&mut self.inner).poll_next(cx) {
124
+ Poll::Ready(Some(Ok(stream))) => stream,
125
+ Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))),
126
+ Poll::Ready(None) => return Poll::Ready(None),
127
+ Poll::Pending => return Poll::Pending,
128
+ };
129
+ match self.semaphore.clone().try_acquire_owned() {
130
+ Ok(permit) => {
131
+ configure_accepted(&stream);
132
+ self.active.fetch_add(1, Ordering::Relaxed);
133
+ return Poll::Ready(Some(Ok(LimitedConn {
134
+ inner: stream,
135
+ _permit: permit,
136
+ active: self.active.clone(),
137
+ })));
138
+ }
139
+ Err(_) => {
140
+ let active = self.active.load(Ordering::Relaxed);
141
+ eprintln!("tonic-rails: connection refused: max_connections reached ({active} active)");
142
+ }
143
+ }
144
+ }
145
+ }
146
+ }
@@ -0,0 +1,507 @@
1
+ //! `TonicRails::Native::Server` — owns the tokio runtime and the tonic
2
+ //! server, and runs Ruby dispatch workers.
3
+ //!
4
+ //! Lifecycle (driven by `Tonic::Rails::Server`):
5
+ //!
6
+ //! ```text
7
+ //! new(addr, max_message_size) → set_* options → on_request(proc)
8
+ //! start_server bind, build tonic router, spawn on tokio
9
+ //! serve_worker × N (Ruby threads) dispatch RPCs until stopped
10
+ //! begin_shutdown health → NOT_SERVING, stop accepting, GOAWAY
11
+ //! wait_for_drain(secs) wait for in-flight RPCs (workers keep running)
12
+ //! stop_workers workers return; Ruby joins them
13
+ //! stop_server drop the runtime, closing anything left
14
+ //! ```
15
+
16
+ use std::collections::HashSet;
17
+ use std::net::SocketAddr;
18
+ use std::sync::Arc;
19
+ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
20
+ use std::time::{Duration, Instant};
21
+
22
+ use magnus::{
23
+ Error, RArray, RHash, RString, Ruby, TryConvert, Value, gc, typed_data::DataTypeFunctions, value::Opaque,
24
+ };
25
+ use parking_lot::Mutex;
26
+ use prost::Message;
27
+ use tokio::runtime::Runtime;
28
+ use tokio::sync::{Semaphore, watch};
29
+ use tonic_health::ServingStatus;
30
+ use tonic_health::server::HealthReporter;
31
+
32
+ use crate::body::ServerStats;
33
+ use crate::dispatch::{DispatchInner, DispatchService};
34
+ use crate::gvl::{self, Cancel, DispatchContext};
35
+ use crate::net::{LimitedIncoming, bind};
36
+
37
+ /// The overall-server health entry (empty service name).
38
+ const OVERALL: &str = "";
39
+
40
+ struct Options {
41
+ tls: Option<(String, String)>,
42
+ tls_client_ca: Option<String>,
43
+ max_connections: usize,
44
+ concurrency_limit_per_connection: usize,
45
+ worker_threads: Option<usize>,
46
+ keepalive_interval: Duration,
47
+ keepalive_timeout: Duration,
48
+ reflection: bool,
49
+ dispatch_queue_size: usize,
50
+ routes: Vec<String>,
51
+ descriptors: Vec<Vec<u8>>,
52
+ }
53
+
54
+ impl Default for Options {
55
+ fn default() -> Self {
56
+ Self {
57
+ tls: None,
58
+ tls_client_ca: None,
59
+ max_connections: 64,
60
+ concurrency_limit_per_connection: 32,
61
+ worker_threads: None,
62
+ keepalive_interval: Duration::from_secs(15),
63
+ keepalive_timeout: Duration::from_secs(5),
64
+ reflection: false,
65
+ dispatch_queue_size: 128,
66
+ routes: Vec::new(),
67
+ descriptors: Vec::new(),
68
+ }
69
+ }
70
+ }
71
+
72
+ /// State that exists only while the server is running.
73
+ struct Running {
74
+ runtime: Runtime,
75
+ request_rx: crossbeam_channel::Receiver<Box<DispatchContext>>,
76
+ stop_workers: Arc<AtomicBool>,
77
+ shutdown_tx: watch::Sender<bool>,
78
+ done_rx: crossbeam_channel::Receiver<()>,
79
+ health: HealthReporter,
80
+ services: Vec<String>,
81
+ }
82
+
83
+ #[derive(magnus::TypedData)]
84
+ #[magnus(class = "TonicRails::Native::Server", mark)]
85
+ pub struct NativeServer {
86
+ bind_address: String,
87
+ max_message_size: usize,
88
+ callback: Mutex<Option<Opaque<Value>>>,
89
+ options: Mutex<Options>,
90
+ running: Mutex<Option<Running>>,
91
+ stats: Arc<ServerStats>,
92
+ compression: Arc<AtomicBool>,
93
+ active_connections: Arc<AtomicUsize>,
94
+ }
95
+
96
+ impl DataTypeFunctions for NativeServer {
97
+ fn mark(&self, marker: &gc::Marker) {
98
+ if let Some(callback) = *self.callback.lock() {
99
+ marker.mark(callback);
100
+ }
101
+ }
102
+ }
103
+
104
+ fn runtime_error(ruby: &Ruby, message: impl Into<std::borrow::Cow<'static, str>>) -> Error {
105
+ Error::new(ruby.exception_runtime_error(), message)
106
+ }
107
+
108
+ fn arg_error(ruby: &Ruby, message: impl Into<std::borrow::Cow<'static, str>>) -> Error {
109
+ Error::new(ruby.exception_arg_error(), message)
110
+ }
111
+
112
+ fn read_file(ruby: &Ruby, what: &str, path: &str) -> Result<Vec<u8>, Error> {
113
+ std::fs::read(path).map_err(|e| runtime_error(ruby, format!("cannot read TLS {what} {path}: {e}")))
114
+ }
115
+
116
+ fn secs(value: f64) -> Duration {
117
+ Duration::from_secs_f64(value.max(0.0))
118
+ }
119
+
120
+ impl NativeServer {
121
+ pub fn new(ruby: &Ruby, bind_address: String, max_message_size: usize) -> Result<Self, Error> {
122
+ if max_message_size == 0 || max_message_size > u32::MAX as usize {
123
+ return Err(arg_error(
124
+ ruby,
125
+ "max_message_size must be between 1 and 4294967295",
126
+ ));
127
+ }
128
+ bind_address
129
+ .parse::<SocketAddr>()
130
+ .map_err(|e| arg_error(ruby, format!("invalid bind address {bind_address:?}: {e}")))?;
131
+ Ok(Self {
132
+ bind_address,
133
+ max_message_size,
134
+ callback: Mutex::new(None),
135
+ options: Mutex::new(Options::default()),
136
+ running: Mutex::new(None),
137
+ stats: Arc::new(ServerStats::default()),
138
+ compression: Arc::new(AtomicBool::new(false)),
139
+ active_connections: Arc::new(AtomicUsize::new(0)),
140
+ })
141
+ }
142
+
143
+ pub fn on_request(&self, callback: Value) {
144
+ *self.callback.lock() = Some(callback.into());
145
+ }
146
+
147
+ pub fn set_tls(ruby: &Ruby, this: &Self, cert_path: String, key_path: String) -> Result<(), Error> {
148
+ for (what, path) in [("certificate", &cert_path), ("key", &key_path)] {
149
+ if !std::path::Path::new(path).is_file() {
150
+ return Err(runtime_error(ruby, format!("TLS {what} file not found: {path}")));
151
+ }
152
+ }
153
+ this.options.lock().tls = Some((cert_path, key_path));
154
+ Ok(())
155
+ }
156
+
157
+ pub fn set_tls_client_ca(ruby: &Ruby, this: &Self, ca_path: String) -> Result<(), Error> {
158
+ if !std::path::Path::new(&ca_path).is_file() {
159
+ return Err(runtime_error(
160
+ ruby,
161
+ format!("TLS client CA file not found: {ca_path}"),
162
+ ));
163
+ }
164
+ this.options.lock().tls_client_ca = Some(ca_path);
165
+ Ok(())
166
+ }
167
+
168
+ pub fn set_max_connections(&self, max: usize) {
169
+ self.options.lock().max_connections = max;
170
+ }
171
+
172
+ pub fn set_concurrency_limit_per_connection(&self, limit: usize) {
173
+ self.options.lock().concurrency_limit_per_connection = limit;
174
+ }
175
+
176
+ pub fn set_worker_threads(&self, n: usize) {
177
+ self.options.lock().worker_threads = (n > 0).then_some(n);
178
+ }
179
+
180
+ pub fn set_keepalive(&self, interval: f64, timeout: f64) {
181
+ let mut options = self.options.lock();
182
+ options.keepalive_interval = secs(interval);
183
+ options.keepalive_timeout = secs(timeout);
184
+ }
185
+
186
+ pub fn set_dispatch_queue_size(&self, size: usize) {
187
+ self.options.lock().dispatch_queue_size = size.max(1);
188
+ }
189
+
190
+ pub fn set_log_requests(&self, enabled: bool) {
191
+ self.stats.log_requests.store(enabled, Ordering::Relaxed);
192
+ }
193
+
194
+ pub fn set_compression(&self, enabled: bool) {
195
+ self.compression.store(enabled, Ordering::Relaxed);
196
+ }
197
+
198
+ pub fn set_reflection(&self, enabled: bool) {
199
+ self.options.lock().reflection = enabled;
200
+ }
201
+
202
+ /// gRPC paths (`/package.Service/Method`) served by Ruby handlers.
203
+ pub fn set_routes(&self, routes: Vec<String>) {
204
+ self.options.lock().routes = routes;
205
+ }
206
+
207
+ /// Serialized `google.protobuf.FileDescriptorProto`s for reflection.
208
+ pub fn set_reflection_descriptors(&self, descriptors: RArray) -> Result<(), Error> {
209
+ let mut out = Vec::with_capacity(descriptors.len());
210
+ for value in descriptors.into_iter() {
211
+ let s = RString::try_convert(value)?;
212
+ out.push(unsafe { s.as_slice() }.to_vec());
213
+ }
214
+ self.options.lock().descriptors = out;
215
+ Ok(())
216
+ }
217
+
218
+ pub fn start_server(ruby: &Ruby, this: &Self) -> Result<(), Error> {
219
+ let mut running = this.running.lock();
220
+ if running.is_some() {
221
+ return Err(runtime_error(ruby, "server already running"));
222
+ }
223
+ let callback = this
224
+ .callback
225
+ .lock()
226
+ .ok_or_else(|| runtime_error(ruby, "no dispatch callback set: call #on_request first"))?;
227
+
228
+ let options = this.options.lock();
229
+ let addr: SocketAddr = this.bind_address.parse().expect("validated in new");
230
+
231
+ let routes: HashSet<String> = options.routes.iter().cloned().collect();
232
+ let mut services: Vec<String> = routes
233
+ .iter()
234
+ .filter_map(|p| {
235
+ p.strip_prefix('/')?
236
+ .split_once('/')
237
+ .map(|(svc, _)| svc.to_owned())
238
+ })
239
+ .collect();
240
+ services.sort();
241
+ services.dedup();
242
+
243
+ // TLS problems must fail startup — never fall back to plaintext.
244
+ let tls = match &options.tls {
245
+ Some((cert, key)) => {
246
+ let identity = tonic::transport::Identity::from_pem(
247
+ read_file(ruby, "certificate", cert)?,
248
+ read_file(ruby, "key", key)?,
249
+ );
250
+ let mut config = tonic::transport::ServerTlsConfig::new().identity(identity);
251
+ if let Some(ca) = &options.tls_client_ca {
252
+ config = config.client_ca_root(tonic::transport::Certificate::from_pem(read_file(
253
+ ruby,
254
+ "client CA",
255
+ ca,
256
+ )?));
257
+ }
258
+ Some(config)
259
+ }
260
+ None => None,
261
+ };
262
+
263
+ let reflection = if options.reflection {
264
+ let mut files = Vec::with_capacity(options.descriptors.len());
265
+ for bytes in &options.descriptors {
266
+ files.push(
267
+ prost_types::FileDescriptorProto::decode(bytes.as_slice())
268
+ .map_err(|e| arg_error(ruby, format!("invalid reflection descriptor: {e}")))?,
269
+ );
270
+ }
271
+ let mut builder = tonic_reflection::server::Builder::configure()
272
+ .register_file_descriptor_set(prost_types::FileDescriptorSet { file: files })
273
+ .register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
274
+ .with_service_name("grpc.health.v1.Health")
275
+ .with_service_name("grpc.reflection.v1.ServerReflection");
276
+ for service in &services {
277
+ builder = builder.with_service_name(service.clone());
278
+ }
279
+ Some(
280
+ builder
281
+ .build_v1()
282
+ .map_err(|e| runtime_error(ruby, format!("failed to build reflection service: {e}")))?,
283
+ )
284
+ } else {
285
+ None
286
+ };
287
+
288
+ let mut rt = tokio::runtime::Builder::new_multi_thread();
289
+ rt.enable_all().thread_name("tonic-rails");
290
+ if let Some(n) = options.worker_threads {
291
+ rt.worker_threads(n);
292
+ }
293
+ let runtime = rt
294
+ .build()
295
+ .map_err(|e| runtime_error(ruby, format!("failed to create tokio runtime: {e}")))?;
296
+
297
+ let listener = {
298
+ let _guard = runtime.enter();
299
+ bind(addr, 1024).map_err(|e| runtime_error(ruby, format!("failed to bind {addr}: {e}")))?
300
+ };
301
+
302
+ let (health, health_service) = tonic_health::server::health_reporter();
303
+ runtime.block_on(async {
304
+ health.set_service_status(OVERALL, ServingStatus::Serving).await;
305
+ for service in &services {
306
+ health.set_service_status(service, ServingStatus::Serving).await;
307
+ }
308
+ });
309
+
310
+ let (request_tx, request_rx) = crossbeam_channel::bounded(options.dispatch_queue_size);
311
+ let dispatch = DispatchService(Arc::new(DispatchInner {
312
+ callback,
313
+ routes,
314
+ request_tx,
315
+ stats: this.stats.clone(),
316
+ max_message_size: this.max_message_size,
317
+ compression: this.compression.clone(),
318
+ }));
319
+
320
+ let mut grpc_routes = tonic::service::Routes::new(health_service);
321
+ if let Some(reflection) = reflection {
322
+ grpc_routes = grpc_routes.add_service(reflection);
323
+ }
324
+ let grpc_routes =
325
+ tonic::service::Routes::from(grpc_routes.into_axum_router().fallback_service(dispatch));
326
+
327
+ let window = this.max_message_size.min(u32::MAX as usize) as u32;
328
+ let mut builder = tonic::transport::Server::builder()
329
+ .http2_keepalive_interval(Some(options.keepalive_interval))
330
+ .http2_keepalive_timeout(Some(options.keepalive_timeout))
331
+ .http2_max_header_list_size(16 * 1024)
332
+ .http2_max_pending_accept_reset_streams(Some(64))
333
+ .initial_stream_window_size(window)
334
+ .initial_connection_window_size(window.saturating_mul(2));
335
+ if options.concurrency_limit_per_connection > 0 {
336
+ let limit = options.concurrency_limit_per_connection;
337
+ builder = builder
338
+ .concurrency_limit_per_connection(limit)
339
+ .max_concurrent_streams(u32::try_from(limit).unwrap_or(u32::MAX));
340
+ }
341
+ if let Some(tls) = tls {
342
+ builder = builder
343
+ .tls_config(tls)
344
+ .map_err(|e| runtime_error(ruby, format!("invalid TLS configuration: {e}")))?;
345
+ }
346
+ let router = builder.add_routes(grpc_routes);
347
+
348
+ let permits = match options.max_connections {
349
+ 0 => Semaphore::MAX_PERMITS,
350
+ n => n,
351
+ };
352
+ let incoming = LimitedIncoming {
353
+ inner: tokio_stream::wrappers::TcpListenerStream::new(listener),
354
+ semaphore: Arc::new(Semaphore::new(permits)),
355
+ active: this.active_connections.clone(),
356
+ };
357
+
358
+ let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
359
+ let (done_tx, done_rx) = crossbeam_channel::bounded(1);
360
+ runtime.spawn(async move {
361
+ let signal = async move {
362
+ let _ = shutdown_rx.wait_for(|stop| *stop).await;
363
+ };
364
+ if let Err(e) = router.serve_with_incoming_shutdown(incoming, signal).await {
365
+ eprintln!("tonic-rails: server error: {e}");
366
+ }
367
+ let _ = done_tx.send(());
368
+ });
369
+
370
+ *running = Some(Running {
371
+ runtime,
372
+ request_rx,
373
+ stop_workers: Arc::new(AtomicBool::new(false)),
374
+ shutdown_tx,
375
+ done_rx,
376
+ health,
377
+ services,
378
+ });
379
+ Ok(())
380
+ }
381
+
382
+ /// Run one dispatch worker. Blocks (GVL released while idle) until
383
+ /// `stop_workers` is called. Interrupts are serviced in between.
384
+ pub fn serve_worker(ruby: &Ruby, this: &Self) -> Result<(), Error> {
385
+ let (request_rx, stop) = match this.running.lock().as_ref() {
386
+ Some(r) => (r.request_rx.clone(), r.stop_workers.clone()),
387
+ None => return Err(runtime_error(ruby, "server not started")),
388
+ };
389
+
390
+ while !stop.load(Ordering::Acquire) {
391
+ let cancel = Cancel::default();
392
+ gvl::without_gvl(&cancel, || {
393
+ while !stop.load(Ordering::Acquire) && !cancel.is_cancelled() {
394
+ match request_rx.recv_timeout(Duration::from_millis(100)) {
395
+ Ok(ctx) => gvl::dispatch(ctx, &request_rx),
396
+ Err(crossbeam_channel::RecvTimeoutError::Timeout) => {}
397
+ Err(crossbeam_channel::RecvTimeoutError::Disconnected) => return,
398
+ }
399
+ }
400
+ });
401
+ ruby.thread_check_ints()?;
402
+ }
403
+ Ok(())
404
+ }
405
+
406
+ /// Mark everything NOT_SERVING and stop accepting new connections.
407
+ /// Existing connections receive GOAWAY; in-flight RPCs continue.
408
+ pub fn begin_shutdown(&self) {
409
+ if let Some(r) = self.running.lock().as_ref() {
410
+ r.runtime.block_on(async {
411
+ r.health
412
+ .set_service_status(OVERALL, ServingStatus::NotServing)
413
+ .await;
414
+ for service in &r.services {
415
+ r.health
416
+ .set_service_status(service, ServingStatus::NotServing)
417
+ .await;
418
+ }
419
+ });
420
+ let _ = r.shutdown_tx.send(true);
421
+ }
422
+ }
423
+
424
+ /// Wait (GVL released, so workers can finish in-flight RPCs) until
425
+ /// all connections have closed or `timeout` seconds pass.
426
+ /// Returns true if the server drained completely.
427
+ pub fn wait_for_drain(ruby: &Ruby, this: &Self, timeout: f64) -> Result<bool, Error> {
428
+ let Some(done_rx) = this.running.lock().as_ref().map(|r| r.done_rx.clone()) else {
429
+ return Ok(true);
430
+ };
431
+ let deadline = Instant::now() + secs(timeout);
432
+ loop {
433
+ let cancel = Cancel::default();
434
+ let drained = gvl::without_gvl(&cancel, || {
435
+ while !cancel.is_cancelled() {
436
+ let remaining = deadline.saturating_duration_since(Instant::now());
437
+ if remaining.is_zero() {
438
+ return Some(false);
439
+ }
440
+ match done_rx.recv_timeout(remaining.min(Duration::from_millis(100))) {
441
+ Ok(()) | Err(crossbeam_channel::RecvTimeoutError::Disconnected) => return Some(true),
442
+ Err(crossbeam_channel::RecvTimeoutError::Timeout) => {}
443
+ }
444
+ }
445
+ None
446
+ });
447
+ match drained.flatten() {
448
+ Some(result) => return Ok(result),
449
+ None => ruby.thread_check_ints()?,
450
+ }
451
+ }
452
+ }
453
+
454
+ pub fn stop_workers(&self) {
455
+ if let Some(r) = self.running.lock().as_ref() {
456
+ r.stop_workers.store(true, Ordering::Release);
457
+ }
458
+ }
459
+
460
+ /// Tear down the runtime. Connections still open are closed.
461
+ pub fn stop_server(&self) {
462
+ let Some(running) = self.running.lock().take() else {
463
+ return;
464
+ };
465
+ running.stop_workers.store(true, Ordering::Release);
466
+ let _ = running.shutdown_tx.send(true);
467
+ let mut runtime = Some(running.runtime);
468
+ // Shutdown joins tokio worker threads; never do that holding the GVL.
469
+ let cancel = Cancel::default();
470
+ gvl::without_gvl(&cancel, || {
471
+ if let Some(rt) = runtime.take() {
472
+ rt.shutdown_timeout(Duration::from_secs(1));
473
+ }
474
+ });
475
+ // An interrupt was pending, so the closure never ran.
476
+ if let Some(rt) = runtime {
477
+ rt.shutdown_background();
478
+ }
479
+ }
480
+
481
+ pub fn is_running(&self) -> bool {
482
+ self.running.lock().is_some()
483
+ }
484
+
485
+ pub fn stats(ruby: &Ruby, this: &Self) -> Result<RHash, Error> {
486
+ let s = &this.stats;
487
+ let queued = this.running.lock().as_ref().map_or(0, |r| r.request_rx.len());
488
+ let hash = ruby.hash_new();
489
+ for (key, value) in [
490
+ ("total", s.total.load(Ordering::Relaxed)),
491
+ ("successful", s.successful.load(Ordering::Relaxed)),
492
+ ("errors", s.errors.load(Ordering::Relaxed)),
493
+ ("cancelled", s.cancelled.load(Ordering::Relaxed)),
494
+ ("deadline_exceeded", s.deadline_exceeded.load(Ordering::Relaxed)),
495
+ ("resource_exhausted", s.resource_exhausted.load(Ordering::Relaxed)),
496
+ ("unimplemented", s.unimplemented.load(Ordering::Relaxed)),
497
+ (
498
+ "active_connections",
499
+ this.active_connections.load(Ordering::Relaxed),
500
+ ),
501
+ ("queued", queued),
502
+ ] {
503
+ hash.aset(ruby.to_symbol(key), value)?;
504
+ }
505
+ Ok(hash)
506
+ }
507
+ }