From 8e5ef1fb31e9d63b735c089b21cbc82c50acce46 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 16 Jul 2026 11:30:42 -0400 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20perf:=20Halve=20Per-Delta=20Redis?= =?UTF-8?q?=20Round=20Trips=20in=20Resumable=20Streams=20(#14313)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `RedisEventTransport.emitChunk` awaited `INCR` (sequence allocation) and then `PUBLISH` as two sequential round trips, per streamed delta. Fold both into one Lua `EVAL` that allocates the sequence and publishes server-side. Measured 50.6% reduction on the emit path (0.219ms -> 0.108ms per delta, 2000 deltas, loopback Redis). The saving multiplies by the token count of every response, and scales with RTT: on setups where Redis sits behind a network boundary (WSL2 loopback, cross-host, cross-AZ) at ~1-2ms/RTT this is ~0.5-1s on a 500-token response. The sequence is spliced into the payload server-side rather than round-tripped through cjson, which would coerce empty arrays to objects and alter float precision. The channel is passed as ARGV rather than KEYS: ioredis applies `keyPrefix` to EVAL keys but never to a pub/sub channel, so keying it would publish to a prefixed channel no subscriber listens on. PUBLISH is broadcast cluster-wide rather than slot-routed, so it needs no key for Cluster correctness. Also parallelize `getResumeState`'s three independent job-store reads (`getContentParts` / `getRunSteps` / `peekSteers`), collapsing 3 round trips into 1 on every resume. Safe despite `readCachedGraph`'s cache-drop side effect: each call catches its own unusable-graph throw and falls back to durable reconstruction, so ordering cannot change the result. `createMockPublisher` gains an `eval` that delegates to its own incr/publish mocks, keeping the existing error-propagation tests meaningful now that sequence allocation and publish are one operation. --- .../api/src/stream/GenerationJobManager.ts | 12 ++- ...sEventTransport.stream_integration.spec.ts | 84 ++++++++++++++++- .../src/stream/__tests__/helpers/publisher.ts | 31 +++++- .../implementations/RedisEventTransport.ts | 94 +++++++++++++------ 4 files changed, 186 insertions(+), 35 deletions(-) diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index f6e304d989..46f819effd 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -1691,9 +1691,15 @@ class GenerationJobManagerClass { return null; } - const result = await this.jobStore.getContentParts(streamId); + /** Independent reads (streamId-only): parallel to collapse 3 Redis round trips into 1. + * Safe despite readCachedGraph's cache-drop side effect — each call catches its own + * unusable-graph throw and falls back to reconstruction, so ordering cannot change the result. */ + const [result, runSteps, queuedSteers] = await Promise.all([ + this.jobStore.getContentParts(streamId), + this.jobStore.getRunSteps(streamId), + this.jobStore.peekSteers(streamId), + ]); const aggregatedContent = result?.content ?? []; - const runSteps = await this.jobStore.getRunSteps(streamId); let titleEvent: t.ResumeState['titleEvent']; if (jobData.titleEvent) { try { @@ -1733,7 +1739,7 @@ class GenerationJobManagerClass { } /** Steers still queued (not yet injected); injected ones are already in aggregatedContent. */ - const pendingSteers = (await this.jobStore.peekSteers(streamId)).map(toPendingSteer); + const pendingSteers = queuedSteers.map(toPendingSteer); logger.debug(`[GenerationJobManager] getResumeState:`, { streamId, diff --git a/packages/api/src/stream/__tests__/RedisEventTransport.stream_integration.spec.ts b/packages/api/src/stream/__tests__/RedisEventTransport.stream_integration.spec.ts index e254cf66ba..f20da80e91 100644 --- a/packages/api/src/stream/__tests__/RedisEventTransport.stream_integration.spec.ts +++ b/packages/api/src/stream/__tests__/RedisEventTransport.stream_integration.spec.ts @@ -1,5 +1,5 @@ -import type { Redis, Cluster } from 'ioredis'; import { logger } from '@librechat/data-schemas'; +import type { Redis, Cluster } from 'ioredis'; import { createMockPublisher } from './helpers/publisher'; logger.silent = true; @@ -182,6 +182,88 @@ describe('RedisEventTransport Integration Tests', () => { }); }); + describe('Payload fidelity through server-side seq splicing', () => { + /** The seq is spliced into the payload by Lua rather than encoded with it, so the + * fragments either side must reassemble to exactly what JSON.stringify would emit. + * Shapes here are the ones a naive cjson round-trip would corrupt. */ + test('should round-trip payload shapes that cjson would coerce', async () => { + if (!ioredisClient) { + console.warn('Redis not available, skipping test'); + return; + } + + const { RedisEventTransport } = await import('../implementations/RedisEventTransport'); + + const subscriber = (ioredisClient as Redis).duplicate(); + const transport = new RedisEventTransport(ioredisClient, subscriber); + + const streamId = `payload-fidelity-${Date.now()}`; + const received: unknown[] = []; + + transport.subscribe(streamId, { + onChunk: (event) => received.push(event), + }); + + await new Promise((resolve) => setTimeout(resolve, 100)); + + const payloads: unknown[] = [ + { empty: [], nestedEmpty: { inner: [] } }, + { float: 0.1234567890123, negative: -273.15, zero: 0 }, + { unicode: 'héllo 🌍 "quoted" \\ backslash\n newline' }, + { nullish: null, emptyString: '', emptyObject: {} }, + { text: 'ordinary delta' }, + ]; + + for (const payload of payloads) { + await transport.emitChunk(streamId, payload); + } + + await new Promise((resolve) => setTimeout(resolve, 300)); + + expect(received).toEqual(payloads); + + transport.destroy(); + subscriber.disconnect(); + }); + + test('should assign 0-indexed sequences and set a TTL on the counter only once', async () => { + if (!ioredisClient) { + console.warn('Redis not available, skipping test'); + return; + } + + const { RedisEventTransport } = await import('../implementations/RedisEventTransport'); + + const subscriber = (ioredisClient as Redis).duplicate(); + const transport = new RedisEventTransport(ioredisClient, subscriber); + + const streamId = `seq-alloc-${Date.now()}`; + /** Bare key: the client applies REDIS_KEY_PREFIX itself, and ioredis prefixes EVAL keys + * the same way, so both sides land on the same prefixed key. */ + const seqKey = `stream:{${streamId}}:seq`; + + transport.subscribe(streamId, { onChunk: () => {} }); + await new Promise((resolve) => setTimeout(resolve, 100)); + + await transport.emitChunk(streamId, { index: 0 }); + /** First INCR arms the TTL; it must never be refreshed, or a long stream could + * have its counter reset mid-generation. */ + const ttlAfterFirst = await (ioredisClient as Redis).ttl(seqKey); + expect(ttlAfterFirst).toBeGreaterThan(0); + + for (let i = 1; i < 5; i++) { + await transport.emitChunk(streamId, { index: i }); + } + + /** Counter is 1-based in Redis; seq is 0-based, so 5 emits => counter 5, last seq 4. */ + expect(await (ioredisClient as Redis).get(seqKey)).toBe('5'); + expect(await (ioredisClient as Redis).ttl(seqKey)).toBeLessThanOrEqual(ttlAfterFirst); + + transport.destroy(); + subscriber.disconnect(); + }); + }); + describe('Sequential Event Ordering', () => { test('should maintain strict order when emitChunk is awaited', async () => { if (!ioredisClient) { diff --git a/packages/api/src/stream/__tests__/helpers/publisher.ts b/packages/api/src/stream/__tests__/helpers/publisher.ts index 3e1799d373..be15c67535 100644 --- a/packages/api/src/stream/__tests__/helpers/publisher.ts +++ b/packages/api/src/stream/__tests__/helpers/publisher.ts @@ -4,12 +4,13 @@ export interface MockPublisher { expire: jest.Mock; get: jest.Mock; del: jest.Mock; + eval: jest.Mock; } /** Mock publisher with Redis command simulation for atomic sequence counters */ export function createMockPublisher(): MockPublisher { const counters = new Map(); - return { + const publisher: MockPublisher = { publish: jest.fn().mockResolvedValue(1), incr: jest.fn().mockImplementation((key: string) => { const current = (counters.get(key) ?? 0) + 1; @@ -27,5 +28,33 @@ export function createMockPublisher(): MockPublisher { } return Promise.resolve(keys.length); }), + eval: jest.fn(), }; + + /** + * Stands in for PUBLISH_SEQ_LUA, which allocates the sequence and publishes in one server-side + * round trip. Delegates to the incr/publish mocks rather than reimplementing them, so a test + * can still fail either half independently and observe the ordering between them. + */ + publisher.eval.mockImplementation( + async ( + _script: string, + _numKeys: number, + seqKey: string, + channel: string, + prefix: string, + suffix: string, + ttlSeconds: string, + ) => { + const val = (await publisher.incr(seqKey)) as number; + if (val === 1) { + await publisher.expire(seqKey, Number(ttlSeconds)); + } + const seq = val - 1; + await publisher.publish(channel, `${prefix}${seq}${suffix}`); + return seq; + }, + ); + + return publisher; } diff --git a/packages/api/src/stream/implementations/RedisEventTransport.ts b/packages/api/src/stream/implementations/RedisEventTransport.ts index d0eaebad47..eb332ece04 100644 --- a/packages/api/src/stream/implementations/RedisEventTransport.ts +++ b/packages/api/src/stream/implementations/RedisEventTransport.ts @@ -50,6 +50,32 @@ interface ReorderBuffer { flushTimeout: ReturnType | null; } +/** + * Allocate a sequence number and publish the event in a single round trip. + * + * The payload is spliced server-side rather than round-tripped through `cjson`: decoding and + * re-encoding arbitrary event data would coerce empty arrays to objects and alter float + * precision. The caller pre-serializes everything around the seq, so this only concatenates. + * + * The TTL is set once on the first INCR and never refreshed, so an active stream cannot have + * its counter reset mid-generation. + * + * The channel is passed as ARGV, not KEYS: ioredis applies `keyPrefix` to EVAL keys but never + * to a pub/sub channel, so keying it here would publish to a prefixed channel that no + * subscriber listens on. PUBLISH is broadcast cluster-wide rather than slot-routed, so it does + * not need to be a key for Cluster correctness. + * + * KEYS: [sequence] + * ARGV: [channel, payloadPrefix, payloadSuffix, sequenceTtlSeconds] + * RETURNS: the 0-indexed seq assigned to this event + */ +const PUBLISH_SEQ_LUA = + 'local val = redis.call("INCR", KEYS[1]) ' + + 'if val == 1 then redis.call("EXPIRE", KEYS[1], tonumber(ARGV[4])) end ' + + 'local seq = val - 1 ' + + 'redis.call("PUBLISH", ARGV[1], ARGV[2] .. string.format("%d", seq) .. ARGV[3]) ' + + 'return seq'; + /** Max time (ms) to wait for out-of-order messages before force-flushing */ const REORDER_TIMEOUT_MS = 500; /** Max messages to buffer before force-flushing (prevents memory issues) */ @@ -126,21 +152,40 @@ export class RedisEventTransport implements IEventTransport { private static readonly SEQUENCE_TTL_SECONDS = 86400; /** - * Get next sequence number for a stream (0-indexed, backed by Redis INCR). - * A 24-hour TTL is set on the first INCR only (val === 1) as a safety net - * for orphaned keys from crashed processes. It is never refreshed, so an - * active stream cannot have its counter reset mid-generation. - * Keys are also deleted explicitly by cleanup() on normal stream teardown. + * Split a seq-less message into the JSON fragments surrounding its `seq`, so the sequence + * can be spliced in by {@link PUBLISH_SEQ_LUA} without re-encoding the payload. + * + * Omitting a field (e.g. `data: undefined`) yields an empty tail, matching what + * `JSON.stringify` would have dropped from the whole-object encoding. */ - private async getNextSequence(streamId: string): Promise { - const key = KEYS.sequence(streamId); - const val = await this.publisher.incr(key); - if (val === 1) { - this.publisher.expire(key, RedisEventTransport.SEQUENCE_TTL_SECONDS).catch((err) => { - logger.warn(`[RedisEventTransport] Failed to set TTL on sequence key ${key}:`, err); - }); - } - return val - 1; + private static buildPayloadParts(message: Omit): [string, string] { + const { type, ...rest } = message; + const encodedRest = JSON.stringify(rest); + const inner = encodedRest.slice(1, -1); + return [`{"type":${JSON.stringify(type)},"seq":`, inner.length > 0 ? `,${inner}}` : '}']; + } + + /** + * Allocate a sequence number and publish, in one Redis round trip. + * + * Keys are deleted explicitly by cleanup() on normal stream teardown; the TTL is a safety + * net for orphaned keys from crashed processes. + */ + private async publishWithSequence( + streamId: string, + message: Omit, + ): Promise { + const [prefix, suffix] = RedisEventTransport.buildPayloadParts(message); + const seq = await this.publisher.eval( + PUBLISH_SEQ_LUA, + 1, + KEYS.sequence(streamId), + CHANNELS.events(streamId), + prefix, + suffix, + String(RedisEventTransport.SEQUENCE_TTL_SECONDS), + ); + return seq as number; } /** Reset subscriber reorder buffer state to initial values */ @@ -516,15 +561,12 @@ export class RedisEventTransport implements IEventTransport { * Publish a chunk event to all subscribers across all instances. * Includes sequence number for ordered delivery in Redis Cluster mode. * - * Performance: each emit requires two sequential Redis round-trips (INCR + PUBLISH). - * This is the unavoidable cost of cross-replica sequence coordination. + * Performance: sequence allocation and publish share one round trip. This runs per streamed + * delta, so the saved round trip is multiplied by the token count of every response. */ async emitChunk(streamId: string, event: unknown): Promise { try { - const channel = CHANNELS.events(streamId); - const seq = await this.getNextSequence(streamId); - const message: PubSubMessage = { type: EventTypes.CHUNK, seq, data: event }; - await this.publisher.publish(channel, JSON.stringify(message)); + await this.publishWithSequence(streamId, { type: EventTypes.CHUNK, data: event }); } catch (err) { logger.error(`[RedisEventTransport] Failed to publish chunk:`, err); } @@ -535,12 +577,8 @@ export class RedisEventTransport implements IEventTransport { * Includes sequence number to ensure delivery after all chunks. */ async emitDone(streamId: string, event: unknown): Promise { - const channel = CHANNELS.events(streamId); - const seq = await this.getNextSequence(streamId); - const message: PubSubMessage = { type: EventTypes.DONE, seq, data: event }; - try { - await this.publisher.publish(channel, JSON.stringify(message)); + await this.publishWithSequence(streamId, { type: EventTypes.DONE, data: event }); } catch (err) { logger.error(`[RedisEventTransport] Failed to publish done:`, err); throw err; @@ -552,12 +590,8 @@ export class RedisEventTransport implements IEventTransport { * Includes sequence number to ensure delivery after all chunks. */ async emitError(streamId: string, error: string): Promise { - const channel = CHANNELS.events(streamId); - const seq = await this.getNextSequence(streamId); - const message: PubSubMessage = { type: EventTypes.ERROR, seq, error }; - try { - await this.publisher.publish(channel, JSON.stringify(message)); + await this.publishWithSequence(streamId, { type: EventTypes.ERROR, error }); } catch (err) { logger.error(`[RedisEventTransport] Failed to publish error:`, err); throw err;