omq-backend-rust 0.1.7 → 0.2.0

data/lib/omq/rust/engine.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "async"
+require_relative "fd_watcher"
 
 module OMQ
   module Rust
@@ -14,18 +15,18 @@ module OMQ
 
 
       def initialize(socket_type, options)
-        @socket_type    = socket_type
-        @options        = options
-        @peer_connected = Async::Promise.new
-        @all_peers_gone = Async::Promise.new
-        @subscriber_joined = Async::Promise.new
-        @connections    = {}
-        @closed         = false
-        @parent_task    = nil
-        @on_io_thread   = false
-        @materialized   = false
-        @recv_sentinels = 0
-        @compression_options = {}
+        @socket_type          = socket_type
+        @options              = options
+        @peer_connected       = Async::Promise.new
+        @all_peers_gone       = Async::Promise.new
+        @subscriber_joined    = Async::Promise.new
+        @connections          = {}
+        @closed               = false
+        @parent_task          = nil
+        @on_io_thread         = false
+        @materialized         = false
+        @recv_sentinels       = 0
+        @compression_options  = {}
 
         @native = Native::RustSocket.new(socket_type.to_s)
 
@@ -40,6 +41,8 @@ module OMQ
           @parent_task = parent
         elsif Async::Task.current?
           @parent_task = Async::Task.current
+        elsif !Reactor.native_fiber_scheduler?
+          @parent_task = nil
         else
           @parent_task  = Reactor.root_task
           @on_io_thread = true
@@ -81,7 +84,7 @@ module OMQ
         result = @native.enqueue_send(parts)
         return if result == :ok
 
-        @send_signal_r ||= IO.for_fd(@native.send_fd, autoclose: false)
+        @send_signal_r ||= io_for_native_fd(@native.send_fd)
         loop do
           result = @native.enqueue_send(parts)
           return if result == :ok
@@ -130,7 +133,10 @@ module OMQ
         @peer_connected.resolve(nil) unless @peer_connected.resolved?
         @all_peers_gone.resolve(nil) unless @all_peers_gone.resolved?
         @subscriber_joined.resolve(nil) unless @subscriber_joined.resolved?
+        FdWatcher.unwatch_owner(self) unless Reactor.native_fiber_scheduler?
         @native.close
+        close_io_wrapper(@recv_signal_r)
+        close_io_wrapper(@send_signal_r)
       end
 
 
@@ -176,7 +182,7 @@ module OMQ
         Native.send(:io_threads=, OMQ::Rust.io_threads)
         @native.set_options(extract_options)
         @native.materialize
-        @recv_signal_r = IO.for_fd(@native.recv_fd, autoclose: false)
+        @recv_signal_r = io_for_native_fd(@native.recv_fd)
         @materialized  = true
 
         @routing.replay_pending(@native)
@@ -206,21 +212,42 @@ module OMQ
 
 
       def spawn_lifecycle_watcher(fd, promise)
-        io = IO.for_fd(fd, autoclose: false)
-        @parent_task.async(transient: true) do
-          io.wait_readable
-          promise.resolve(true) unless promise.resolved? || @closed
-        rescue IOError, Errno::EBADF
+        if Reactor.native_fiber_scheduler?
+          io = io_for_native_fd(fd)
+          @parent_task.async(transient: true) do
+            io.wait_readable
+            promise.resolve(true) unless promise.resolved? || @closed
+          rescue IOError, Errno::EBADF
+          ensure
+            close_io_wrapper(io)
+          end
+        else
+          FdWatcher.watch_once(fd, owner: self) do
+            promise.resolve(true) unless promise.resolved? || @closed
+          end
         end
       end
 
 
       def start_monitor_forwarder
-        monitor_io = IO.for_fd(@native.monitor_fd, autoclose: false)
-        @parent_task.async(transient: true, annotation: "rust-monitor") do
-          loop do
-            monitor_io.wait_readable
-            monitor_io.read_nonblock(256, exception: false)
+        if Reactor.native_fiber_scheduler?
+          monitor_io = io_for_native_fd(@native.monitor_fd)
+          @parent_task.async(transient: true, annotation: "rust-monitor") do
+            until @closed
+              monitor_io.wait_readable
+              monitor_io.read_nonblock(256, exception: false)
+              while (data = @native.try_recv_monitor)
+                track_connection_event(data)
+                @monitor_queue.enqueue(MonitorEvent.new(**data))
+              end
+            end
+          rescue IOError, Errno::EBADF
+          ensure
+            close_io_wrapper(monitor_io)
+          end
+        else
+          FdWatcher.watch_loop(@native.monitor_fd, owner: self) do |io|
+            io.read_nonblock(256, exception: false)
             while (data = @native.try_recv_monitor)
               track_connection_event(data)
               @monitor_queue.enqueue(MonitorEvent.new(**data))
@@ -230,6 +257,19 @@ module OMQ
       end
 
 
+      def io_for_native_fd(fd)
+        IO.for_fd(fd, autoclose: false)
+      end
+
+
+      def close_io_wrapper(io)
+        return unless io && !io.closed?
+
+        io.close
+      rescue IOError, SystemCallError
+      end
+
+
       def track_connection_event(data)
         detail = data[:detail] || {}
         connection_id = detail[:connection_id]

data/lib/omq/rust/fd_watcher.rb

@@ -0,0 +1,184 @@
+# frozen_string_literal: true
+
+module OMQ
+  module Rust
+    class FdWatcher
+      Watch = Struct.new(:id, :io, :owner, :once, :callback, keyword_init: true)
+
+      class << self
+        def watch_once(fd, owner:, &block)
+          watch(fd, owner: owner, once: true, &block)
+        end
+
+
+        def watch_loop(fd, owner:, &block)
+          watch(fd, owner: owner, once: false, &block)
+        end
+
+
+        def unwatch_owner(owner)
+          removed = nil
+
+          mutex.synchronize do
+            reset_after_fork
+            removed = watchers.values.select { |watch| watch.owner.equal?(owner) }
+            removed.each { |watch| watchers.delete(watch.id) }
+          end
+
+          removed.each { |watch| close_io(watch.io) }
+          wake
+        end
+
+
+        private
+
+
+        def watch(fd, owner:, once:, &block)
+          raise ArgumentError, "block required" unless block
+
+          io = IO.for_fd(fd, autoclose: false)
+          mutex.synchronize do
+            reset_after_fork
+            ensure_started
+            id = next_id
+            watchers[id] = Watch.new(id: id, io: io, owner: owner, once: once, callback: block)
+            id
+          end
+        rescue StandardError
+          close_io(io)
+          raise
+        ensure
+          wake if io
+        end
+
+
+        def reset_after_fork
+          pid = Process.pid
+          return if @pid == pid
+
+          watchers.each_value { |watch| close_io(watch.io) }
+          close_io(@wake_r)
+          close_io(@wake_w)
+          @watchers = {}
+          @wake_r   = nil
+          @wake_w   = nil
+          @thread   = nil
+          @next_id  = 0
+          @pid      = pid
+        end
+
+
+        def ensure_started
+          return if @thread&.alive?
+
+          @wake_r, @wake_w = IO.pipe
+          @thread = Thread.new do
+            Thread.current.name = "omq-rust-watch" if Thread.current.respond_to?(:name=)
+            run
+          end
+        end
+
+
+        def run
+          loop do
+            wake_io, current = snapshot
+            readable = [wake_io, *current.map(&:io)].reject(&:closed?)
+            ready = IO.select(readable)&.first || []
+
+            drain_wake(wake_io) if ready.include?(wake_io)
+
+            current.each do |watch|
+              next unless ready.include?(watch.io)
+
+              dispatch(watch)
+            end
+          rescue IOError, SystemCallError
+            sweep_closed
+          end
+        end
+
+
+        def dispatch(watch)
+          active = watch.once ? delete_if_current(watch) : current?(watch)
+          return unless active
+
+          watch.callback.call(watch.io)
+        rescue StandardError
+          delete_if_current(watch)
+        ensure
+          close_io(watch.io) if watch.once
+        end
+
+
+        def snapshot
+          mutex.synchronize { [@wake_r, watchers.values] }
+        end
+
+
+        def current?(watch)
+          mutex.synchronize { watchers[watch.id].equal?(watch) }
+        end
+
+
+        def delete_if_current(watch)
+          mutex.synchronize do
+            if watchers[watch.id].equal?(watch)
+              watchers.delete(watch.id)
+              true
+            else
+              false
+            end
+          end
+        end
+
+
+        def sweep_closed
+          removed = nil
+          mutex.synchronize do
+            removed = watchers.values.select { |watch| watch.io.closed? }
+            removed.each { |watch| watchers.delete(watch.id) }
+          end
+          removed.each { |watch| close_io(watch.io) }
+        end
+
+
+        def drain_wake(io)
+          loop do
+            result = io.read_nonblock(256, exception: false)
+            break if result == :wait_readable || result.nil? || result.empty?
+          end
+        end
+
+
+        def wake
+          return unless @wake_w && !@wake_w.closed?
+
+          @wake_w.write_nonblock(".", exception: false)
+        rescue IOError, SystemCallError
+        end
+
+
+        def close_io(io)
+          io.close if io && !io.closed?
+        rescue IOError, SystemCallError
+        end
+
+
+        def watchers
+          @watchers ||= {}
+        end
+
+
+        def mutex
+          @mutex ||= Mutex.new
+        end
+
+
+        def next_id
+          @next_id ||= 0
+          @next_id += 1
+        end
+      end
+    end
+  end
+end

data/lib/omq/rust/version.rb

@@ -2,6 +2,6 @@
 
 module OMQ
   module Rust
-    VERSION = "0.1.7"
+    VERSION = "0.2.0"
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: omq-backend-rust
 version: !ruby/object:Gem::Version
-  version: 0.1.7
+  version: 0.2.0
 platform: ruby
 authors:
 - Patrik Wenger
@@ -53,15 +53,18 @@ files:
 - LICENSE
 - README.md
 - ext/omq_backend_rust/Cargo.toml
+- ext/omq_backend_rust/build.rs
 - ext/omq_backend_rust/extconf.rb
 - ext/omq_backend_rust/src/error.rs
 - ext/omq_backend_rust/src/lib.rs
 - ext/omq_backend_rust/src/notify.rs
 - ext/omq_backend_rust/src/options.rs
+- ext/omq_backend_rust/src/rb.rs
 - ext/omq_backend_rust/src/runtime.rs
 - ext/omq_backend_rust/src/socket.rs
 - lib/omq/backend/rust.rb
 - lib/omq/rust/engine.rb
+- lib/omq/rust/fd_watcher.rb
 - lib/omq/rust/version.rb
 homepage: https://github.com/zeromq/omq.rb/tree/main/gems/omq-backend-rust
 licenses: