@typeberry/lib 0.5.1-aff8b4f → 0.5.1-eb84786

package/packages/jam/jamnp-s/stream-manager.js

@@ -0,0 +1,215 @@
+import { BytesBlob } from "#@typeberry/bytes";
+import { Logger } from "#@typeberry/logger";
+import { encodeMessageLength, handleMessageFragmentation, StreamErrorKind, } from "#@typeberry/networking";
+import { tryAsStreamId, tryAsStreamKind, } from "./protocol/stream.js";
+import { handleAsyncErrors } from "./utils.js";
+const logger = Logger.new(import.meta.filename, "stream");
+export class StreamManager {
+    /** A map of handlers for incoming stream kinds. */
+    incomingHandlers = new Map();
+    /** A map of handlers for outgoing stream kinds. */
+    outgoingHandlers = new Map();
+    /** A collection of open streams, peers and their handlers. */
+    streams = new Map();
+    /** Promises for stream background tasks (reading data). */
+    backgroundTasks = new Map();
+    /** Add supported incoming handlers. */
+    registerIncomingHandlers(...handlers) {
+        for (const handler of handlers) {
+            this.incomingHandlers.set(handler.kind, handler);
+        }
+    }
+    /** Add supported outgoing handlers. */
+    registerOutgoingHandlers(...handlers) {
+        for (const handler of handlers) {
+            this.outgoingHandlers.set(handler.kind, handler);
+        }
+    }
+    /** Get peer associated with a stream. */
+    getPeer(streamId) {
+        return this.streams.get(streamId)?.peer ?? null;
+    }
+    /** Wait until all of the streams are closed. */
+    async waitForFinish() {
+        for (const task of this.backgroundTasks.values()) {
+            await task;
+        }
+    }
+    /** Attempt to find an already open stream of given kind. */
+    withStreamOfKind(peerId, kind, work) {
+        // TODO [ToDr] That might not be super performant, perhaps we should
+        // maintain a mapping of Peer->open streams as well?
+        for (const streamData of this.streams.values()) {
+            if (streamData.handler.kind === kind && streamData.peer.id === peerId) {
+                work(streamData.handler, streamData.streamSender);
+                return;
+            }
+        }
+    }
+    /** Open a new stream of given kind, with the peer given. */
+    withNewStream(peer, kind, work) {
+        const handler = this.outgoingHandlers.get(kind);
+        if (handler === undefined) {
+            throw new Error(`Unsupported outgoing stream kind: ${kind}`);
+        }
+        const stream = peer.openStream();
+        const quicStream = this.registerStream(peer, handler, stream, BytesBlob.empty());
+        // send the initial byte with stream kind
+        quicStream.bufferAndSend(BytesBlob.blobFromNumbers([kind]), false);
+        work(handler, quicStream);
+    }
+    /** Handle an incoming stream. */
+    async onIncomingStream(peer, stream) {
+        const { readable, streamId } = stream;
+        const reader = readable.getReader();
+        let bytes = BytesBlob.empty();
+        try {
+            // We expect a one-byte identifier first.
+            const data = await reader.read();
+            bytes = BytesBlob.blobFrom(data.value !== undefined ? data.value : new Uint8Array());
+            logger.trace `🚰 --> [${peer.id}:${streamId}] Initial data: ${bytes}`;
+        }
+        finally {
+            reader.releaseLock();
+        }
+        if (bytes.raw.length < 1) {
+            throw new Error(`Expected 1-byte stream identifier, got: ${bytes}`);
+        }
+        // stream kind
+        const kind = tryAsStreamKind(bytes.raw[0]);
+        const handler = this.incomingHandlers.get(kind);
+        if (handler === undefined) {
+            throw new Error(`Unsupported stream kind: ${kind}`);
+        }
+        logger.log `🚰 --> [${peer.id}:${stream.streamId}] Stream identified as: ${kind}`;
+        this.registerStream(peer, handler, stream, BytesBlob.blobFrom(bytes.raw.subarray(1)));
+    }
+    registerStream(peer, handler, stream, initialData) {
+        const streamId = tryAsStreamId(stream.streamId);
+        // NOTE: `onError` callback may be called multiple times.
+        const onError = (e, kind) => {
+            this.streams.delete(streamId);
+            this.backgroundTasks.delete(streamId);
+            if (kind === StreamErrorKind.Exception) {
+                logger.error `🚰 --- [${peer.id}:${streamId}] Stream error: ${e}. Disconnecting peer.`;
+            }
+            if (kind !== StreamErrorKind.LocalClose) {
+                // whenever we have an error, we are going to inform the handler
+                // and close the stream,
+                handler.onClose(streamId, true);
+                // but also disconnect from the peer.
+                peer.disconnect();
+            }
+        };
+        stream.addOnError(onError);
+        const quicStream = new QuicStreamSender(streamId, stream, onError);
+        this.streams.set(streamId, {
+            handler,
+            streamSender: quicStream,
+            peer,
+        });
+        // when we start reading the stream, there might be some data there already,
+        // so we need to make sure that the stream is fully initialized.
+        const readStreamPromise = handleAsyncErrors(() => readStreamForever(peer, handler, quicStream, initialData, stream.readable.getReader()), (e) => onError(e, StreamErrorKind.Exception));
+        // there could be an error already during the first read, so
+        // only insert the background task when it's still active.
+        if (this.streams.has(streamId)) {
+            this.backgroundTasks.set(streamId, readStreamPromise);
+        }
+        return quicStream;
+    }
+}
+async function readStreamForever(peer, handler, quicStream, initialData, reader) {
+    // finally start listening for more data.
+    let bytes = initialData;
+    let isDone = false;
+    const callback = handleMessageFragmentation((data) => {
+        const bytes = BytesBlob.blobFrom(new Uint8Array(data));
+        logger.trace `🚰 --> [${peer.id}:${quicStream.streamId}] ${bytes}`;
+        handler.onStreamMessage(quicStream, bytes);
+    }, () => {
+        logger.error `🚰 --> [${peer.id}:${quicStream.streamId}] got too much data. Disconnecting.`;
+        peer.disconnect();
+    });
+    for (;;) {
+        // TODO [ToDr] We are going to read messages from the socket as fast as we can,
+        // yet it doesn't mean we are able to handle them as fast. This should rather
+        // be a promise, so that we can make back pressure here.
+        callback(bytes.raw);
+        if (isDone) {
+            logger.log `🚰 --> [${peer.id}:${quicStream.streamId}] remote finished.`;
+            return;
+        }
+        // await for more data
+        const data = await reader.read();
+        isDone = data.done;
+        bytes = BytesBlob.blobFrom(data.value !== undefined ? data.value : new Uint8Array());
+    }
+}
+const MAX_OUTGOING_BUFFER_BYTES = 16384;
+class QuicStreamSender {
+    streamId;
+    internal;
+    onError;
+    bufferedLength = 0;
+    bufferedData = [];
+    currentWriterPromise = null;
+    constructor(streamId, internal, onError) {
+        this.streamId = streamId;
+        this.internal = internal;
+        this.onError = onError;
+    }
+    /** Send given piece of data to the other end. */
+    bufferAndSend(data, prefixWithLength = true) {
+        if (this.bufferedLength > MAX_OUTGOING_BUFFER_BYTES) {
+            return false;
+        }
+        this.bufferedData.push({ data, addPrefix: prefixWithLength });
+        this.bufferedLength += data.length;
+        // some other async task already has a lock, so it will keep writing.
+        if (this.currentWriterPromise !== null) {
+            return true;
+        }
+        // let's start an async task to write the buffer
+        this.currentWriterPromise = handleAsyncErrors(async () => {
+            const writer = this.internal.writable.getWriter();
+            try {
+                for (;;) {
+                    const chunk = this.bufferedData.shift();
+                    if (chunk === undefined) {
+                        return;
+                    }
+                    const { data, addPrefix } = chunk;
+                    logger.trace `🚰 <-- [${this.streamId}] write: ${data}`;
+                    if (addPrefix) {
+                        await writer.write(encodeMessageLength(data.raw));
+                    }
+                    await writer.write(data.raw);
+                    this.bufferedLength -= data.length;
+                }
+            }
+            finally {
+                writer.releaseLock();
+                this.currentWriterPromise = null;
+            }
+        }, (e) => this.onError(e, StreamErrorKind.Exception));
+        return true;
+    }
+    close() {
+        handleAsyncErrors(async () => {
+            logger.trace `🚰 <-- [${this.streamId}] closing`;
+            if (this.currentWriterPromise !== null) {
+                await this.currentWriterPromise;
+            }
+            await this.internal.writable.close();
+        }, (e) => this.onError(e, StreamErrorKind.Exception));
+    }
+    /**
+     * Wait for the data to finish writing.
+     *
+     * Usually should not be used, but may be useful for tests sometimes.
+     */
+    flush() {
+        return this.currentWriterPromise ?? Promise.resolve();
+    }
+}

package/packages/jam/jamnp-s/stream-manager.test.d.ts

@@ -0,0 +1,2 @@
+export {};
+//# sourceMappingURL=stream-manager.test.d.ts.map

package/packages/jam/jamnp-s/stream-manager.test.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"stream-manager.test.d.ts","sourceRoot":"","sources":["../../../../../packages/jam/jamnp-s/stream-manager.test.ts"],"names":[],"mappings":""}

package/packages/jam/jamnp-s/stream-manager.test.js

@@ -0,0 +1,188 @@
+import assert from "node:assert";
+import { describe, it } from "node:test";
+import { createDisconnectedPeer, TestManualStream } from "#@typeberry/networking/testing.js";
+import { OK } from "#@typeberry/utils";
+import { StreamManager } from "./stream-manager.js";
+// Test StreamHandler implementation
+class TestStreamHandler {
+    kind;
+    messages = [];
+    closeCalls = [];
+    constructor(kind) {
+        this.kind = kind;
+    }
+    onStreamMessage(streamSender, message) {
+        this.messages.push({ streamSender, message });
+    }
+    onClose(streamId, isError) {
+        this.closeCalls.push({ streamId, isError });
+    }
+}
+function createTestHandler(kind) {
+    return new TestStreamHandler(kind);
+}
+describe("StreamManager", () => {
+    describe("stream management", () => {
+        it("should return null for unknown stream ID", () => {
+            const manager = new StreamManager();
+            const peer = manager.getPeer(123);
+            assert.strictEqual(peer, null);
+        });
+        it("should open new stream and call work function", async () => {
+            const manager = new StreamManager();
+            const peer = createDisconnectedPeer("peer1");
+            const handler = createTestHandler(1);
+            manager.registerOutgoingHandlers(handler);
+            let workCalled = false;
+            let senderFlush = Promise.resolve();
+            manager.withNewStream(peer, 1, (h, sender) => {
+                workCalled = true;
+                assert.strictEqual(h, handler);
+                assert.strictEqual(typeof sender.streamId, "number");
+                senderFlush = sender.flush();
+                return OK;
+            });
+            assert.strictEqual(workCalled, true);
+            assert.strictEqual(peer._openedStreams.length, 1);
+            await senderFlush;
+            // Check that stream kind was sent
+            const stream = peer._openedStreams[0];
+            const written = stream._writtenData;
+            assert.strictEqual(written.length, 1);
+            assert.deepStrictEqual(written[0], new Uint8Array([1]));
+        });
+        it("should throw error for unsupported outgoing stream kind", () => {
+            const manager = new StreamManager();
+            const peer = createDisconnectedPeer("peer1");
+            assert.throws(() => {
+                manager.withNewStream(peer, 99, () => OK);
+            }, /Unsupported outgoing stream kind: 99/);
+        });
+        it("should find existing stream of given kind", () => {
+            const manager = new StreamManager();
+            const peer = createDisconnectedPeer("peer1");
+            const handler = createTestHandler(1);
+            manager.registerOutgoingHandlers(handler);
+            // Open a stream first
+            manager.withNewStream(peer, 1, () => OK);
+            // Now try to find it
+            let foundStream = false;
+            manager.withStreamOfKind(peer.id, 1, (h) => {
+                foundStream = true;
+                assert.strictEqual(h, handler);
+                return OK;
+            });
+            assert.strictEqual(foundStream, true);
+        });
+        it("should not call work function if stream kind not found", () => {
+            const manager = new StreamManager();
+            const peer = createDisconnectedPeer("peer1");
+            let workCalled = false;
+            manager.withStreamOfKind(peer.id, 1, () => {
+                workCalled = true;
+                return OK;
+            });
+            assert.strictEqual(workCalled, false);
+        });
+    });
+    describe("incoming streams", () => {
+        it("should handle incoming stream with valid kind", async () => {
+            const manager = new StreamManager();
+            const peer = createDisconnectedPeer("peer1");
+            const handler = createTestHandler(1);
+            manager.registerIncomingHandlers(handler);
+            const stream = new TestManualStream(42);
+            // Simulate incoming stream with kind byte
+            const kindData = new Uint8Array([1, 0x41, 0x42]); // kind=1, followed by some data
+            stream._simulateIncomingData(kindData);
+            stream._incomingData.close();
+            await manager.onIncomingStream(peer, stream);
+            // Check that peer is tracked
+            const retrievedPeer = manager.getPeer(42);
+            assert.strictEqual(retrievedPeer, peer);
+        });
+        it("should throw error for stream without kind byte", async () => {
+            const manager = new StreamManager();
+            const peer = createDisconnectedPeer("peer1");
+            const stream = new TestManualStream(42);
+            // Simulate empty stream
+            stream._incomingData.close();
+            await assert.rejects(() => manager.onIncomingStream(peer, stream), /Expected 1-byte stream identifier/);
+        });
+        it("should throw error for unsupported incoming stream kind", async () => {
+            const manager = new StreamManager();
+            const peer = createDisconnectedPeer("peer2");
+            const stream = new TestManualStream(42);
+            // Simulate incoming stream with unsupported kind
+            stream._simulateIncomingData(new Uint8Array([99])); // unsupported kind
+            stream._incomingData.close();
+            await assert.rejects(() => manager.onIncomingStream(peer, stream), /Unsupported stream kind: 99/);
+        });
+    });
+    describe("message handling", () => {
+        it("should handle fragmented messages", async () => {
+            const manager = new StreamManager();
+            const peer = createDisconnectedPeer("peer3");
+            const handler = createTestHandler(1);
+            manager.registerIncomingHandlers(handler);
+            const stream = new TestManualStream(42);
+            // Send initial kind byte
+            stream._simulateIncomingData(new Uint8Array([1]));
+            // Start handling the stream
+            const handlePromise = manager.onIncomingStream(peer, stream);
+            // Send a complete message with length prefix
+            const message = new Uint8Array([0x41, 0x42, 0x43]);
+            const lengthPrefix = new Uint8Array([3, 0, 0, 0]); // length = 3
+            stream._simulateIncomingData(new Uint8Array([...lengthPrefix, ...message]));
+            // Close the stream
+            stream._incomingData.close();
+            await handlePromise;
+            await manager.waitForFinish();
+            // Check that handler received the message
+            assert.strictEqual(handler.messages.length, 1);
+            assert.deepStrictEqual(handler.messages[0].message.raw, message);
+        });
+    });
+    describe("error handling", () => {
+        it("should handle stream errors and disconnect peer", async () => {
+            const manager = new StreamManager();
+            const peer = createDisconnectedPeer("peer2");
+            const handler = createTestHandler(1);
+            manager.registerIncomingHandlers(handler);
+            const stream = new TestManualStream(42);
+            // Send initial kind byte
+            stream._simulateIncomingData(new Uint8Array([1]));
+            // Start handling the stream
+            await manager.onIncomingStream(peer, stream);
+            // Simulate stream error
+            await stream._incomingData.abort(new Error("Test error"));
+            // Check that peer was disconnected
+            assert.strictEqual(peer._disconnectCalled, true);
+            // Check that handler's onClose was called with error=true
+            assert.strictEqual(handler.closeCalls.length, 1);
+            assert.strictEqual(handler.closeCalls[0].isError, true);
+        });
+    });
+    describe("lifecycle", () => {
+        it("should wait for all streams to finish", async () => {
+            const manager = new StreamManager();
+            const peer = createDisconnectedPeer("1");
+            const handler = createTestHandler(1);
+            manager.registerIncomingHandlers(handler);
+            const stream1 = new TestManualStream(42);
+            const stream2 = new TestManualStream(43);
+            // Start handling both streams
+            stream1._simulateIncomingData(new Uint8Array([1]));
+            stream2._simulateIncomingData(new Uint8Array([1]));
+            // Close streams immediately
+            stream1._incomingData.close();
+            stream2._incomingData.close();
+            const handle1 = manager.onIncomingStream(peer, stream1);
+            const handle2 = manager.onIncomingStream(peer, stream2);
+            await Promise.all([handle1, handle2]);
+            await manager.waitForFinish();
+            // If we reach here, waitForFinish worked correctly
+            assert.ok(true);
+        });
+    });
+});

