diff --git a/discojs/package.json b/discojs/package.json index cca3fc989..a2cc8002d 100644 --- a/discojs/package.json +++ b/discojs/package.json @@ -102,11 +102,10 @@ "debug": "catalog:", "fast-check": "3.23.2", "immutable": "catalog:", - "simple-peer": "9.11.1", "zod": "catalog:" }, "devDependencies": { - "@types/simple-peer": "9.11.9", + "@types/node": "22.19.15", "nodemon": "3.1.14", "vitest": "catalog:" } diff --git a/discojs/src/client/decentralized/packet.spec.ts b/discojs/src/client/decentralized/packet.spec.ts new file mode 100644 index 000000000..bfea49612 --- /dev/null +++ b/discojs/src/client/decentralized/packet.spec.ts @@ -0,0 +1,43 @@ +import { assert, describe, it } from "vitest"; +import { Packet } from "#client/decentralized/packet"; + +describe("packet", () => { + it("encodes and decodes an ack packet", () => { + const ack = Packet.ack(42); + const encoded = ack.encode(); + const decoded = Packet.decode(encoded); + + assert.isTrue(decoded.isAck); + assert.strictEqual(decoded.messageID, 42); + }); + + it("encodes and decodes a first chunk packet with 32-bit chunk counter and total count", () => { + const payload = Buffer.from("hello world"); + const firstChunk = Packet.firstChunk(101, 500, payload); + const encoded = firstChunk.encode(); + const decoded = Packet.decode(encoded); + + assert.isTrue(decoded.isFirstChunk); + assert.strictEqual(decoded.messageID, 101); + assert.strictEqual(decoded.chunkID, 0); + assert.strictEqual(decoded.totalChunks, 500); + assert.strictEqual(decoded.payload.toString(), "hello world"); + }); + + it("encodes and decodes a subsequent chunk packet", () => { + const payload = Buffer.from("subsequent chunk"); + const chunk = Packet.chunk(101, 299, payload); + const encoded = chunk.encode(); + const decoded = Packet.decode(encoded); + + assert.isFalse(decoded.isAck); + assert.isFalse(decoded.isFirstChunk); + assert.strictEqual(decoded.messageID, 101); + assert.strictEqual(decoded.chunkID, 299); + assert.strictEqual(decoded.payload.toString(), "subsequent chunk"); + }); + + it("rejects invalid buffer sizes on decode", () => { + assert.throws(() => Packet.decode(Buffer.alloc(2))); + }); +}); diff --git a/discojs/src/client/decentralized/packet.ts b/discojs/src/client/decentralized/packet.ts new file mode 100644 index 000000000..f18255608 --- /dev/null +++ b/discojs/src/client/decentralized/packet.ts @@ -0,0 +1,112 @@ +export type MessageID = number; +export type ChunkID = number; + +// helper function to make it evident that the number passed is a chunk id +const asChunkID = (n: number): ChunkID => n; + +export class Packet { + constructor( + public readonly messageID: MessageID, + public readonly chunkID: ChunkID, + public readonly totalChunks?: number, + public readonly payload: Buffer = Buffer.alloc(0), + ) {} + + // message id size in bytes (uint16) + static MESSAGE_ID_BYTES = 2; + // chunk counter size in bytes (uint32) + static CHUNK_ID_BYTES = 4; + // total chunk count size in bytes (uint32) + static TOTAL_CHUNKS_BYTES = 4; + + // message id (2 bytes) + chunk counter (4 bytes) + static HEADER_SIZE = Packet.MESSAGE_ID_BYTES + Packet.CHUNK_ID_BYTES; + + // message id (2 bytes) + chunk counter (4 bytes) + total chunk count (4 bytes) + static FIRST_HEADER_SIZE = Packet.HEADER_SIZE + Packet.TOTAL_CHUNKS_BYTES; + + // chunk counter of an acknowledgement, unused by messages + static ACK_CHUNK_ID: ChunkID = 0xffff_ffff; + + get isAck(): boolean { + return this.chunkID === Packet.ACK_CHUNK_ID; + } + + get isFirstChunk(): boolean { + return this.chunkID === asChunkID(0); + } + + static ack(messageID: MessageID): Packet { + return new Packet(messageID, Packet.ACK_CHUNK_ID); + } + + static firstChunk( + messageID: MessageID, + totalChunks: number, + payload: Buffer, + ): Packet { + return new Packet(messageID, asChunkID(0), totalChunks, payload); + } + + static chunk( + messageID: MessageID, + chunkID: ChunkID, + payload: Buffer, + ): Packet { + return new Packet(messageID, chunkID, undefined, payload); + } + + encode(): Buffer { + if (this.isAck) { + const buf = Buffer.alloc(Packet.HEADER_SIZE); + buf.writeUInt16BE(this.messageID, 0); + buf.writeUInt32BE(Packet.ACK_CHUNK_ID, Packet.MESSAGE_ID_BYTES); + return buf; + } + + if (this.isFirstChunk) { + if (this.totalChunks === undefined) { + throw new Error("first chunk packet requires totalChunks"); + } + const buf = Buffer.alloc(Packet.FIRST_HEADER_SIZE + this.payload.length); + buf.writeUInt16BE(this.messageID, 0); + buf.writeUInt32BE(asChunkID(0), Packet.MESSAGE_ID_BYTES); + buf.writeUInt32BE(this.totalChunks, Packet.HEADER_SIZE); + this.payload.copy(buf, Packet.FIRST_HEADER_SIZE, 0); + return buf; + } + + const buf = Buffer.alloc(Packet.HEADER_SIZE + this.payload.length); + buf.writeUInt16BE(this.messageID, 0); + buf.writeUInt32BE(this.chunkID, Packet.MESSAGE_ID_BYTES); + this.payload.copy(buf, Packet.HEADER_SIZE, 0); + return buf; + } + + static decode(data: Buffer): Packet { + if (!Buffer.isBuffer(data) || data.length < Packet.HEADER_SIZE) { + throw new Error("received invalid message type"); + } + + const messageID: MessageID = data.readUInt16BE(0); + const chunkID: ChunkID = data.readUInt32BE(Packet.MESSAGE_ID_BYTES); + + if (chunkID === Packet.ACK_CHUNK_ID) { + return Packet.ack(messageID); + } + + if (chunkID === asChunkID(0)) { + if (data.length < Packet.FIRST_HEADER_SIZE) { + throw new Error("received invalid message type"); + } + const totalChunks = data.readUInt32BE(Packet.HEADER_SIZE); + const payload = Buffer.alloc(data.length - Packet.FIRST_HEADER_SIZE); + data.copy(payload, 0, Packet.FIRST_HEADER_SIZE); + return Packet.firstChunk(messageID, totalChunks, payload); + } + + const payload = Buffer.alloc(data.length - Packet.HEADER_SIZE); + data.copy(payload, 0, Packet.HEADER_SIZE); + return Packet.chunk(messageID, chunkID, payload); + } +} diff --git a/discojs/src/client/decentralized/peer.spec.ts b/discojs/src/client/decentralized/peer.spec.ts index 4608b4f8a..67916c5b4 100644 --- a/discojs/src/client/decentralized/peer.spec.ts +++ b/discojs/src/client/decentralized/peer.spec.ts @@ -78,4 +78,48 @@ describe("peer", () => { assert.isTrue(received); }); + + it("transfers GPT-sized payload with byte-for-byte reassembly", async () => { + // 10 MB buffer simulating GPT weights + const size = 10 * 1024 * 1024; + const message = Buffer.alloc(size); + // Fill with non-trivial pattern to thoroughly verify byte ordering + for (let i = 0; i < size; i += 4) { + message.writeUInt32LE((i * 1103515245 + 12345) >>> 0, i); + } + + const receivedPromise = new Promise((resolve) => { + peer2.on("data", (data) => resolve(data)); + }); + + peer1.send(message); + await peer1.drain(); + const received = await receivedPromise; + + assert.strictEqual(received.length, message.length); + assert.strictEqual(Buffer.compare(received, message), 0); + }); + + it("handles payloads with more than 255 chunks", async () => { + // Override maxChunkSize to small size (1024 bytes) to test >255 chunk count boundary + Object.defineProperty(peer1, "maxChunkSize", { get: () => 1024 }); + + // 300 chunks + const size = 300 * 1024; + const message = Buffer.alloc(size); + for (let i = 0; i < size; i += 4) { + message.writeUInt32LE((i * 1664525 + 1013904223) >>> 0, i); + } + + const receivedPromise = new Promise((resolve) => { + peer2.on("data", (data) => resolve(data)); + }); + + peer1.send(message); + await peer1.drain(); + const received = await receivedPromise; + + assert.strictEqual(received.length, message.length); + assert.strictEqual(Buffer.compare(received, message), 0); + }); }); diff --git a/discojs/src/client/decentralized/peer.ts b/discojs/src/client/decentralized/peer.ts index ecedbb3b8..3ce8659f6 100644 --- a/discojs/src/client/decentralized/peer.ts +++ b/discojs/src/client/decentralized/peer.ts @@ -1,21 +1,13 @@ import { List, Map, Range, Seq, Set } from "immutable"; import wrtc from "@epfml/isomorphic-wrtc"; -import SimplePeer from "simple-peer"; import type { NodeID } from "#client/types"; - -type MessageID = number; -type ChunkID = number; - -// message id + (chunk counter == 0) + chunk count -const FIRST_HEADER_SIZE = 2 + 1 + 1; - -// message id + chunk counter -const HEADER_SIZE = 2 + 1; - -// chunk counter of an acknowledgement, unused by messages -// as they have at most 0xff chunks, numbered from 0 -const ACK_CHUNK_ID = 0xff; +import { EventEmitter } from "#utils/event_emitter"; +import { + type MessageID, + type ChunkID, + Packet, +} from "#client/decentralized/packet"; // at which interval to poll const TICK = 10; @@ -23,45 +15,41 @@ const TICK = 10; // how long to wait for sent messages to be received before giving up const MAX_DRAIN_TIME = 10_000; -// we can't use the definition in DOM as we're platform independent -export type SignalData = - | { type: "answer" | "offer" | "pranswer" | "rollback"; sdp?: string } - | { type: "transceiverRequest"; transceiverRequest: { kind: string } } - | { type: "renegotiate"; renegotiate: true } - | { type: "candidate"; candidate: RTCIceCandidate }; - -interface Events { - close: () => void; - connect: () => void; - error: (err: Error) => void; - signal: (signal: SignalData) => void; - data: (data: Buffer) => void; -} +// Fallback SCTP message size per RFC 8841 / WebRTC spec when transport is not yet connected +const SAFE_MAX_MESSAGE_SIZE = 65_536; -// Peer wraps a SimplePeer, adding message fragmentation -// -// WebRTC implementations have various maximum message size -// but with huge models, our messages might be bigger. -// We split messages by chunks and reconstruct these -// on the other side. -// -// As the WebRTC's DataChannel is not a stream, we need -// reorder messages, so we use a header on each chunk -// with a message id and chunk counter. The first chunk -// (chunk counter == 0), also add the total number of chunk. +export type SignalData = + | RTCSessionDescriptionInit + | { type: "candidate"; candidate: RTCIceCandidateInit }; + +type EventPayloads = { + close: void; + connect: void; + error: Error; + signal: SignalData; + data: Buffer; + [key: string]: unknown; +}; + +// Peer adds message fragmentation, flow control, and receipt acknowledgements +// over standard WebRTC RTCPeerConnection and RTCDataChannel. // -// Once a message is fully received, its id is sent back -// with ACK_CHUNK_ID as chunk counter, so that the sender -// knows when it can close the connection without losing it. +// WebRTC implementations have various maximum message sizes (e.g. 256 KiB in Chromium). +// For large models (megabytes to gigabytes), messages are split into chunks +// and reconstructed on the receiving peer. // -// see feross/simple-peer#393 for more info -export class Peer { - private readonly peer: SimplePeer.Instance; - - private bufferSize?: number; +// As the WebRTC DataChannel is not a stream, we frame each chunk with a message ID, +// chunk counter, and total chunk count. Once a message is fully received, its ID is sent back +// with ACK_CHUNK_ID so the sender can drain before closing the connection. +export class Peer extends EventEmitter { + private readonly pc: RTCPeerConnection; + private channel?: RTCDataChannel; + private isDestroyed = false; + private signalQueue = Promise.resolve(); + private queuedCandidates: RTCIceCandidateInit[] = []; private sendCounter: MessageID = 0; - private sendQueue = List(); + private sendQueue = List(); private unacknowledged = Set(); private receiving = Map< @@ -74,18 +62,145 @@ export class Peer { constructor( public readonly id: NodeID, - initiator: boolean = false, + isInitiator: boolean = false, ) { - this.peer = new SimplePeer({ wrtc, initiator }); + super(); - this.peer.on("data", (data: unknown) => { - if ( - Buffer.isBuffer(data) && - data.length === HEADER_SIZE && - data.readUInt8(2) === ACK_CHUNK_ID - ) - this.unacknowledged = this.unacknowledged.delete(data.readUInt16BE()); - }); + this.pc = new wrtc.RTCPeerConnection(); + + this.pc.onicecandidate = (event) => { + if (event.candidate) { + const candidate: RTCIceCandidateInit = event.candidate.toJSON(); + this.emit("signal", { type: "candidate", candidate }); + } + }; + + this.pc.onconnectionstatechange = () => { + const state = this.pc.connectionState; + if (state === "failed" || state === "closed") { + this.emit("close", undefined); + } + }; + + if (isInitiator) { + const channel = this.pc.createDataChannel("disco"); + this.setupChannel(channel); + queueMicrotask(() => { + void this.makeOffer(); + }); + } else { + this.pc.ondatachannel = (event) => { + this.setupChannel(event.channel); + }; + } + } + + private setupChannel(channel: RTCDataChannel): void { + this.channel = channel; + this.channel.binaryType = "arraybuffer"; + try { + this.channel.bufferedAmountLowThreshold = this.maxChunkSize; + } catch { + // Some WebRTC implementations may not support bufferedAmountLowThreshold + } + + this.channel.onopen = () => { + this.emit("connect", undefined); + this.flush(); + }; + + this.channel.onmessage = (event: MessageEvent) => { + const data = Buffer.isBuffer(event.data) + ? event.data + : Buffer.from(event.data as ArrayBuffer); + this.handleData(data); + }; + + this.channel.onbufferedamountlow = () => { + this.flush(); + }; + + this.channel.onclose = () => { + this.emit("close", undefined); + }; + + this.channel.onerror = (event) => { + const message = + "error" in event && event.error instanceof Error + ? event.error.message + : "DataChannel error"; + this.emit("error", new Error(message)); + }; + } + + private async makeOffer(): Promise { + if (this.isDestroyed) return; + try { + const offer = await this.pc.createOffer(); + await this.pc.setLocalDescription(offer); + this.emit("signal", { + type: "offer", + sdp: offer.sdp, + }); + } catch (err) { + this.emit("error", err instanceof Error ? err : new Error(String(err))); + } + } + + signal(signal: SignalData): void { + this.signalQueue = this.signalQueue + .then(() => this.handleSignal(signal)) + .catch((err) => { + this.emit("error", err instanceof Error ? err : new Error(String(err))); + }); + } + + private async handleSignal(signal: SignalData): Promise { + if (this.isDestroyed) return; + + if (signal.type === "offer") { + if (signal.sdp === undefined) { + throw new Error("signal offer without session description"); + } + await this.pc.setRemoteDescription( + new wrtc.RTCSessionDescription({ type: "offer", sdp: signal.sdp }), + ); + await this.drainQueuedCandidates(); + const answer = await this.pc.createAnswer(); + await this.pc.setLocalDescription(answer); + this.emit("signal", { + type: "answer", + sdp: answer.sdp, + }); + } else if (signal.type === "answer") { + if (signal.sdp === undefined) { + throw new Error("signal answer without session description"); + } + await this.pc.setRemoteDescription( + new wrtc.RTCSessionDescription({ type: "answer", sdp: signal.sdp }), + ); + await this.drainQueuedCandidates(); + } else if (signal.type === "candidate") { + if (!signal.candidate) return; + if (this.pc.remoteDescription && this.pc.remoteDescription.type) { + await this.pc.addIceCandidate( + new wrtc.RTCIceCandidate(signal.candidate), + ); + } else { + this.queuedCandidates.push(signal.candidate); + } + } + } + + private async drainQueuedCandidates(): Promise { + for (const c of this.queuedCandidates) { + try { + await this.pc.addIceCandidate(new wrtc.RTCIceCandidate(c)); + } catch { + // Ignore candidate errors if transport already connected + } + } + this.queuedCandidates = []; } send(msg: Buffer): void { @@ -97,222 +212,183 @@ export class Peer { } private acknowledge(messageID: MessageID): void { - const ack = Buffer.alloc(HEADER_SIZE); - ack.writeUint16BE(messageID); - ack.writeUint8(ACK_CHUNK_ID, 2); + const ack = Packet.ack(messageID); this.sendQueue = this.sendQueue.push(ack); this.flush(); } private flush(): void { - if (this.bufferSize === undefined) { - throw new Error("flush without known buffer size"); - } - - const chunk = this.sendQueue.first(); - if (chunk === undefined) { - return; // nothing to flush - } - - const remainingBufferSize = this.bufferSize - this.peer.bufferSize; - if (chunk.length > remainingBufferSize) { - setTimeout(() => { - this.flush(); - }, TICK); + if ( + this.isDestroyed || + !this.channel || + this.channel.readyState !== "open" + ) { return; } - this.sendQueue = this.sendQueue.shift(); - this.peer.send(chunk); + const HIGH_WATER_MARK = 1024 * 1024; // 1 MB queue backpressure cap - // and loop - this.flush(); - } + while (!this.sendQueue.isEmpty()) { + if (this.channel.bufferedAmount >= HIGH_WATER_MARK) { + setTimeout(() => this.flush(), TICK); + return; + } - get maxChunkSize(): number { - if (this.bufferSize === undefined) { - throw new Error("chunk without known buffer size"); + const chunk = this.sendQueue.first(); + if (chunk === undefined) break; + + try { + this.channel.send(chunk.encode() as ArrayBufferView); + this.sendQueue = this.sendQueue.shift(); + } catch (err) { + this.emit( + "error", + err instanceof Error ? err : new Error("Failed to send chunk"), + ); + return; + } } - return this.bufferSize; } - private chunk(b: Buffer): Seq.Indexed { + private get maxChunkSize(): number { + return this.pc.sctp?.maxMessageSize || SAFE_MAX_MESSAGE_SIZE; + } + + private chunk(b: Buffer): Seq.Indexed { const messageID = this.sendCounter; this.sendCounter++; if (this.sendCounter > 0xffff) { - throw new Error("too much messages sent to this peer"); + throw new Error("too many messages sent to this peer"); } - // special case as Range(1, 0) yields a value let tail = Seq.Indexed([]); - if (b.length > this.maxChunkSize) { + if (b.length > this.maxChunkSize - Packet.FIRST_HEADER_SIZE) { tail = Range( - this.maxChunkSize - FIRST_HEADER_SIZE, + this.maxChunkSize - Packet.FIRST_HEADER_SIZE, b.length, - this.maxChunkSize - HEADER_SIZE, + this.maxChunkSize - Packet.HEADER_SIZE, ).map((offset) => - b.subarray(offset, offset + this.maxChunkSize - HEADER_SIZE), + b.subarray(offset, offset + this.maxChunkSize - Packet.HEADER_SIZE), ); } const totalChunkCount = 1 + tail.count(); - if (totalChunkCount > 0xff) { + if (totalChunkCount >= Packet.ACK_CHUNK_ID) { throw new Error( - `The payload is too big: ${totalChunkCount * this.maxChunkSize} bytes > 255,` + - " consider reducing the model size or increasing the chunk size", + `The payload is too big: ${totalChunkCount} chunks exceeds max chunks (${Packet.ACK_CHUNK_ID})`, ); } - const firstChunk = Buffer.alloc( - b.length > this.maxChunkSize - FIRST_HEADER_SIZE - ? this.maxChunkSize - : FIRST_HEADER_SIZE + b.length, - ); - firstChunk.writeUint16BE(messageID); - const asChunkID = (n: number): ChunkID => n; - firstChunk.writeUint8(asChunkID(0), 2); - firstChunk.writeUint8(totalChunkCount, 3); - b.copy( - firstChunk, - FIRST_HEADER_SIZE, + const firstPayload = b.subarray( 0, - this.maxChunkSize - FIRST_HEADER_SIZE, + Math.min(b.length, this.maxChunkSize - Packet.FIRST_HEADER_SIZE), + ); + const headerChunk = Packet.firstChunk( + messageID, + totalChunkCount, + firstPayload, ); - return Seq.Indexed([firstChunk]).concat( - Range(asChunkID(1), Number.POSITIVE_INFINITY) + return Seq.Indexed([headerChunk]).concat( + Range(1, Number.POSITIVE_INFINITY) .zip(tail) - .map(([id, raw]) => { - const chunk = Buffer.alloc(HEADER_SIZE + raw.length); - chunk.writeUint16BE(messageID); - chunk.writeUint8(id, 2); - raw.copy(chunk, HEADER_SIZE, 0); - return chunk; - }), + .map(([id, raw]) => Packet.chunk(messageID, id, raw)), ); } // Resolves once every message sent so far has been received, // as destroying the peer drops whatever is still in transit. // Gives up after a while in case the other peer is gone. - async drain(): Promise { - const deadline = Date.now() + MAX_DRAIN_TIME; + public async drain(timeout: number = MAX_DRAIN_TIME): Promise { + const deadline = Date.now() + timeout; while ( - !this.peer.destroyed && + !this.isDestroyed && + this.channel?.readyState === "open" && !this.unacknowledged.isEmpty() && Date.now() < deadline - ) + ) { await new Promise((resolve) => setTimeout(resolve, TICK)); + } } async destroy(): Promise { - return new Promise((resolve, reject) => { - this.peer.once("error", reject); - this.peer.once("close", resolve); - this.peer.destroy(); - }); - } + if (this.isDestroyed) return; + this.isDestroyed = true; - signal(signal: SignalData): void { - // extract max buffer size from the signal - if (signal.type === "offer" || signal.type === "answer") { - if (signal.sdp === undefined) { - throw new Error("signal answer|offer without session description"); + try { + if (this.channel) { + this.channel.close(); } - if (this.bufferSize !== undefined) { - throw new Error("buffer size set twice"); - } - - const match = signal.sdp.match(/a=max-message-size:(\d+)/); - if (match === null) { - // TODO default value instead? - throw new Error("no max-message-size found in signal"); - } - const max = parseInt(match[1], 10); - if (isNaN(max)) { - throw new Error(`unable to parse max-message-size as int: ${match[1]}`); - } - - this.bufferSize = max; + this.pc.close(); + } catch { + // Ignore closing errors } - this.peer.signal(signal); + this.emit("close", undefined); + return Promise.resolve(); } - on(event: K, listener: Events[K]): void { - if (event !== "data") { - this.peer.on(event, listener); + private handleData(data: Buffer): void { + const packet = Packet.decode(data); + if (packet.isAck) { + this.unacknowledged = this.unacknowledged.delete(packet.messageID); return; } - // gotta help typescript here - const dataListener = listener as Events["data"]; - this.peer.on("data", (data: unknown) => { - if (!Buffer.isBuffer(data) || data.length < HEADER_SIZE) { - throw new Error("received invalid message type"); - } - const messageID: MessageID = data.readUInt16BE(); //readUint16BE (case sensitive) fails at runtime - const chunkID: ChunkID = data.readUInt8(2); // same for readUint8 - if (chunkID === ACK_CHUNK_ID) return; // handled in the constructor + const received = this.receiving.get(packet.messageID, { + total: undefined, + chunks: Map(), + }); + let total = received.total; + const chunks = received.chunks; - const received = this.receiving.get(messageID, { - total: undefined, - chunks: Map(), - }); - let total = received.total; - const chunks = received.chunks; + if (chunks.has(packet.chunkID)) { + throw new Error( + `chunk ${packet.messageID}:${packet.chunkID} already received`, + ); + } - if (chunks.has(chunkID)) { - throw new Error(`chunk ${messageID}:${chunkID} already received`); + // Receive the header message + if (packet.isFirstChunk) { + if (total !== undefined) { + throw new Error("first header received twice"); } - let chunk: Buffer; - if (chunkID !== 0) { - chunk = Buffer.alloc(data.length - HEADER_SIZE); - data.copy(chunk, 0, HEADER_SIZE); - } else { - if (data.length < FIRST_HEADER_SIZE) { - throw new Error("received invalid message type"); - } - if (total !== undefined) { - throw new Error("first header received twice"); - } - - const readTotal = data.readUInt8(3); - total = readTotal; - chunk = Buffer.alloc(data.length - FIRST_HEADER_SIZE); - data.copy(chunk, 0, FIRST_HEADER_SIZE); - - if (chunks.keySeq().some((id) => id > readTotal)) { - throw new Error( - "received total of chunk but got now-out-of-bound chunks", - ); - } + const totalChunks = packet.totalChunks; + if (totalChunks === undefined) { + throw new Error("first chunk packet missing totalChunks"); + } else if (chunks.keySeq().some((id) => id >= totalChunks)) { + throw new Error( + "received total of chunk but got now-out-of-bound chunks", + ); } - this.receiving = this.receiving.set(messageID, { - total, - chunks: chunks.set(chunkID, chunk), - }); - const readyMessages = this.receiving - .filter( - ({ total, chunks }) => total !== undefined && chunks.size === total, - ) - .sort() - .map(({ chunks }) => - chunks - .entrySeq() - .toList() - .sortBy(([id, _]) => id), - ) - .map((chunks) => Buffer.concat(chunks.map(([_, b]) => b).toArray())); - this.receiving = this.receiving.deleteAll(readyMessages.keys()); - - readyMessages.forEach((message, messageID) => { - this.acknowledge(messageID); - // TODO debug - dataListener(message); - }); + // Set the total number of chunks expected for this message ID + total = totalChunks; + } + + this.receiving = this.receiving.set(packet.messageID, { + total, + chunks: chunks.set(packet.chunkID, packet.payload), + }); + + const readyMessages = this.receiving + .filter( + ({ total, chunks }) => total !== undefined && chunks.size === total, + ) + .sort() + .map(({ chunks }) => + chunks + .entrySeq() + .toList() + .sortBy(([id]) => id), + ) + .map((chunks) => Buffer.concat(chunks.map(([_, b]) => b).toArray())); + + readyMessages.forEach((message, messageID) => { + this.acknowledge(messageID); + this.emit("data", message); + this.receiving = this.receiving.delete(messageID); }); } } diff --git a/discojs/src/client/decentralized/peer_pool.ts b/discojs/src/client/decentralized/peer_pool.ts index b208f7c45..7e1e60573 100644 --- a/discojs/src/client/decentralized/peer_pool.ts +++ b/discojs/src/client/decentralized/peer_pool.ts @@ -26,6 +26,9 @@ async function createStaggered(create: () => T): Promise { return create(); } +// Maximum time in ms to wait for pending messages to be delivered and acknowledged during shutdown +const SHUTDOWN_GRACE_PERIOD = 10_000; + export class PeerPool { private peers = Map(); @@ -41,15 +44,12 @@ export class PeerPool { async shutdown(): Promise { debug(`[${this.id}] is shutting down all its connections`); - // Let the messages still being sent, such as our last round's weights, - // reach the other peers before closing the connections - await Promise.all(this.peers.valueSeq().map((peer) => peer.drain())); - - // Add a timeout o.w. the promise hangs forever if the other peer is already disconnected - await Promise.race([ - Promise.all(this.peers.valueSeq().map((peer) => peer.disconnect())), - new Promise((res, _) => setTimeout(res, 1000)), // Wait for other peers to finish - ]); + // Let messages still being sent reach the other peers before closing + await Promise.all( + this.peers + .valueSeq() + .map((peer) => peer.disconnect(SHUTDOWN_GRACE_PERIOD)), + ); this.peers = Map(); this.pendingSignals = Map(); } diff --git a/discojs/src/client/event_connection.ts b/discojs/src/client/event_connection.ts index d5ddd787d..0cb22358d 100644 --- a/discojs/src/client/event_connection.ts +++ b/discojs/src/client/event_connection.ts @@ -34,7 +34,7 @@ export interface EventConnection { handler: (event: NarrowMessage) => void, ) => void; send: (msg: T) => void; - disconnect: () => Promise; + disconnect: (gracePeriod?: number) => Promise; } /** @@ -199,12 +199,10 @@ export class PeerConnection this.#peer.send(Buffer.from(msgpack.encode(msg))); } - // Resolves once the messages sent so far have left the connection - async drain(): Promise { - await this.#peer.drain(); - } - - async disconnect(): Promise { + async disconnect(gracePeriod = 0): Promise { + if (gracePeriod > 0) { + await this.#peer.drain(gracePeriod); + } await this.#peer.destroy(); } } diff --git a/discojs/tsconfig.lib.json b/discojs/tsconfig.lib.json index 6d137af39..8981b36e3 100644 --- a/discojs/tsconfig.lib.json +++ b/discojs/tsconfig.lib.json @@ -1,6 +1,7 @@ { "extends": "../tsconfig.base.lib.json", "compilerOptions": { + "types": ["node"], // Copied from `package.json` since some tools don't work without it "paths": { "#root/*": ["./src/*"], diff --git a/discojs/tsconfig.vitest.json b/discojs/tsconfig.vitest.json index f33d5a873..d07d9971e 100644 --- a/discojs/tsconfig.vitest.json +++ b/discojs/tsconfig.vitest.json @@ -1,6 +1,7 @@ { "extends": "../tsconfig.base.json", "compilerOptions": { + "types": ["node"], "noEmit": true }, "include": ["src"] diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 95b04c3a1..958674a38 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -350,16 +350,13 @@ importers: immutable: specifier: 'catalog:' version: 5.1.6 - simple-peer: - specifier: 9.11.1 - version: 9.11.1(supports-color@8.1.1) zod: specifier: 'catalog:' version: 4.3.6 devDependencies: - '@types/simple-peer': - specifier: 9.11.9 - version: 9.11.9 + '@types/node': + specifier: 22.19.15 + version: 22.19.15 nodemon: specifier: 3.1.14 version: 3.1.14 @@ -2017,9 +2014,6 @@ packages: '@types/serve-static@2.2.0': resolution: {integrity: sha512-8mam4H1NHLtu7nmtalF7eyBH14QyOASmcxHhSfEoRyr0nP/YdoesEtU+uSRvMe96TW/HPTtkoKqQLl53N7UXMQ==} - '@types/simple-peer@9.11.9': - resolution: {integrity: sha512-6Gdl7TSS5oh9nuwKD4Pl8cSmaxWycYeZz9HLnJBNvIwWjZuGVsmHe9RwW3+9RxfhC1aIR9Z83DvaJoMw6rhkbg==} - '@types/sinonjs__fake-timers@8.1.1': resolution: {integrity: sha512-0kSuKjAS0TrGLJ0M/+8MaFkGsQhZpB6pxOmvS3K8FYI72K//YmdfoW9X2qPsAKh1mkwxGD5zib9s1FIFed6E8g==} @@ -2927,9 +2921,6 @@ packages: resolution: {integrity: sha512-xUtoPkMggbz0MPyPiIWr1Kp4aeWJjDZ6SMvURhimjdZgsRuDplF5/s9hcgGhyXMhs+6vpnuoiZ2kFiu3FMnS8Q==} engines: {node: '>=18'} - err-code@3.0.1: - resolution: {integrity: sha512-GiaH0KJUewYok+eeY05IIgjtAe4Yltygk9Wqp1V5yVWLdhf0hYZchRjNIT9bb0mSwRcIusT3cx7PJUf3zEIfUA==} - es-define-property@1.0.1: resolution: {integrity: sha512-e3nRfgfUZ4rNGL232gUgX06QNyyez04KdjFrF+LTRoOXmrOgFKDg4BCdsjW8EnT69eqdYGmRpJwiPVYNrCaW3g==} engines: {node: '>= 0.4'} @@ -4146,9 +4137,6 @@ packages: queue-microtask@1.2.3: resolution: {integrity: sha512-NuaNSa6flKT5JaSYQzJok04JzTL1CA6aGhv5rfLW3PgqA+M2ChpZQnAC8h8i4ZFkBS8X5RqkDBHA7r4hej3K9A==} - randombytes@2.1.0: - resolution: {integrity: sha512-vYl3iOX+4CKUWuxGi9Ukhie6fsqXqS9FE2Zaic4tNFD2N2QQaXOMFbuKK4QmDHC0JO6B1Zp41J0LpT0oR68amQ==} - range-parser@1.2.1: resolution: {integrity: sha512-Hrgsx+orqoygnmhFbKaHE6c296J+HTAQXoxEF6gNupROmmGJRoyzfG3ccAveqCBrwr/2yxQ5BVd/GTl5agOwSg==} engines: {node: '>= 0.6'} @@ -4336,9 +4324,6 @@ packages: simple-get@4.0.1: resolution: {integrity: sha512-brv7p5WgH0jmQJr1ZDDfKDOSeWWg+OVypG99A/5vYGPqJ6pxiaHLy8nxtFjBA7oMa01ebA9gfh1uMCFqOuXxvA==} - simple-peer@9.11.1: - resolution: {integrity: sha512-D1SaWpOW8afq1CZGWB8xTfrT3FekjQmPValrqncJMX7QFl8YwhrPTZvMCANLtgBwwdS+7zURyqxDDEmY558tTw==} - simple-update-notifier@2.0.0: resolution: {integrity: sha512-a2B9Y0KlNXl9u/vsW6sTIu9vGEpfKu2wRV6l1H3XEas/0gUIzGzBoP/IouTcUQbm9JWZLH3COxyn03TYlFax6w==} engines: {node: '>=10'} @@ -6163,10 +6148,6 @@ snapshots: '@types/http-errors': 2.0.5 '@types/node': 22.19.15 - '@types/simple-peer@9.11.9': - dependencies: - '@types/node': 22.19.15 - '@types/sinonjs__fake-timers@8.1.1': {} '@types/sizzle@2.3.10': {} @@ -7215,8 +7196,6 @@ snapshots: environment@1.1.0: {} - err-code@3.0.1: {} - es-define-property@1.0.1: {} es-errors@1.3.0: {} @@ -8480,10 +8459,6 @@ snapshots: queue-microtask@1.2.3: {} - randombytes@2.1.0: - dependencies: - safe-buffer: 5.2.1 - range-parser@1.2.1: {} raw-body@3.0.2: @@ -8747,18 +8722,6 @@ snapshots: once: 1.4.0 simple-concat: 1.0.1 - simple-peer@9.11.1(supports-color@8.1.1): - dependencies: - buffer: 6.0.3 - debug: 4.4.3(supports-color@8.1.1) - err-code: 3.0.1 - get-browser-rtc: 1.1.0 - queue-microtask: 1.2.3 - randombytes: 2.1.0 - readable-stream: 3.6.2 - transitivePeerDependencies: - - supports-color - simple-update-notifier@2.0.0: dependencies: semver: 7.8.5