@dxos/teleport 0.1.55-main.ff36828 → 0.1.55

package/src/muxing/balancer.test.ts

@@ -0,0 +1,50 @@
+//
+// Copyright 2023 DXOS.org
+//
+
+import { expect } from 'chai';
+import * as varint from 'varint';
+
+import { describe, test } from '@dxos/test';
+
+import { encodeChunk, decodeChunk } from './balancer';
+
+describe('Balancer', () => {
+  test('varints', () => {
+    const values = [0, 1, 5, 127, 128, 255, 256, 257, 1024, 1024 * 1024];
+    for (const value of values) {
+      const encoded = varint.encode(value, Buffer.allocUnsafe(4)).slice(0, varint.encode.bytes);
+      const length = varint.encode.bytes;
+      expect(encoded.length).to.eq(length);
+
+      const decoded = varint.decode(encoded);
+      expect(decoded).to.equal(value);
+      expect(varint.decode.bytes).to.equal(length);
+    }
+  });
+
+  it('should correctly encode and decode a chunk without dataLength', () => {
+    const channelId = 1;
+    const chunk = Uint8Array.from([0x11, 0x22, 0x33]);
+
+    const encoded = encodeChunk({ chunk, channelId });
+    const decoded = decodeChunk(encoded, () => false);
+
+    expect(decoded.channelId).to.equal(channelId);
+    expect(decoded.dataLength).to.equal(undefined);
+    expect(decoded.chunk).to.deep.equal(chunk);
+  });
+
+  it('should correctly encode and decode a chunk with dataLength', () => {
+    const channelId = 2;
+    const chunk = Uint8Array.from([0x44, 0x55, 0x66]);
+    const dataLength = chunk.length;
+
+    const encoded = encodeChunk({ chunk, channelId, dataLength });
+    const decoded = decodeChunk(encoded, (channelId) => channelId === 2);
+
+    expect(decoded.channelId).to.equal(channelId);
+    expect(decoded.dataLength).to.equal(dataLength);
+    expect(decoded.chunk).to.deep.equal(chunk);
+  });
+});

package/src/muxing/balancer.ts

@@ -2,14 +2,29 @@
 // Copyright 2023 DXOS.org
 //
 
-import { Trigger } from '@dxos/async';
+import * as varint from 'varint';
 
-// TODO(egorgripasov): Is BinaryPort a better name?
-import { RpcPort } from './rpc-port';
+import { Trigger, Event } from '@dxos/async';
+import { log } from '@dxos/log';
+
+import { Framer } from './framer';
+
+const MAX_CHUNK_SIZE = 8192;
 
 type Chunk = {
-  msg: Uint8Array;
-  trigger: Trigger;
+  chunk: Uint8Array;
+  channelId: number;
+  dataLength?: number;
+};
+
+type ChunkEnvelope = {
+  msg: Buffer;
+  trigger?: Trigger;
+};
+
+type ChannelBuffer = {
+  buffer: Buffer;
+  msgLength: number;
 };
 
 /**
@@ -22,45 +37,96 @@ export class Balancer {
   private _lastCallerIndex = 0;
   private _channels: number[] = [];
 
+  private readonly _framer = new Framer();
   // TODO(egorgripasov): Will cause a memory leak if channels do not appreciate the backpressure.
-  private readonly _calls: Map<number, Chunk[]> = new Map();
+  private readonly _sendBuffers: Map<number, ChunkEnvelope[]> = new Map();
+  private readonly _receiveBuffers = new Map<number, ChannelBuffer>();
 
-  constructor(private readonly _port: RpcPort, private readonly _sysChannelId: number) {
+  public incomingData = new Event<Uint8Array>();
+  public readonly stream = this._framer.stream;
+
+  constructor(private readonly _sysChannelId: number) {
     this._channels.push(_sysChannelId);
+
+    // Handle incoming messages.
+    this._framer.port.subscribe(this._processIncomingMessage.bind(this));
+  }
+
+  get bytesSent() {
+    return this._framer.bytesSent;
+  }
+
+  get bytesReceived() {
+    return this._framer.bytesReceived;
   }
 
   addChannel(channel: number) {
     this._channels.push(channel);
   }
 
-  pushChunk(msg: Uint8Array, trigger: Trigger, channelId: number) {
-    const noCalls = this._calls.size === 0;
+  pushData(data: Uint8Array, trigger: Trigger, channelId: number) {
+    const noCalls = this._sendBuffers.size === 0;
 
     if (!this._channels.includes(channelId)) {
       throw new Error(`Unknown channel ${channelId}`);
     }
 
-    if (!this._calls.has(channelId)) {
-      this._calls.set(channelId, []);
+    if (!this._sendBuffers.has(channelId)) {
+      this._sendBuffers.set(channelId, []);
     }
 
-    const channelCalls = this._calls.get(channelId)!;
-    channelCalls.push({ msg, trigger });
+    const sendBuffer = this._sendBuffers.get(channelId)!;
+
+    const chunks = [];
+    for (let idx = 0; idx < data.length; idx += MAX_CHUNK_SIZE) {
+      chunks.push(data.subarray(idx, idx + MAX_CHUNK_SIZE));
+    }
+
+    chunks.forEach((chunk, index) => {
+      const msg = encodeChunk({
+        chunk,
+        channelId,
+        dataLength: index === 0 ? data.length : undefined,
+      });
+      sendBuffer.push({ msg, trigger: index === chunks.length - 1 ? trigger : undefined });
+    });
 
     // Start processing calls if this is the first call.
     if (noCalls) {
-      process.nextTick(async () => {
-        await this._processCalls();
-      });
+      this._sendChunks().catch((err) => log.catch(err));
     }
   }
 
   destroy() {
-    this._calls.clear();
+    this._sendBuffers.clear();
+    this._framer.destroy();
+  }
+
+  private _processIncomingMessage(msg: Uint8Array) {
+    const { channelId, dataLength, chunk } = decodeChunk(msg, (channelId) => !this._receiveBuffers.has(channelId));
+    if (!this._receiveBuffers.has(channelId)) {
+      if (chunk.length < dataLength!) {
+        this._receiveBuffers.set(channelId, {
+          buffer: Buffer.from(chunk),
+          msgLength: dataLength!,
+        });
+      } else {
+        this.incomingData.emit(chunk);
+      }
+    } else {
+      const channelBuffer = this._receiveBuffers.get(channelId)!;
+      channelBuffer.buffer = Buffer.concat([channelBuffer.buffer, chunk]);
+      if (channelBuffer.buffer.length < channelBuffer.msgLength) {
+        return;
+      }
+      const msg = channelBuffer.buffer;
+      this._receiveBuffers.delete(channelId);
+      this.incomingData.emit(msg);
+    }
   }
 
   private _getNextCallerId() {
-    if (this._calls.has(this._sysChannelId)) {
+    if (this._sendBuffers.has(this._sysChannelId)) {
       return this._sysChannelId;
     }
 
@@ -70,37 +136,64 @@ export class Balancer {
     return this._channels[index];
   }
 
-  private _getNextCall(): Chunk {
-    let call;
-    while (!call) {
+  private _getNextChunk(): ChunkEnvelope {
+    let chunk;
+    while (!chunk) {
       const channelId = this._getNextCallerId();
-      const channelCalls = this._calls.get(channelId);
-      if (!channelCalls) {
+      const sendBuffer = this._sendBuffers.get(channelId);
+      if (!sendBuffer) {
         continue;
       }
 
-      call = channelCalls.shift();
-      if (channelCalls.length === 0) {
-        this._calls.delete(channelId);
+      chunk = sendBuffer.shift();
+      if (sendBuffer.length === 0) {
+        this._sendBuffers.delete(channelId);
       }
     }
-    return call;
+    return chunk;
   }
 
-  private async _processCalls() {
-    if (this._calls.size === 0) {
+  private async _sendChunks() {
+    if (this._sendBuffers.size === 0) {
       return;
     }
 
-    const call = this._getNextCall();
+    const chunk = this._getNextChunk();
 
     try {
-      await this._port.send(call.msg);
-      call.trigger.wake();
+      await this._framer.port.send(chunk.msg);
+      chunk.trigger?.wake();
     } catch (err: any) {
-      call.trigger.throw(err);
+      chunk.trigger?.throw(err);
     }
 
-    await this._processCalls();
+    await this._sendChunks();
   }
 }
+
+export const encodeChunk = ({ channelId, dataLength, chunk }: Chunk): Buffer => {
+  const channelTagLength = varint.encodingLength(channelId);
+  const dataLengthLength = dataLength ? varint.encodingLength(dataLength) : 0;
+  const message = Buffer.allocUnsafe(channelTagLength + dataLengthLength + chunk.length);
+  varint.encode(channelId, message);
+  if (dataLength) {
+    varint.encode(dataLength, message, channelTagLength);
+  }
+  message.set(chunk, channelTagLength + dataLengthLength);
+  return message;
+};
+
+export const decodeChunk = (data: Uint8Array, withLength: (channelId: number) => boolean): Chunk => {
+  const channelId = varint.decode(data);
+  let dataLength: number | undefined;
+  let offset = varint.decode.bytes;
+
+  if (withLength(channelId)) {
+    dataLength = varint.decode(data, offset);
+    offset += varint.decode.bytes;
+  }
+
+  const chunk = data.subarray(offset);
+
+  return { channelId, dataLength, chunk };
+};

package/src/muxing/framer.test.ts

@@ -5,7 +5,6 @@
 import { expect } from 'chai';
 import { pipeline } from 'node:stream';
 import randomBytes from 'randombytes';
-import * as varint from 'varint';
 import waitForExpect from 'wait-for-expect';
 
 import { sleep } from '@dxos/async';
@@ -50,21 +49,8 @@ const pipe = (a: NodeJS.ReadWriteStream, b: NodeJS.ReadWriteStream): (() => void
 };
 
 describe('Framer', () => {
-  test('varints', () => {
-    const values = [0, 1, 5, 127, 128, 255, 256, 257, 1024, 1024 * 1024];
-    for (const value of values) {
-      const encoded = varint.encode(value, Buffer.allocUnsafe(4)).slice(0, varint.encode.bytes);
-      const length = varint.encode.bytes;
-      expect(encoded.length).to.eq(length);
-
-      const decoded = varint.decode(encoded);
-      expect(decoded).to.equal(value);
-      expect(varint.decode.bytes).to.equal(length);
-    }
-  });
-
   test('frame encoding', () => {
-    const sizes = [0, 1, 5, 127, 128, 255, 256, 257, 1024, 1024 * 1024];
+    const sizes = [0, 1, 5, 127, 128, 255, 256, 257, 1024, 1024 * 60];
     for (const size of sizes) {
       const payload = randomBytes(size);
       const frame = encodeFrame(payload);

package/src/muxing/framer.ts

@@ -4,10 +4,11 @@
 
 import { Duplex } from 'node:stream';
 import invariant from 'tiny-invariant';
-import * as varint from 'varint';
 
 import { RpcPort } from './rpc-port';
 
+const FRAME_LENGTH_SIZE = 2;
+
 /**
  * Framer that turns a stream of binary messages into a framed RpcPort.
  *
@@ -18,14 +19,21 @@ export class Framer {
   private _messageCb?: (msg: Uint8Array) => void;
   private _subscribeCb?: () => void;
   private _buffer?: Buffer; // The rest of the bytes from the previous write call.
-  private _responseQueue: (() => void)[] = [];
+  private _sendCallbacks: (() => void)[] = [];
+
+  private _bytesSent = 0;
+  private _bytesReceived = 0;
 
   private readonly _stream = new Duplex({
     objectMode: false,
-    read: () => {},
+    read: () => {
+      this._processResponseQueue();
+    },
     write: (chunk, encoding, callback) => {
       invariant(!this._subscribeCb, 'Internal Framer bug. Concurrent writes detected.');
 
+      this._bytesReceived += chunk.length;
+
       if (this._buffer && this._buffer.length > 0) {
         this._buffer = Buffer.concat([this._buffer, chunk]);
       } else {
@@ -50,11 +58,13 @@ export class Framer {
     send: (message) => {
       // log('write', { len: message.length, frame: Buffer.from(message).toString('hex') })
       return new Promise<void>((resolve) => {
-        const canContinue = this._stream.push(encodeFrame(message));
+        const frame = encodeFrame(message);
+        this._bytesSent += frame.length;
+        const canContinue = this._stream.push(frame);
         if (!canContinue) {
-          this._responseQueue.push(resolve);
+          this._sendCallbacks.push(resolve);
         } else {
-          process.nextTick(resolve);
+          resolve();
         }
       });
     },
@@ -72,13 +82,17 @@ export class Framer {
     return this._stream;
   }
 
-  constructor() {
-    this.stream.on('drain', this._processResponseQueue.bind(this));
+  get bytesSent() {
+    return this._bytesSent;
+  }
+
+  get bytesReceived() {
+    return this._bytesReceived;
   }
 
   private _processResponseQueue() {
-    const responseQueue = this._responseQueue;
-    this._responseQueue = [];
+    const responseQueue = this._sendCallbacks;
+    this._sendCallbacks = [];
     responseQueue.forEach((cb) => cb());
   }
 
@@ -109,7 +123,6 @@ export class Framer {
 
   destroy() {
     // TODO(dmaretskyi): Call stream.end() instead?
-    this._stream.removeAllListeners('drain');
     this._stream.destroy();
   }
 }
@@ -118,35 +131,30 @@ export class Framer {
  * Attempts to read a frame from the input buffer.
  */
 export const decodeFrame = (buffer: Buffer, offset: number): { payload: Buffer; bytesConsumed: number } | undefined => {
-  try {
-    const frameLength = varint.decode(buffer, offset);
-    const tagLength = varint.decode.bytes;
-
-    if (buffer.length < offset + tagLength + frameLength) {
-      // Not enough bytes to read the frame.
-      return undefined;
-    }
+  if (buffer.length < offset + FRAME_LENGTH_SIZE) {
+    // Not enough bytes to read the frame length.
+    return undefined;
+  }
 
-    const payload = buffer.subarray(offset + tagLength, offset + tagLength + frameLength);
+  const frameLength = buffer.readUInt16BE(offset);
+  const bytesConsumed = FRAME_LENGTH_SIZE + frameLength;
 
-    return {
-      payload,
-      bytesConsumed: tagLength + frameLength,
-    };
-  } catch (err) {
-    if (err instanceof RangeError) {
-      // Not enough bytes to read the tag.
-      return undefined;
-    } else {
-      throw err;
-    }
+  if (buffer.length < offset + bytesConsumed) {
+    // Not enough bytes to read the frame.
+    return undefined;
   }
+
+  const payload = buffer.subarray(offset + FRAME_LENGTH_SIZE, offset + bytesConsumed);
+
+  return {
+    payload,
+    bytesConsumed,
+  };
 };
 
 export const encodeFrame = (payload: Uint8Array): Buffer => {
-  const tagLength = varint.encodingLength(payload.length);
-  const frame = Buffer.allocUnsafe(tagLength + payload.length);
-  varint.encode(payload.length, frame);
-  frame.set(payload, tagLength);
+  const frame = Buffer.allocUnsafe(FRAME_LENGTH_SIZE + payload.length);
+  frame.writeUInt16BE(payload.length, 0);
+  frame.set(payload, FRAME_LENGTH_SIZE);
   return frame;
 };

package/src/muxing/muxer.ts

@@ -14,7 +14,6 @@ import { ConnectionInfo } from '@dxos/protocols/proto/dxos/devtools/swarm';
 import { Command } from '@dxos/protocols/proto/dxos/mesh/muxer';
 
 import { Balancer } from './balancer';
-import { Framer } from './framer';
 import { RpcPort } from './rpc-port';
 
 const Command = schema.getCodecForType('dxos.mesh.muxer.Command');
@@ -32,9 +31,17 @@ export type CreateChannelOpts = {
   contentType?: string;
 };
 
+export type MuxerStats = {
+  channels: ConnectionInfo.StreamStats[];
+  bytesSent: number;
+  bytesReceived: number;
+};
+
 const STATS_INTERVAL = 1000;
 
-const SYSTEM_CHANNEL_ID = -1;
+const MAX_SAFE_FRAME_SIZE = 1_000_000;
+
+const SYSTEM_CHANNEL_ID = 0;
 
 /**
  * Channel based multiplexer.
@@ -47,26 +54,23 @@ const SYSTEM_CHANNEL_ID = -1;
  * A higher level API (could be build on top of this muxer) for channel discovery is required.
  */
 export class Muxer {
-  private readonly _framer = new Framer();
-  private readonly _balancer = new Balancer(this._framer.port, SYSTEM_CHANNEL_ID);
-  public readonly stream = this._framer.stream;
-
+  private readonly _balancer = new Balancer(SYSTEM_CHANNEL_ID);
   private readonly _channelsByLocalId = new Map<number, Channel>();
   private readonly _channelsByTag = new Map<string, Channel>();
+  private readonly _ctx = new Context();
 
-  private _nextId = 0;
+  private _nextId = 1;
   private _destroyed = false;
   private _destroying = false;
 
   public close = new Event<Error | undefined>();
+  public statsUpdated = new Event<MuxerStats>();
 
-  public statsUpdated = new Event<ConnectionInfo.StreamStats[]>();
-
-  private readonly _ctx = new Context();
+  public readonly stream = this._balancer.stream;
 
   constructor() {
     // Add a channel for control messages.
-    this._framer.port.subscribe(async (msg) => {
+    this._balancer.incomingData.on(async (msg) => {
       await this._handleCommand(Command.decode(msg));
     });
 
@@ -220,7 +224,6 @@ export class Muxer {
 
     this._destroyed = true;
     this._balancer.destroy();
-    this._framer.destroy();
 
     for (const channel of this._channelsByTag.values()) {
       channel.destroy?.(err);
@@ -276,7 +279,7 @@ export class Muxer {
   private async _sendCommand(cmd: Command, channelId = -1) {
     try {
       const trigger = new Trigger<void>();
-      this._balancer.pushChunk(Command.encode(cmd), trigger, channelId);
+      this._balancer.pushData(Command.encode(cmd), trigger, channelId);
       await trigger.wait();
     } catch (err: any) {
       this.destroy(err);
@@ -307,6 +310,10 @@ export class Muxer {
   }
 
   private async _sendData(channel: Channel, data: Uint8Array): Promise<void> {
+    if (data.length > MAX_SAFE_FRAME_SIZE) {
+      log.warn('frame size exceeds maximum safe value', { size: data.length, threshold: MAX_SAFE_FRAME_SIZE });
+    }
+
     channel.stats.bytesSent += data.length;
     if (channel.remoteId === null) {
       // Remote side has not opened the channel yet.
@@ -337,14 +344,20 @@ export class Muxer {
       return;
     }
 
-    this.statsUpdated.emit(
-      Array.from(this._channelsByTag.values()).map((channel) => ({
+    const bytesSent = this._balancer.bytesSent;
+    const bytesReceived = this._balancer.bytesReceived;
+
+    this.statsUpdated.emit({
+      channels: Array.from(this._channelsByTag.values()).map((channel) => ({
         id: channel.id,
         tag: channel.tag,
+        contentType: channel.contentType,
         bytesSent: channel.stats.bytesSent,
         bytesReceived: channel.stats.bytesReceived,
       })),
-    );
+      bytesSent,
+      bytesReceived,
+    });
   }
 }
 

package/src/teleport.ts

@@ -11,12 +11,11 @@ import { failUndefined } from '@dxos/debug';
 import { PublicKey } from '@dxos/keys';
 import { log } from '@dxos/log';
 import { schema, RpcClosedError } from '@dxos/protocols';
-import { ConnectionInfo } from '@dxos/protocols/proto/dxos/devtools/swarm';
 import { ControlService } from '@dxos/protocols/proto/dxos/mesh/teleport/control';
 import { createProtoRpcPeer, ProtoRpcPeer } from '@dxos/rpc';
 import { Callback } from '@dxos/util';
 
-import { CreateChannelOpts, Muxer, RpcPort } from './muxing';
+import { CreateChannelOpts, Muxer, MuxerStats, RpcPort } from './muxing';
 
 export type TeleportParams = {
   initiator: boolean;
@@ -84,13 +83,23 @@ export class Teleport {
         await this.destroy(err);
       });
     }
+
+    this._muxer.statsUpdated.on((stats) => {
+      log.trace('dxos.mesh.teleport.stats', {
+        localPeerId,
+        remotePeerId,
+        bytesSent: stats.bytesSent,
+        bytesReceived: stats.bytesReceived,
+        channels: stats.channels,
+      });
+    });
   }
 
   get stream(): Duplex {
     return this._muxer.stream;
   }
 
-  get stats(): Event<ConnectionInfo.StreamStats[]> {
+  get stats(): Event<MuxerStats> {
     return this._muxer.statsUpdated;
   }
 
@@ -105,7 +114,6 @@ export class Teleport {
 
   async close(err?: Error) {
     // TODO(dmaretskyi): Try soft close.
-
     await this.destroy(err);
   }
 
@@ -207,6 +215,10 @@ export interface TeleportExtension {
   onClose(err?: Error): Promise<void>;
 }
 
+type ControlRpcBundle = {
+  Control: ControlService;
+};
+
 type ControlExtensionOpts = {
   heartbeatInterval: number;
   heartbeatTimeout: number;
@@ -220,13 +232,17 @@ class ControlExtension implements TeleportExtension {
     },
   });
 
+  public readonly onExtensionRegistered = new Callback<(extensionName: string) => void>();
+
   private _extensionContext!: ExtensionContext;
   private _rpc!: ProtoRpcPeer<{ Control: ControlService }>;
 
-  public readonly onExtensionRegistered = new Callback<(extensionName: string) => void>();
-
   constructor(private readonly opts: ControlExtensionOpts) {}
 
+  async registerExtension(name: string) {
+    await this._rpc.rpc.Control.registerExtension({ name });
+  }
+
   async onOpen(extensionContext: ExtensionContext): Promise<void> {
     this._extensionContext = extensionContext;
 
@@ -273,12 +289,4 @@ class ControlExtension implements TeleportExtension {
     await this._ctx.dispose();
     await this._rpc.close();
   }
-
-  async registerExtension(name: string) {
-    await this._rpc.rpc.Control.registerExtension({ name });
-  }
 }
-
-type ControlRpcBundle = {
-  Control: ControlService;
-};

package/src/testing/test-extension-with-streams.ts

@@ -84,6 +84,20 @@ export class TestExtensionWithStreams implements TeleportExtension {
     networkStream.on('close', () => {
       networkStream.removeAllListeners();
     });
+
+    streamEntry.reportingTimer = setInterval(() => {
+      const { bytesSent, bytesReceived, sendErrors, receiveErrors } = streamEntry;
+      // log.info('stream stats', { streamTag, bytesSent, bytesReceived, sendErrors, receiveErrors });
+      log.trace('dxos.test.stream-stats', {
+        streamTag,
+        bytesSent,
+        bytesReceived,
+        sendErrors,
+        receiveErrors,
+        from: this.extensionContext?.localPeerId,
+        to: this.extensionContext?.remotePeerId,
+      });
+    }, 100);
   }
 
   private _closeStream(streamTag: string): Stats {
@@ -92,6 +106,7 @@ export class TestExtensionWithStreams implements TeleportExtension {
     const stream = this._streams.get(streamTag)!;
 
     clearTimeout(stream.timer);
+    clearTimeout(stream.reportingTimer);
 
     const { bytesSent, bytesReceived, sendErrors, receiveErrors, startTimestamp } = stream;
 
@@ -244,4 +259,5 @@ type TestStream = {
   receiveErrors: number;
   timer?: NodeJS.Timer;
   startTimestamp?: number;
+  reportingTimer?: NodeJS.Timer;
 };