Skip to content

Commit ef0ebe3

Browse files
kettanaitokettanaito
authored
fix(WebSocketClientManager): use localStorage for clients persistence (#2127)
1 parent 74215a3 commit ef0ebe3

5 files changed

Lines changed: 364 additions & 91 deletions

File tree

‎src/browser/setupWorker/stop/createStop.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { devUtils } from '~/core/utils/internal/devUtils'
2+
import { MSW_WEBSOCKET_CLIENTS_KEY } from '~/core/ws/WebSocketClientManager'
23
import { SetupWorkerInternalContext, StopHandler } from '../glossary'
34
import { printStopMessage } from './utils/printStopMessage'
45

@@ -24,6 +25,9 @@ export const createStop = (
2425
context.isMockingEnabled = false
2526
window.clearInterval(context.keepAliveInterval)
2627

28+
// Clear the WebSocket clients from the shared storage.
29+
localStorage.removeItem(MSW_WEBSOCKET_CLIENTS_KEY)
30+
2731
printStopMessage({ quiet: context.startOptions?.quiet })
2832
}
2933
}

‎src/core/ws.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -72,10 +72,12 @@ function createWebSocketLinkHandler(url: Path): WebSocketLink {
7272
typeof url,
7373
)
7474

75-
const clientManager = new WebSocketClientManager(wsBroadcastChannel)
75+
const clientManager = new WebSocketClientManager(wsBroadcastChannel, url)
7676

