emb 0.3.0 → 0.4.0.pre3

data/lib/emb/connection_router.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+require_relative 'round_robin_pool'
+require 'redis_client'
+
+module Emb
+  # Owns a client's instance fan-out: one RoundRobinPool per configured url,
+  # instance-level round-robin on top of each pool's connection rotation, and
+  # the pre-send retry across instances.
+  #
+  # Retry is deliberately limited to PRESEND_CONNECTION_ERROR: a connection
+  # that was never established means the command was not written, so moving it
+  # to another instance is safe. Errors after a command may have been sent
+  # (timeouts, mid-flight connection loss) are never re-dispatched — retrying
+  # them would duplicate inference on the server.
+  class ConnectionRouter
+    # Captured at load so `rescue` keeps working when tests or apps replace the
+    # RedisClient constant (stub_const) — a runtime constant lookup in the
+    # rescue clause would otherwise resolve against the replacement.
+    PRESEND_CONNECTION_ERROR = RedisClient::CannotConnectError
+
+    attr_reader :pools
+
+    def initialize(size, urls, redis_options)
+      @pools = urls.map do |url|
+        RoundRobinPool.new(size) do
+          RedisClient.new(url: url, **redis_options)
+        end
+      end
+      @next_instance = 0
+      @instance_mutex = Mutex.new
+    end
+
+    def call(*args)
+      attempts = @pools.size
+      idx = pick_instance
+      last_error = nil
+      attempts.times do
+        return perform_command(@pools[idx], args)
+      rescue PRESEND_CONNECTION_ERROR => e
+        last_error = e
+        idx = (idx + 1) % @pools.size if @pools.size > 1
+      end
+      raise last_error
+    end
+
+    private
+
+    def pick_instance
+      return 0 if @pools.size == 1
+
+      @instance_mutex.synchronize do
+        idx = @next_instance % @pools.size
+        @next_instance += 1
+        idx
+      end
+    end
+
+    def perform_command(pool, args)
+      debug = Emb.debug?
+      started = Process.clock_gettime(Process::CLOCK_MONOTONIC) if debug
+      result = pool.with { |r| r.call(*args) }
+      log_command(args, started) if debug
+      result
+    end
+
+    def log_command(args, started)
+      elapsed = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000
+      $stdout.puts "[EMB] #{args.map(&:inspect).join(' ')} (#{format('%.2f', elapsed)}ms)"
+    end
+  end
+end

data/lib/emb/errors.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Emb
+  # Raised when an EMB.MULTI batch fails — after redis-client exhausted its
+  # configured re-sends (`reconnect_attempts`) on a transient error, or on the
+  # first attempt for an operation error (server error reply) or the default
+  # `reconnect_attempts: 0` configuration. The underlying redis error is
+  # preserved as `cause`; `attempts` counts the wire sends that were made.
+  class ServerError < StandardError
+    attr_reader :attempts
+
+    def initialize(message, attempts: 1)
+      super(message)
+      @attempts = attempts
+    end
+  end
+end

data/lib/emb/job_middleware.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+require 'batch_loader'
+require_relative 'batch_scope'
+
+module Emb
+  # Job middleware: clears the per-thread batch scope after each job execution,
+  # even when the job raises. Registered as a Sidekiq/Shoryuken server
+  # middleware (positional args differ per framework and are unused) and used by
+  # the ActiveJob perform callback.
+  class JobMiddleware
+    def call(*_args, &)
+      BatchScope.wrap(&)
+    end
+  end
+end

data/lib/emb/middleware.rb

@@ -1,20 +1,18 @@
 # frozen_string_literal: true
 
 require 'batch_loader'
+require_relative 'batch_scope'
 
 module Emb
-  # Rack middleware that clears the per-thread batch scope at the end of each
-  # request, bounding cache growth in long-lived request threads. The scope is
-  # cleared even when the app raises; a fresh scope starts with the next request.
+  # Rack middleware: clears the per-thread batch scope after each request,
+  # even when the app raises.
   class Middleware
     def initialize(app)
       @app = app
     end
 
     def call(env)
-      @app.call(env)
-    ensure
-      BatchLoader::Executor.clear_current
+      BatchScope.wrap { @app.call(env) }
     end
   end
 end

data/lib/emb/proxy.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative 'values_reply'
+
 module Emb
   class Proxy
     attr_reader :name
@@ -9,8 +11,20 @@ module Emb
       @name = name
     end
 
