async-nats 0.1.0

data/lib/async/nats/client.rb

@@ -0,0 +1,332 @@
+# frozen_string_literal: true
+
+# Copyright 2025 Tinco Andringa
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require "nats/io/client"
+require "async"
+require "async/queue"
+require_relative "socket"
+
+module Async
+  module NATS
+    # Async-based NATS client that inherits from NATS::IO::Client
+    # and overrides connection and threading methods to use async operations
+    class Client < ::NATS::IO::Client
+      # Override create_socket to use Async::NATS::Socket
+      def create_socket
+        socket_class = case @uri.scheme
+        when "nats", "tls"
+          Async::NATS::Socket
+        when "ws", "wss"
+          # TODO: Implement async websocket support
+          raise NotImplementedError, "WebSocket support not yet implemented for async-nats"
+        else
+          raise NotImplementedError, "#{@uri.scheme} protocol is not supported"
+        end
+
+        socket_class.new(
+          uri: @uri,
+          tls: {context: tls_context, hostname: @hostname},
+          connect_timeout: ::NATS::IO::DEFAULT_CONNECT_TIMEOUT
+        )
+      end
+
+      # Override establish_connection! to use async operations
+      private def establish_connection!
+        @ruby_pid = Process.pid # For fork detection
+
+        srv = nil
+        begin
+          srv = select_next_server
+
+          # Use the hostname from the server for TLS hostname verification.
+          if client_using_secure_connection? && 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 using async socket
+          @io = create_socket
+          @io.connect
+
+          # Capture state that we have had a TCP connection established against
+          # this server and could potentially be used for reconnecting.
+          srv[:was_connected] = true
+
+          # Connection established and now in process of sending CONNECT to NATS
+          @status = CONNECTING
+
+          # Established TCP connection successfully so can start connect
+          process_connect_init
+
+          # Reset reconnection attempts if connection is valid
+          srv[:reconnect_attempts] = 0
+          srv[:auth_required] ||= true if @server_info[:auth_required]
+
+          # Add back to rotation since successfully connected
+          server_pool << srv
+        rescue ::NATS::IO::NoServersError => e
+          @disconnect_cb.call(e) if @disconnect_cb
+          raise @last_err || e
+        rescue => e
+          # Capture sticky error
+          synchronize do
+            @last_err = e
+            srv[:auth_required] ||= true if @server_info[:auth_required]
+            server_pool << srv if can_reuse_server?(srv)
+          end
+
+          err_cb_call(self, e, nil) if @err_cb
+
+          if should_not_reconnect?
+            @disconnect_cb.call(e) if @disconnect_cb
+            raise e
+          end
+
+          # Clean up any connecting state and close connection without
+          # 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.
+          if @options[:reconnect_time_wait]
+            Async::Task.current.sleep(@options[:reconnect_time_wait])
+          end
+
+          # Continue retrying until there are no options left in the server pool
+          retry
+        end
+
+        # Initialize queues using async-compatible queues
+        @flush_queue = create_async_queue(::NATS::IO::MAX_FLUSH_KICK_SIZE)
+        @pending_queue = create_async_queue(::NATS::IO::MAX_PENDING_SIZE)
+        @pings_outstanding = 0
+        @pongs_received = 0
+        @pending_size = 0
+
+        # Server roundtrip went ok so consider to be connected at this point
+        @status = CONNECTED
+
+        # Connected to NATS so Ready to start parser loop, flusher and ping interval
+        start_async_tasks!
+
+        self
+      end
+
+      # Create an async-compatible queue
+      private def create_async_queue(size)
+        # Use a simple array with mutex for now, can be optimized with Async::Queue
+        queue = SizedQueue.new(size)
+        queue
+      end
+
+      # Override start_threads! to use async tasks instead of threads
+      private def start_async_tasks!
+        # Start the read loop as an async task
+        @read_loop_task = Async::Task.current.async do
+          read_loop
+        end
+
+        # Start the flusher as an async task
+        @flusher_task = Async::Task.current.async do
+          flusher_loop
+        end
+
+        # Start the ping interval as an async task
+        @ping_interval_task = Async::Task.current.async do
+          ping_interval_loop
+        end
+
+        # Use async-based executor instead of Concurrent::ThreadPoolExecutor
+        @subscription_executor = Async::NATS::Executor.new(
+          name: "nats:subscription",
+          max_threads: ::NATS::IO::DEFAULT_TOTAL_SUB_CONCURRENCY,
+          min_threads: defined?(JRUBY_VERSION) ? 2 : 0
+        )
+      end
+
+      # Override read_loop to use async operations
+      private def read_loop
+        loop do
+          should_bail = synchronize do
+            @status == CLOSED or @status == RECONNECTING
+          end
+          if !@io || @io.closed? || should_bail
+            return
+          end
+
+          # Read data from socket with async I/O
+          data = @io.read(::NATS::IO::MAX_SOCKET_READ_BYTES)
+          @parser.parse(data) if data
+        rescue Errno::ETIMEDOUT
+          # Retry on timeout
+          retry
+        rescue => e
+          # In case of reading/parser errors, trigger reconnection logic
+          process_op_error(e)
+        end
+      end
+
+      # Override flusher_loop to use async operations
+      private def flusher_loop
+        loop do
+          # Blocks waiting for the flusher to be kicked...
+          @flush_queue.pop
+
+          should_bail = synchronize do
+            (@status != CONNECTED && !draining?) || @status == CONNECTING
+          end
+          return if should_bail
+
+          # Skip in case nothing remains pending already.
+          next if @pending_queue.empty?
+
+          force_flush!
+
+          synchronize do
+            @pending_size = 0
+          end
+        end
+      end
+
+      # Override ping_interval_loop to use async operations
+      private def ping_interval_loop
+        loop do
+          Async::Task.current.sleep(@options[:ping_interval])
+
+          # Skip ping interval until connected
+          next if !connected?
+
+          if @pings_outstanding >= @options[:max_outstanding_pings]
+            process_op_error(::NATS::IO::StaleConnectionError.new("nats: stale connection"))
+            return
+          end
+
+          @pings_outstanding += 1
+          send_command(PING_REQUEST)
+          @flush_queue << :ping
+        end
+      rescue => e
+        process_op_error(e)
+      end
+
+      # Override drain to use async primitives when in async context
+      def drain
+        return if draining?
+
+        # Check if we're in an async context
+        if Async::Task.current?
+          # Async version - returns a task that can be awaited
+          synchronize do
+            @drain_task ||= Async::Task.current.async do
+              do_async_drain
+            end
+          end
+        else
+          # Fall back to parent thread-based implementation
+          super
+        end
+      end
+
+      # Override close to stop async tasks
+      def close
+        # Stop async tasks if they exist
+        # Use safe stop that works from both async and thread contexts
+        begin
+          @read_loop_task&.stop if @read_loop_task&.alive?
+        rescue => e
+          # Ignore errors during task stop
+        end
+        
+        begin
+          @flusher_task&.stop if @flusher_task&.alive?
+        rescue => e
+          # Ignore errors during task stop
+        end
+        
+        begin
+          @ping_interval_task&.stop if @ping_interval_task&.alive?
+        rescue => e
+          # Ignore errors during task stop
+        end
+
+        # Call parent close
+        super
+      end
+
+      private
+
+      # Async-aware drain implementation using Async::Barrier
+      def do_async_drain
+        synchronize { @status = DRAINING_SUBS }
+
+        # Create a barrier for all subscriptions
+        barrier = Async::Barrier.new
+        subs_to_drain = []
+
+        # Send UNSUB for all subscriptions and add barrier tasks
+        @subs.each do |_, sub|
+          next if sub == @resp_sub
+          
+          drain_sub(sub)
+          subs_to_drain << sub
+          
+          # Add a task to the barrier that waits for this sub to drain
+          barrier.async do
+            wait_for_sub_drained(sub)
+          end
+        end
+        
+        force_flush!
+
+        # Wait for all subscriptions to drain with timeout
+        drain_timeout = @options[:drain_timeout] || 30
+        begin
+          Async::Task.current.with_timeout(drain_timeout) do
+            barrier.wait  # Waits for all tasks in the barrier
+          end
+        rescue Async::TimeoutError
+          barrier.stop  # Stop all waiting tasks
+          e = ::NATS::IO::DrainTimeoutError.new("nats: draining connection timed out")
+          err_cb_call(self, e, nil) if @err_cb
+        ensure
+          barrier.stop rescue nil
+        end
+
+        # Shutdown subscription executor
+        @subscription_executor.shutdown
+        @subscription_executor.wait_for_termination(@options[:drain_timeout] || 30)
+
+        # Set status to DRAINING_PUBS and close
+        synchronize { @status = DRAINING_PUBS }
+        unsubscribe(@resp_sub) if @resp_sub
+        close
+      end
+
+      # Wait for a single subscription to drain (all messages processed)
+      def wait_for_sub_drained(sub)
+        loop do
+          # Check if subscription queue is empty
+          queue_size = sub.synchronize { sub.pending_queue.size }
+          break if queue_size == 0
+          
+          # Yield to other tasks briefly
+          Async::Task.current.sleep(0.01)
+        end
+      end
+    end
+  end
+end

