tonic-rails 0.4.0 → 0.6.0

data/ext/tonic_rails_native/src/gvl.rs

@@ -20,7 +20,7 @@ use magnus::{
     Error, RHash, RString, Ruby, Value, prelude::*, r_hash::ForEach, typed_data::DataTypeFunctions,
     value::Opaque,
 };
-use parking_lot::Mutex;
+use parking_lot::{Condvar, Mutex};
 use rb_sys::{RB_NOGVL_INTR_FAIL, rb_nogvl, rb_thread_call_with_gvl};
 use tokio::sync::{Notify, mpsc};
 
@@ -100,25 +100,40 @@ where
 
 /// Minimal single-future executor that parks the current (non-tokio)
 /// thread.  tokio's channel and `Notify` futures are runtime-agnostic.
+///
+/// It parks on a pthread condition variable rather than
+/// `std::thread::park`, which on macOS uses libdispatch: libdispatch
+/// traps in the child of a multi-threaded process, and the client blocks
+/// here in processes forked after it made calls.
 pub fn block_on<F: Future>(fut: F) -> F::Output {
-    struct ThreadWaker(std::thread::Thread);
-    impl Wake for ThreadWaker {
+    #[derive(Default)]
+    struct Signal {
+        woken: Mutex<bool>,
+        condvar: Condvar,
+    }
+    impl Wake for Signal {
         fn wake(self: Arc<Self>) {
-            self.0.unpark();
+            self.wake_by_ref();
         }
         fn wake_by_ref(self: &Arc<Self>) {
-            self.0.unpark();
+            *self.woken.lock() = true;
+            self.condvar.notify_one();
         }
     }
 
-    let waker = Waker::from(Arc::new(ThreadWaker(std::thread::current())));
+    let signal = Arc::new(Signal::default());
+    let waker = Waker::from(signal.clone());
     let mut cx = Context::from_waker(&waker);
     let mut fut = std::pin::pin!(fut);
     loop {
         if let Poll::Ready(value) = fut.as_mut().poll(&mut cx) {
             return value;
         }
-        std::thread::park();
+        let mut woken = signal.woken.lock();
+        while !*woken {
+            signal.condvar.wait(&mut woken);
+        }
+        *woken = false;
     }
 }
 

data/ext/tonic_rails_native/src/lib.rs

@@ -20,6 +20,7 @@
 //! Modules:
 //! - [`framing`]  — gRPC wire format (frames, gzip, timeouts, trailers)
 //! - [`body`]     — response bodies and request stats
+//! - [`client`]   — `TonicRails::Native::Channel` and `ResponseStream`, for outgoing calls
 //! - [`dispatch`] — the fallback tower service that queues RPCs for Ruby
 //! - [`gvl`]      — GVL release/acquire and the Ruby stream objects
 //! - [`log`]      — transport log events for Ruby's logger
@@ -27,6 +28,7 @@
 //! - [`server`]   — `TonicRails::Native::Server`
 
 mod body;
+mod client;
 mod dispatch;
 mod framing;
 mod gvl;
@@ -36,6 +38,7 @@ mod server;
 
 use magnus::{Error, Ruby, function, method, prelude::*};
 
+use client::{NativeChannel, ResponseStream};
 use gvl::{RequestReader, StreamSender};
 use server::NativeServer;
 
@@ -43,6 +46,8 @@ use server::NativeServer;
 fn init(ruby: &Ruby) -> Result<(), Error> {
     let native = ruby.define_module("TonicRails")?.define_module("Native")?;
 
+    native.define_module_function("bind_listener", function!(server::bind_listener, 1))?;
+
     let class = native.define_class("Server", ruby.class_object())?;
     class.define_singleton_method("new", function!(NativeServer::new, 2))?;
     class.define_method("on_request", method!(NativeServer::on_request, 1))?;
@@ -66,6 +71,7 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     class.define_method("set_reflection", method!(NativeServer::set_reflection, 1))?;
     class.define_method("set_grpc_web", method!(NativeServer::set_grpc_web, 3))?;
     class.define_method("set_routes", method!(NativeServer::set_routes, 1))?;
+    class.define_method("set_listener_fd", method!(NativeServer::set_listener_fd, 1))?;
     class.define_method(
         "set_reflection_descriptors",
         method!(NativeServer::set_reflection_descriptors, 1),
@@ -92,5 +98,16 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     reader.define_method("next", method!(RequestReader::next, 0))?;
     reader.define_method("error", method!(RequestReader::error, 0))?;
 
+    let channel = native.define_class("Channel", ruby.class_object())?;
+    channel.define_singleton_method("new", function!(NativeChannel::new, 4))?;
+    channel.define_method("unary", method!(NativeChannel::unary, 4))?;
+    channel.define_method("server_streaming", method!(NativeChannel::server_streaming, 4))?;
+
+    let stream = native.define_class("ResponseStream", ruby.class_object())?;
+    stream.undef_default_alloc_func();
+    stream.define_method("next", method!(ResponseStream::next, 0))?;
+    stream.define_method("result", method!(ResponseStream::result, 0))?;
+    stream.define_method("close", method!(ResponseStream::close, 0))?;
+
     Ok(())
 }

data/ext/tonic_rails_native/src/net.rs

@@ -9,17 +9,44 @@ use std::time::Duration;
 
 use socket2::{Domain, Protocol, Socket, Type};
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
-use tokio::net::{TcpListener, TcpStream};
+use tokio::net::TcpStream;
 use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 use tokio_stream::Stream;
 use tonic::transport::server::{Connected, TcpConnectInfo};
 
 use crate::log::{Level, LogSink};
 
-/// Bind with SO_REUSEADDR + SO_REUSEPORT so several worker processes can
-/// share a port (the kernel load-balances accepts) and restarts don't
-/// fail on TIME_WAIT.  Must be called inside a tokio runtime context.
-pub fn bind(addr: SocketAddr, backlog: i32) -> std::io::Result<TcpListener> {
+/// A stream of accepted connections.
+pub type Accept = Pin<Box<dyn Stream<Item = std::io::Result<TcpStream>> + Send>>;
+
+/// Bind `addr` and accept on it.  Must be called inside a tokio runtime
+/// context.
+pub fn listen(addr: SocketAddr, backlog: i32) -> std::io::Result<Accept> {
+    let listener = tokio::net::TcpListener::from_std(bind(addr, backlog)?)?;
+    Ok(Box::pin(tokio_stream::wrappers::TcpListenerStream::new(listener)))
+}
+
+/// Accept on a listening socket inherited from a parent process, shared
+/// with sibling workers (multi-process mode).  The fd is duplicated, so
+/// the caller keeps ownership of the original.  `active` is this
+/// process's open connection count.  Must be called inside a tokio
+/// runtime context.
+#[cfg(unix)]
+pub fn listen_inherited(fd: i32, active: Arc<AtomicUsize>) -> std::io::Result<Accept> {
+    Ok(Box::pin(SharedAccept::new(inherited(fd)?, active)?))
+}
+
+#[cfg(not(unix))]
+pub fn listen_inherited(_fd: i32, _active: Arc<AtomicUsize>) -> std::io::Result<Accept> {
+    Err(std::io::Error::new(
+        std::io::ErrorKind::Unsupported,
+        "inherited listeners need fork(), which this platform lacks",
+    ))
+}
+
+/// Bind with SO_REUSEADDR + SO_REUSEPORT so restarts don't fail on
+/// TIME_WAIT.  The listener is non-blocking, ready for tokio.
+pub fn bind(addr: SocketAddr, backlog: i32) -> std::io::Result<std::net::TcpListener> {
     let domain = if addr.is_ipv4() {
         Domain::IPV4
     } else {
@@ -35,7 +62,95 @@ pub fn bind(addr: SocketAddr, backlog: i32) -> std::io::Result<TcpListener> {
     socket.set_nonblocking(true)?;
     socket.bind(&addr.into())?;
     socket.listen(backlog)?;
-    TcpListener::from_std(socket.into())
+    Ok(socket.into())
+}
+
+#[cfg(unix)]
+fn inherited(fd: i32) -> std::io::Result<std::net::TcpListener> {
+    use std::os::fd::BorrowedFd;
+
+    // SAFETY: the caller guarantees `fd` is open for the duration of
+    // this call; we only duplicate it.
+    let owned = unsafe { BorrowedFd::borrow_raw(fd) }.try_clone_to_owned()?;
+    let socket = Socket::from(owned);
+    socket.set_nonblocking(true)?;
+    Ok(socket.into())
+}
+
+/// Accepts on a listening socket shared with sibling worker processes.
+///
+/// Every process is woken for each new connection and the first to
+/// accept wins, so gRPC's few long-lived connections tend to pile onto
+/// whichever process wakes fastest.  Instead, when a connection is
+/// waiting, a process holds back for `backoff` per connection it already
+/// has (up to `max_backoff`) before trying to accept, so the least busy
+/// process usually gets it, as with Puma's `wait_for_less_busy_worker`.
+#[cfg(unix)]
+struct SharedAccept {
+    listener: tokio::io::unix::AsyncFd<std::net::TcpListener>,
+    active: Arc<AtomicUsize>,
+    backoff: Duration,
+    max_backoff: Duration,
+    pause: Option<Pin<Box<tokio::time::Sleep>>>,
+    /// The pause for the waiting connection is over: accept right away.
+    paused: bool,
+}
+
+#[cfg(unix)]
+impl SharedAccept {
+    fn new(listener: std::net::TcpListener, active: Arc<AtomicUsize>) -> std::io::Result<Self> {
+        Ok(Self {
+            listener: tokio::io::unix::AsyncFd::new(listener)?,
+            active,
+            backoff: Duration::from_millis(1),
+            max_backoff: Duration::from_millis(10),
+            pause: None,
+            paused: false,
+        })
+    }
+}
+
+#[cfg(unix)]
+impl Stream for SharedAccept {
+    type Item = std::io::Result<TcpStream>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = &mut *self;
+        loop {
+            if let Some(pause) = this.pause.as_mut() {
+                if pause.as_mut().poll(cx).is_pending() {
+                    return Poll::Pending;
+                }
+                this.pause = None;
+                this.paused = true;
+            }
+            let mut guard = match this.listener.poll_read_ready(cx) {
+                Poll::Ready(Ok(guard)) => guard,
+                Poll::Ready(Err(e)) => return Poll::Ready(Some(Err(e))),
+                Poll::Pending => return Poll::Pending,
+            };
+            let active = this.active.load(Ordering::Relaxed);
+            if !this.paused && active > 0 {
+                let pause = this.backoff.saturating_mul(active as u32).min(this.max_backoff);
+                this.pause = Some(Box::pin(tokio::time::sleep(pause)));
+                continue;
+            }
+            // Whether we accept or a sibling beat us to it, the next
+            // connection gets a fresh pause.
+            this.paused = false;
+            match guard.try_io(|listener| listener.get_ref().accept()) {
+                Ok(Ok((stream, _))) => {
+                    return Poll::Ready(Some(
+                        stream
+                            .set_nonblocking(true)
+                            .and_then(|()| TcpStream::from_std(stream)),
+                    ));
+                }
+                Ok(Err(e)) => return Poll::Ready(Some(Err(e))),
+                Err(_would_block) => continue,
+            }
+        }
+    }
 }
 
 /// TCP_NODELAY for latency; TCP keepalive to reap half-open peers that

data/ext/tonic_rails_native/src/server.rs

@@ -36,7 +36,7 @@ use crate::body::ServerStats;
 use crate::dispatch::{DispatchInner, DispatchService};
 use crate::gvl::{self, Cancel, DispatchContext};
 use crate::log::{Level, LogEvent, LogSink};
-use crate::net::{LimitedIncoming, bind};
+use crate::net::{LimitedIncoming, bind, listen, listen_inherited};
 
 /// The overall-server health entry (empty service name).
 const OVERALL: &str = "";
@@ -54,6 +54,8 @@ struct Options {
     dispatch_queue_size: usize,
     routes: Vec<String>,
     descriptors: Vec<Vec<u8>>,
+    /// Accept on this inherited listening socket instead of binding.
+    listener_fd: Option<i32>,
 }
 
 impl Default for Options {
@@ -71,6 +73,7 @@ impl Default for Options {
             dispatch_queue_size: 128,
             routes: Vec::new(),
             descriptors: Vec::new(),
+            listener_fd: None,
         }
     }
 }
@@ -191,6 +194,29 @@ fn read_file(ruby: &Ruby, what: &str, path: &str) -> Result<Vec<u8>, Error> {
     std::fs::read(path).map_err(|e| runtime_error(ruby, format!("cannot read TLS {what} {path}: {e}")))
 }
 
+/// `TonicRails::Native.bind_listener(address)`: bind and listen the way
+/// the server does, returning the socket's fd for worker processes to
+/// inherit and accept on (`Server#set_listener_fd`).  The caller owns it.
+#[cfg(unix)]
+pub fn bind_listener(ruby: &Ruby, address: String) -> Result<i32, Error> {
+    use std::os::fd::IntoRawFd;
+
+    let addr: SocketAddr = address
+        .parse()
+        .map_err(|e| arg_error(ruby, format!("invalid bind address {address:?}: {e}")))?;
+    let listener =
+        bind(addr, 1024).map_err(|e| runtime_error(ruby, format!("failed to bind {addr}: {e}")))?;
+    Ok(listener.into_raw_fd())
+}
+
+#[cfg(not(unix))]
+pub fn bind_listener(ruby: &Ruby, _address: String) -> Result<i32, Error> {
+    Err(Error::new(
+        ruby.exception_not_imp_error(),
+        "multi-process mode needs fork(), which this platform lacks",
+    ))
+}
+
 fn secs(value: f64) -> Duration {
     Duration::from_secs_f64(value.max(0.0))
 }
@@ -292,6 +318,12 @@ impl NativeServer {
         self.options.lock().reflection = enabled;
     }
 
+    /// Accept on a listening socket bound by a parent process (see
+    /// `bind_listener`) instead of binding `bind_address`.
+    pub fn set_listener_fd(&self, fd: i32) {
+        self.options.lock().listener_fd = Some(fd);
+    }
+
     /// gRPC paths (`/package.Service/Method`) served by Ruby handlers.
     pub fn set_routes(&self, routes: Vec<String>) {
         self.options.lock().routes = routes;
@@ -387,9 +419,15 @@ impl NativeServer {
             .build()
             .map_err(|e| runtime_error(ruby, format!("failed to create tokio runtime: {e}")))?;
 
-        let listener = {
+        let accept = {
             let _guard = runtime.enter();
-            bind(addr, 1024).map_err(|e| runtime_error(ruby, format!("failed to bind {addr}: {e}")))?
+            match options.listener_fd {
+                Some(fd) => listen_inherited(fd, this.active_connections.clone()).map_err(|e| {
+                    runtime_error(ruby, format!("cannot use inherited listener (fd {fd}): {e}"))
+                })?,
+                None => listen(addr, 1024)
+                    .map_err(|e| runtime_error(ruby, format!("failed to bind {addr}: {e}")))?,
+            }
         };
 
         let (health, health_service) = tonic_health::server::health_reporter();
@@ -453,7 +491,7 @@ impl NativeServer {
             n => n,
         };
         let incoming = LimitedIncoming {
-            inner: tokio_stream::wrappers::TcpListenerStream::new(listener),
+            inner: accept,
             semaphore: Arc::new(Semaphore::new(permits)),
             active: this.active_connections.clone(),
             log: this.stats.log.clone(),

data/lib/puma/plugin/tonic_rails.rb

@@ -9,7 +9,8 @@ require "puma/plugin"
 #
 # The gRPC server runs as one child process of the Puma master, in single
 # and cluster mode alike, so it has its own GVL and never competes with web
-# workers.  It starts once Puma has booted, drains gracefully when Puma
+# workers.  With `config.processes > 1`, that child supervises the gRPC
+# worker processes.  It starts once Puma has booted, drains gracefully when Puma
 # stops or restarts, stops itself if Puma goes away, and if it dies
 # unexpectedly Puma shuts down too, so a supervisor can restart both.
 Puma::Plugin.create do
@@ -62,7 +63,8 @@ Puma::Plugin.create do
   end
 
   # Puma fires its stop hooks from its TERM trap handler, so this must
-  # not block indefinitely: allow the server's drain timeout plus a margin,
+  # not block indefinitely: allow the server's drain timeout plus a margin
+  # (longer than the one a multi-process server gives its own workers),
   # then kill it.
   def stop_grpc
     pid = @grpc_pid
@@ -87,7 +89,7 @@ Puma::Plugin.create do
 
   def stop_timeout
     drain = defined?(Tonic::Rails) ? Tonic::Rails.configuration.drain_timeout.to_f : 10.0
-    drain + 5
+    drain + 10
   end
 
   def monitor_grpc

data/lib/tonic/rails/base.rb

@@ -248,7 +248,10 @@ module Tonic
 
         @request = build_request(reader, spec)
         response = nil
-        run_callbacks(:rpc) { response = respond(public_send(spec[:handler], @request), spec, sender) }
+        # Calls this handler makes with Tonic::Rails::Client default to its deadline.
+        Client.with_deadline(@grpc_deadline_at) do
+          run_callbacks(:rpc) { response = respond(public_send(spec[:handler], @request), spec, sender) }
+        end
         response
       rescue RpcError
         raise