tonic-rails 0.4.0 → 0.6.0

data/lib/tonic/rails/client.rb

@@ -0,0 +1,371 @@
+# frozen_string_literal: true
+
+require "active_support/inflector"
+
+module Tonic
+  module Rails
+    # Calls other gRPC services, without the grpc gem.
+    #
+    #   ITEMS = Tonic::Rails::Client.new("items.internal:50051", deadline: 2)
+    #
+    #   items = ITEMS.stub("demo.v1.ItemService")
+    #   item = items.get_item(Demo::V1::GetItemRequest.new(id: 1))
+    #   item = items.get_item(id: 1)
+    #   item = items.get_item({ id: 1 }, deadline: 0.5, metadata: { "x-request-id" => request_id })
+    #   items.list_items(category: "books").each { |item| ... }
+    #
+    # Stubs are built from the protobuf descriptors that
+    # `tonic_rails:compile_protos` (protoc --ruby_out) generates; no
+    # `_services_pb.rb` is needed.  A failed call raises
+    # Tonic::Rails::RpcError.
+    #
+    # A client holds one HTTP/2 connection and is safe to share between
+    # threads, so create it once (a constant or an initializer).  It keeps
+    # working across `fork`: each process connects on its own, with no
+    # hooks to call.
+    class Client
+      # A response with its headers and trailers; see `return_call:`.
+      Call = Data.define(:response, :headers, :trailers)
+
+      DEFAULT_MAX_MESSAGE_SIZE = 4 * 1024 * 1024
+
+      # @return [String] "host:port"
+      attr_reader :address
+
+      # @return [Numeric, nil] default deadline in seconds
+      attr_reader :deadline
+
+      # @param address [String] "host:port"
+      # @param tls [Boolean] false for plaintext, which must be explicit
+      # @param tls_ca_path [String, nil] CA certificates (PEM) to trust instead
+      #   of the system's root certificates
+      # @param tls_cert_path [String, nil] client certificate (PEM), for mutual TLS
+      # @param tls_key_path [String, nil] client private key (PEM), for mutual TLS
+      # @param tls_server_name [String, nil] name to verify the server's
+      #   certificate against, if not the host in `address`
+      # @param deadline [Numeric, nil] seconds each call may take, unless the
+      #   call sets its own
+      # @param metadata [Hash] sent with every call (see Tonic::Rails::Metadata)
+      # @param max_message_size [Integer] largest message sent or received, in bytes
+      # @param connect_timeout [Numeric] seconds to wait for a connection
+      def initialize(address, tls: true, tls_ca_path: nil, tls_cert_path: nil, tls_key_path: nil,
+                     tls_server_name: nil, deadline: nil, metadata: {}, max_message_size: DEFAULT_MAX_MESSAGE_SIZE,
+                     connect_timeout: 5)
+        Tonic::Rails.native_extension!
+        validate_address!(address)
+        validate_tls!(tls, tls_cert_path, tls_key_path)
+        positive!(:deadline, deadline) unless deadline.nil?
+        positive!(:connect_timeout, connect_timeout)
+        unless max_message_size.is_a?(Integer) && max_message_size.positive?
+          raise ArgumentError, "max_message_size must be a positive Integer"
+        end
+
+        @address = address
+        @deadline = deadline
+        @metadata = Metadata.encode(metadata)
+        @stubs = {}
+        tls_options = [read_pem(tls_ca_path, "CA"), read_pem(tls_cert_path, "certificate"),
+                       read_pem(tls_key_path, "key"), tls_server_name]
+        @native = TonicRails::Native::Channel.new(
+          "#{tls ? 'https' : 'http'}://#{address}", max_message_size, (connect_timeout * 1000).ceil,
+          (tls_options if tls)
+        )
+      end
+
+      # A stub with one method per RPC of the service, named in snake_case
+      # (`GetItem` becomes `get_item`).
+      #
+      # @param service [String, Google::Protobuf::ServiceDescriptor, #service_name]
+      #   the fully-qualified service name, its descriptor, or a grpc gem
+      #   service class (`Demo::V1::ItemService::Service`)
+      # @return [Tonic::Rails::Client::Stub]
+      def stub(service)
+        descriptor = service_descriptor(service)
+        @stubs[descriptor.name] ||= Stub.for(descriptor).new(self)
+      end
+
+      def inspect
+        "#<#{self.class.name} #{address}>"
+      end
+
+      # The deadline of the RPC this thread is handling, as a monotonic
+      # clock time; outgoing calls default to it.
+      # @api private
+      def self.current_deadline
+        Thread.current[:tonic_rails_deadline]
+      end
+
+      # @api private
+      def self.with_deadline(deadline)
+        previous = Thread.current[:tonic_rails_deadline]
+        Thread.current[:tonic_rails_deadline] = deadline
+        yield
+      ensure
+        Thread.current[:tonic_rails_deadline] = previous
+      end
+
+      # Make a unary call.  Stubs call this; use them instead.
+      # @api private
+      def unary(path, input, output, request, deadline: nil, metadata: {}, return_call: false)
+        message = request_message(path, input, request)
+        metadata = call_metadata(metadata)
+        timeout_ms = remaining_ms(deadline_at(deadline))
+
+        Instrumentation.instrument_request(path, address, metadata) do
+          code, error_message, response, headers, trailers, details =
+            @native.unary(path, input.encode(message), metadata, timeout_ms)
+          headers = decode_metadata(headers)
+          trailers = decode_metadata(trailers)
+          check_status!(code, error_message, trailers, details)
+
+          response = decode_response(output, response, path)
+          return_call ? Call.new(response: response, headers: headers, trailers: trailers) : response
+        end
+      end
+
+      # Start a server-streaming call.  Stubs call this; use them instead.
+      # @api private
+      def server_streaming(path, input, output, request, deadline: nil, metadata: {})
+        message = request_message(path, input, request)
+        metadata = call_metadata(metadata)
+        # Fixed now, so the stream keeps the deadline of the RPC this
+        # thread is handling even if it's read later.
+        deadline_at = deadline_at(deadline)
+
+        ServerStream.new(path) do |stream, block|
+          timeout_ms = remaining_ms(deadline_at)
+          Instrumentation.instrument_request(path, address, metadata) do
+            native = @native.server_streaming(path, input.encode(message), metadata, timeout_ms)
+            begin
+              while (bytes = native.next)
+                block.call(decode_response(output, bytes, path))
+              end
+              code, error_message, _, headers, trailers, details = native.result
+              trailers = decode_metadata(trailers)
+              stream.finish(decode_metadata(headers), trailers)
+              check_status!(code, error_message, trailers, details)
+            ensure
+              native.close
+            end
+          end
+        end
+      end
+
+      private
+
+      def validate_address!(address)
+        return if address.is_a?(String) && address.match?(%r{\A[^\s/]+:\d+\z})
+
+        raise ArgumentError, "address must be \"host:port\", got #{address.inspect}"
+      end
+
+      def validate_tls!(tls, cert_path, key_path)
+        raise ArgumentError, "tls must be true or false" unless [true, false].include?(tls)
+        raise ArgumentError, "tls_cert_path and tls_key_path must be set together" if cert_path.nil? != key_path.nil?
+        raise ArgumentError, "tls_cert_path and tls_key_path need tls: true" if cert_path && !tls
+      end
+
+      def positive!(name, value)
+        raise ArgumentError, "#{name} must be a positive number of seconds" unless value.is_a?(Numeric) && value.positive?
+      end
+
+      def read_pem(path, what)
+        return unless path
+        raise ArgumentError, "TLS #{what} file not found: #{path}" unless File.file?(path)
+
+        File.read(path)
+      end
+
+      def service_descriptor(service)
+        return service if service.is_a?(Google::Protobuf::ServiceDescriptor)
+
+        name = service.respond_to?(:service_name) ? service.service_name : service.to_s
+        descriptor = Google::Protobuf::DescriptorPool.generated_pool.lookup(name)
+        return descriptor if descriptor.is_a?(Google::Protobuf::ServiceDescriptor)
+
+        raise NotConfiguredError,
+              "#{name} is not a known protobuf service. Require its generated _pb.rb file (protoc --ruby_out) first."
+      end
+
+      def request_message(path, input, request)
+        case request
+        when input then request
+        when Hash then input.new(request)
+        else raise TypeError, "#{path} expects a #{input} or Hash, got #{request.class}"
+        end
+      end
+
+      def call_metadata(metadata)
+        metadata.empty? ? @metadata : @metadata.merge(Metadata.encode(metadata))
+      end
+
+      # When the call must end, as a monotonic clock time: the call's own
+      # deadline, or else the earlier of the client's default and the
+      # deadline of the RPC being handled.
+      def deadline_at(deadline)
+        now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        case deadline
+        when nil then [(now + @deadline if @deadline), self.class.current_deadline].compact.min
+        when Numeric then now + deadline
+        when Time then now + (deadline - Time.now)
+        else raise ArgumentError, "deadline must be seconds (Numeric) or a Time, got #{deadline.class}"
+        end
+      end
+
+      # Milliseconds left until `deadline_at`, or nil for no deadline.
+      def remaining_ms(deadline_at)
+        return unless deadline_at
+
+        remaining = deadline_at - Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        raise RpcError.new("deadline exceeded", code: :deadline_exceeded) unless remaining.positive?
+
+        (remaining * 1000).ceil
+      end
+
+      def check_status!(code, message, trailers, details)
+        return if code.zero?
+
+        raise RpcError.new(message.empty? ? nil : message, code: code, metadata: trailers,
+                                                           details: RpcError.decode_details(details))
+      end
+
+      # Wire-form [name, value] pairs to a Hash: `-bin` values decoded,
+      # repeated names collected into an Array, and transport headers
+      # (content-type, grpc-*) dropped.
+      def decode_metadata(pairs)
+        pairs.each_with_object({}) do |(name, value), metadata|
+          next if name.start_with?("grpc-") || Metadata::RESERVED.include?(name)
+
+          value = value.unpack1("m") if name.end_with?("-bin")
+          metadata[name] = metadata.key?(name) ? [*metadata[name], value] : value
+        end
+      end
+
+      def decode_response(output, bytes, path)
+        output.decode(bytes)
+      rescue Google::Protobuf::ParseError
+        raise RpcError.new("invalid protobuf response from #{path}", code: :internal)
+      end
+
+      # Calls an RPC's methods.  Built per service from its descriptor; get
+      # one with Client#stub.
+      class Stub
+        class << self
+          # @api private
+          attr_reader :service_name
+
+          # @api private
+          def for(descriptor)
+            Class.new(self) do
+              @service_name = descriptor.name
+              descriptor.each { |method| define_rpc(method) }
+            end
+          end
+
+          def inspect
+            service_name ? "#{Stub.name}(#{service_name})" : super
+          end
+
+          private
+
+          def define_rpc(method)
+            path = "/#{service_name}/#{method.name}"
+            name = ActiveSupport::Inflector.underscore(method.name)
+            if method.client_streaming
+              define_method(name) do |*, **|
+                raise NotImplementedError,
+                      "#{path} is a client-streaming RPC; the tonic-rails client makes unary and server-streaming calls"
+              end
+              return
+            end
+
+            input = method.input_type.msgclass
+            output = method.output_type.msgclass
+            if method.server_streaming
+              define_method(name) do |request = nil, deadline: nil, metadata: {}, **fields|
+                request = rpc_request(name, request, fields)
+                @client.server_streaming(path, input, output, request, deadline: deadline, metadata: metadata)
+              end
+            else
+              # `get_item(id: 1)` works as well as `get_item({ id: 1 })`.
+              define_method(name) do |request = nil, deadline: nil, metadata: {}, return_call: false, **fields|
+                request = rpc_request(name, request, fields)
+                @client.unary(path, input, output, request, deadline: deadline, metadata: metadata,
+                                                            return_call: return_call)
+              end
+            end
+          end
+        end
+
+        # @return [Tonic::Rails::Client]
+        attr_reader :client
+
+        def initialize(client)
+          @client = client
+        end
+
+        def inspect
+          "#<#{self.class.inspect} #{client.address}>"
+        end
+
+        private
+
+        def rpc_request(name, request, fields)
+          unless request.nil? || fields.empty?
+            raise ArgumentError, "#{name}: pass the request or its fields, not both (#{fields.keys.join(', ')})"
+          end
+
+          request || fields
+        end
+      end
+
+      # The responses of a server-streaming call.  Nothing is sent until
+      # it's read; then messages are yielded as they arrive, a few at most
+      # buffered ahead of the reader.
+      #
+      #   items.list_items(category: "books").each { |item| ... }
+      #   items.list_items(category: "books").first(10)
+      #
+      # Stopping early (`break`, `first(n)`, an exception) cancels the call.
+      # An error from the server raises Tonic::Rails::RpcError after the
+      # messages that came before it.  A stream can be read once.
+      class ServerStream
+        include Enumerable
+
+        # @return [Hash, nil] response headers, once the stream has ended
+        attr_reader :headers
+
+        # @return [Hash, nil] trailing metadata, once the stream has ended
+        attr_reader :trailers
+
+        # @api private
+        def initialize(path, &read)
+          @path = path
+          @read = read
+        end
+
+        # @yieldparam response [Google::Protobuf::AbstractMessage]
+        # @return [self, Enumerator] an Enumerator without a block
+        def each(&block)
+          return enum_for(:each) unless block
+          raise Error, "the #{@path} stream has already been read" unless @read
+
+          read = @read
+          @read = nil
+          read.call(self, block)
+          self
+        end
+
+        # @api private
+        def finish(headers, trailers)
+          @headers = headers
+          @trailers = trailers
+        end
+
+        def inspect
+          "#<#{self.class.name} #{@path}>"
+        end
+      end
+    end
+  end
+end

data/lib/tonic/rails/cluster.rb

@@ -0,0 +1,222 @@
+# frozen_string_literal: true
+
+require "socket"
+
+module Tonic
+  module Rails
+    # Multi-process mode (`config.processes > 1`).
+    #
+    # The process that calls Server#start becomes a supervisor: it binds the
+    # listening socket, then forks one worker per process.  Workers inherit
+    # the loaded app (copy-on-write) and the socket, and each runs a full
+    # server (its own tokio runtime, created after the fork) accepting on
+    # it.  The supervisor never starts a runtime itself, so nothing
+    # multi-threaded crosses a fork.
+    #
+    # Sharing one socket rather than binding with SO_REUSEPORT per worker
+    # means connections spread across workers on every Unix (macOS doesn't
+    # load-balance SO_REUSEPORT), port 0 works, a bind error surfaces once,
+    # and connections waiting in the backlog survive a worker's death.
+    #
+    # Lifecycle:
+    # - startup fails if a worker exits before it's serving (usually a
+    #   configuration or boot error, which a retry wouldn't fix)
+    # - a worker that dies later is logged and replaced, with backoff if
+    #   replacements keep failing to boot
+    # - stopping (TERM/INT or Server#stop) closes the supervisor's copy of
+    #   the socket and sends TERM to every worker, which flips its health to
+    #   NOT_SERVING and drains; workers still running after `drain_timeout`
+    #   plus a margin are killed
+    # - a worker whose supervisor disappears stops itself
+    class Cluster
+      # Seconds to wait beyond drain_timeout before killing a worker.
+      STOP_MARGIN = 5
+      MAX_BACKOFF = 30
+
+      Worker = Struct.new(:index, :pid, :ready_io, :ready)
+
+      def initialize(server, configuration, stop_signal)
+        @server = server
+        @configuration = configuration
+        @stop_signal = stop_signal
+        @workers = []
+        @failures = Hash.new(0)
+        @respawn_at = {}
+      end
+
+      # Run the workers until stopped.
+      # @yield once every worker is serving
+      def run(install_signal_handlers:)
+        @listener = Socket.for_fd(TonicRails::Native.bind_listener(@configuration.bind_address))
+        @listener.autoclose = true
+        previous_handlers = install_signal_handlers ? trap_signals : {}
+        @configuration.processes.times { |index| @workers << spawn(index) }
+        return unless booted?
+
+        @started = true
+        yield if block_given?
+        supervise
+      ensure
+        @listener&.close
+        stop_workers
+        previous_handlers&.each { |signal, handler| Signal.trap(signal, handler || "DEFAULT") }
+      end
+
+      private
+
+      def logger = @configuration.logger
+
+      def trap_signals
+        server = @server
+        %w[TERM INT].to_h { |signal| [signal, Signal.trap(signal) { server.stop }] }
+      end
+
+      def spawn(index)
+        reader, writer = IO.pipe
+        pid = fork do
+          reader.close
+          @workers.each { |worker| worker&.ready_io&.close }
+          run_worker(index, writer)
+        end
+        writer.close
+        Worker.new(index, pid, reader, false)
+      end
+
+      # In the worker process.
+      def run_worker(index, ready)
+        Process.setproctitle("tonic-rails worker #{index}")
+        supervisor = Process.ppid
+        Thread.new do
+          sleep 1 while Process.ppid == supervisor
+          @server.stop
+        end
+        @server.run_worker(@listener, index) do
+          ready.write("1")
+          ready.close
+        end
+      rescue StandardError => e
+        Tonic::Rails.log_exception(e, "worker #{index}")
+        exit 1
+      end
+
+      # Wait until every worker is serving: true once they are, false if a
+      # stop was requested first.  Raises if a worker exited before serving.
+      def booted?
+        until @workers.all?(&:ready)
+          return false if @stop_signal.pop(timeout: 0.05)
+
+          @workers.each do |worker|
+            next if ready?(worker)
+
+            status = reap(worker)
+            raise Error, "tonic-rails worker #{worker.index} exited during startup (#{describe(status)})" if status
+          end
+        end
+        true
+      end
+
+      def supervise
+        loop do
+          break if @stop_signal.pop(timeout: 0.5)
+
+          @workers.each_index { |index| supervise_worker(index) }
+        end
+      end
+
+      def supervise_worker(index)
+        worker = @workers[index]
+        if worker.nil?
+          @workers[index] = spawn(index) if monotonic >= @respawn_at[index]
+          return
+        end
+
+        ready?(worker)
+        status = reap(worker)
+        return unless status
+
+        replace(worker, status)
+      end
+
+      # A worker died after startup.  Replace it right away if it had been
+      # serving; if replacements keep dying before they serve, back off.
+      def replace(worker, status)
+        index = worker.index
+        @failures[index] = worker.ready ? 0 : @failures[index] + 1
+        delay = @failures[index].zero? ? 0 : [2**(@failures[index] - 1), MAX_BACKOFF].min
+        logger&.error(
+          "tonic-rails: worker #{index} (pid #{worker.pid}) exited unexpectedly (#{describe(status)}); " \
+          "starting a replacement#{" in #{delay}s" if delay.positive?}",
+        )
+        @workers[index] = nil
+        @respawn_at[index] = monotonic + delay
+      end
+
+      def ready?(worker)
+        return true if worker.ready
+
+        case worker.ready_io.read_nonblock(1, exception: false)
+        when String
+          worker.ready = true
+          worker.ready_io.close
+          logger&.info("tonic-rails: worker #{worker.index} (pid #{worker.pid}) serving") if @started
+          true
+        else
+          false
+        end
+      end
+
+      # The worker's exit status if it has exited, otherwise nil.
+      def reap(worker)
+        _, status = Process.waitpid2(worker.pid, Process::WNOHANG)
+        worker.ready_io.close if status
+        status
+      rescue Errno::ECHILD
+        worker.ready_io.close
+        :unknown
+      end
+
+      def stop_workers
+        live = @workers.compact
+        return if live.empty?
+
+        logger&.info(
+          "tonic-rails: shutting down #{live.size} workers, draining for up to #{@configuration.drain_timeout}s",
+        )
+        live.each { |worker| signal(worker, :TERM) }
+        deadline = monotonic + @configuration.drain_timeout.to_f + STOP_MARGIN
+        until (live = live.reject { |worker| reap(worker) }).empty?
+          if monotonic > deadline
+            logger&.warn("tonic-rails: #{live.size} worker(s) still running after the drain timeout; killing them")
+            live.each { |worker| kill(worker) }
+            break
+          end
+          sleep 0.05
+        end
+        @workers = []
+        logger&.info("tonic-rails: stopped")
+      end
+
+      def kill(worker)
+        signal(worker, :KILL)
+        Process.wait(worker.pid)
+      rescue Errno::ECHILD
+        nil
+      end
+
+      def signal(worker, name)
+        Process.kill(name, worker.pid)
+      rescue Errno::ESRCH
+        nil
+      end
+
+      def describe(status)
+        return "status unknown" unless status.is_a?(Process::Status)
+        return "signal #{Signal.signame(status.termsig)}" if status.signaled?
+
+        "exit status #{status.exitstatus}"
+      end
+
+      def monotonic = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    end
+  end
+end

data/lib/tonic/rails/configuration.rb

@@ -26,6 +26,22 @@ module Tonic
       # (RAILS_MAX_THREADS).  Default 4.
       attr_accessor :dispatch_threads
 
+      # Server processes.  Ruby runs one thread at a time per process, so a
+      # CPU-bound service uses one core however many dispatch threads it
+      # has; N processes use N.  Above 1, the app is loaded once and forked
+      # into N workers that share the listening socket, supervised by the
+      # process that started the server.  Unix only.  Default 1.
+      attr_accessor :processes
+
+      # Called in each worker process after it's forked, with the worker's
+      # index (multi-process mode only), to reopen connections that can't
+      # be shared across fork:
+      #
+      #   config.on_worker_boot = ->(index) { Redis.current.disconnect! }
+      #
+      # Active Record reconnects by itself.  Default nil.
+      attr_accessor :on_worker_boot
+
       # RPCs that may wait for a free dispatch thread.  Beyond this the
       # server answers RESOURCE_EXHAUSTED immediately.  Default 128.
       attr_accessor :dispatch_queue_size
@@ -110,6 +126,8 @@ module Tonic
         @logger = Logger.new($stdout)
         @worker_threads = nil
         @dispatch_threads = 4
+        @processes = 1
+        @on_worker_boot = nil
         @dispatch_queue_size = 128
         @tls_cert_path = nil
         @tls_key_path = nil
@@ -132,6 +150,10 @@ module Tonic
         @health_check_interval = 5
       end
 
+      def multi_process?
+        processes.is_a?(Integer) && processes > 1
+      end
+
       def tls?
         !tls_cert_path.nil? && !tls_key_path.nil?
       end
@@ -140,6 +162,7 @@ module Tonic
       def validate!
         positive_integer!(:dispatch_threads)
         positive_integer!(:dispatch_queue_size)
+        positive_integer!(:processes)
         positive_integer!(:max_message_size)
         non_negative_integer!(:max_connections)
         non_negative_integer!(:concurrency_limit_per_connection)
@@ -152,12 +175,20 @@ module Tonic
 
         validate_callables!
         validate_grpc_web!
+        validate_processes!
 
         self
       end
 
       private
 
+      def validate_processes!
+        return unless multi_process?
+        return if Process.respond_to?(:fork)
+
+        raise ArgumentError, "processes > 1 needs fork(), which this platform lacks; run one server process per core instead"
+      end
+
       def validate_grpc_web!
         %i[grpc_web_origins grpc_web_allowed_headers grpc_web_exposed_headers].each do |name|
           value = public_send(name)
@@ -169,6 +200,7 @@ module Tonic
 
       def validate_callables!
         raise ArgumentError, "health_check must respond to #call" if health_check && !health_check.respond_to?(:call)
+        raise ArgumentError, "on_worker_boot must respond to #call" if on_worker_boot && !on_worker_boot.respond_to?(:call)
         unless health_check_interval.is_a?(Numeric) && health_check_interval.positive?
           raise ArgumentError, "health_check_interval must be a positive number, got #{health_check_interval.inspect}"
         end