7777
return {
78-
clients: clientManager.clients,
78+
get clients() {
79+
return clientManager.clients
80+
},
7981
on(event, listener) {
8082
const handler = new WebSocketHandler(url)
8183

‎src/core/ws/WebSocketClientManager.test.ts‎

Lines changed: 43 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,79 +1,73 @@
11
/**
22
* @vitest-environment node-websocket
33
*/
4-
import { randomUUID } from 'node:crypto'
54
import {
65
WebSocketClientConnection,
6+
WebSocketData,
77
WebSocketTransport,
88
} from '@mswjs/interceptors/WebSocket'
99
import {
1010
WebSocketClientManager,
1111
WebSocketBroadcastChannelMessage,
12-
WebSocketRemoteClientConnection,
1312
} from './WebSocketClientManager'
1413

1514
const channel = new BroadcastChannel('test:channel')
1615
vi.spyOn(channel, 'postMessage')
1716

1817
const socket = new WebSocket('ws://localhost')
19-
const transport = {
20-
onOutgoing: vi.fn(),
21-
onIncoming: vi.fn(),
22-
onClose: vi.fn(),
23-
send: vi.fn(),
24-
close: vi.fn(),
25-
} satisfies WebSocketTransport
18+
19+
class TestWebSocketTransport extends EventTarget implements WebSocketTransport {
20+
send(_data: WebSocketData): void {}
21+
close(_code?: number | undefined, _reason?: string | undefined): void {}
22+
}
2623

2724
afterEach(() => {
2825
vi.resetAllMocks()
2926
})
3027

3128
it('adds a client from this runtime to the list of clients', () => {
32-
const manager = new WebSocketClientManager(channel)
33-
const connection = new WebSocketClientConnection(socket, transport)
29+
const manager = new WebSocketClientManager(channel, '*')
30+
const connection = new WebSocketClientConnection(
31+
socket,
32+
new TestWebSocketTransport(),
33+
)
3434

3535
manager.addConnection(connection)
3636

3737
// Must add the client to the list of clients.
3838
expect(Array.from(manager.clients.values())).toEqual([connection])
39-
40-
// Must emit the connection open event to notify other runtimes.
41-
expect(channel.postMessage).toHaveBeenCalledWith({
42-
type: 'connection:open',
43-
payload: {
44-
clientId: connection.id,
45-
url: socket.url,
46-
},
47-
} satisfies WebSocketBroadcastChannelMessage)
4839
})
4940

50-
it('adds a client from another runtime to the list of clients', async () => {
51-
const clientId = randomUUID()
52-
const url = new URL('ws://localhost')
53-
const manager = new WebSocketClientManager(channel)
41+
it('adds multiple clients from this runtime to the list of clients', () => {
42+
const manager = new WebSocketClientManager(channel, '*')
43+
const connectionOne = new WebSocketClientConnection(
44+
socket,
45+
new TestWebSocketTransport(),
46+
)
47+
manager.addConnection(connectionOne)
5448

55-
channel.dispatchEvent(
56-
new MessageEvent<WebSocketBroadcastChannelMessage>('message', {
57-
data: {
58-
type: 'connection:open',
59-
payload: {
60-
clientId,
61-
url: url.href,
62-
},
63-
},
64-
}),
49+
// Must add the client to the list of clients.
50+
expect(Array.from(manager.clients.values())).toEqual([connectionOne])
51+
52+
const connectionTwo = new WebSocketClientConnection(
53+
socket,
54+
new TestWebSocketTransport(),
6555
)
56+
manager.addConnection(connectionTwo)
6657

67-
await vi.waitFor(() => {
68-
expect(Array.from(manager.clients.values())).toEqual([
69-
new WebSocketRemoteClientConnection(clientId, url, channel),
70-
])
71-
})
58+
// Must add the new cilent to the list as well.
59+
expect(Array.from(manager.clients.values())).toEqual([
60+
connectionOne,
61+
connectionTwo,
62+
])
7263
})
7364

7465
it('replays a "send" event coming from another runtime', async () => {
75-
const manager = new WebSocketClientManager(channel)
76-
const connection = new WebSocketClientConnection(socket, transport)
66+
const manager = new WebSocketClientManager(channel, '*')
67+
const connection = new WebSocketClientConnection(
68+
socket,
69+
new TestWebSocketTransport(),
70+
)
7771
manager.addConnection(connection)
7872
vi.spyOn(connection, 'send')
7973

@@ -98,8 +92,11 @@ it('replays a "send" event coming from another runtime', async () => {
9892
})
9993

10094
it('replays a "close" event coming from another runtime', async () => {
101-
const manager = new WebSocketClientManager(channel)
102-
const connection = new WebSocketClientConnection(socket, transport)
95+
const manager = new WebSocketClientManager(channel, '*')
96+
const connection = new WebSocketClientConnection(
97+
socket,
98+
new TestWebSocketTransport(),
99+
)
103100
manager.addConnection(connection)
104101
vi.spyOn(connection, 'close')
105102

@@ -125,7 +122,8 @@ it('replays a "close" event coming from another runtime', async () => {
125122
})
126123

127124
it('removes the extraneous message listener when the connection closes', async () => {
128-
const manager = new WebSocketClientManager(channel)
125+
const manager = new WebSocketClientManager(channel, '*')
126+
const transport = new TestWebSocketTransport()
129127
const connection = new WebSocketClientConnection(socket, transport)
130128
vi.spyOn(connection, 'close').mockImplementationOnce(() => {
131129
/**
@@ -135,7 +133,7 @@ it('removes the extraneous message listener when the connection closes', async (
135133
* All we care here is that closing the connection triggers
136134
* the transport closure, which it always does.
137135
*/
138-
connection['transport'].onClose()
136+
transport.dispatchEvent(new Event('close'))
139137
})
140138
vi.spyOn(connection, 'send')
141139

‎src/core/ws/WebSocketClientManager.ts‎

Lines changed: 108 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,14 @@
1+
import { invariant } from 'outvariant'
12
import type {
23
WebSocketData,
34
WebSocketClientConnection,
45
WebSocketClientConnectionProtocol,
56
} from '@mswjs/interceptors/WebSocket'
7+
import { matchRequestUrl, type Path } from '../utils/matching/matchRequestUrl'
8+
9+
export const MSW_WEBSOCKET_CLIENTS_KEY = 'msw:ws:clients'
610

711
export type WebSocketBroadcastChannelMessage =
8-
| {
9-
type: 'connection:open'
10-
payload: {
11-
clientId: string
12-
url: string
13-
}
14-
}
1512
| {
1613
type: 'extraneous:send'
1714
payload: {
@@ -28,33 +25,122 @@ export type WebSocketBroadcastChannelMessage =
2825
}
2926
}
3027

31-
export const kAddByClientId = Symbol('kAddByClientId')
28+
type SerializedClient = {
29+
clientId: string
30+
url: string
31+
}
3232

3333
/**
3434
* A manager responsible for accumulating WebSocket client
3535
* connections across different browser runtimes.
3636
*/
3737
export class WebSocketClientManager {
38+
private inMemoryClients: Set<WebSocketClientConnectionProtocol>
39+
40+
constructor(
41+
private channel: BroadcastChannel,
42+
private url: Path,
43+
) {
44+
this.inMemoryClients = new Set()
45+
46+
if (typeof localStorage !== 'undefined') {
47+
// When the worker clears the local storage key in "worker.stop()",
48+
// also clear the in-memory clients map.
49+
localStorage.removeItem = new Proxy(localStorage.removeItem, {
50+
apply: (target, thisArg, args) => {
51+
const [key] = args
52+
53+
if (key === MSW_WEBSOCKET_CLIENTS_KEY) {
54+
this.inMemoryClients.clear()
55+
}
56+
57+
return Reflect.apply(target, thisArg, args)
58+
},
59+
})
60+
}
61+
}
62+
3863
/**
3964
* All active WebSocket client connections.
4065
*/
41-
public clients: Set<WebSocketClientConnectionProtocol>
66+
get clients(): Set<WebSocketClientConnectionProtocol> {
67+
// In the browser, different runtimes use "localStorage"
68+
// as the shared source of all the clients.
69+
if (typeof localStorage !== 'undefined') {
70+
const inMemoryClients = Array.from(this.inMemoryClients)
71+
72+
console.log('get clients()', inMemoryClients, this.getSerializedClients())
73+
74+
return new Set(
75+
inMemoryClients.concat(
76+
this.getSerializedClients()
77+
// Filter out the serialized clients that are already present
78+
// in this runtime in-memory. This is crucial because a remote client
79+
// wrapper CANNOT send a message to the client in THIS runtime
80+
// (the "message" event on broadcast channel won't trigger).
81+
.filter((serializedClient) => {
82+
if (
83+
inMemoryClients.every(
84+
(client) => client.id !== serializedClient.clientId,
85+
)
86+
) {
87+
return serializedClient
88+
}
89+
})
90+
.map((serializedClient) => {
91+
return new WebSocketRemoteClientConnection(
92+
serializedClient.clientId,
93+
new URL(serializedClient.url),
94+
this.channel,
95+
)
96+
}),
97+
),
98+
)
99+
}
100+
101+
// In Node.js, the manager acts as a singleton, and all clients
102+
// are kept in-memory.
103+
return this.inMemoryClients
104+
}
42105

43-
constructor(private channel: BroadcastChannel) {
44-
this.clients = new Set()
106+
private getSerializedClients(): Array<SerializedClient> {
107+
invariant(
108+
typeof localStorage !== 'undefined',
109+
'Failed to call WebSocketClientManager#getSerializedClients() in a non-browser environment. This is likely a bug in MSW. Please, report it on GitHub: https://github.com/mswjs/msw',
110+
)
45111

46-
this.channel.addEventListener('message', (message) => {
47-
const { type, payload } = message.data as WebSocketBroadcastChannelMessage
112+
const clientsJson = localStorage.getItem(MSW_WEBSOCKET_CLIENTS_KEY)
48113

49-
switch (type) {
50-
case 'connection:open': {
51-
// When another runtime notifies about a new connection,
52-
// create a connection wrapper class and add it to the set.
53-
this.onRemoteConnection(payload.clientId, new URL(payload.url))
54-
break
55-
}
56-
}
114+
if (!clientsJson) {
115+
return []
116+
}
117+
118+
const allClients = JSON.parse(clientsJson) as Array<SerializedClient>
119+
const matchingClients = allClients.filter((client) => {
120+
return matchRequestUrl(new URL(client.url), this.url).matches
57121
})
122+
123+
return matchingClients
124+
}
125+
126+
private addClient(client: WebSocketClientConnection): void {
127+
this.inMemoryClients.add(client)
128+
129+
if (typeof localStorage !== 'undefined') {
130+
const serializedClients = this.getSerializedClients()
131+
132+
// Serialize the current client for other runtimes to create
133+
// a remote wrapper over it. This has no effect on the current runtime.
134+
const nextSerializedClients = serializedClients.concat({
135+
clientId: client.id,
136+
url: client.url.href,
137+
} as SerializedClient)
138+
139+
localStorage.setItem(
140+
MSW_WEBSOCKET_CLIENTS_KEY,
141+
JSON.stringify(nextSerializedClients),
142+
)
143+
}
58144
}
59145

60146
/**
@@ -64,16 +150,7 @@ export class WebSocketClientManager {
64150
* for the opened connections in the same runtime.
65151
*/
66152
public addConnection(client: WebSocketClientConnection): void {
67-
this.clients.add(client)
68-
69-
// Signal to other runtimes about this connection.
70-
this.channel.postMessage({
71-
type: 'connection:open',
72-
payload: {
73-
clientId: client.id,
74-
url: client.url.toString(),
75-
},
76-
} as WebSocketBroadcastChannelMessage)
153+
this.addClient(client)
77154

78155
// Instruct the current client how to handle events
79156
// coming from other runtimes (e.g. when calling `.broadcast()`).
@@ -116,19 +193,6 @@ export class WebSocketClientManager {
116193
once: true,
117194
})
118195
}
119-
120-
/**
121-
* Adds a client connection wrapper to operate with
122-
* WebSocket client connections in other runtimes.
123-
*/
124-
private onRemoteConnection(id: string, url: URL): void {
125-
this.clients.add(
126-
// Create a connection-compatible instance that can
127-
// operate with this client from a different runtime
128-
// using the BroadcastChannel messages.
129-
new WebSocketRemoteClientConnection(id, url, this.channel),
130-
)
131-
}
132196
}
133197

134198
/**

0 commit comments

Comments
 (0)