emb 0.2.6 → 0.4.0.pre1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
checksums.yaml CHANGED
@@ -1,7 +1,7 @@
1
1
  ---
2
2
  SHA256:
3
- metadata.gz: 0465face51bde07116c17303587651ae8c947875b9702a587654ead5256c71da
4
- data.tar.gz: b0e8cf2d155a51c9b65a2bb006eec6be4d7a60e932301bea426b2ae986a91015
3
+ metadata.gz: 646506ae908e1f9db100c67c6d4e20dc6c6fc84e9d2d2a62c7954385ff632d73
4
+ data.tar.gz: 07e42cf06b265483d14899ad8607909ad8caa99243e28bd7fe66dfad6eef1496
5
5
  SHA512:
6
- metadata.gz: 843e3cdb952a4a1ece48413aad778939d155f4e345c9ebbce3ccd18ce703e53183a099dcae3b815856d062a0d7a80b3d028587704c736402ce659133eb1d7b8a
7
- data.tar.gz: 786734c2ac2765bde4810313f998f83550e0823cd47887e16ebf5b2014fe2c865e22b021286187cedaf105d2d5b818c8885d85acceab72ead5e07915da17d9f3
6
+ metadata.gz: 0ea5a3f7613be7280da47b7e52ebcff8b7773ef33ebbdcdbaa231e28084b8e4b3b12755be415b965b82ce395ea252c1dbf7592a4981740e2ba45b67e97338d49
7
+ data.tar.gz: 869f1f0df61b44fc09aa6ce95b4b04bbcf0de1542f2697594160e16baf977dc6707df7a7a6da77a0e68be3d1f199c77f2cac171e0a41cabdf467c8958a352e89
data/Gemfile CHANGED
@@ -4,7 +4,6 @@ source 'https://rubygems.org'
4
4
 
5
5
  gemspec
6
6
 
7
- gem 'connection_pool'
8
7
  gem 'redis-client'
9
8
 
10
9
  gem 'rake', require: false
data/README.md CHANGED
@@ -55,7 +55,7 @@ resolve in this order — the first one wins:
55
55
  ```ruby
56
56
  Emb.configure do |c|
57
57
  c.pool = 8
58
- c.batch = false # opt out of lazy batching app-wide
58
+ c.lazy = :multi # defer and coalesce embed calls app-wide
59
59
  end
60
60
 
61
61
  Emb.configuration # => the shared Emb::Configuration
@@ -67,11 +67,28 @@ Emb::Client.new(pool: 20) # per-call still wins
67
67
 
68
68
  ### Out-of-the-box defaults
69
69
 
70
- The shipped defaults are benchmark-derived (see `BENCHMARK.md`): **lazy batching is on by
71
- default** (`batch: true` — each embed coalesces into one `EMB.MULTI`), pool `5`, pure-Ruby
72
- RESP driver, `protocol: 2`, `reconnect_attempts: 3`. To keep the eager behavior (immediate
73
- `EMB` per call), opt out globally via `Emb.configure { |c| c.batch = false }` or per client
74
- with `Emb.new(batch: false)`.
70
+ The shipped defaults are benchmark-derived (see `BENCHMARK.md`) with secure-by-default
71
+ network behavior: **eager execution by default** (`lazy: false` — every `Emb[:model][t]`
72
+ sends one `EMB` round trip), pool `5`, pure-Ruby RESP driver, `protocol: 2`,
73
+ `read_timeout`/`write_timeout` both **10s**, and `reconnect_attempts: 0`. Opt into
74
+ coalescing with `lazy: :multi` or concurrent fan-out with `lazy: :batch` — globally via
75
+ `Emb.configure { |c| c.lazy = :batch }` or per client with `Emb.new(lazy: :batch)`.
76
+
77
+ **Why the timeout and reconnect defaults matter:** batched `EMB.MULTI` (up to 512 pairs)
78
+ can take over a second of inference on a shared CPU, and redis-client's silent default is
79
+ 1.0s — a slower reply times out. The gem therefore defaults to an explicit 10s timeout
80
+ and `reconnect_attempts: 0`: a failing batch fails closed after one attempt and raises
81
+ `Emb::ServerError` (see [Lazy execution modes](#lazy-execution-modes)). Set
82
+ `Emb.configure { |c| c.reconnect_attempts = 2 }` and redis-client re-sends
83
+ **connection and protocol failures** up to that many extra times before the batch fails
84
+ closed — each re-send re-runs server inference, so keep the budget small. Operation
85
+ errors (unknown model, auth, bad pairs) are never retried, and neither are **read
86
+ timeouts**: redis-client treats `ReadTimeoutError` as terminal, and re-sending a timed-out
87
+ command could duplicate inference the server already did. `lazy: :batch` shares
88
+ additionally retry a refused connection on the next configured instance. Raise the
89
+ timeouts if you raise `batch_size`. `reconnect_attempts` also accepts an Array of
90
+ per-retry delays (a redis-client passthrough); each entry grants one retry and counts
91
+ toward `Emb::ServerError#attempts`.
75
92
 
76
93
  ### Connection pool
77
94
 
@@ -79,11 +96,48 @@ with `Emb.new(batch: false)`.
79
96
  Emb.setup(url: "redis://localhost:6379", pool: 10)
80
97
  ```
81
98
 
82
- The default pool size is **5**. The pool is usually not the bottleneck for
83
- inference-bound workloads (small pools are fine); it becomes a knob only at high
84
- concurrency on a multi-model box see [Performance](#performance). If a pool checkout
85
- would wait too long, `RedisClient`'s `connect_timeout`/`read_timeout` (above) bound the
86
- wait.
99
+ The default pool size is **5**. The client keeps `pool` persistent connections per
100
+ emb **instance** and routes commands through them **in round-robin order**
101
+ consecutive commands use different connections instead of reusing one hot connection.
102
+
103
+ `url` also accepts an **array of instances** (interchangeable replicas serving the same
104
+ model set); commands then round-robin across instances first, then across connections
105
+ within the selected instance:
106
+
107
+ ```ruby
108
+ Emb.setup(url: ["redis://emb-a:6379", "redis://emb-b:6379", "redis://emb-c:6379"], pool: 5)
109
+ ```
110
+
111
+ Each url gets its own pool of `pool` connections, so `pool: 5` with three urls maintains
112
+ 15 connections total. If an instance refuses a connection **before a command is sent**
113
+ (connection refused/unreachable), the command is retried on the next instance — this is
114
+ why `reconnect_attempts` stays `0`: a retry after a *sent* command could duplicate
115
+ inference, but a never-sent command is safe to move elsewhere.
116
+
117
+ This matters when emb runs as a pool of instances behind a **connection-level load
118
+ balancer** (AWS Service Connect, an NLB, or Envoy): the balancer assigns each
119
+ keep-alive connection to exactly one instance for its lifetime, so relying on a
120
+ single connection would pin all of a process's embedding traffic to one instance —
121
+ and queue it there even while other instances sit idle. Round-robin spreading across
122
+ `pool` connections spreads traffic across up to `pool` instances, so set `pool` to
123
+ **at least your expected emb instance count** (e.g. `pool: 10` for 10 instances).
124
+
125
+ Commands beyond the pool's parallelism do **not** time out: when `pool` commands are
126
+ already in flight, later commands wait on the shared connections until one frees
127
+ (there is no checkout timeout — unlike the previous pool's 5s
128
+ `ConnectionPool::TimeoutError`). Only the wait for a free pool connection is
129
+ unbounded — once a command runs, `RedisClient`'s `connect_timeout`,
130
+ `read_timeout`, and `write_timeout` still apply. For inference workloads this
131
+ lets commands ride out token-budget batching instead of erroring.
132
+
133
+ Locally (a single server) only the backend selection is identical whatever the
134
+ pool size — every command reaches the same server. The pool size still controls
135
+ how many persistent connections the client holds and how many commands run in
136
+ parallel. The pool is usually not the bottleneck for
137
+ inference-bound workloads (small pools are fine); it becomes a knob at high
138
+ concurrency on a multi-model box — see [Performance](#performance). With
139
+ round-robin selection up to `pool` commands run in parallel; beyond that, commands
140
+ share connections and serialize on them.
87
141
 
88
142
  ### Authentication
89
143
 
@@ -124,8 +178,8 @@ Emb.setup(
124
178
  ```