package/packages/jam/jamnp-s/tasks/sync.d.ts

@@ -0,0 +1,72 @@
+import { type BlockView, Header, type HeaderHash, type HeaderView, type TimeSlot } from "#@typeberry/block";
+import type { ChainSpec } from "#@typeberry/config";
+import type { BlocksDb } from "#@typeberry/database";
+import { type Blake2b, WithHash } from "#@typeberry/hash";
+import type { Peer, PeerId } from "#@typeberry/networking";
+import type { AuxData, Connections } from "../peers.js";
+import type { StreamManager } from "../stream-manager.js";
+export declare const SYNC_AUX: AuxData<SyncAux>;
+type SyncAux = {
+    finalBlockHash: HeaderHash;
+    finalBlockSlot: TimeSlot;
+    bestHeader: WithHash<HeaderHash, Header> | null;
+};
+export declare class SyncTask {
+    private readonly spec;
+    private readonly blake2b;
+    private readonly streamManager;
+    private readonly connections;
+    private readonly blocks;
+    private readonly onNewBlocks;
+    static start(spec: ChainSpec, blake2b: Blake2b, streamManager: StreamManager, connections: Connections, blocks: BlocksDb, onNewBlocks: (blocks: BlockView[], peerId: PeerId) => void): SyncTask;
+    private othersBest;
+    private constructor();
+    private onUp0Handshake;
+    private onUp0Annoucement;
+    private getUp0Handshake;
+    /**
+     * Open a UP0 stream with given peer.
+     *
+     * This will automatically send a handshake as well.
+     */
+    openUp0(peer: Peer): void;
+    /** Broadcast header we have seen or produced to our peers. */
+    broadcastHeader(header: WithHash<HeaderHash, HeaderView>): void;
+    private handleGetBlockSequence;
+    /** Should be called periodically to request best seen blocks from other peers. */
+    maintainSync(): {
+        kind: SyncResult;
+        ours?: undefined;
+        theirs?: undefined;
+        requested?: undefined;
+    } | {
+        kind: SyncResult;
+        ours: number & import("@typeberry/numbers").WithBytesRepresentation<4> & import("@typeberry/utils").WithOpaque<"TimeSlot[u32]">;
+        theirs: number & import("@typeberry/numbers").WithBytesRepresentation<4> & import("@typeberry/utils").WithOpaque<"TimeSlot[u32]">;
+        requested?: undefined;
+    } | {
+        kind: SyncResult;
+        ours: number & import("@typeberry/numbers").WithBytesRepresentation<4> & import("@typeberry/utils").WithOpaque<"TimeSlot[u32]">;
+        requested: RequestedBlocks[];
+        theirs?: undefined;
+    };
+}
+/** Some extra details about how maintaining sync went. */
+export declare enum SyncResult {
+    /** We didn't find our best header? */
+    OurBestHeaderMissing = 1,
+    /** There is no new blocks that we can sync. */
+    NoNewBlocks = 2,
+    /** Sent request to some peers. */
+    BlocksRequested = 3
+}
+/** Information about blocks requested from other peers. */
+export type RequestedBlocks = {
+    /** Peer id we sent the request to. */
+    peerId: PeerId;
+    theirs: TimeSlot;
+    /** Number of blocks requested. */
+    count: number;
+};
+export {};
+//# sourceMappingURL=sync.d.ts.map

package/packages/jam/jamnp-s/tasks/sync.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"sync.d.ts","sourceRoot":"","sources":["../../../../../../packages/jam/jamnp-s/tasks/sync.ts"],"names":[],"mappings":"AAAA,OAAO,EACL,KAAK,SAAS,EACd,MAAM,EACN,KAAK,UAAU,EACf,KAAK,UAAU,EACf,KAAK,QAAQ,EAEd,MAAM,kBAAkB,CAAC;AAE1B,OAAO,KAAK,EAAE,SAAS,EAAE,MAAM,mBAAmB,CAAC;AACnD,OAAO,KAAK,EAAE,QAAQ,EAAE,MAAM,qBAAqB,CAAC;AACpD,OAAO,EAAE,KAAK,OAAO,EAAE,QAAQ,EAAE,MAAM,iBAAiB,CAAC;AAEzD,OAAO,KAAK,EAAE,IAAI,EAAE,MAAM,EAAE,MAAM,uBAAuB,CAAC;AAG1D,OAAO,KAAK,EAAE,OAAO,EAAE,WAAW,EAAE,MAAM,aAAa,CAAC;AAGxD,OAAO,KAAK,EAAE,aAAa,EAAE,MAAM,sBAAsB,CAAC;AAG1D,eAAO,MAAM,QAAQ,EAAE,OAAO,CAAC,OAAO,CAErC,CAAC;AAEF,KAAK,OAAO,GAAG;IACb,cAAc,EAAE,UAAU,CAAC;IAC3B,cAAc,EAAE,QAAQ,CAAC;IACzB,UAAU,EAAE,QAAQ,CAAC,UAAU,EAAE,MAAM,CAAC,GAAG,IAAI,CAAC;CACjD,CAAC;AASF,qBAAa,QAAQ;IA8DjB,OAAO,CAAC,QAAQ,CAAC,IAAI;IACrB,OAAO,CAAC,QAAQ,CAAC,OAAO;IACxB,OAAO,CAAC,QAAQ,CAAC,aAAa;IAC9B,OAAO,CAAC,QAAQ,CAAC,WAAW;IAC5B,OAAO,CAAC,QAAQ,CAAC,MAAM;IACvB,OAAO,CAAC,QAAQ,CAAC,WAAW;IAlE9B,MAAM,CAAC,KAAK,CACV,IAAI,EAAE,SAAS,EACf,OAAO,EAAE,OAAO,EAChB,aAAa,EAAE,aAAa,EAC5B,WAAW,EAAE,WAAW,EACxB,MAAM,EAAE,QAAQ,EAEhB,WAAW,EAAE,CAAC,MAAM,EAAE,SAAS,EAAE,EAAE,MAAM,EAAE,MAAM,KAAK,IAAI;IAmD5D,OAAO,CAAC,UAAU,CAAkB;IAEpC,OAAO;IAoBP,OAAO,CAAC,cAAc;IAsBtB,OAAO,CAAC,gBAAgB;IAwCxB,OAAO,CAAC,eAAe;IAiBvB;;;;OAIG;IACH,OAAO,CAAC,IAAI,EAAE,IAAI;IAOlB,8DAA8D;IAE9D,eAAe,CAAC,MAAM,EAAE,QAAQ,CAAC,UAAU,EAAE,UAAU,CAAC;IAsBxD,OAAO,CAAC,sBAAsB;IA4B9B,kFAAkF;IAClF,YAAY;;;;;;;;;;;;;;;;CAmFb;AAED,0DAA0D;AAC1D,oBAAY,UAAU;IACpB,sCAAsC;IACtC,oBAAoB,IAAI;IACxB,+CAA+C;IAC/C,WAAW,IAAI;IACf,kCAAkC;IAClC,eAAe,IAAI;CACpB;AAED,2DAA2D;AAC3D,MAAM,MAAM,eAAe,GAAG;IAC5B,sCAAsC;IACtC,MAAM,EAAE,MAAM,CAAC;IAEf,MAAM,EAAE,QAAQ,CAAC;IACjB,kCAAkC;IAClC,KAAK,EAAE,MAAM,CAAC;CACf,CAAC"}