omq-zstd 0.4.0

data/lib/omq/transport/zstd_tcp/connection.rb

@@ -0,0 +1,162 @@
+# frozen_string_literal: true
+
+require "delegate"
+
+module OMQ
+  module Transport
+    module ZstdTcp
+      class ZstdConnection < SimpleDelegator
+        # @return [Integer, nil] wire bytesize of the last received message
+        #
+        attr_reader :last_wire_size_in
+
+
+        def initialize(conn, codec)
+          super(conn)
+          @codec              = codec
+          @dict_shipped       = false
+          @recv_dict          = nil
+          @last_wire_size_in  = nil
+        end
+
+
+        def send_message(parts)
+          compressed = @codec.compress_parts(parts)
+          ship_dict!
+          __getobj__.send_message(compressed)
+        end
+
+
+        def write_message(parts)
+          compressed = @codec.compress_parts(parts)
+          ship_dict!
+          __getobj__.write_message(compressed)
+        end
+
+
+        def write_messages(messages)
+          compressed = messages.map { |parts| @codec.compress_parts(parts) }
+          ship_dict!
+          __getobj__.write_messages(compressed)
+        end
+
+
+        def receive_message
+          loop do
+            parts   = __getobj__.receive_message
+            decoded = decode_parts(parts)
+            if decoded
+              @last_wire_size_in = parts.sum { |p| p.bytesize }
+              return decoded
+            end
+          end
+        end
+
+
+        def respond_to?(name, include_private = false)
+          return false if name == :write_wire
+          super
+        end
+
+
+        private
+
+
+        def ship_dict!
+          return if @dict_shipped
+
+          dict_bytes = @codec.send_dict_bytes
+          return unless dict_bytes
+
+          __getobj__.write_message([dict_bytes])
+          @dict_shipped = true
+        end
+
+
+        def decode_parts(parts)
+          budget    = @codec.max_message_size
+          decoded   = []
+          all_dicts = true
+
+          parts.each do |wire|
+            plaintext = decode_part(wire, budget)
+            if plaintext
+              all_dicts = false
+              budget -= plaintext.bytesize if budget
+              decoded << plaintext
+            end
+          end
+
+          all_dicts ? nil : decoded
+        end
+
+
+        def decode_part(wire, budget)
+          raise ProtocolError, "short frame" if wire.bytesize < 4
+
+          head = wire.byteslice(0, 4)
+
+          case head
+          when Codec::NUL_PREAMBLE
+            plaintext = wire.byteslice(4, wire.bytesize - 4) || "".b
+            enforce_budget!(plaintext.bytesize, budget)
+            plaintext
+          when Codec::ZSTD_MAGIC
+            decode_zstd_frame(wire, budget)
+          when Codec::ZDICT_MAGIC
+            install_recv_dict(wire)
+            nil
+          else
+            raise ProtocolError, "unrecognized preamble: #{head.unpack1('H*')}"
+          end
+        end
+
+
+        def decode_zstd_frame(wire, budget)
+          fcs = @codec.parse_frame_content_size(wire)
+          raise ProtocolError, "Zstd frame missing Frame_Content_Size" if fcs.nil?
+
+          if budget && fcs > budget
+            raise ProtocolError, "declared FCS #{fcs} exceeds limit #{budget}"
+          end
+
+          decompress_opts = budget ? { max_output_size: budget } : {}
+
+          if @recv_dict
+            @recv_dict.decompress(wire, **decompress_opts)
+          else
+            RZstd.decompress(wire, **decompress_opts)
+          end
+        rescue RZstd::DecompressError => e
+          raise ProtocolError, "decompression failed: #{e.message}"
+        rescue RZstd::MissingContentSizeError => e
+          raise ProtocolError, "Zstd frame missing Frame_Content_Size (#{e.message})"
+        rescue RZstd::OutputSizeLimitError => e
+          raise ProtocolError, "declared FCS exceeds limit (#{e.message})"
+        end
+
+
+        def install_recv_dict(wire)
+          if wire.bytesize < 8
+            raise ProtocolError, "dict frame too short"
+          end
+
+          if wire.bytesize > Codec::MAX_DICT_SIZE
+            raise ProtocolError, "dict exceeds #{Codec::MAX_DICT_SIZE} bytes"
+          end
+
+          @recv_dict = RZstd::Dictionary.new(wire.b)
+        end
+
+
+        def enforce_budget!(size, budget)
+          return if budget.nil?
+          return if size <= budget
+
+          raise ProtocolError, "decompressed message size exceeds maximum"
+        end
+
+      end
+    end
+  end
+end

data/lib/omq/transport/zstd_tcp/transport.rb