125
179
 
126
180
  See the [redis-client documentation](https://github.com/redis-rb/redis-client) for
127
- all available options. Only `pool` and `batch` are handled by the gem — everything
128
- else passes through to `RedisClient.new`.
181
+ all available options. Only `pool`, `lazy`, and `batch_size` are handled by the
182
+ gem — everything else passes through to `RedisClient.new`.
129
183
 
130
184
  ## Instance-based clients
131
185
 
@@ -169,6 +223,12 @@ automatically.
169
223
 
170
224
  ## Server info & config
171
225
 
226
+ > **Per-instance data.** `Emb.server_info`, `stats`, `info`, and `config` read
227
+ > from ONE instance, and with round-robin connection selection that instance
228
+ > may change between calls (with `pool: 1` it never does) — each call is a
229
+ > sample. `Emb.config[key] = value` mutates one (arbitrary) instance. Use
230
+ > per-connection queries if you need a specific or aggregated view.
231
+
172
232
  The server exposes Redis-style `INFO` and `CONFIG` commands; the gem wraps them.
173
233
 
174
234
  ### `Emb.stats` — server statistics as a hash
@@ -190,15 +250,18 @@ Emb.stats
190
250
  ### `Emb.server_info` — sectioned INFO, parsed
191
251
 
192
252
  The Redis-style `INFO` reply is parsed into a nested Hash. **No arguments = all
193
- sections**; pass section names to filter (`:server`, `:cache`, `:keyspace`, `:stats`, `:clients`):
253
+ sections**; pass section names to filter (`:server`, `:cache`, `:keyspace`, `:stats`, `:memory`, `:cpu`, `:clients`):
194
254
 
195
255
  ```ruby
196
256
  Emb.server_info
197
257
  # => {Server: {redis_version: "0.2.4", emb_version: "0.2.4", uptime_secs: "7", ...},
198
258
  # Cache: {cache_hits: 0, cache_misses: 0, cache_hit_rate: "0.0%", ...},
199
- # Keyspace: {db0: "model=minilm,keys=0,hits=0,misses=0,hit_rate=0.0%"}, ...}
259
+ # Keyspace: {db0: "model=minilm,keys=0,hits=0,misses=0,hit_rate=0.0%"},
260
+ # Memory: {used_memory_rss_bytes: 262438912, used_memory_heap_bytes: 154960,
261
+ # goroutines: 4, total_system_memory_bytes: 25769803776},
262
+ # CPU: {used_cpu_user_usec: 188900, used_cpu_sys_usec: 50462, gomaxprocs: 10}, ...}
200
263
 
201
- Emb.server_info(:server, :cache) # only those two sections
264
+ Emb.server_info(:memory, :cpu) # live process resources only
202
265
  ```
203
266
 
204
267
  ### `Emb.config` — hot config read & change
@@ -273,38 +336,73 @@ client.multi do |m|
273
336
  end
274
337
  ```
275
338
 
276
- ### Lazy batching (`Emb.batch`)
339
+ ### Lazy execution modes
340
+
341
+ Embed-call behavior is governed by a single `lazy` mode — `false` (default, eager),
342
+ `:multi` (defer and coalesce into one `EMB` for a single model / one `EMB.MULTI` for mixed scopes, serial), or `:batch` (defer and execute
343
+ the coalesced chunk shares **concurrently**). The three are mutually exclusive.
344
+
345
+ | mode | `Emb[:model][text]` | execution |
346
+ |---|---|---|
347
+ | `false` (default) | immediate `EMB` round trip | serial, per call |
348
+ | `:multi` | deferred → coalesces into one `EMB` (single model) or `EMB.MULTI` (mixed) | serial, one command at a time |
349
+ | `:batch` | deferred → coalesces into `EMB`/`EMB.MULTI` chunks | **concurrent** — chunk shares run in parallel |
277
350
 
278
- Instead of collecting pairs by hand, `Emb.batch` returns lazy embeddings that all
279
- coalesce into a single `EMB.MULTI` round trip when the first one is used. This is
280
- powered by the [batch-loader](https://github.com/exAspArk/batch-loader) gem.
351
+ In `:batch` mode the shares fan out across the configured instances (one share per
352
+ instance when the share count allows) or across the instance's pool connections when a
353
+ single url is configured — a batch whose pieces take 60ms and 10ms completes in roughly
354
+ the slowest share (~60ms) instead of the sum (~70ms). Deferred work is powered by the
355
+ [batch-loader](https://github.com/exAspArk/batch-loader) gem:
281
356
 
282
357
  ```ruby
358
+ Emb.setup(lazy: :multi) # or :batch — Emb[:minilm] now defers
359
+
283
360
  users = User.all # some application objects
284
361
 
285
362
  # Create loaders first...
286
- l1 = Emb.batch[:minilm]["hello"]
287
- l2 = Emb.batch[:minilm]["world"]
288
- l3 = Emb.batch[:bge]["bonjour"]
363
+ l1 = Emb[:minilm]["hello"]
364
+ l2 = Emb[:minilm]["world"]
365
+ l3 = Emb[:bge]["bonjour"]
289
366
 
290
- # ...then consume them. The first use sends ONE EMB.MULTI for all three.
367
+ # ...then consume them. The first use sends ONE EMB (single-model scope)
368
+ # or ONE EMB.MULTI (mixed-model scope) for all three.
291
369
  l1.sum # => 12.345
292
370
  l2.sum # => -0.678
293
371
  l3.sum # => 3.141
294
372
  ```
295
373
 
296
- Instance clients expose the same API:
297
-
298
- ```ruby
299
- client.batch[:minilm]["hello"].sum
300
- ```
301
-
302
374
  Each lazy value materializes to the same shape as the eager API: a single text
303
- yields an `Array<Float>`, multiple texts yield `Array<Array<Float>>`:
375
+ yields an `Array<Float>`, multiple texts yield `Array<Array<Float>>`. For explicit
376
+ composition that always sends immediately regardless of mode, use `Emb.multi { }`
377
+ (see [Multi-model queries](#multi-model-queries)).
378
+
379
+ The batch executes in the current thread's scope; the Rack/ActiveJob middleware
380
+ clears that scope after every request/job (`Emb::BatchScope`), so deferred work
381
+ can never accumulate across requests (in eager mode there is nothing to clear).
382
+
383
+ **Fail-closed batches.** If a batch command fails — a timeout, an operation error, or in
384
+ `:batch` mode a share failure after pre-send retries are exhausted — the batch raises
385
+ **`Emb::ServerError`** to the code that first used it, and every deferred item of that
386
+ batch is removed from the scope — retrying (or using other items of the failed batch)
387
+ **does not re-send the batch** and resolves to `[]` instead. The error's `cause` is the
388
+ underlying redis error (`RedisClient::ReadTimeoutError`, `RedisClient::CommandError`, ...)
389
+ and its message includes the model(s), text count, and attempt count. Transient failures
390
+ (timeout, connection error, protocol error) re-send up to `reconnect_attempts` extra
391
+ times when that option is set above 0 (default 0 = a single attempt); operation errors
392
+ are never retried. A non-redis error raised by the batch call itself (e.g. a `TypeError`
393
+ from bad arguments) is a local bug, not a server failure: the batch's pending items are
394
+ still cleared, but the original error is re-raised. This
395
+ prevents a slow server from turning one failed batch into endless duplicate work
396
+ (retries re-running the whole batch) or growth of the pending set across retries.
397
+ Pair-level failures the server reports as `null` (MGET semantics) are unaffected.
398
+
399
+ > **Breaking change (gem ≥ next release):** a failed batch raises `Emb::ServerError`
400
+ > instead of the raw `RedisClient::*` error. Rescue `Emb::ServerError` and read `cause`
401
+ > for the original error. Eager `Emb.multi` and the `lazy: false` path are unchanged.
304
402
 
305
403
  ```ruby
306
- vec = Emb.batch[:minilm]["hello"] # use -> Array of Float
307
- vecs = Emb.batch[:minilm]["hello", "world"] # use -> Array of Array of Float
404
+ vec = Emb[:minilm]["hello"] # use -> Array of Float
405
+ vecs = Emb[:minilm]["hello", "world"] # use -> Array of Array of Float
308
406
  ```
309
407
 
310
408
  Embeddings are cached per thread, so reusing a lazy value (or creating an
@@ -318,43 +416,59 @@ Loaders only fire when a value is **used**. Create all loaders *first*, then
318
416
  consume them, so they share one round trip:
319
417
 
320
418
  ```ruby
321
- texts.each { |t| process(Emb.batch[:minilm][t]) } # wrong: one MULTI per item
322
- loaders = texts.map { |t| Emb.batch[:minilm][t] } # right: ONE MULTI for all
419
+ texts.each { |t| process(Emb[:minilm][t]) } # wrong: one EMB per item
420
+ loaders = texts.map { |t| Emb[:minilm][t] } # right: ONE EMB for all
323
421
  loaders.each { |l| process(l) }
324
422
  ```
325
423
 
326
424
  A loader that is created but never used **never embeds** (unless a sibling batch
327
425
  fires first) and is silently dropped when the thread's scope ends. Duration and
328
- scope: batching is per-thread — a multithreaded app issues one `EMB.MULTI` per
426
+ scope: batching is per-thread — a multithreaded app issues one `EMB`/`EMB.MULTI` per
329
427
  thread per flush.
330
428
 
331
- ### `batch` configuration option
429
+ ### `lazy` configuration option
332
430
 
333
- Setting `batch: true` makes the standard proxy API lazy, so existing call sites
334
- batch automatically without restructuring:
431
+ Setting `lazy: :multi` or `lazy: :batch` makes the standard proxy API defer, so
432
+ existing call sites batch automatically without restructuring:
335
433
 
336
434
  ```ruby
337
- Emb.setup(url: "redis://localhost:6379", batch: true)
435
+ Emb.setup(url: "redis://localhost:6379", lazy: :multi)
338
436
  # or
339
- Emb.new(url: "redis://localhost:6379", batch: true)
437
+ Emb.new(url: "redis://localhost:6379", lazy: :batch)
340
438
  ```
341
439
 
342
- With `batch: true`, `Emb[:minilm]["hello"]` returns a lazy embedding that sends
343
- `EMB.MULTI` on first use. The default is `false` the proxy API stays eager,
344
- sending `EMB` immediately. `Emb.batch` works regardless of the option, and
345
- `Emb.multi` remains the explicit, eager, deterministic batching API.
440
+ Under `lazy: :multi`, `Emb[:minilm]["hello"]` returns a lazy embedding that sends
441
+ `EMB` on first use (serial chunks; `EMB.MULTI` only for mixed-model scopes). Under `lazy: :batch`, the chunk shares
442
+ execute concurrently with multiple `url`s they fan out across instances. The
443
+ default is eager (`lazy: false`). `Emb.multi` remains the explicit, eager,
444
+ deterministic composition API in every mode.
346
445
 
347
446
  ### Clearing the cache per request
348
447
 
349
- The per-thread batch scope holds cached embeddings for the life of the thread. In
350
- request-shaped processes (Rails, Rack apps, Sidekiq) mount `Emb::Middleware` to
351
- clear the scope at the end of each request:
448
+ The per-thread batch scope holds cached embeddings for the life of the thread.
449
+ In a Rails application the bundled Railtie mounts `Emb::Middleware`
450
+ automatically, so the scope is cleared at the end of every request with no
451
+ configuration:
352
452
 
353
453
  ```ruby
354
- # config/application.rb (Rails)
355
- config.middleware.use Emb::Middleware
454
+ # nothing to do — Emb::Railtie inserts Emb::Middleware for you
455
+ ```
456
+
457
+ Opt out if you want to manage the stack yourself:
458
+
459
+ ```ruby
460
+ # config/application.rb
461
+ config.emb.middleware = false
462
+ ```
463
+
464
+ If your Gemfile loads `emb` before Rails is required (non-standard boot order), the
465
+ guarded railtie require in `emb.rb` is skipped — add `require "emb/railtie"` in
466
+ `config/application.rb` right after `require "rails/all"` (or in an initializer).
356
467
 
357
- # Any Rack app
468
+ The middleware is also safe to mount manually in any Rack app (the Railtie
469
+ skips insertion when it is already present):
470
+
471
+ ```ruby
358
472
  use Emb::Middleware
359
473
  ```
360
474
 
@@ -362,6 +476,46 @@ The scope is cleared even when the app raises, and a fresh scope starts
362
476
  automatically with the next request. Loaders created but never used within a
363
477
  request are dropped — the create-then-consume contract applies per request.
364
478
 
479
+ ### Job-scoped cache clearing
480
+
481
+ The Railtie also registers `Emb::JobMiddleware` — the same per-scope clearing
482
+ for background work — for every job framework present. Each job execution
483
+ starts with a fresh batch scope: cached embeddings and loaders created but
484
+ never used are dropped at the end of the job (even when it raises), so worker
485
+ threads never leak scope state between jobs.
486
+
487
+ - **ActiveJob** (SolidQueue, Sidekiq, Shoryuken, async, test adapters): an
488
+ `around_perform` callback on `ActiveJob::Base` registered by the Railtie.
489
+ - **Plain Sidekiq workers** (non-ActiveJob): a server middleware added by the
490
+ Railtie.
491
+ - **Plain Shoryuken workers** (non-ActiveJob): a server middleware added by the
492
+ Railtie.
493
+
494
+ ```ruby
495
+ # opt out of all job-scope protection
496
+ config.emb.job_middleware = false
497
+ ```
498
+
499
+ In a non-Rails process, register the middleware manually:
500
+
501
+ ```ruby
502
+ # sidekiq.rb / an initializer
503
+ Sidekiq.configure_server do |config|
504
+ config.server_middleware { |chain| chain.add Emb::JobMiddleware }
505
+ end
506
+
507
+ # Shoryuken
508
+ Shoryuken.configure_server do |config|
509
+ config.server_middleware { |chain| chain.add Emb::JobMiddleware }
510
+ end
511
+ ```
512
+
513
+ The Railtie also registers `Emb::JobMiddleware` into the `Sidekiq::Testing` middleware
514
+ chain when `Sidekiq::Testing` is loaded *before the app finishes booting* (e.g. a dev
515
+ `INLINE_SIDEKIQ` initializer). For RSpec test modes, `sidekiq/testing` typically loads
516
+ after the app boots, so require it before your environment to get per-job clearing in
517
+ fake/inline tests — e.g. `gem "sidekiq", require: ["sidekiq", "sidekiq/testing"]`.
518
+
365
519
  ## Performance
366
520
 
367
521
  ### Eager-burst pipelining (no new API)
@@ -372,7 +526,7 @@ batching doesn't fit), coalesce them into one packet with `RedisClient#pipelined
372
526
  ```ruby
373
527
  client = Emb.new(url: "redis://localhost:6379")
374
528
 
375
- a = client.pool.with do |conn|
529
+ a = client.pools.first.with do |conn|
376
530
  conn.pipelined do |pipe|
377
531
  texts.each { |t| pipe.call("EMB", "minilm", t) }
378
532
  end
@@ -380,9 +534,22 @@ end
380
534
  # a -> Array of Float32-binary replies; unpack each with r.unpack("e*")
381
535
  ```
382
536
 
383
- `Client#pool` exposes the pool's `RedisClient` connections. This measurably beats
384
- plain per-call round trips; use it for bursts, and `Emb.batch`/`Emb.multi` when you want
385
- server-side coalescing into one `EMB.MULTI`.
537
+ `Client#pools` exposes the per-instance pools' `RedisClient` connections (the
538
+ first pool for a single-instance client). This measurably beats
539
+ plain per-call round trips; use it for bursts, and `Emb.multi` or a deferred `lazy`
540
+ mode when you want
541
+ server-side coalescing into one packed inference.
542
+
543
+ ### Benchmark harness coverage
544
+
545
+ The client benchmark harness (`gems/emb/bench/bench.rb`, `just bench-ruby`)
546
+ reports one row per execution mechanism: `eager` (one `EMB` per call), `multi`
547
+ (coalesced `EMB`, or `EMB.MULTI` for mixed models), `batch` (concurrent `EMB` chunk shares),
548
+ `pipelined` (raw RESP pipelining), and `threaded` (eager across threads).
549
+ `just bench-ruby-multi` starts two partitioned emb instances and adds the
550
+ url-array rows `eager-2node` (round-robin distribution) and `batch-2node`
551
+ (concurrent fan-out across instances); the two-node rows are skipped unless a
552
+ second instance is reachable (`EMB_BENCH_PORT2`).
386
553
 