-    def [](text, *texts)
-      return Emb.build_batch_loader(@client, @name, texts.empty? ? text : [text, *texts]) if @client.batch?
+    # Queries the model for one or more texts. format: :binary (default)
+    # returns the float32-unpacked vectors from the compact BLOB wire;
+    # format: :values sends the VALUES keyword and returns the decoded
+    # envelope — a Hash with dtype:/shape:/values: (Flats, flat for a single
+    # text, grouped per text when several are requested).
+    def [](text, *texts, format: :binary)
+      format = normalize_format(format)
+      if @client.lazy?
+        return Emb.build_batch_loader(@client, @name, texts.empty? ? text : [text, *texts], format: format)
+      end
+
+      if format == :values
+        return values_query(text, texts)
+      end
 
       set = Array(@client.send_command('EMB', @name.to_s, text, *texts))
       result = set.map { |entry| entry.unpack('e*') }
@@ -23,5 +37,23 @@ module Emb
     def inspect
       "#<Emb::Proxy #{@name}>"
     end
+
+    private
+
+    def normalize_format(format)
+      unless %i[binary values].include?(format)
+        raise ArgumentError, "unknown format #{format.inspect} (expected :binary or :values)"
+      end
+
+      format
+    end
+
+    def values_query(text, texts)
+      reply = @client.send_command('EMB', @name.to_s, 'VALUES', text, *texts)
+      envelope = Emb::ValuesReply.parse(reply)
+      return envelope if texts.empty?
+
+      envelope.merge(values: Emb::ValuesReply.rows(envelope[:shape], envelope[:values]))
+    end
   end
 end

data/lib/emb/railtie.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+if defined?(Rails::Railtie)
+  require_relative '../emb'
+
+  module Emb
+    # Rails integration, loaded by emb.rb when Rails is present.
+    #
+    #   config.emb.middleware = false      # skip Emb::Middleware in the stack
+    #   config.emb.job_middleware = false  # skip job-scope protection
+    class Railtie < Rails::Railtie
+      config.emb = ActiveSupport::OrderedOptions.new
+      config.emb.middleware = true
+      config.emb.job_middleware = true
+
+      initializer 'emb.middleware' do |app|
+        next if config.emb.middleware == false
+        next if app.middleware.include?(Emb::Middleware)
+
+        app.middleware.use Emb::Middleware
+      end
+
+      config.after_initialize do
+        next if config.emb.job_middleware == false
+
+        ActiveSupport.on_load(:active_job) do |base|
+          base.around_perform { |_job, block| Emb::BatchScope.wrap { block.call } }
+        end
+
+        if defined?(Sidekiq)
+          Sidekiq.configure_server do |sidekiq_config|
+            sidekiq_config.server_middleware { |chain| chain.add Emb::JobMiddleware }
+          end
+        end
+
+        if defined?(Sidekiq::Testing)
+          # Testing modes use their own middleware chain.
+          Sidekiq::Testing.server_middleware { |chain| chain.add Emb::JobMiddleware }
+        end
+
+        if defined?(Shoryuken)
+          Shoryuken.configure_server do |shoryuken_config|
+            shoryuken_config.server_middleware { |chain| chain.add Emb::JobMiddleware }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/emb/round_robin_pool.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+module Emb
+  # A thread-safe pool of N RedisClient connections with round-robin selection.
+  # Behind connection-level load balancers (AWS Service Connect, an NLB, or an
+  # Envoy TCP proxy) each keep-alive connection is pinned to one upstream
+  # instance, so rotating commands across the pool spreads traffic across every
+  # instance — even single-threaded at zero concurrency. Connections are created
+  # up front but connect lazily on first use.
+  #
+  # Two behaviors deliberately match the connection_pool gem it replaces: a
+  # nested `with` from the same thread re-enters the held connection, and after
+  # `fork` (Puma preload_app, unicorn, resque) the pool closes inherited sockets
+  # and rebuilds its mutexes in the child so parent and child never share a
+  # connection.
+  class RoundRobinPool
+    # Pools are tracked only to reset them in forked children. WeakMap so a
+    # pool is reclaimed together with its client.
+    INSTANCES = Process.respond_to?(:fork) ? ObjectSpace::WeakMap.new : nil
+    private_constant :INSTANCES
+
+    THREAD_KEY = :emb_round_robin_pool_held
+    private_constant :THREAD_KEY
+
+    attr_reader :size, :connections
+
+    def self.after_fork
+      INSTANCES&.each_value(&:reload_after_fork!)
+    end
+
+    def initialize(size, &)
+      raise ArgumentError, "pool size must be >= 1 (got #{size})" if size < 1
+
+      @size = size
+      @connections = Array.new(size, &)
+      @locks = Array.new(size) { Mutex.new }
+      @next = 0
+      @index_mutex = Mutex.new
+      INSTANCES&.[]=(self, self)
+    end
+
+    # Yields the next connection in rotation order. Safe from multiple threads:
+    # up to `size` commands run in parallel, each on its own connection. A
+    # nested `with` from the same thread re-enters the connection this pool
+    # already holds without re-locking; other pools are unaffected.
+    def with(&)
+      held = Thread.current[THREAD_KEY]
+      if held&.key?(self)
+        yield @connections[held[self]]
+      else
+        take(held, &)
+      end
+    end
+
+    # Child side of after_fork(): drop inherited sockets and sync state.
+    def reload_after_fork!
+      @connections.each { |conn| conn.close if conn.respond_to?(:close) }
+      @locks = Array.new(@size) { Mutex.new }
+      @next = 0
+      @index_mutex = Mutex.new
+    end
+
+    if Process.respond_to?(:fork)
+      # Hooks Process._fork (MRI 3.1+) so registered pools reset in the child.
+      module ForkTracker
+        def _fork
+          pid = super
+          RoundRobinPool.after_fork if pid.zero?
+          pid
+        end
+      end
+      Process.singleton_class.prepend(ForkTracker)
+    end
+
+    private
+
+    # Acquires the next connection, records it as held by this thread/pool, and
+    # releases both on exit — even when the block raises.
+    def take(held)
+      idx, connection = pick
+      held ||= {}
+      Thread.current[THREAD_KEY] = held
+      held[self] = idx
+      begin
+        @locks[idx].synchronize { yield connection }
+      ensure
+        held.delete(self)
+        Thread.current[THREAD_KEY] = nil if held.empty?
+      end
+    end
+
+    def pick
+      @index_mutex.synchronize do
+        idx = @next % @size
+        @next += 1
+        [idx, @connections[idx]]
+      end
+    end
+  end
+end

data/lib/emb/script_reply_decode.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+module Emb
+  # Opt-in float decoding for scripted replies — the decode: keyword on
+  # Emb::Commands#eval / #evalsha. Kept out of the Commands module so the
+  # command surface stays small. decode: is validated (normalize_decode) before
+  # any command is sent, then applied after reply parsing.
+  module ScriptReplyDecode
+    private
+
+    # Validates an (un-normalized) decode option and canonicalizes it. Raises
+    # ArgumentError for unknown modes so callers fail before any command is
+    # sent. nil and an empty hash both mean "no decoding".
+    def normalize_decode(decode)
+      case decode
+      when nil, :f32
+        decode
+      when Hash
+        schema = decode.each_with_object({}) do |(field, mode), acc|
+          unless mode == :f32
+            raise ArgumentError, "unsupported decode mode #{mode.inspect} for field #{field.inspect} (supported: :f32)"
+          end
+
+          acc[field.to_s] = :f32
+        end
+        schema.empty? ? nil : schema
+      else
+        raise ArgumentError, "unsupported decode mode #{decode.inspect} (supported: :f32 or {field => :f32})"
+      end
+    end
+
+    # Applies a normalized decode to the parsed (script_hash'ed) reply array.
+    def apply_decode(parsed, decode)
+      return parsed if decode.nil?
+
+      case decode
+      when :f32
+        parsed.each_with_index.map { |v, i| decode_f32(v, position(i)) }
+      when Hash
+        parsed.each_with_index.map { |v, i| decode_hash_fields(v, decode, position(i)) }
+      end
+    end
+
+    def position(index)
+      index.zero? ? 'the reply' : "element #{index + 1}"
+    end
+
+    # decode: :f32 — a float32-packed bulk unpacks to floats; a numeric array
+    # (the pre-float32_bytes reply style) passes through unchanged; anything
+    # else is a contract mismatch and raises.
+    def decode_f32(value, where)
+      case value
+      when String
+        unless (value.bytesize % 4).zero?
+          raise ArgumentError,
+                "decode :f32: #{where} is a #{value.bytesize}-byte bulk, not a multiple of 4 — not a float32 vector"
+        end
+
+        value.unpack('e*')
+      when Array
+        return value if value.all?(Numeric)
+
+        raise ArgumentError, "decode :f32: #{where} is an array of non-numbers, not a float vector"
+      else
+        raise ArgumentError,
+              "decode :f32: #{where} is a #{value.class}, expected a float32-packed bulk or a numeric array"
+      end
+    end
+
+    # decode: {field => :f32} — the reply (or each per-text reply) must be a
+    # hash; present fields decode as :f32, absent fields stay untouched.
+    def decode_hash_fields(value, schema, where)
+      unless value.is_a?(Hash)
+        raise ArgumentError, "decode #{schema.inspect}: #{where} is a #{value.class}, expected a hash reply"
+      end
+
+      schema.each_key do |field|
+        next unless value.key?(field)
+
+        value[field] = decode_f32(value[field], "field #{field.inspect} of #{where}")
+      end
+      value
+    end
+  end
+end