@@ -0,0 +1,253 @@
+# frozen_string_literal: true
+
+require "socket"
+require "uri"
+require "io/stream"
+
+module OMQ
+  module Transport
+    module ZstdTcp
+      SCHEME = "zstd+tcp"
+
+
+      # WeakKeyMap keyed by Engine — the Codec for a socket lives as long
+      # as the Engine (and therefore the socket) lives.  When the socket is
+      # closed and the Engine is GC'd, the entry disappears automatically.
+      #
+      @codecs = ObjectSpace::WeakKeyMap.new
+
+
+      class << self
+        # Creates a bound zstd+tcp listener.
+        #
+        # @param endpoint [String] e.g. "zstd+tcp://127.0.0.1:5555"
+        # @param engine [Engine]
+        # @param level [Integer] Zstd compression level
+        # @param dict [String, nil] user-supplied dictionary bytes
+        # @return [Listener]
+        #
+        def listener(endpoint, engine, level: -3, dict: nil, **)
+          codec = codec_for(engine, level: level, dict: dict)
+
+          host, port = parse_endpoint(endpoint)
+          host       = normalize_bind_host(host)
+          servers    = ::Socket.tcp_server_sockets(host, port)
+
+          if servers.empty?
+            raise ::Socket::ResolutionError, "no addresses for #{host.inspect}"
+          end
+
+          actual_port  = servers.first.local_address.ip_port
+          display_host = host || "*"
+          host_part    = display_host.include?(":") ? "[#{display_host}]" : display_host
+          resolved     = "#{SCHEME}://#{host_part}:#{actual_port}"
+
+          Listener.new(resolved, servers, actual_port, engine, codec)
+        end
+
+
+        # Creates a zstd+tcp dialer for an endpoint.
+        #
+        # @param endpoint [String] e.g. "zstd+tcp://127.0.0.1:5555"
+        # @param engine [Engine]
+        # @param level [Integer] Zstd compression level
+        # @param dict [String, nil] user-supplied dictionary bytes
+        # @return [Dialer]
+        #
+        def dialer(endpoint, engine, level: -3, dict: nil, **)
+          codec = codec_for(engine, level: level, dict: dict)
+          Dialer.new(endpoint, engine, codec)
+        end
+
+
+        def validate_endpoint!(endpoint)
+          host, _port = parse_endpoint(endpoint)
+          host = normalize_connect_host(host)
+          Addrinfo.getaddrinfo(host, nil, nil, :STREAM) if host
+        end
+
+
+        def parse_endpoint(endpoint)
+          uri = URI.parse(endpoint)
+          [uri.hostname, uri.port]
+        end
+
+
+        def normalize_bind_host(host)
+          return nil if host == "*"
+          host
+        end
+
+
+        def normalize_connect_host(host)
+          host == "*" ? "127.0.0.1" : host
+        end
+
+
+        def connect_timeout(options)
+          ri = options.reconnect_interval
+          ri.is_a?(Range) ? ri.end : [ri * 10, 30].min
+        end
+
+
+        private
+
+
+        # Returns the shared Codec for this engine (socket), creating one
+        # on first call.  Stored in a WeakKeyMap keyed by engine — the
+        # entry is automatically removed when the engine is GC'd.
+        #
+        # Inside a Ractor the module-level @codecs ivar is inaccessible,
+        # so we fall back to creating a fresh Codec per call.  This is
+        # fine: each Ractor worker owns a single socket with one endpoint,
+        # so there is nothing to share.
+        #
+        # @param engine [Engine]
+        # @param level [Integer]
+        # @param dict [String, nil]
+        # @return [Codec]
+        #
+        def codec_for(engine, level:, dict:)
+          @codecs[engine] ||= Codec.new level: level, dict: dict,
+            max_message_size: engine.options.max_message_size
+
+        rescue Ractor::IsolationError
+          Codec.new level: level, dict: dict,
+            max_message_size: engine.options.max_message_size
+        end
+      end
+
+
+      # A zstd+tcp dialer — stateful factory for outgoing connections.
+      #
+      # Holds the Codec (compression cache, training state, dict) and
+      # wraps new connections with {ZstdConnection}.
+      #
+      class Dialer
+        # @return [String] the endpoint this dialer connects to
+        #
+        attr_reader :endpoint
+
+
+        # @param endpoint [String]
+        # @param engine [Engine]
+        # @param codec [Codec]
+        #
+        def initialize(endpoint, engine, codec)
+          @endpoint = endpoint
+          @engine   = engine
+          @codec    = codec
+        end
+
+
+        # Establishes a TCP connection to the endpoint.
+        #
+        # @return [void]
+        #
+        def connect
+          host, port = ZstdTcp.parse_endpoint(@endpoint)
+          host       = ZstdTcp.normalize_connect_host(host)
+          sock       = ::Socket.tcp host, port, connect_timeout: ZstdTcp.connect_timeout(@engine.options)
+
+          TCP.apply_buffer_sizes sock, @engine.options
+
+          @engine.handle_connected IO::Stream::Buffered.wrap(sock), endpoint: @endpoint
+        rescue
+          sock&.close
+          raise
+        end
+
+
+        # Wraps a raw ZMTP connection with Zstd compression.
+        #
+        # @param conn [Protocol::ZMTP::Connection]
+        # @return [ZstdConnection]
+        #
+        def wrap_connection(conn)
+          ZstdConnection.new(conn, @codec)
+        end
+
+      end
+
+
+      # A bound zstd+tcp listener.
+      #
+      class Listener
+        # @return [String] resolved endpoint with actual port
+        #
+        attr_reader :endpoint
+
+
+        # @return [Integer] bound port
+        #
+        attr_reader :port
+
+
+        # @param endpoint [String] resolved endpoint URI
+        # @param servers [Array<Socket>]
+        # @param port [Integer] bound port number
+        # @param engine [Engine]
+        # @param codec [Codec]
+        #
+        def initialize(endpoint, servers, port, engine, codec)
+          @endpoint = endpoint
+          @servers  = servers
+          @port     = port
+          @engine   = engine
+          @codec    = codec
+          @tasks    = []
+        end
+
+
+        # Wraps a raw ZMTP connection with Zstd compression.
+        #
+        # @param conn [Protocol::ZMTP::Connection]
+        # @return [ZstdConnection]
+        #
+        def wrap_connection(conn)
+          ZstdConnection.new(conn, @codec)
+        end
+
+
+        # Spawns accept loop tasks under +parent_task+.
+        #
+        # @param parent_task [Async::Task]
+        # @yieldparam io [IO::Stream::Buffered]
+        #
+        def start_accept_loops(parent_task, &on_accepted)
+          @tasks = @servers.map do |server|
+            parent_task.async(transient: true, annotation: "zstd+tcp accept #{@endpoint}") do
+              loop do
+                client, _addr = server.accept
+
+                Async::Task.current.defer_stop do
+                  TCP.apply_buffer_sizes(client, @engine.options)
+
+                  stream = IO::Stream::Buffered.wrap(client)
+
+                  on_accepted.call stream
+                end
+              end
+            rescue Async::Stop
+            rescue IOError
+            ensure
+              server.close rescue nil
+            end
+          end
+        end
+
+
+        # Stops the listener and closes all server sockets.
+        #
+        # @return [void]
+        #
+        def stop
+          @tasks.each(&:stop)
+          @servers.each { |s| s.close rescue nil }
+        end
+
+      end
+
+    end
+  end
+end

