|
| 1 | +import type { ChannelOptions } from 'birpc' |
| 2 | +import type { RpcClientEvents } from 'devframe/client' |
| 3 | +import type { DevframeRpcClientFunctions, DevframeRpcServerFunctions, RpcFunctionsHost } from 'devframe/types' |
| 4 | +import type { MessagePort } from 'node:worker_threads' |
| 5 | +import { MessageChannel } from 'node:worker_threads' |
| 6 | +import { createHostContext } from 'devframe/node' |
| 7 | +import { RpcFunctionsCollectorBase } from 'devframe/rpc' |
| 8 | +import { createRpcClient } from 'devframe/rpc/client' |
| 9 | +import { createRpcServer } from 'devframe/rpc/server' |
| 10 | +import { createRpcSharedStateClientHost, createRpcSharedStateServerHost } from 'devframe/rpc/shared-state' |
| 11 | +import { createEventEmitter } from 'devframe/utils/events' |
| 12 | +import { structuredCloneDeserialize, structuredCloneSerialize } from 'devframe/utils/structured-clone' |
| 13 | +import { expect, it } from 'vitest' |
| 14 | +import { createContextRpcServer } from '../node/rpc-core' |
| 15 | + |
| 16 | +/** JSON records also travel through transports that cannot clone native values. */ |
| 17 | +function channelFor(port: MessagePort): ChannelOptions { |
| 18 | + return { |
| 19 | + post: message => port.postMessage(message), |
| 20 | + on: (handler) => { |
| 21 | + port.on('message', handler) |
| 22 | + }, |
| 23 | + off: (handler) => { |
| 24 | + port.off('message', handler) |
| 25 | + }, |
| 26 | + serialize: value => JSON.stringify(structuredCloneSerialize(value)), |
| 27 | + deserialize: value => structuredCloneDeserialize(JSON.parse(value)), |
| 28 | + } |
| 29 | +} |
| 30 | + |
| 31 | +async function createStateServer(mode: string) { |
| 32 | + if (mode === 'node context') { |
| 33 | + const context = await createHostContext({ |
| 34 | + cwd: process.cwd(), |
| 35 | + mode: 'dev', |
| 36 | + host: { |
| 37 | + mountStatic() {}, |
| 38 | + resolveOrigin: () => 'http://localhost', |
| 39 | + getStorageDir: () => process.cwd(), |
| 40 | + }, |
| 41 | + }) |
| 42 | + const { rpcGroup } = createContextRpcServer({ context, auth: false }) |
| 43 | + return { group: rpcGroup, sharedState: context.rpc.sharedState } |
| 44 | + } |
| 45 | + |
| 46 | + const collector = new RpcFunctionsCollectorBase<DevframeRpcServerFunctions, undefined>(undefined) |
| 47 | + const group = createRpcServer<DevframeRpcClientFunctions, DevframeRpcServerFunctions>(collector.functions) |
| 48 | + const broadcast: RpcFunctionsHost['broadcast'] = async (options) => { |
| 49 | + await Promise.all(group.clients |
| 50 | + .filter(client => options.filter?.(client) !== false) |
| 51 | + .map(client => client.$callRaw({ ...options, optional: true, event: true }))) |
| 52 | + } |
| 53 | + const sharedState = createRpcSharedStateServerHost({ register: collector.register.bind(collector), broadcast }) |
| 54 | + return { group, sharedState } |
| 55 | +} |
| 56 | + |
| 57 | +it.each(['custom channels', 'node context'])('shares state across %s with per-connection subscriptions', async (mode) => { |
| 58 | + expect.assertions(13) |
| 59 | + const { group, sharedState } = await createStateServer(mode) |
| 60 | + const counter = await sharedState.get('counter', { initialValue: { count: 1 } }) |
| 61 | + const channels = [new MessageChannel(), new MessageChannel()] |
| 62 | + const peers = channels.map((channel, index) => { |
| 63 | + const meta = { id: index, subscribedStates: new Set<string>() } |
| 64 | + const serverChannel = { ...channelFor(channel.port1), meta } |
| 65 | + group.updateChannels(current => current.push(serverChannel)) |
| 66 | + const client = new RpcFunctionsCollectorBase<DevframeRpcClientFunctions, undefined>(undefined) |
| 67 | + const rpc = createRpcClient<DevframeRpcServerFunctions, DevframeRpcClientFunctions>(client.functions, { |
| 68 | + channel: channelFor(channel.port2), |
| 69 | + }) |
| 70 | + const state = createRpcSharedStateClientHost({ |
| 71 | + call: rpc.$call, |
| 72 | + callEvent: rpc.$callEvent, |
| 73 | + client, |
| 74 | + isTrusted: true, |
| 75 | + events: createEventEmitter<RpcClientEvents>(), |
| 76 | + connectionMeta: { backend: 'none' }, |
| 77 | + }) |
| 78 | + return { meta, rpc, state, serverChannel } |
| 79 | + }) |
| 80 | + const [first, second] = peers |
| 81 | + try { |
| 82 | + const firstCounter = await first.state.get<{ count: number }>('counter') |
| 83 | + expect(firstCounter.value()).toEqual({ count: 1 }) |
| 84 | + expect(first.meta.subscribedStates.has('counter')).toBe(true) |
| 85 | + expect(second.meta.subscribedStates.has('counter')).toBe(false) |
| 86 | + |
| 87 | + const secondCounter = await second.state.get<{ count: number }>('counter') |
| 88 | + expect(second.meta.subscribedStates.has('counter')).toBe(true) |
| 89 | + firstCounter.mutate((draft) => { |
| 90 | + draft.count = 2 |
| 91 | + }) |
| 92 | + await expect.poll(() => counter.value().count).toBe(2) |
| 93 | + await expect.poll(() => secondCounter.value().count).toBe(2) |
| 94 | + |
| 95 | + group.clients.find(client => client.$meta === first.meta)?.$close() |
| 96 | + group.updateChannels(current => current.splice(current.indexOf(first.serverChannel), 1)) |
| 97 | + first.rpc.$close() |
| 98 | + await expect(first.rpc.$call('devframe:rpc:server-state:get', 'counter')).rejects.toThrow() |
| 99 | + secondCounter.mutate((draft) => { |
| 100 | + draft.count = 3 |
| 101 | + }) |
| 102 | + await expect.poll(() => counter.value().count).toBe(3) |
| 103 | + expect(firstCounter.value().count).toBe(2) |
| 104 | + expect(sharedState.keys()).toEqual(['counter']) |
| 105 | + expect(first.state.delete('counter')).toBe(true) |
| 106 | + expect(sharedState.delete('counter')).toBe(true) |
| 107 | + expect(sharedState.keys()).toEqual([]) |
| 108 | + } |
| 109 | + finally { |
| 110 | + for (const peer of peers) peer.rpc.$close() |
| 111 | + for (const client of group.clients) client.$close() |
| 112 | + group.updateChannels(current => current.splice(0)) |
| 113 | + for (const channel of channels) { |
| 114 | + channel.port1.close() |
| 115 | + channel.port2.close() |
| 116 | + } |
| 117 | + } |
| 118 | +}) |
| 119 | + |
| 120 | +it('rejects an initial snapshot when its RPC connection closes', async () => { |
| 121 | + expect.assertions(1) |
| 122 | + const client = new RpcFunctionsCollectorBase<DevframeRpcClientFunctions, undefined>(undefined) |
| 123 | + const channel = new MessageChannel() |
| 124 | + const rpc = createRpcClient<DevframeRpcServerFunctions, DevframeRpcClientFunctions>(client.functions, { |
| 125 | + channel: channelFor(channel.port1), |
| 126 | + }) |
| 127 | + const state = createRpcSharedStateClientHost({ |
| 128 | + call: rpc.$call, |
| 129 | + callEvent: rpc.$callEvent, |
| 130 | + client, |
| 131 | + isTrusted: true, |
| 132 | + events: createEventEmitter<RpcClientEvents>(), |
| 133 | + connectionMeta: { backend: 'none' }, |
| 134 | + }) |
| 135 | + try { |
| 136 | + const pending = state.get('counter') |
| 137 | + rpc.$close() |
| 138 | + await expect(pending).rejects.toThrow('closed') |
| 139 | + } |
| 140 | + finally { |
| 141 | + channel.port1.close() |
| 142 | + channel.port2.close() |
| 143 | + } |
| 144 | +}, 1000) |
0 commit comments