387
554
  ### RESP driver: pure-Ruby default, `hiredis` on demand
388
555
 
@@ -409,7 +576,7 @@ instance. Scale out without a cluster client:
409
576
  (optionally seed with `-cache` and a warmup pass).
410
577
 
411
578
  Lazy batching stays per server (a thread's loaders target one client), so batching and
412
- horizontal scale compose: each instance receives one `EMB.MULTI` per request.
579
+ horizontal scale compose: each instance receives one `EMB` (single model) or `EMB.MULTI` (mixed) share per request.
413
580
 
414
581
  ### Commands
415
582
 
data/lib/emb/batch.rb CHANGED
@@ -1,69 +1,116 @@
1
1
  # frozen_string_literal: true
2
2
 
3
3
  require 'batch_loader'
4
+ require 'redis_client'
5
+ require_relative 'batch_dispatch'
4
6
 
5
7
  module Emb
8
+ extend BatchDispatch
9
+
6
10
  BATCH_KEY = :emb
7
11
 
12
+ # Raised by resolve_slice when the server reply carries fewer entries than
13
+ # the slice's texts. Subclasses RedisClient::ProtocolError so the existing
14
+ # fail-closed rescue/wrap paths treat it as a client error, but stays
15
+ # distinct so transient_error? does not count it as a transport retry: the
16
+ # reply shape was wrong on the single send — nothing was re-sent.
17
+ ShortReplyError = Class.new(RedisClient::ProtocolError)
18
+
8
19
  BATCH_BLOCK = lambda do |items, loader, _args|
9
20
  items.group_by(&:first).each do |client, client_items|
10
21
  chunk = client.respond_to?(:batch_size) && client.batch_size ? client.batch_size : Emb.configuration.batch_size
11
22
 
12
- # Chunk at batch_size pairs per EMB.MULTI so a single command stays well
13
- # under the server's max_pairs cap and within client read timeouts.
14
- client_items.each_slice(chunk) do |slice|
15
- pairs = slice.flat_map { |_, model, text| Array(text).flat_map { |t| [model.to_s, t] } }
16
- results = Array(client.send_command('EMB.MULTI', *pairs))
17
-
18
- offset = 0
19
- slice.each do |item|
20
- _, _, text = item
21
- texts = Array(text)
22
- values = results[offset, texts.size].map { |entry| entry&.unpack('e*') }
23
- offset += texts.size
24
-
25
- # eager Proxy#[] shape: vector for a single text, vectors for many
26
- loader.call(item, values.size == 1 ? values.first : values)
23
+ slices = pack_slices(client_items, chunk)
24
+ if client.respond_to?(:parallel_batch?) && client.parallel_batch? && slices.size > 1
25
+ dispatch_parallel(client, slices, loader)
26
+ else
27
+ # Serial dispatch: one share in flight at a time. Redis errors fail
28
+ # closed with context; anything else is a local bug and is re-raised
29
+ # unchanged after the pending set is dropped.
30
+ slices.each do |slice|
31
+ resolve_slice(loader, slice, dispatch_slice(client, slice))
32
+ rescue RedisClient::Error => e
33
+ fail_batch!(e, slice: slice, budget: retry_budget(client))
34
+ rescue StandardError
35
+ clear_batch_pending!
36
+ raise
27
37
  end
28
38
  end
29
39
  end
30
40
  end
31
41
 
32
42
  class << self
33
- def build_batch_loader(client, model, text)
34
- BatchLoader.for([client, model, text]).batch(key: BATCH_KEY, &BATCH_BLOCK)
35
- end
36
- end
43
+ # BatchDispatch mechanics (pack_slices, dispatch_parallel, resolve_slice,
44
+ # ...) are internal to BATCH_BLOCK; nothing outside Emb calls them with an
45
+ # explicit receiver.
46
+ private(*BatchDispatch.instance_methods(false))
37
47
 
38
- class BatchProxy
39
- def initialize(client)
40
- @client = client
41
- @models = {}
48
+ def build_batch_loader(client, model, text)
49
+ # default_value []: an item whose batch failed (fail-closed) resolves to
50
+ # an empty vector collection instead of nil, so resolver methods like
51
+ # `loader.sum` do not blow up with NoMethodError-on-nil.
52
+ BatchLoader.for([client, model, text]).batch(default_value: [], key: BATCH_KEY, &BATCH_BLOCK)
42
53
  end
43
54
 
44
- def [](name)
45
- @models[name.to_sym] ||= BatchModelProxy.new(@client, name.to_sym)
55
+ # Removes every pending item of the batch scope. batch-loader prunes
56
+ # pending items only after a successful batch block, so failed batches
57
+ # would otherwise stay queued: retries would re-run the whole batch and
58
+ # stale items would be re-sent by later batches in the same scope. Guarded
59
+ # for a nil executor (no batch scope).
60
+ def clear_batch_pending!
61
+ key = [BATCH_BLOCK.source_location, BATCH_KEY]
62
+ BatchLoader::Executor.current&.items_by_block&.delete(key)
46
63
  end
47
64
 
48
- def inspect
49
- "#<Emb::BatchProxy client=#{@client.inspect}>"
65
+ # Fail-closed tail for a failed batch: clear the pending set, then raise
66
+ # Emb::ServerError carrying the cause and the models/texts/attempts
67
+ # context. The raise happens inside a rescue of the original error so Ruby
68
+ # attaches it as `cause` regardless of whether this runs while a rescue is
69
+ # active (serial path) or from the forcing thread after parallel workers
70
+ # captured the error (parallel path). `attempts` counts the error's retry
71
+ # class: connection/protocol errors (the ones redis-client actually
72
+ # re-sends) report `budget + 1`; operation errors and read timeouts (never
73
+ # re-sent — a timeout may already have executed server work) report 1. A
74
+ # pre-send connection refusal is retried across instances by the
75
+ # connection router before it can reach here as a terminal error.
76
+ def fail_batch!(error, slice:, budget:)
77
+ clear_batch_pending!
78
+ attempts = transient_error?(error) ? budget + 1 : 1
79
+ models = slice.map { |_, model, _| model }.uniq.join(', ')
80
+ texts = slice.sum { |_, _, text| Array(text).size }
81
+ message = "batch failed after #{attempts} attempt(s) " \
82
+ "(models: #{models}, #{texts} text(s)) #{error.class}: #{error.message}"
83
+ begin
84
+ raise error
85
+ rescue StandardError
86
+ raise ServerError.new(message, attempts: attempts)
87
+ end
50
88
  end
51
- end
52
89
 
53
- class BatchModelProxy
54
- attr_reader :name
90
+ # How many additional re-sends redis-client performs for transient
91
+ # failures: the per-client option when set, else the global default.
92
+ # Normalized to an Integer (redis-client also accepts a delay Array, whose
93
+ # truthy slots each grant one retry).
94
+ def retry_budget(client)
95
+ value = client.reconnect_attempts if client.respond_to?(:reconnect_attempts)
96
+ value = Emb.configuration.reconnect_attempts if value.nil?
97
+ return value if value.is_a?(Integer)
55
98
 
56
- def initialize(client, name)
57
- @client = client
58
- @name = name
99
+ value.is_a?(Array) ? value.count(&:itself) : 0
59
100
  end
60
101
 
61
- def [](text, *texts)
62
- Emb.build_batch_loader(@client, @name, texts.empty? ? text : [text, *texts])
102
+ # The error classes redis-client actually re-dispatches under
103
+ # reconnect_attempts: ConnectionError (connect/transport breaks) and
104
+ # ProtocolError. ReadTimeoutError is intercepted before the retry loop
105
+ # (a timed-out command may already have executed server work), so it is
106
+ # terminal and counts a single attempt. Emb::ShortReplyError is raised
107
+ # locally by resolve_slice (reply shape wrong on the single send), so it
108
+ # is likewise terminal and counts one attempt.
109
+ def transient_error?(error)
110
+ (error.is_a?(RedisClient::ConnectionError) && !error.is_a?(RedisClient::ReadTimeoutError)) ||
111
+ (error.is_a?(RedisClient::ProtocolError) && !error.is_a?(ShortReplyError))
63
112
  end
64
113
 
65
- def inspect
66
- "#<Emb::BatchModelProxy #{@name}>"
67
- end
114
+ private :clear_batch_pending!, :fail_batch!, :retry_budget, :transient_error?
68
115
  end
69
116
  end
@@ -0,0 +1,147 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Emb
4
+ # Share-dispatch mechanics for the deferred batch path: EMB/MULTI wire
5
+ # shaping, per-slice result mapping, and the bounded concurrent fan-out used
6
+ # by `lazy: :batch`. Extended into Emb by batch.rb.
7
+ module BatchDispatch
8
+ # Sends one chunk share and returns the raw reply entries. Single-model
9
+ # shares use plain `EMB <model> <text>...` (one inference, model once);
10
+ # mixed-model shares keep EMB.MULTI per-pair nil semantics. Errors after
11
+ # the command may have been sent are terminal and propagate so the forcing
12
+ # thread can fail closed.
13
+ def dispatch_slice(client, slice)
14
+ models = slice.map { |_, model, _| model }.uniq
15
+ args = models.size == 1 ? same_model_args(slice, models.first) : mixed_model_args(slice)
16
+ Array(client.send_command(*args))
17
+ end
18
+
19
+ def same_model_args(slice, model)
20
+ ['EMB', model.to_s, *slice.flat_map { |_, _, text| Array(text) }]
21
+ end
22
+
23
+ def mixed_model_args(slice)
24
+ ['EMB.MULTI', *slice.flat_map { |_, model, text| Array(text).flat_map { |t| [model.to_s, t] } }]
25
+ end
26
+
27
+ # Maps a slice's reply entries onto its items in deferral order. Runs on
28
+ # the forcing thread (batch-loader's executor is per-thread), so workers
29
+ # never resolve loaders. A reply shorter than the slice's texts is a
30
+ # protocol violation: fail the batch via Emb::ShortReplyError (distinct
31
+ # from a client-raised ProtocolError so it is not counted as a transport
32
+ # retry).
33
+ def resolve_slice(loader, slice, results)
34
+ expected = slice.sum { |_, _, text| Array(text).size }
35
+ unless results.size >= expected
36
+ raise ShortReplyError, "expected #{expected} reply entries, got #{results.size}"
37
+ end
38
+
39
+ offset = 0
40
+ slice.each do |item|
41
+ _, _, text = item
42
+ texts = Array(text)
43
+ values = entry_values(results, offset, texts)
44
+ offset += texts.size
45
+
46
+ # eager Proxy#[] shape: vector for a single text, vectors for many
47
+ loader.call(item, values)
48
+ end
49
+ end
50
+
51
+ def entry_values(results, offset, texts)
52
+ values = results[offset, texts.size].map { |entry| entry&.unpack('e*') }
53
+ values.size == 1 ? values.first : values
54
+ end
55
+
56
+ # The worker captures failures as outcomes; only the forcing thread
57
+ # resolves loaders (see resolve_slice).
58
+ def dispatch_share(client, slice)
59
+ [:ok, slice, dispatch_slice(client, slice)]
60
+ rescue StandardError => e
61
+ [:error, slice, e]
62
+ end
63
+
64
+ # Dispatches all shares concurrently over at most the client's connection
65
+ # capacity workers, then fails closed on the first terminal error.
66
+ def dispatch_parallel(client, slices, loader)
67
+ workers = slices.size.clamp(1, worker_capacity(client))
68
+ queue = share_queue(slices, workers)
69
+ outcomes = Array.new(slices.size)
70
+ run_worker_threads(client, queue, outcomes, workers)
71
+ resolve_outcomes(client, outcomes, loader)
72
+ end
73
+
74
+ def share_queue(slices, workers)
75
+ queue = Queue.new
76
+ slices.each_with_index { |slice, index| queue << [index, slice] }
77
+ workers.times { queue << nil }
78
+ queue
79
+ end
80
+
81
+ def worker_capacity(client)
82
+ return Float::INFINITY unless client.respond_to?(:pools)
83
+
84
+ client.pools.sum { |pool| pool.respond_to?(:size) ? pool.size : 1 }
85
+ end
86
+
87
+ def run_worker_threads(client, queue, outcomes, workers)
88
+ workers.times.map do
89
+ Thread.new do
90
+ while (job = queue.pop)
91
+ index, slice = job
92
+ outcomes[index] = dispatch_share(client, slice)
93
+ end
94
+ end
95
+ end.each(&:join)
96
+ end
97
+
98
+ # Redis errors fail closed with context (Emb::ServerError, matching the
99
+ # serial path); non-redis errors are local bugs and re-raise unchanged.
100
+ def resolve_outcomes(client, outcomes, loader)
101
+ first_error, failed_slice = collect_outcomes(outcomes, loader)
102
+ return unless first_error
103
+
104
+ raise first_error unless first_error.is_a?(RedisClient::Error)
105
+
106
+ fail_batch!(first_error, slice: failed_slice, budget: retry_budget(client))
107
+ rescue StandardError
108
+ clear_batch_pending!
109
+ raise
110
+ end
111
+
112
+ def collect_outcomes(outcomes, loader)
113
+ first_error = nil
114
+ failed_slice = nil
115
+ outcomes.each do |status, slice, result|
116
+ if status == :ok
117
+ resolve_slice(loader, slice, result)
118
+ else
119
+ first_error ||= result
120
+ failed_slice ||= slice
121
+ end
122
+ end
123
+ [first_error, failed_slice]
124
+ end
125
+
126
+ # Packs items into shares by accumulated text count so one command stays
127
+ # within `chunk` texts. An item larger than the chunk goes alone; the
128
+ # server truncates it with null reply slots, as in the eager path. `used`
129
+ # carries the running text count of the current slice instead of
130
+ # re-summing it for every item (O(n) per item would be O(n²)).
131
+ def pack_slices(items, chunk)
132
+ slices = []
133
+ used = 0
134
+ items.each do |item|
135
+ size = Array(item[2]).size
136
+ if slices.empty? || used + size > chunk
137
+ slices << [item]
138
+ used = size
139
+ else
140
+ slices.last << item
141
+ used += size
142
+ end
143
+ end
144
+ slices
145
+ end
146
+ end
147
+ end
@@ -0,0 +1,20 @@
1
+ # frozen_string_literal: true
2
+
3
+ require 'batch_loader'
4
+
5
+ module Emb
6
+ # Runs a block of work and clears the per-thread batch scope afterwards,
7
+ # even when the block raises. Clearing is unconditional: batch-loader scopes
8
+ # are thread-global and shared by every client, so a deferred-mode client
9
+ # (global `lazy` config or a per-call override) inside the block would
10
+ # otherwise leak cached values and pending loaders into the next request/job.
11
+ # Under the eager default (`lazy: false`) embed calls never create a scope,
12
+ # so the clear is a no-op — the middleware is observably inert there.
13
+ module BatchScope
14
+ def self.wrap
15
+ yield
16
+ ensure
17
+ BatchLoader::Executor.clear_current
18
+ end
19
+ end
20
+ end
data/lib/emb/client.rb CHANGED
@@ -1,52 +1,41 @@
1
1
  # frozen_string_literal: true
2
2
 
3
- require 'connection_pool'
3
+ require_relative 'connection_router'
4
4
  require 'redis_client'
5
5
 
6
6
  module Emb
7
7
  class Client
8
8
  include Commands
9
9
 
10
- attr_reader :pool, :batch_size
11
-
12
- def initialize(pool: nil, batch: nil, **redis_options)
10
+ # Accepts a single redis URL (String) or several (Array of Strings) naming
11
+ # interchangeable emb instances serving the same model set; each url gets
12
+ # its own pool of `pool` connections (see ConnectionRouter).
13
+ def initialize(pool: nil, lazy: nil, **redis_options)
13
14
  cfg = Emb.configuration
14
- @batch_enabled = batch.nil? ? cfg.batch : batch
15
+ @lazy_mode = lazy.nil? ? cfg.lazy : validate_lazy!(lazy)
15
16
  @batch_size = redis_options.delete(:batch_size) || cfg.batch_size
16
- size = pool.nil? ? cfg.pool : pool
17
17
  url = extract_url!(redis_options, cfg)
18
+ # Captured before ConnectionRouter consumes the merged options, so
19
+ # fail-closed batches can report the retry budget (Emb::ServerError).
18
20
  redis_options = merged_redis_options(redis_options, cfg, url)
19
-
20
- @pool = ConnectionPool.new(size: size) do
21
- RedisClient.new(url: url, **redis_options)
22
- end
23
-
21
+ @reconnect_attempts = redis_options.fetch(:reconnect_attempts, cfg.reconnect_attempts)
22
+ @router = ConnectionRouter.new(pool || cfg.pool, instance_urls(url), redis_options)
24
23
  @registry = {}
25
24
  end
26
25
 
27
- def send_command(*args)
28
- return @pool.with { |r| r.call(*args) } unless Emb.debug?
26
+ def send_command(...) = @router.call(...)
29
27
 
30
- start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
31
- result = @pool.with { |r| r.call(*args) }
32
- elapsed = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - start) * 1000
28
+ def pools = @router.pools
33
29
 