data/lib/omq/transport/zstd_tcp.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+# OMQ Zstd+TCP transport — adds zstd+tcp:// endpoint support.
+#
+# Usage:
+#   require "omq/zstd"
+#
+#   push = OMQ::PUSH.new
+#   push.connect("zstd+tcp://127.0.0.1:5555", level: -3)
+
+require "omq"
+require "rzstd"
+
+require_relative "zstd_tcp/codec"
+require_relative "zstd_tcp/connection"
+require_relative "zstd_tcp/transport"
+
+module OMQ
+  module Transport
+    module ZstdTcp
+      class ProtocolError < StandardError; end
+    end
+  end
+end
+
+OMQ::Engine.transports["zstd+tcp"] = OMQ::Transport::ZstdTcp

data/lib/omq/zstd/version.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module OMQ
+  module Zstd
+    VERSION = "0.4.0"
+  end
+end

data/lib/omq/zstd.rb

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+require_relative "zstd/version"
+require_relative "transport/zstd_tcp"

metadata

@@ -0,0 +1,79 @@
+--- !ruby/object:Gem::Specification
+name: omq-zstd
+version: !ruby/object:Gem::Version
+  version: 0.4.0
+platform: ruby
+authors:
+- Patrik Wenger
+bindir: bin
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: omq
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.23'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.23'
+- !ruby/object:Gem::Dependency
+  name: rzstd
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.2'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.2'
+description: Adds zstd+tcp:// endpoint support to OMQ with per-frame Zstd compression,
+  bounded decompression, in-band dictionary shipping, and sender-side dictionary training.
+email:
+- paddor@gmail.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- DESIGN.md
+- LICENSE
+- README.md
+- RFC.md
+- lib/omq/transport/zstd_tcp.rb
+- lib/omq/transport/zstd_tcp/codec.rb
+- lib/omq/transport/zstd_tcp/connection.rb
+- lib/omq/transport/zstd_tcp/transport.rb
+- lib/omq/zstd.rb
+- lib/omq/zstd/version.rb
+homepage: https://github.com/paddor/omq-zstd
+licenses:
+- ISC
+metadata: {}
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '3.3'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 4.0.6
+specification_version: 4
+summary: Zstd+TCP transport for OMQ
+test_files: []