data/lib/emb/values_batch.rb

@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+
+require_relative 'values_reply'
+
+module Emb
+  # VALUES-format batch resolution for the deferred path. Split out of
+  # BatchDispatch so both modules stay under the Metrics rubric.
+  module ValuesBatch
+    module_function
+
+    # Resolves a VALUES reply onto the slice's items. Single-model slices get
+    # ONE envelope reply carrying every text's flattened values (row-major),
+    # so the rows are re-sliced per item; mixed-model slices get per-pair
+    # envelopes and map one row each. Failures stay nil, mirroring the BLOB
+    # path's MGET semantics.
+    def resolve(loader, slice, results)
+      if single_model?(slice)
+        rows = envelope_rows(results)
+        expect_rows!(slice, rows.size)
+        assign_rows(loader, slice, rows)
+      else
+        resolve_entries(loader, slice, results)
+      end
+    end
+
+    def single_model?(slice)
+      slice.map { |item| item[1] }.uniq.size == 1
+    end
+
+    def envelope_rows(results)
+      envelope = Emb::ValuesReply.parse(rehydrate_envelope(results))
+      Emb::ValuesReply.rows(envelope[:shape], envelope[:values])
+    end
+
+    # Under RESP3, redis-client decodes the server's single-model VALUES map
+    # into a Ruby Hash, but dispatch_slice wraps every reply in Array(), which
+    # turns that Hash into an array of [key, value] pairs. Rehydrate those
+    # pairs back into a Hash so parse sees a real envelope; RESP2 flat pair
+    # arrays (top-level strings) pass through untouched.
+    def rehydrate_envelope(results)
+      return results.to_h if results.all? { |e| e.is_a?(Array) && e.size == 2 }
+
+      results
+    end
+
+    def expect_rows!(slice, got)
+      expected = slice.sum { |_, _, text, _| Array(text).size }
+      return if got >= expected
+
+      raise ShortReplyError, "expected #{expected} VALUES rows, got #{got}"
+    end
+
+    def assign_rows(loader, slice, rows)
+      offset = 0
+      slice.each do |item|
+        _, _, text, = item
+        texts = Array(text)
+        values = rows[offset, texts.size]
+        offset += texts.size
+        loader.call(item, values.size == 1 ? values.first : values)
+      end
+    end
+
+    def resolve_entries(loader, slice, results)
+      expect_entries!(slice, results.size)
+      offset = 0
+      slice.each do |item|
+        _, _, text, = item
+        texts = Array(text)
+        values = entry_values(results, offset, texts)
+        offset += texts.size
+        loader.call(item, values.size == 1 ? values.first : values)
+      end
+    end
+
+    # A short multi-model reply is a protocol violation: an explicit null slot
+    # is the only legal way to signal a missing pair (see the BLOB path).
+    def expect_entries!(slice, got)
+      expected = slice.sum { |_, _, text, _| Array(text).size }
+      return if got >= expected
+
+      raise ShortReplyError, "expected #{expected} VALUES reply entries, got #{got}"
+    end
+
+    def entry_values(results, offset, texts)
+      results[offset, texts.size].map { |entry| entry && row_of(entry) }
+    end
+
+    def row_of(entry)
+      envelope = Emb::ValuesReply.parse(entry)
+      Emb::ValuesReply.rows(envelope[:shape], envelope[:values]).first
+    end
+  end
+end

data/lib/emb/values_reply.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module Emb
+  # Decoding for EMB VALUES envelopes (the RedisAI META+VALUES shape). The
+  # server replies with flat alternating key/value pairs under RESP2 — dtype,
+  # shape, values — where values are decimal bulk strings carrying the float64
+  # widening of each float32 dimension.
+  module ValuesReply
+    module_function
+
+    # Parses one envelope into a Hash. RESP2 replies arrive as a flat 6-element
+    # pair array; under `protocol: 3` redis-client already decoded the server's
+    # RESP3 map into a Ruby Hash (keys frozen strings, values already Floats),
+    # which is used directly. values are converted to Ruby Floats; an
+    # unparsable entry raises ArgumentError naming the offending value.
+    def parse(reply)
+      pairs = reply.is_a?(Hash) ? reply : reply.each_slice(2).to_h
+      shape = pairs.fetch('shape')
+      values = Array(pairs.fetch('values')).map { |v| float_value(v) }
+      { dtype: pairs.fetch('dtype'), shape: shape, values: values }
+    end
+
+    # rows reshapes flat values into per-text rows using shape [m, dim].
+    def rows(shape, values)
+      dim = shape[1].to_i
+      values.each_slice(dim).to_a
+    end
+
+    def float_value(v)
+      Float(v)
+    rescue ArgumentError, TypeError
+      raise ArgumentError, "EMB VALUES: unparsable decimal value #{v.inspect} (expected a numeric bulk string)"
+    end
+  end
+end

data/lib/emb.rb

@@ -2,13 +2,17 @@
 
 require_relative 'emb/version'
 require_relative 'emb/configuration'
+require_relative 'emb/errors'
 require_relative 'emb/commands'
 require_relative 'emb/runtime_config'
 require_relative 'emb/client'
 require_relative 'emb/proxy'
+require_relative 'emb/batch_scope'
+require_relative 'emb/middleware'
+require_relative 'emb/job_middleware'
 require_relative 'emb/multi'
 require_relative 'emb/batch'
-require_relative 'emb/middleware'
+require_relative 'emb/railtie' if defined?(Rails::Railtie)
 
 module Emb
   class << self
@@ -28,19 +32,33 @@ module Emb
     end
 
     def [](name)    = default_client[name]
-    def batch       = default_client.batch
     def models      = default_client.models
     def info(name)  = default_client.info(name)
     def stats       = default_client.stats
     def server_info(*sections)  = default_client.server_info(*sections)
+    def cache_flush(model = nil) = default_client.cache_flush(model)
+    def save_cache               = default_client.save_cache
     def config                  = default_client.config
     def help        = default_client.help
     def ping        = default_client.ping
     def ready       = default_client.ready
     def ready?      = default_client.ready?
     def multi(&)    = default_client.multi(&)
+
+    def eval(model, script, texts, args = [], decode: nil)
+      default_client.eval(model, script, texts, args, decode: decode)
+    end
+
+    def evalsha(model, sha, texts, args = [], decode: nil)
+      default_client.evalsha(model, sha, texts, args, decode: decode)
+    end
+
+    def script = default_client.script
+
     def reset_registry! = default_client.reset_registry!
+
     def debug? = @debug
+
     def send_command(*) = default_client.send_command(*)
 
     def debug!

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: emb
 version: !ruby/object:Gem::Version
-  version: 0.3.0
+  version: 0.4.0.pre3
 platform: ruby
 authors:
 - elcuervo
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2026-09-01 00:00:00.000000000 Z
+date: 2026-09-10 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: batch-loader
@@ -24,20 +24,6 @@ dependencies:
     - - "~>"
       - !ruby/object:Gem::Version
         version: '2.0'
-- !ruby/object:Gem::Dependency
-  name: connection_pool
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: '2.5'
-  type: :runtime
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: '2.5'
 - !ruby/object:Gem::Dependency
   name: redis-client
   requirement: !ruby/object:Gem::Requirement
@@ -63,13 +49,23 @@ files:
 - README.md
 - lib/emb.rb
 - lib/emb/batch.rb
+- lib/emb/batch_dispatch.rb
+- lib/emb/batch_scope.rb
 - lib/emb/client.rb
 - lib/emb/commands.rb
 - lib/emb/configuration.rb
+- lib/emb/connection_router.rb
+- lib/emb/errors.rb
+- lib/emb/job_middleware.rb
 - lib/emb/middleware.rb
 - lib/emb/multi.rb
 - lib/emb/proxy.rb
+- lib/emb/railtie.rb
+- lib/emb/round_robin_pool.rb
 - lib/emb/runtime_config.rb
+- lib/emb/script_reply_decode.rb
+- lib/emb/values_batch.rb
+- lib/emb/values_reply.rb
 - lib/emb/version.rb
 homepage: https://github.com/elcuervo/emb
 licenses: