fix: ensure stable WebSocket object identity across lifecycle hooks#1729
fix: ensure stable WebSocket object identity across lifecycle hooks#1729raunak-rpm wants to merge 1 commit intoelysiajs:mainfrom
Conversation
WalkthroughIntroduces a per-WebSocket WeakMap cache and helper to ensure a single ElysiaWS instance is reused per connection across open/message/drain/close hooks, and adds tests validating identity, per-connection state, and per-message body updates. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant BunWS as "Bun WebSocket"
participant Cache as "getElysiaWS / elysiaWSCache"
participant ElysiaWS
Client->>BunWS: Upgrade -> open(ws)
BunWS->>Cache: getElysiaWS(ws)
Cache-->>ElysiaWS: create/store ElysiaWS(ws, context)
Cache-->>BunWS: return ElysiaWS
BunWS->>ElysiaWS: open handler(ElysiaWS)
Client->>BunWS: send message (body)
BunWS->>Cache: getElysiaWS(ws, body)
Cache-->>ElysiaWS: return existing ElysiaWS (update body)
BunWS->>ElysiaWS: message handler(ElysiaWS, body)
Client->>BunWS: close
BunWS->>Cache: getElysiaWS(ws)
Cache-->>ElysiaWS: return existing ElysiaWS
BunWS->>ElysiaWS: close handler(ElysiaWS)
BunWS->>Cache: optional cleanup (WeakMap allows GC)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@src/adapter/bun/index.ts`:
- Around line 505-527: The WeakMap and cached wrapper use ElysiaWS with its
default generic parameters which narrows body to never; change the type to a
widened generic (e.g., ElysiaWS<any, any>) so the cache and returned value from
getElysiaWS are assignable to the open/drain/close hooks. Update the declaration
of elysiaWSCache (WeakMap<ServerWebSocket<any>, ElysiaWS>) and any related
types/usages in getElysiaWS to use ElysiaWS<any, any>, ensuring the cached.body
assignment and passing cached into options.open/options.drain/options.close
compile cleanly.
🧹 Nitpick comments (1)
src/adapter/bun/index.ts (1)
505-527: Consider simplifying: the WeakMap only ever holds one entry per closureThe
elysiaWSCacheWeakMap andgetElysiaWSare created inside the per-upgrade route handler. Since each upgrade creates a fresh closure, the WeakMap will only ever cache oneServerWebSocket → ElysiaWSmapping. A plainletvariable would achieve the same result with less overhead:let cachedElysiaWS: ElysiaWS<any, any> | undefined const getElysiaWS = (ws: ServerWebSocket<any>, body?: unknown) => { if (!cachedElysiaWS) { cachedElysiaWS = new ElysiaWS(ws, context as any, body) } else if (body !== undefined) { cachedElysiaWS.body = body } return cachedElysiaWS }That said, the current approach is correct and the WeakMap makes the "per-socket cache" intent very clear, so this is entirely optional.
commit: |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@src/adapter/bun/index.ts`:
- Around line 522-524: The current logic lets cached.body persist across hooks
because the guard "else if (body !== undefined) { cached.body = body }" skips
resetting on non-message events; to restore previous behavior, ensure
cached.body is explicitly cleared when building a fresh ElysiaWS for hooks that
shouldn't inherit the last message (e.g., in the drain/close paths or when
creating a new ElysiaWS instance). Update the code that constructs
ElysiaWS/cached (referencing ElysiaWS and cached.body) to set cached.body =
undefined for non-message hooks or always initialize a new cached.body when
creating the hook-specific ElysiaWS so message values do not leak into
subsequent drain/close handlers.
🧹 Nitpick comments (3)
src/adapter/bun/index.ts (1)
505-527: The WeakMap will only ever hold one entry per connection — a plain variable is simpler.The
elysiaWSCacheandgetElysiaWSare created inside the per-upgrade route handler. Each invocation (i.e. each connection) gets its own WeakMap instance, which will only ever store a single key-value pair. A simple captured variable achieves the same result without the overhead of a WeakMap lookup:♻️ Suggested simplification
- const elysiaWSCache = new WeakMap< - ServerWebSocket<any>, - ElysiaWS<any, any> - >() - - const getElysiaWS = ( - ws: ServerWebSocket<any>, - body?: unknown - ): ElysiaWS<any, any> => { - let cached = elysiaWSCache.get(ws) - - if (!cached) { - cached = new ElysiaWS(ws, context as any, body) - elysiaWSCache.set(ws, cached) - } else if (body !== undefined) { - cached.body = body - } - - return cached - } + let cachedElysiaWS: ElysiaWS<any, any> | undefined + + const getElysiaWS = ( + ws: ServerWebSocket<any>, + body?: unknown + ): ElysiaWS<any, any> => { + if (!cachedElysiaWS) { + cachedElysiaWS = new ElysiaWS(ws, context as any, body) + } else if (body !== undefined) { + cachedElysiaWS.body = body + } + + return cachedElysiaWS + }If future changes move the cache to a broader scope (e.g. shared across routes), the WeakMap would make sense, but in the current per-connection closure it's unnecessary indirection.
test/ws/identity.test.ts (2)
35-37:app.stop()is not awaited in any test — may cause port/resource leaks between tests.
BunAdapter.stopisasync. Not awaiting it means the server may not be fully shut down before the next test starts. While.listen(0)mitigates port conflicts, unawaited cleanup can cause flaky failures under parallel execution or resource pressure. This applies to all 8 tests in this file.♻️ Suggested fix (apply to all tests)
- await wsClosed(ws) - app.stop() + await wsClosed(ws) + await app.stop()
262-295: Assertion inside the WebSocket handler may silently swallow failures.Line 273 (
expect(ws).toBe(stableWs)) throws inside the server-sidemessagehandler if it fails. The adapter'stry/catchwould catch it and route it tohandleErrors, but since there's no custom error handler in this test, the failure won't propagate to the test runner. Consider moving identity assertions outside the handler (as other tests do) for reliable failure reporting.♻️ Suggested approach
.ws('/ws', { open(ws) { stableWs = ws }, message(ws, message) { bodies.push(ws.body) - expect(ws).toBe(stableWs) + // capture for external assertion ws.send('ok') } }) .listen(0) ... expect(bodies[0]).toBe('first') expect(bodies[1]).toBe('second') + expect(stableWs).toBeDefined() + // Identity already guaranteed by other tests; body update is the focus here
| } else if (body !== undefined) { | ||
| cached.body = body | ||
| } |
There was a problem hiding this comment.
Subtle behavioral change: body now persists across hooks.
Previously, each hook created a fresh ElysiaWS with body = undefined (except message which passed the parsed message). Now, after a message hook fires, body retains the last message value in subsequent drain/close calls since the body !== undefined guard skips the reset.
This is unlikely to cause issues in practice (close/drain handlers rarely inspect body), but it's a change from the prior behavior worth being aware of.
🤖 Prompt for AI Agents
In `@src/adapter/bun/index.ts` around lines 522 - 524, The current logic lets
cached.body persist across hooks because the guard "else if (body !== undefined)
{ cached.body = body }" skips resetting on non-message events; to restore
previous behavior, ensure cached.body is explicitly cleared when building a
fresh ElysiaWS for hooks that shouldn't inherit the last message (e.g., in the
drain/close paths or when creating a new ElysiaWS instance). Update the code
that constructs ElysiaWS/cached (referencing ElysiaWS and cached.body) to set
cached.body = undefined for non-message hooks or always initialize a new
cached.body when creating the hook-specific ElysiaWS so message values do not
leak into subsequent drain/close handlers.
3 participants
Summary
Fixes #1716
The
wsobject passed to WebSocket lifecycle hooks (open,message,drain,close) was a newElysiaWSwrapper each time, breaking JavaScript identity (===). This made common patterns unreliable:Map/SetcloseRoot Cause
In the Bun adapter (
src/adapter/bun/index.ts), every lifecycle hook callednew ElysiaWS(ws, context), creating a fresh wrapper despite the underlying BunServerWebSocketbeing the same native object across all hooks.Fix
Introduced a per-connection
WeakMap<ServerWebSocket, ElysiaWS>cache with agetElysiaWS()helper. The first hook to fire (typicallyopen) creates and caches the wrapper; subsequent hooks (message,drain,close) reuse the same instance.WeakMapensures automatic garbage collection when the raw socket is collected — no memory leaksmessagehook updatesws.bodyon each call to reflect the current message payloadTests
Added 8 targeted tests in
test/ws/identity.test.ts:openandmessageopenandcloseopen,message, andcloseopenpersist tomessageandcloseSet-based connection tracking with cleanup oncloseMap-based per-socket state accumulationbodyupdates per message while preserving identityAll 1508 tests pass (1500 existing + 8 new), verified stable across multiple runs.
Summary by CodeRabbit