nats-pure 2.2.1 → 2.3.0

data/lib/nats/io/client.rb

@@ -18,6 +18,7 @@ require_relative 'errors'
 require_relative 'msg'
 require_relative 'subscription'
 require_relative 'jetstream'
+require_relative "rails" if defined?(::Rails::Engine)
 
 require 'nats/nuid'
 require 'thread'
@@ -26,6 +27,7 @@ require 'json'
 require 'monitor'
 require 'uri'
 require 'securerandom'
+require 'concurrent'
 
 begin
   require "openssl"
@@ -81,15 +83,28 @@ module NATS
     DRAINING_PUBS = 6
   end
 
+  # Fork Detection handling
+  # Based from similar approach as mperham/connection_pool: https://github.com/mperham/connection_pool/pull/166
+  if Process.respond_to?(:fork) && Process.respond_to?(:_fork) # MRI 3.1+
+    module ForkTracker
+      def _fork
+        super.tap do |pid|
+          Client.after_fork if pid.zero? # in the child process
+        end
+      end
+    end
+    Process.singleton_class.prepend(ForkTracker)
+  end
+
   # Client creates a connection to the NATS Server.
   class Client
     include MonitorMixin
     include Status
 
-    attr_reader :status, :server_info, :server_pool, :options, :connected_server, :stats, :uri
+    attr_reader :status, :server_info, :server_pool, :options, :connected_server, :stats, :uri, :subscription_executor, :reloader
 
-    DEFAULT_PORT = 4222
-    DEFAULT_URI = ("nats://localhost:#{DEFAULT_PORT}".freeze)
+    DEFAULT_PORT = { nats: 4222, ws: 80, wss: 443 }.freeze
+    DEFAULT_URI = ("nats://localhost:#{DEFAULT_PORT[:nats]}".freeze)
 
     CR_LF = ("\r\n".freeze)
     CR_LF_SIZE = (CR_LF.bytesize)
@@ -106,9 +121,39 @@ module NATS
     SUB_OP = ('SUB'.freeze)
     EMPTY_MSG = (''.freeze)
 
-    def initialize
-      super # required to initialize monitor
-      @options = nil
+    INSTANCES = ObjectSpace::WeakMap.new # tracks all alive client instances
+    private_constant :INSTANCES
+
+    class << self
+      # Reloader should free resources managed by external framework
+      # that were implicitly acquired in subscription callbacks.
+      attr_writer :default_reloader
+
+      def default_reloader
+        @default_reloader ||= proc { |&block| block.call }.tap { |r| Ractor.make_shareable(r) if defined? Ractor }
+      end
+
+      # Re-establish connection in a new process after forking to start new threads.
+      def after_fork
+        INSTANCES.each do |client|
+          if client.options[:reconnect]
+            was_connected = !client.disconnected?
+            client.send(:close_connection, Status::DISCONNECTED, true)
+            client.connect if was_connected
+          else
+            client.send(:err_cb_call, self, NATS::IO::ForkDetectedError, nil)
+            client.close
+          end
+        rescue => e
+          warn "nats: Error during handling after_fork callback: #{e}" # TODO: Report as async error via error callback?
+        end
+      end
+    end
+
+    def initialize(uri = nil, opts = {})
+      super() # required to initialize monitor
+      @initial_uri = uri
+      @initial_options = opts
 
       # Read/Write IO
       @io = nil
@@ -132,7 +177,7 @@ module NATS
       @uri = nil
       @server_pool = []
 
-      @status = DISCONNECTED
+      @status = nil
 
       # Subscriptions
       @subs = { }
@@ -194,29 +239,63 @@ module NATS
 
       # Draining
       @drain_t = nil
+
+      # Prepare for calling connect or automatic delayed connection
+      parse_and_validate_options if uri || opts.any?
+
+      # Keep track of all client instances to handle them after process forking in Ruby 3.1+
+      INSTANCES[self] = self if !defined?(Ractor) || Ractor.current == Ractor.main # Ractors doesn't work in forked processes
+
+      @reloader = opts.fetch(:reloader, self.class.default_reloader)
     end
 
-    # Establishes a connection to NATS.
+    # Prepare connecting to NATS, but postpone real connection until first usage.
     def connect(uri=nil, opts={})