34
- $stdout.puts "[EMB] #{args.map(&:inspect).join(' ')} (#{format('%.2f', elapsed)}ms)"
35
-
36
- result
37
- end
30
+ attr_reader :batch_size, :lazy_mode, :reconnect_attempts
38
31
 
39
32
  def [](name)
40
33
  @registry[name] ||= Proxy.new(self, name.to_sym)
41
34
  end
42
35
 
43
- def batch
44
- @batch ||= BatchProxy.new(self)
45
- end
36
+ def lazy? = @lazy_mode != false
46
37
 
47
- def batch?
48
- @batch_enabled
49
- end
38
+ def parallel_batch? = @lazy_mode == :batch
50
39
 
51
40
  # Live view of the server's runtime configuration (CONFIG GET/SET).
52
41
  def config
@@ -103,9 +92,23 @@ module Emb
103
92
 
104
93
  private
105
94
 
95
+ def validate_lazy!(value)
96
+ unless Configuration::LAZY_MODES.include?(value)
97
+ raise ArgumentError, "lazy must be false, :multi, or :batch (got #{value.inspect})"
98
+ end
99
+
100
+ value
101
+ end
102
+
103
+ def instance_urls(url)
104
+ raise ArgumentError, 'url array must not be empty' if url.is_a?(Array) && url.empty?
105
+
106
+ url.nil? ? [nil] : Array(url)
107
+ end
108
+
106
109
  def merged_redis_options(opts, cfg, url)
107
110
  defaults = cfg.to_h
108
- keys = defaults.keys - %i[url pool batch batch_size]
111
+ keys = defaults.keys - %i[url pool lazy batch_size]
109
112
  keys -= %i[host port] if url
110
113
 
111
114
  keys.each do |key|
@@ -3,25 +3,51 @@
3
3
  module Emb
4
4
  class Configuration
5
5
  OPTIONS = %i[
6
- host port url pool batch batch_size driver protocol
6
+ host port url pool lazy batch_size driver protocol
7
7
  connect_timeout read_timeout write_timeout reconnect_attempts
8
8
  ].freeze
9
9
 
10
+ # Execution modes for embed calls. false = eager (default, one EMB round
11
+ # trip per call); :multi = defer and coalesce into EMB.MULTI, serial;
12
+ # :batch = defer and execute chunk shares concurrently. Mutually exclusive
13
+ # by construction.
14
+ LAZY_MODES = [false, :multi, :batch].freeze
15
+
10
16
  attr_accessor(*OPTIONS)
11
17
 
18
+ def lazy=(value)
19
+ unless LAZY_MODES.include?(value)
20
+ raise ArgumentError, "lazy must be false, :multi, or :batch (got #{value.inspect})"
21
+ end
22
+
23
+ @lazy = value
24
+ end
25
+
12
26
  def initialize
13
27
  self.host = 'localhost'
14
28
  self.port = 6379
15
29
  self.url = nil
16
30
  self.pool = 5
17
- self.batch = true
31
+ self.lazy = false
18
32
  self.batch_size = 512
19
33
  self.driver = nil
20
34
  self.protocol = 2
21
35
  self.connect_timeout = nil
22
- self.read_timeout = nil
23
- self.write_timeout = nil
24
- self.reconnect_attempts = 3
36
+ # Read/write timeouts are explicit, NOT nil: nil forwards nothing and
37
+ # redis-client silently applies its 1.0s default, which makes 512-pair
38
+ # EMB.MULTI batches fail under load. 10s covers worst-case inference on
39
+ # shared CPUs; scale up if you raise batch_size.
40
+ self.read_timeout = 10
41
+ self.write_timeout = 10
42
+ # 0 = default: a failing batch fails closed after one attempt and raises
43
+ # Emb::ServerError. Set > 0 to opt into bounded re-sends: redis-client
44
+ # retries connection/protocol failures (never read timeouts) up to that
45
+ # many extra times — EMB.MULTI is not idempotent, so each re-send
46
+ # duplicates inference — and the batch still terminates in
47
+ # Emb::ServerError. An Array of per-retry delays is also accepted (one
48
+ # retry per entry). Operation errors (server error replies) are never
49
+ # retried.
50
+ self.reconnect_attempts = 0
25
51
  end
26
52
 
27
53
  def to_h
@@ -0,0 +1,72 @@
1
+ # frozen_string_literal: true
2
+
3
+ require_relative 'round_robin_pool'
4
+ require 'redis_client'
5
+
6
+ module Emb
7
+ # Owns a client's instance fan-out: one RoundRobinPool per configured url,
8
+ # instance-level round-robin on top of each pool's connection rotation, and
9
+ # the pre-send retry across instances.
10
+ #
11
+ # Retry is deliberately limited to PRESEND_CONNECTION_ERROR: a connection
12
+ # that was never established means the command was not written, so moving it
13
+ # to another instance is safe. Errors after a command may have been sent
14
+ # (timeouts, mid-flight connection loss) are never re-dispatched — retrying
15
+ # them would duplicate inference on the server.
16
+ class ConnectionRouter
17
+ # Captured at load so `rescue` keeps working when tests or apps replace the
18
+ # RedisClient constant (stub_const) — a runtime constant lookup in the
19
+ # rescue clause would otherwise resolve against the replacement.
20
+ PRESEND_CONNECTION_ERROR = RedisClient::CannotConnectError
21
+
22
+ attr_reader :pools
23
+
24
+ def initialize(size, urls, redis_options)
25
+ @pools = urls.map do |url|
26
+ RoundRobinPool.new(size) do
27
+ RedisClient.new(url: url, **redis_options)
28
+ end
29
+ end
30
+ @next_instance = 0
31
+ @instance_mutex = Mutex.new
32
+ end
33
+
34
+ def call(*args)
35
+ attempts = @pools.size
36
+ idx = pick_instance
37
+ last_error = nil
38
+ attempts.times do
39
+ return perform_command(@pools[idx], args)
40
+ rescue PRESEND_CONNECTION_ERROR => e
41
+ last_error = e
42
+ idx = (idx + 1) % @pools.size if @pools.size > 1
43
+ end
44
+ raise last_error
45
+ end
46
+
47
+ private
48
+
49
+ def pick_instance
50
+ return 0 if @pools.size == 1
51
+
52
+ @instance_mutex.synchronize do
53
+ idx = @next_instance % @pools.size
54
+ @next_instance += 1
55
+ idx
56
+ end
57
+ end
58
+
59
+ def perform_command(pool, args)
60
+ debug = Emb.debug?
61
+ started = Process.clock_gettime(Process::CLOCK_MONOTONIC) if debug
62
+ result = pool.with { |r| r.call(*args) }
63
+ log_command(args, started) if debug
64
+ result
65
+ end
66
+
67
+ def log_command(args, started)
68
+ elapsed = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000
69
+ $stdout.puts "[EMB] #{args.map(&:inspect).join(' ')} (#{format('%.2f', elapsed)}ms)"
70
+ end
71
+ end
72
+ end
data/lib/emb/errors.rb ADDED
@@ -0,0 +1,17 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Emb
4
+ # Raised when an EMB.MULTI batch fails — after redis-client exhausted its
5
+ # configured re-sends (`reconnect_attempts`) on a transient error, or on the
6
+ # first attempt for an operation error (server error reply) or the default
7
+ # `reconnect_attempts: 0` configuration. The underlying redis error is
8
+ # preserved as `cause`; `attempts` counts the wire sends that were made.
9
+ class ServerError < StandardError
10
+ attr_reader :attempts
11
+
12
+ def initialize(message, attempts: 1)
13
+ super(message)
14
+ @attempts = attempts
15
+ end
16
+ end
17
+ end
@@ -0,0 +1,16 @@
1
+ # frozen_string_literal: true
2
+
3
+ require 'batch_loader'
4
+ require_relative 'batch_scope'
5
+
6
+ module Emb
7
+ # Job middleware: clears the per-thread batch scope after each job execution,
8
+ # even when the job raises. Registered as a Sidekiq/Shoryuken server
9
+ # middleware (positional args differ per framework and are unused) and used by
10
+ # the ActiveJob perform callback.
11
+ class JobMiddleware
12
+ def call(*_args, &)
13
+ BatchScope.wrap(&)
14
+ end
15
+ end
16
+ end
@@ -1,20 +1,18 @@
1
1
  # frozen_string_literal: true
2
2
 
3
3
  require 'batch_loader'
4
+ require_relative 'batch_scope'
4
5
 
5
6
  module Emb
6
- # Rack middleware that clears the per-thread batch scope at the end of each
7
- # request, bounding cache growth in long-lived request threads. The scope is
8
- # cleared even when the app raises; a fresh scope starts with the next request.
7
+ # Rack middleware: clears the per-thread batch scope after each request,
8
+ # even when the app raises.
9
9
  class Middleware
10
10
  def initialize(app)
11
11
  @app = app
12
12
  end
13
13
 
14
14
  def call(env)
15
- @app.call(env)
16
- ensure
17
- BatchLoader::Executor.clear_current
15
+ BatchScope.wrap { @app.call(env) }
18
16
  end
19
17
  end
20
18
  end
data/lib/emb/proxy.rb CHANGED
@@ -10,7 +10,7 @@ module Emb
10
10
  end
11
11
 
12
12
  def [](text, *texts)
13
- return Emb.build_batch_loader(@client, @name, texts.empty? ? text : [text, *texts]) if @client.batch?
13
+ return Emb.build_batch_loader(@client, @name, texts.empty? ? text : [text, *texts]) if @client.lazy?
14
14
 
15
15
  set = Array(@client.send_command('EMB', @name.to_s, text, *texts))
16
16
  result = set.map { |entry| entry.unpack('e*') }
@@ -0,0 +1,49 @@
1
+ # frozen_string_literal: true
2
+
3
+ if defined?(Rails::Railtie)
4
+ require_relative '../emb'
5
+
6
+ module Emb
7
+ # Rails integration, loaded by emb.rb when Rails is present.
8
+ #
9
+ # config.emb.middleware = false # skip Emb::Middleware in the stack
10
+ # config.emb.job_middleware = false # skip job-scope protection
11
+ class Railtie < Rails::Railtie
12
+ config.emb = ActiveSupport::OrderedOptions.new
13
+ config.emb.middleware = true
14
+ config.emb.job_middleware = true
15
+
16
+ initializer 'emb.middleware' do |app|
17
+ next if config.emb.middleware == false
18
+ next if app.middleware.include?(Emb::Middleware)
19
+
20
+ app.middleware.use Emb::Middleware
21
+ end
22
+
23
+ config.after_initialize do
24
+ next if config.emb.job_middleware == false
25
+
26
+ ActiveSupport.on_load(:active_job) do |base|
27
+ base.around_perform { |_job, block| Emb::BatchScope.wrap { block.call } }
28
+ end
29
+
30
+ if defined?(Sidekiq)
31
+ Sidekiq.configure_server do |sidekiq_config|
32
+ sidekiq_config.server_middleware { |chain| chain.add Emb::JobMiddleware }
33
+ end
34
+ end
35
+
36
+ if defined?(Sidekiq::Testing)
37
+ # Testing modes use their own middleware chain.
38
+ Sidekiq::Testing.server_middleware { |chain| chain.add Emb::JobMiddleware }
39
+ end
40
+
41
+ if defined?(Shoryuken)
42
+ Shoryuken.configure_server do |shoryuken_config|
43
+ shoryuken_config.server_middleware { |chain| chain.add Emb::JobMiddleware }
44
+ end
45
+ end
46
+ end
47
+ end
48
+ end
49
+ end
@@ -0,0 +1,100 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Emb
4
+ # A thread-safe pool of N RedisClient connections with round-robin selection.
5
+ # Behind connection-level load balancers (AWS Service Connect, an NLB, or an
6
+ # Envoy TCP proxy) each keep-alive connection is pinned to one upstream
7
+ # instance, so rotating commands across the pool spreads traffic across every
8
+ # instance — even single-threaded at zero concurrency. Connections are created
9
+ # up front but connect lazily on first use.
10
+ #
11
+ # Two behaviors deliberately match the connection_pool gem it replaces: a
12
+ # nested `with` from the same thread re-enters the held connection, and after
13
+ # `fork` (Puma preload_app, unicorn, resque) the pool closes inherited sockets
14
+ # and rebuilds its mutexes in the child so parent and child never share a
15
+ # connection.
16
+ class RoundRobinPool
17
+ # Pools are tracked only to reset them in forked children. WeakMap so a
18
+ # pool is reclaimed together with its client.
19
+ INSTANCES = Process.respond_to?(:fork) ? ObjectSpace::WeakMap.new : nil
20
+ private_constant :INSTANCES
21
+
22
+ THREAD_KEY = :emb_round_robin_pool_held
23
+ private_constant :THREAD_KEY
24
+
25
+ attr_reader :size, :connections
26
+
27
+ def self.after_fork
28
+ INSTANCES&.each_value(&:reload_after_fork!)
29
+ end
30
+
31
+ def initialize(size, &)
32
+ raise ArgumentError, "pool size must be >= 1 (got #{size})" if size < 1
33
+
34
+ @size = size
35
+ @connections = Array.new(size, &)
36
+ @locks = Array.new(size) { Mutex.new }
37
+ @next = 0
38
+ @index_mutex = Mutex.new
39
+ INSTANCES&.[]=(self, self)
40
+ end
41
+
42
+ # Yields the next connection in rotation order. Safe from multiple threads:
43
+ # up to `size` commands run in parallel, each on its own connection. A
44
+ # nested `with` from the same thread re-enters the connection this pool
45
+ # already holds without re-locking; other pools are unaffected.
46
+ def with(&)
47
+ held = Thread.current[THREAD_KEY]
48
+ if held&.key?(self)
49
+ yield @connections[held[self]]
50
+ else
51
+ take(held, &)
52
+ end
53
+ end
54
+
55
+ # Child side of after_fork(): drop inherited sockets and sync state.
56
+ def reload_after_fork!
57
+ @connections.each { |conn| conn.close if conn.respond_to?(:close) }
58
+ @locks = Array.new(@size) { Mutex.new }
59
+ @next = 0
60
+ @index_mutex = Mutex.new
61
+ end
62
+
63
+ if Process.respond_to?(:fork)
64
+ # Hooks Process._fork (MRI 3.1+) so registered pools reset in the child.
65
+ module ForkTracker
66
+ def _fork
67
+ pid = super
68
+ RoundRobinPool.after_fork if pid.zero?
69
+ pid
70
+ end
71
+ end
72
+ Process.singleton_class.prepend(ForkTracker)
73
+ end
74
+
75
+ private
76
+
77
+ # Acquires the next connection, records it as held by this thread/pool, and
78
+ # releases both on exit — even when the block raises.
79
+ def take(held)
80
+ idx, connection = pick
81
+ held ||= {}
82
+ Thread.current[THREAD_KEY] = held
83
+ held[self] = idx
84
+ begin
85
+ @locks[idx].synchronize { yield connection }
86
+ ensure
87
+ held.delete(self)
88
+ Thread.current[THREAD_KEY] = nil if held.empty?
89
+ end
90
+ end
91
+
92
+ def pick
93
+ @index_mutex.synchronize do
94
+ idx = @next % @size
95
+ @next += 1
96
+ [idx, @connections[idx]]
97
+ end
98
+ end
99
+ end
100
+ end
data/lib/emb.rb CHANGED
@@ -2,13 +2,17 @@
2
2
 
3
3
  require_relative 'emb/version'
4
4
  require_relative 'emb/configuration'
5
+ require_relative 'emb/errors'
5
6
  require_relative 'emb/commands'
6
7
  require_relative 'emb/runtime_config'
7
8
  require_relative 'emb/client'
8
9
  require_relative 'emb/proxy'
10
+ require_relative 'emb/batch_scope'
11
+ require_relative 'emb/middleware'
12
+ require_relative 'emb/job_middleware'
9
13
  require_relative 'emb/multi'
10
14
  require_relative 'emb/batch'
11
- require_relative 'emb/middleware'
15
+ require_relative 'emb/railtie' if defined?(Rails::Railtie)
12
16
 
13
17
  module Emb
14
18
  class << self
@@ -28,7 +32,6 @@ module Emb
28
32
  end
29
33
 
30
34
  def [](name) = default_client[name]
31
- def batch = default_client.batch
32
35
  def models = default_client.models
33
36
  def info(name) = default_client.info(name)
34
37
  def stats = default_client.stats
metadata CHANGED
@@ -1,14 +1,14 @@
1
1
  --- !ruby/object:Gem::Specification
2
2
  name: emb
3
3
  version: !ruby/object:Gem::Version
4
- version: 0.2.6
4
+ version: 0.4.0.pre1
5
5
  platform: ruby
6
6
  authors:
7
7
  - elcuervo
8
8
  autorequire:
9
9
  bindir: bin
10
10
  cert_chain: []
11
- date: 2026-09-01 00:00:00.000000000 Z
11
+ date: 2026-09-04 00:00:00.000000000 Z
12
12
  dependencies:
13
13
  - !ruby/object:Gem::Dependency
14
14
  name: batch-loader
@@ -24,20 +24,6 @@ dependencies:
24
24
  - - "~>"
25
25
  - !ruby/object:Gem::Version
26
26
  version: '2.0'
27
- - !ruby/object:Gem::Dependency
28
- name: connection_pool
29
- requirement: !ruby/object:Gem::Requirement
30
- requirements:
31
- - - "~>"
32
- - !ruby/object:Gem::Version
33
- version: '2.5'
34
- type: :runtime
35
- prerelease: false
36
- version_requirements: !ruby/object:Gem::Requirement
37
- requirements:
38
- - - "~>"
39
- - !ruby/object:Gem::Version
40
- version: '2.5'
41
27
  - !ruby/object:Gem::Dependency
42
28
  name: redis-client
43
29
  requirement: !ruby/object:Gem::Requirement
@@ -63,12 +49,19 @@ files:
63
49
  - README.md
64
50
  - lib/emb.rb
65
51
  - lib/emb/batch.rb
52
+ - lib/emb/batch_dispatch.rb
53
+ - lib/emb/batch_scope.rb
66
54
  - lib/emb/client.rb
67
55
  - lib/emb/commands.rb
68
56
  - lib/emb/configuration.rb
57
+ - lib/emb/connection_router.rb
58
+ - lib/emb/errors.rb
59
+ - lib/emb/job_middleware.rb
69
60
  - lib/emb/middleware.rb
70
61
  - lib/emb/multi.rb
71
62
  - lib/emb/proxy.rb
63
+ - lib/emb/railtie.rb
64
+ - lib/emb/round_robin_pool.rb
72
65
  - lib/emb/runtime_config.rb
73
66
  - lib/emb/version.rb
74
67
  homepage: https://github.com/elcuervo/emb