data/lib/async/nats/executor.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+module Async
+  module NATS
+    # Async-based executor using a worker pool pattern with Async::Semaphore
+    # This provides bounded concurrency without creating unlimited nested tasks
+    class Executor
+      def initialize(name:, max_threads:, min_threads: 0)
+        @name = name
+        @max_concurrency = max_threads
+        @running = true
+        @work_queue = Thread::Queue.new  # Thread-safe queue for work items
+        @workers = []
+        @mutex = Mutex.new
+        
+        # Start worker tasks if we're in an async context
+        if Async::Task.current?
+          start_workers
+        end
+      end
+      
+      # Post a block to be executed asynchronously
+      def post(&block)
+        return unless @running
+        
+        # Add work to the queue
+        @work_queue << block
+      end
+      
+      # Shutdown the executor
+      def shutdown
+        @running = false
+        
+        # Signal workers to stop by adding nil items
+        @max_concurrency.times { @work_queue << nil }
+        
+        # Stop all worker tasks
+        @mutex.synchronize do
+          @workers.each do |worker|
+            worker.stop rescue nil
+          end
+          @workers.clear
+        end
+      end
+      
+      # Wait for termination (for compatibility with ThreadPoolExecutor)
+      def wait_for_termination(timeout = nil)
+        deadline = timeout ? Time.now + timeout : nil
+        
+        loop do
+          @mutex.synchronize do
+            return true if @workers.empty? || @workers.all? { |w| !w.alive? }
+          end
+          
+          if deadline && Time.now > deadline
+            return false
+          end
+          
+          sleep(0.01)
+        end
+        
+        true
+      end
+      
+      private
+      
+      def start_workers
+        # Create a pool of worker tasks
+        @max_concurrency.times do |i|
+          worker = Async::Task.current.async do
+            worker_loop
+          end
+          
+          @mutex.synchronize do
+            @workers << worker
+          end
+        end
+      end
+      
+      def worker_loop
+        while @running
+          # Get work from the queue (blocks until work is available)
+          work = @work_queue.pop
+          
+          # nil signals shutdown
+          break if work.nil?
+          
+          begin
+            work.call
+          rescue => e
+            # Log error but don't crash the worker
+            warn "Async::NATS::Executor worker error: #{e.class}: #{e.message}"
+          end
+        end
+      end
+    end
+  end
+end

data/lib/async/nats/socket.rb

@@ -0,0 +1,152 @@
+# frozen_string_literal: true
+
+# Copyright 2025 Tinco Andringa
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require "async"
+require "io/endpoint"
+require "io/endpoint/host_endpoint"
+require "io/stream"
+require "openssl"
+
+module Async
+  module NATS
+    # Async-based socket implementation for NATS connections
+    # This replaces the blocking I/O operations in NATS::IO::Socket with async operations
+    class Socket
+      attr_accessor :socket
+
+      def initialize(options = {})
+        @uri = options[:uri]
+        @connect_timeout = options[:connect_timeout]
+        @write_timeout = options[:write_timeout]
+        @read_timeout = options[:read_timeout]
+        @socket = nil
+        @stream = nil
+        @tls = options[:tls]
+        
+        # Cache async task reference to avoid repeated context checks
+        # This is checked once at initialization to avoid overhead on every I/O operation
+        @async_task = Async::Task.current? rescue nil
+      end
+
+      def connect
+        # Create endpoint for the connection
+        endpoint = IO::Endpoint.tcp(@uri.hostname, @uri.port)
+        
+        # Connect with timeout
+        task = Async::Task.current
+        task.with_timeout(@connect_timeout) do
+          @socket = endpoint.connect
+        end
+
+        # Set TCP no delay by default
+        @socket.setsockopt(::Socket::IPPROTO_TCP, ::Socket::TCP_NODELAY, 1) if @socket.respond_to?(:setsockopt)
+        
+        # Wrap in IO::Stream for buffered I/O
+        @stream = IO::Stream(@socket)
+      rescue Async::TimeoutError
+        raise NATS::IO::SocketTimeoutError
+      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_context = @tls.fetch(:context)
+        
+        # Create SSL socket
+        ssl_socket = OpenSSL::SSL::SSLSocket.new(@socket, tls_context)
+        ssl_socket.sync_close = true
+        ssl_socket.hostname = @tls[:hostname]
+        
+        # Connect with async
+        task = Async::Task.current
+        task.async do
+          ssl_socket.connect
+        end.wait
+        
+        @socket = ssl_socket
+        @stream = IO::Stream(@socket)
+      end
+
+      def read_line(deadline = nil)
+        # deadline is a timeout duration in seconds, not an absolute time
+        # NATS protocol uses \r\n line endings
+        
+        # Use cached task reference (checked once at initialization)
+        if @async_task && deadline
+          @async_task.with_timeout(deadline) do
+            @stream.read_until("\r\n") + "\r\n"
+          end
+        else
+          @stream.read_until("\r\n") + "\r\n"
+        end
+      rescue Async::TimeoutError
+        raise NATS::IO::SocketTimeoutError
+      rescue EOFError
+        raise Errno::ECONNRESET
+      end
+
+      def read(max_bytes, deadline = nil)
+        # deadline is a timeout duration in seconds, not an absolute time
+        
+        # Use cached task reference (checked once at initialization)
+        if @async_task && deadline
+          @async_task.with_timeout(deadline) do
+            @stream.read_partial(max_bytes)
+          end
+        else
+          @stream.read_partial(max_bytes)
+        end
+      rescue Async::TimeoutError
+        raise NATS::IO::SocketTimeoutError
+      rescue EOFError => e
+        if (RUBY_ENGINE == "jruby") && (e.message == "No message available")
+          # FIXME: <EOFError: No message available> can happen in jruby
+          return nil
+        end
+        raise Errno::ECONNRESET
+      end
+
+      def write(data, deadline = nil)
+        # deadline is a timeout duration in seconds, not an absolute time
+        
+        # Use cached task reference (checked once at initialization)
+        if @async_task && deadline
+          @async_task.with_timeout(deadline) do
+            @stream.write(data)
+            @stream.flush
+          end
+        else
+          @stream.write(data)
+          @stream.flush
+        end
+        
+        data.bytesize
+      rescue Async::TimeoutError
+        raise NATS::IO::SocketTimeoutError
+      rescue EOFError
+        raise Errno::ECONNRESET
+      end
+
+      def close
+        @stream&.close
+        @socket&.close
+      end
+
+      def closed?
+        @socket.nil? || @socket.closed?
+      end
+    end
+  end
+end

data/lib/async/nats/version.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module Async
+  module NATS
+    VERSION = "0.1.0"
+  end
+end

data/lib/async/nats.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+require_relative "nats/version"
+require_relative "nats/executor"
+require_relative "nats/socket"
+require_relative "nats/client"
+
+module Async
+  module NATS
+    class Error < StandardError; end
+  end
+end

data/sig/async/nats.rbs

@@ -0,0 +1,6 @@
+module Async
+  module Nats
+    VERSION: String
+    # See the writing guide of rbs: https://github.com/ruby/rbs#guides
+  end
+end