+      if uri || opts.any?
+        @initial_uri = uri
+        @initial_options = opts
+      end
+
       synchronize do
         # In case it has been connected already, then do not need to call this again.
         return if @connect_called
         @connect_called = true
       end
 
+      parse_and_validate_options
+      establish_connection!
+
+      self
+    end
+
+    private def parse_and_validate_options
+      # Reset these in case we have reconnected via fork.
+      @server_pool = []
+      @resp_sub = nil
+      @resp_map = nil
+      @resp_sub_prefix = nil
+      @nuid = NATS::NUID.new
+      @stats = {
+        in_msgs: 0,
+        out_msgs: 0,
+        in_bytes: 0,
+        out_bytes: 0,
+        reconnects: 0
+      }
+      @status = DISCONNECTED
+
       # Convert URI to string if needed.
+      uri = @initial_uri.dup
       uri = uri.to_s if uri.is_a?(URI)
 
+      opts = @initial_options.dup
+
       case uri
       when String
         # Initialize TLS defaults in case any url is using it.
         srvs = opts[:servers] = process_uri(uri)
-        if srvs.any? {|u| u.scheme == 'tls'} and !opts[:tls]
-          tls_context = OpenSSL::SSL::SSLContext.new
-          tls_context.set_params
-          opts[:tls] = {
-            context: tls_context
-          }
+        if srvs.any? {|u| %w[tls wss].include? u.scheme } and !opts[:tls]
+          opts[:tls] = { context: tls_context }
         end
         @single_url_connect_used = true if srvs.size == 1
       when Hash
@@ -287,11 +366,25 @@ module NATS
 
       validate_settings!
 
+      self
+    end
+
+    private def establish_connection!
+      @ruby_pid = Process.pid # For fork detection
+
       srv = nil
       begin
         srv = select_next_server
 
-        # Create TCP socket connection to NATS
+        # Use the hostname from the server for TLS hostname verification.
+        if client_using_secure_connection? and single_url_connect_used?
+          # Always reuse the original hostname used to connect.
+          @hostname ||= srv[:hostname]
+        else
+          @hostname = srv[:hostname]
+        end
+
+        # Create TCP socket connection to NATS.
         @io = create_socket
         @io.connect
 
@@ -302,14 +395,6 @@ module NATS
         # Connection established and now in process of sending CONNECT to NATS
         @status = CONNECTING
 
-        # Use the hostname from the server for TLS hostname verification.
-        if client_using_secure_connection? and single_url_connect_used?
-          # Always reuse the original hostname used to connect.
-          @hostname ||= srv[:hostname]
-        else
-          @hostname = srv[:hostname]
-        end
-
         # Established TCP connection successfully so can start connect
         process_connect_init
 
@@ -341,8 +426,8 @@ module NATS
         # triggering the disconnection/closed callbacks.
         close_connection(DISCONNECTED, false)
 
-        # always sleep here to safe guard against errors before current[:was_connected]
-        # is set for the first time
+        # Always sleep here to safe guard against errors before current[:was_connected]
+        # is set for the first time.
         sleep @options[:reconnect_time_wait] if @options[:reconnect_time_wait]
 
         # Continue retrying until there are no options left in the server pool
@@ -434,6 +519,7 @@ module NATS
       sub.pending_msgs_limit  = opts[:pending_msgs_limit]
       sub.pending_bytes_limit = opts[:pending_bytes_limit]
       sub.pending_queue = SizedQueue.new(sub.pending_msgs_limit)
+      sub.processing_concurrency = opts[:processing_concurrency] if opts.key?(:processing_concurrency)
 
       send_command("SUB #{subject} #{opts[:queue]} #{sid}#{CR_LF}")
       @flush_queue << :sub
@@ -446,41 +532,6 @@ module NATS
         sub.wait_for_msgs_cond = cond
       end
 
-      # Async subscriptions each own a single thread for the
-      # delivery of messages.
-      # FIXME: Support shared thread pool with configurable limits
-      # to better support case of having a lot of subscriptions.
-      sub.wait_for_msgs_t = Thread.new do
-        loop do
-          msg = sub.pending_queue.pop
-
-          cb = nil
-          sub.synchronize do
-
-            # Decrease pending size since consumed already
-            sub.pending_size -= msg.data.size
-            cb = sub.callback
-          end
-
-          begin
-            # Note: Keep some of the alternative arity versions to slightly
-            # improve backwards compatibility.  Eventually fine to deprecate
-            # since recommended version would be arity of 1 to get a NATS::Msg.
-            case cb.arity
-            when 0 then cb.call
-            when 1 then cb.call(msg)
-            when 2 then cb.call(msg.data, msg.reply)
-            when 3 then cb.call(msg.data, msg.reply, msg.subject)
-            else cb.call(msg.data, msg.reply, msg.subject, msg.header)
-            end
-          rescue => e
-            synchronize do
-              err_cb_call(self, e, sub) if @err_cb
-            end
-          end
-        end
-      end if callback
-
       sub
     end
 
@@ -709,6 +760,10 @@ module NATS
       connected? ? @uri : nil
     end
 
+    def disconnected?
+      !@status or @status == DISCONNECTED
+    end
+
     def connected?
       @status == CONNECTED
     end
@@ -812,7 +867,8 @@ module NATS
         if !@options[:ignore_discovered_urls] && connect_urls
           srvs = []
           connect_urls.each do |url|
-            scheme = client_using_secure_connection? ? "tls" : "nats"
+            # Use the same scheme as the currently in use URI.
+            scheme = @uri.scheme
             u = URI.parse("#{scheme}://#{url}")
 
             # Skip in case it is the current server which we already know
@@ -971,12 +1027,8 @@ module NATS
             # Only dispatch message when sure that it would not block
             # the main read loop from the parser.
             msg = Msg.new(subject: subject, reply: reply, data: data, header: hdr, nc: self, sub: sub)
-            sub.pending_queue << msg
 
-            # For sync subscribers, signal that there is a new message.
-            sub.wait_for_msgs_cond.signal if sub.wait_for_msgs_cond
-
-            sub.pending_size += data.size
+            sub.dispatch(msg)
           end
         end
       end
@@ -1016,11 +1068,32 @@ module NATS
       @uri.scheme == "tls" || @tls
     end
 
+    def tls_context
+      return nil unless @tls
+
+      # Allow prepared context and customizations via :tls opts
+      return @tls[:context] if @tls[:context]
+
+      @tls_context ||= OpenSSL::SSL::SSLContext.new.tap do |tls_context|
+        # Use the default verification options from Ruby:
+        # https://github.com/ruby/ruby/blob/96db72ce38b27799dd8e80ca00696e41234db6ba/ext/openssl/lib/openssl/ssl.rb#L19-L29
+        #
+        # Insecure TLS versions not supported already:
+        # https://github.com/ruby/openssl/commit/3e5a009966bd7f806f7180d82cf830a04be28986
+        #
+        tls_context.set_params
+      end
+    end
+
     def single_url_connect_used?
       @single_url_connect_used
     end
 
     def send_command(command)
+      raise NATS::IO::ConnectionClosedError if closed?
+
+      establish_connection! if !status || (disconnected? && should_reconnect?)
+
       @pending_size += command.bytesize
       @pending_queue << command
 
@@ -1047,12 +1120,6 @@ module NATS
       synchronize do
         sub.max = opt_max
         @subs.delete(sid) unless (sub.max && (sub.received < sub.max))
-
-        # Stop messages delivery thread for async subscribers
-        if sub.wait_for_msgs_t && sub.wait_for_msgs_t.alive?
-          sub.wait_for_msgs_t.exit
-          sub.pending_queue.clear
-        end
       end
 
       sub.synchronize do
@@ -1107,11 +1174,6 @@ module NATS
 
         to_delete.each do |sub|
           @subs.delete(sub.sid)
-          # Stop messages delivery thread for async subscribers
-          if sub.wait_for_msgs_t && sub.wait_for_msgs_t.alive?
-            sub.wait_for_msgs_t.exit
-            sub.pending_queue.clear
-          end
         end
         to_delete.clear
 
@@ -1126,6 +1188,9 @@ module NATS
         end
       end
 
+      subscription_executor.shutdown
+      subscription_executor.wait_for_termination(@options[:drain_timeout])
+
       if MonotonicTime::now > drain_timeout
         e = NATS::IO::DrainTimeoutError.new("nats: draining connection timed out")
         err_cb_call(self, e, nil) if @err_cb
@@ -1287,7 +1352,7 @@ module NATS
         @flush_queue.pop
 
         should_bail = synchronize do
-          @status != CONNECTED || @status == CONNECTING
+          (@status != CONNECTED && !draining? ) || @status == CONNECTING
         end
         return if should_bail
 
@@ -1342,6 +1407,7 @@ module NATS
     end
 
     def process_connect_init
+      # FIXME: Can receive PING as well here in recent versions.
       line = @io.read_line(options[:connect_timeout])
       if !line or line.empty?
         raise NATS::IO::ConnectError.new("nats: protocol exception, INFO not received")
@@ -1356,39 +1422,13 @@ module NATS
 
       case
       when (server_using_secure_connection? and client_using_secure_connection?)
-        tls_context = nil
-
-        if @tls
-          # Allow prepared context and customizations via :tls opts
-          tls_context = @tls[:context] if @tls[:context]
-        else
-          # Defaults
-          tls_context = OpenSSL::SSL::SSLContext.new
-
-          # Use the default verification options from Ruby:
-          # https://github.com/ruby/ruby/blob/96db72ce38b27799dd8e80ca00696e41234db6ba/ext/openssl/lib/openssl/ssl.rb#L19-L29
-          #
-          # Insecure TLS versions not supported already:
-          # https://github.com/ruby/openssl/commit/3e5a009966bd7f806f7180d82cf830a04be28986
-          #
-          tls_context.set_params
-        end
-
-        # Setup TLS connection by rewrapping the socket
-        tls_socket = OpenSSL::SSL::SSLSocket.new(@io.socket, tls_context)
-
-        # Close TCP socket after closing TLS socket as well.
-        tls_socket.sync_close = true
-
-        # Required to enable hostname verification if Ruby runtime supports it (>= 2.4):
-        # https://github.com/ruby/openssl/commit/028e495734e9e6aa5dba1a2e130b08f66cf31a21
-        tls_socket.hostname = @hostname
-
-        tls_socket.connect
-        @io.socket = tls_socket
-      when (server_using_secure_connection? and !client_using_secure_connection?)
+        @io.setup_tls!
+      # Server > v2.9.19 returns tls_required regardless of no_tls for WebSocket config being used so need to check URI.
+      when (server_using_secure_connection? and !client_using_secure_connection? and @uri.scheme != "ws")
         raise NATS::IO::ConnectError.new('TLS/SSL required by server')
-      when (client_using_secure_connection? and !server_using_secure_connection?)
+      # Server < v2.9.19 requiring TLS/SSL over websocket but not requiring it over standard protocol
+      # doesn't send `tls_required` in its INFO so we need to check the URI scheme for WebSocket.
+      when (client_using_secure_connection? and !server_using_secure_connection? and @uri.scheme != "wss")
         raise NATS::IO::ConnectError.new('TLS/SSL not supported by server')
       else
         # Otherwise, use a regular connection.
@@ -1433,11 +1473,6 @@ module NATS
       begin
         srv = select_next_server
 
-        # Establish TCP connection with new server
-        @io = create_socket
-        @io.connect
-        @stats[:reconnects] += 1
-
         # Set hostname to use for TLS hostname verification
         if client_using_secure_connection? and single_url_connect_used?
           # Reuse original hostname name in case of using TLS.
@@ -1446,6 +1481,11 @@ module NATS
           @hostname = srv[:hostname]
         end
 
+        # Establish TCP connection with new server
+        @io = create_socket
+        @io.connect
+        @stats[:reconnects] += 1
+
         # Established TCP connection successfully so can start connect
         process_connect_init
 
@@ -1505,6 +1545,7 @@ module NATS
 
     def close_connection(conn_status, do_cbs=true)
       synchronize do
+        @connect_called = false
         if @status == CLOSED
           @status = conn_status
           return
@@ -1553,12 +1594,6 @@ module NATS
         end if should_flush
 
         # Destroy any remaining subscriptions.
-        @subs.each do |_, sub|
-          if sub.wait_for_msgs_t && sub.wait_for_msgs_t.alive?
-            sub.wait_for_msgs_t.exit
-            sub.pending_queue.clear
-          end
-        end
         @subs.clear
 
         if do_cbs
@@ -1580,15 +1615,24 @@ module NATS
     def start_threads!
       # Reading loop for gathering data
       @read_loop_thread = Thread.new { read_loop }
+      @read_loop_thread.name = "nats:read_loop"
       @read_loop_thread.abort_on_exception = true
 
       # Flusher loop for sending commands
       @flusher_thread = Thread.new { flusher_loop }
+      @flusher_thread.name = "nats:flusher_loop"
       @flusher_thread.abort_on_exception = true
 
       # Ping interval handling for keeping alive the connection
       @ping_interval_thread = Thread.new { ping_interval_loop }
+      @ping_interval_thread.name = "nats:ping_loop"
       @ping_interval_thread.abort_on_exception = true
+
+      # Subscription handling thread pool
+      @subscription_executor = Concurrent::ThreadPoolExecutor.new(
+        name: 'nats:subscription', # threads will be given names like nats:subscription-worker-1
+        max_threads: NATS::IO::DEFAULT_TOTAL_SUB_CONCURRENCY,
+      )
     end
 
     # Prepares requests subscription that handles the responses
@@ -1606,29 +1650,25 @@ module NATS
       @resp_sub.pending_msgs_limit = NATS::IO::DEFAULT_SUB_PENDING_MSGS_LIMIT
       @resp_sub.pending_bytes_limit = NATS::IO::DEFAULT_SUB_PENDING_BYTES_LIMIT
       @resp_sub.pending_queue = SizedQueue.new(@resp_sub.pending_msgs_limit)
-      @resp_sub.wait_for_msgs_t = Thread.new do
-        loop do
-          msg = @resp_sub.pending_queue.pop
-          @resp_sub.pending_size -= msg.data.size
-
-          # Pick the token and signal the request under the mutex
-          # from the subscription itself.
-          token = msg.subject.split('.').last
-          future = nil
-          synchronize do
-            future = @resp_map[token][:future]
-            @resp_map[token][:response] = msg
-          end
+      @resp_sub.callback = proc do |msg|
+        # Pick the token and signal the request under the mutex
+        # from the subscription itself.
+        token = msg.subject.split('.').last
+        future = nil
+        synchronize do
+          future = @resp_map[token][:future]
+          @resp_map[token][:response] = msg
+        end
 
-          # Signal back that the response has arrived
-          # in case the future has not been yet delete.
-          @resp_sub.synchronize do
-            future.signal if future
-          end
+        # Signal back that the response has arrived
+        # in case the future has not been yet delete.
+        @resp_sub.synchronize do
+          future.signal if future
         end
       end
 
       sid = (@ssid += 1)
+      @resp_sub.sid = sid
       @subs[sid] = @resp_sub
       send_command("SUB #{@resp_sub.subject} #{sid}#{CR_LF}")
       @flush_queue << :sub
@@ -1661,10 +1701,21 @@ module NATS
     end
 
     def create_socket
-      NATS::IO::Socket.new({
-          uri: @uri,
-          connect_timeout: NATS::IO::DEFAULT_CONNECT_TIMEOUT
-        })
+      socket_class = case @uri.scheme
+        when "nats", "tls"
+          NATS::IO::Socket
+        when "ws", "wss"
+          require_relative 'websocket'
+          NATS::IO::WebSocket
+        else
+          raise NotImplementedError, "#{@uri.scheme} protocol is not supported, check NATS cluster URL spelling"
+        end
+
+      socket_class.new(
+        uri: @uri,
+        tls: { context: tls_context, hostname: @hostname },
+        connect_timeout: NATS::IO::DEFAULT_CONNECT_TIMEOUT,
+      )
     end
 
     def setup_nkeys_connect
@@ -1774,7 +1825,7 @@ module NATS
 
         # Host and Port
         uri_object.hostname ||= "localhost"
-        uri_object.port ||= DEFAULT_PORT
+        uri_object.port ||= DEFAULT_PORT.fetch(uri.scheme.to_sym, DEFAULT_PORT[:nats])
 
         uri_object
       end
@@ -1813,6 +1864,9 @@ module NATS
     DEFAULT_SUB_PENDING_MSGS_LIMIT  = 65536
     DEFAULT_SUB_PENDING_BYTES_LIMIT = 65536 * 1024
 
+    DEFAULT_TOTAL_SUB_CONCURRENCY = 24
+    DEFAULT_SINGLE_SUB_CONCURRENCY = 1
+
     # Implementation adapted from https://github.com/redis/redis-rb
     class Socket
       attr_accessor :socket
@@ -1823,6 +1877,7 @@ module NATS
         @write_timeout = options[:write_timeout]
         @read_timeout = options[:read_timeout]
         @socket = nil
+        @tls = options[:tls]
       end
 
       def connect
@@ -1841,6 +1896,22 @@ module NATS
         @socket.setsockopt(::Socket::IPPROTO_TCP, ::Socket::TCP_NODELAY, 1)
       end
 
+      # (Re-)connect using secure connection if server and client agreed on using it.
+      def setup_tls!
+        # Setup TLS connection by rewrapping the socket
+        tls_socket = OpenSSL::SSL::SSLSocket.new(@socket, @tls.fetch(:context))
+
+        # Close TCP socket after closing TLS socket as well.
+        tls_socket.sync_close = true
+
+        # Required to enable hostname verification if Ruby runtime supports it (>= 2.4):
+        # https://github.com/ruby/openssl/commit/028e495734e9e6aa5dba1a2e130b08f66cf31a21
+        tls_socket.hostname = @tls[:hostname]
+
+        tls_socket.connect
+        @socket = tls_socket
+      end
+
       def read_line(deadline=nil)
         # FIXME: Should accumulate and read in a non blocking way instead
         unless ::IO.select([@socket], nil, nil, deadline)

data/lib/nats/io/errors.rb

@@ -58,6 +58,12 @@ module NATS
 
     # When drain takes too long to complete.
     class DrainTimeoutError < Error; end
+
+    # When a fork is detected, but the client is not configured to re-connect automatically.
+    class ForkDetectedError < Error; end
+
+    # When tried to send command after connection has been closed.
+    class ConnectionClosedError < Error; end
   end
 
   # Timeout is raised when the client gives up waiting for a response from a service.

data/lib/nats/io/jetstream/msg/ack_methods.rb

@@ -41,11 +41,15 @@ module NATS
 
         def nak(**params)
           ensure_is_acked_once!
-
+          payload = if params[:delay]
+                      payload = "#{Ack::Nak} #{{ delay: params[:delay] }.to_json}"
+                    else
+                      Ack::Nak
+                    end
           resp = if params[:timeout]
-                   @nc.request(@reply, Ack::Nak, **params)
+                   @nc.request(@reply, payload, **params)
                  else
-                   @nc.publish(@reply, Ack::Nak)
+                   @nc.publish(@reply, payload)
                  end
           @sub.synchronize { @ackd = true }
 
@@ -104,4 +108,4 @@ module NATS
       end
     end
   end
-end
+end

data/lib/nats/io/msg.rb

@@ -50,7 +50,9 @@ module NATS
 
     def inspect
       hdr = ", header=#{@header}" if @header
-      "#<NATS::Msg(subject: \"#{@subject}\", reply: \"#{@reply}\", data: #{@data.slice(0, 10).inspect}#{hdr})>"
+      dot = '...' if @data.length > 10
+      dat = "#{data.slice(0, 10)}#{dot}"
+      "#<NATS::Msg(subject: \"#{@subject}\", reply: \"#{@reply}\", data: #{dat.inspect}#{hdr})>"
     end
   end
 end

data/lib/nats/io/rails.rb

@@ -0,0 +1,29 @@
+require "rails"
+
+module NATS
+  class Rails < ::Rails::Engine
+    # This class is used to free resources managed by Rails (e.g. database connections)
+    # that were implicitly acquired in subscription callbacks
+    # Implementation is based on https://github.com/sidekiq/sidekiq/blob/5e1a77a6d03193dd977fbfe8961ab78df91bb392/lib/sidekiq/rails.rb
+    class Reloader
+      def initialize(app = ::Rails.application)
+        @app = app
+      end
+
+      def call
+        params = (::Rails::VERSION::STRING >= "7.1") ? {source: "gem.nats"} : {}
+        @app.reloader.wrap(**params) do
+          yield
+        end
+      end
+
+      def inspect
+        "#<NATS::Rails::Reloader @app=#{@app.class.name}>"
+      end
+    end
+
+    config.after_initialize do
+      NATS::Client.default_reloader = NATS::Rails::Reloader.new
+    end
+  end
+end