@dxos/messaging 0.5.8 → 0.5.9-main.0a0e87d

package/src/signal-client/signal-local-state.ts

@@ -0,0 +1,156 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { asyncTimeout, Event } from '@dxos/async';
+import type { Any, Stream } from '@dxos/codec-protobuf';
+import { cancelWithContext, type Context } from '@dxos/context';
+import { PublicKey } from '@dxos/keys';
+import { log } from '@dxos/log';
+import { type Message as SignalMessage, type SwarmEvent } from '@dxos/protocols/proto/dxos/mesh/signal';
+import { ComplexMap, ComplexSet, safeAwaitAll } from '@dxos/util';
+
+import { type SignalRPCClient } from './signal-rpc-client';
+import type { Message } from '../signal-methods';
+
+export class SignalLocalState {
+  /**
+   * Swarm events streams. Keys represent actually joined topic and peerId.
+   */
+  private readonly _swarmStreams = new ComplexMap<{ topic: PublicKey; peerId: PublicKey }, Stream<SwarmEvent>>(
+    ({ topic, peerId }) => topic.toHex() + peerId.toHex(),
+  );
+
+  /**
+   * Represent desired joined topic and peerId.
+   */
+  private readonly _joinedTopics = new ComplexSet<{ topic: PublicKey; peerId: PublicKey }>(
+    ({ topic, peerId }) => topic.toHex() + peerId.toHex(),
+  );
+
+  /**
+   * Represent desired message subscriptions.
+   */
+  private readonly _subscribedMessages = new ComplexSet<{ peerId: PublicKey }>(({ peerId }) => peerId.toHex());
+
+  /**
+   * Message streams. Keys represents actually subscribed peers.
+   * @internal
+   */
+  readonly messageStreams = new ComplexMap<PublicKey, Stream<SignalMessage>>((key) => key.toHex());
+
+  /**
+   * Event to use in tests to wait till subscription is successfully established.
+   * @internal
+   */
+  readonly reconciled = new Event();
+
+  constructor(
+    private readonly _onMessage: (params: { author: PublicKey; recipient: PublicKey; payload: Any }) => Promise<void>,
+    private readonly _onSwarmEvent: (params: { topic: PublicKey; swarmEvent: SwarmEvent }) => Promise<void>,
+  ) {}
+
+  async safeCloseStreams(): Promise<{ failureCount: number }> {
+    const streams = ([...this._swarmStreams.values()] as Stream<any>[]).concat([...this.messageStreams.values()]);
+    this._swarmStreams.clear();
+    this.messageStreams.clear();
+    const failureCount = (await safeAwaitAll(streams, (s) => s.close())).length;
+    return { failureCount };
+  }
+
+  join({ topic, peerId }: { topic: PublicKey; peerId: PublicKey }) {
+    this._joinedTopics.add({ topic, peerId });
+  }
+
+  leave({ topic, peerId }: { topic: PublicKey; peerId: PublicKey }) {
+    void this._swarmStreams.get({ topic, peerId })?.close();
+    this._swarmStreams.delete({ topic, peerId });
+    this._joinedTopics.delete({ topic, peerId });
+  }
+
+  subscribeMessages(peerId: PublicKey) {
+    this._subscribedMessages.add({ peerId });
+  }
+
+  unsubscribeMessages(peerId: PublicKey) {
+    log('unsubscribing from messages', { peerId });
+    this._subscribedMessages.delete({ peerId });
+    void this.messageStreams.get(peerId)?.close();
+    this.messageStreams.delete(peerId);
+  }
+
+  public async reconcile(ctx: Context, client: SignalRPCClient) {
+    await this._reconcileSwarmSubscriptions(ctx, client);
+    await this._reconcileMessageSubscriptions(ctx, client);
+    this.reconciled.emit();
+  }
+
+  private async _reconcileSwarmSubscriptions(ctx: Context, client: SignalRPCClient): Promise<void> {
+    // Unsubscribe from topics that are no longer needed.
+    for (const { topic, peerId } of this._swarmStreams.keys()) {
+      // Join desired topics.
+      if (this._joinedTopics.has({ topic, peerId })) {
+        continue;
+      }
+
+      void this._swarmStreams.get({ topic, peerId })?.close();
+      this._swarmStreams.delete({ topic, peerId });
+    }
+
+    // Subscribe to topics that are needed.
+    for (const { topic, peerId } of this._joinedTopics.values()) {
+      // Join desired topics.
+      if (this._swarmStreams.has({ topic, peerId })) {
+        continue;
+      }
+
+      const swarmStream = await asyncTimeout(cancelWithContext(ctx, client.join({ topic, peerId })), 5_000);
+      // Subscribing to swarm events.
+      // TODO(mykola): What happens when the swarm stream is closed? Maybe send leave event for each peer?
+      swarmStream.subscribe(async (swarmEvent: SwarmEvent) => {
+        if (this._joinedTopics.has({ topic, peerId })) {
+          log('swarm event', { swarmEvent });
+          await this._onSwarmEvent({ topic, swarmEvent });
+        }
+      });
+
+      // Saving swarm stream.
+      this._swarmStreams.set({ topic, peerId }, swarmStream);
+    }
+  }
+
+  private async _reconcileMessageSubscriptions(ctx: Context, client: SignalRPCClient): Promise<void> {
+    // Unsubscribe from messages that are no longer needed.
+    for (const peerId of this.messageStreams.keys()) {
+      // Join desired topics.
+      if (this._subscribedMessages.has({ peerId })) {
+        continue;
+      }
+
+      void this.messageStreams.get(peerId)?.close();
+      this.messageStreams.delete(peerId);
+    }
+
+    // Subscribe to messages that are needed.
+    for (const { peerId } of this._subscribedMessages.values()) {
+      if (this.messageStreams.has(peerId)) {
+        continue;
+      }
+
+      const messageStream = await asyncTimeout(cancelWithContext(ctx, client.receiveMessages(peerId)), 5_000);
+      messageStream.subscribe(async (signalMessage: SignalMessage) => {
+        if (this._subscribedMessages.has({ peerId })) {
+          const message: Message = {
+            author: PublicKey.from(signalMessage.author),
+            recipient: PublicKey.from(signalMessage.recipient),
+            payload: signalMessage.payload,
+          };
+          await this._onMessage(message);
+        }
+      });
+
+      // Saving message stream.
+      this.messageStreams.set(peerId, messageStream);
+    }
+  }
+}

package/src/signal-client/signal-rpc-client-monitor.ts

@@ -0,0 +1,15 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { trace } from '@dxos/tracing';
+
+export class SignalRpcClientMonitor {
+  public recordClientCloseFailure(params: { failureReason: string }) {
+    trace.metrics.increment('mesh.signal.signal-rpc-client.close-failure', 1, {
+      tags: {
+        reason: params.failureReason,
+      },
+    });
+  }
+}

package/src/signal-client/signal-rpc-client.ts

@@ -4,7 +4,7 @@
 
 import WebSocket from 'isomorphic-ws';
 
-import { scheduleTaskInterval, Trigger } from '@dxos/async';
+import { scheduleTaskInterval, TimeoutError, Trigger } from '@dxos/async';
 import { type Any, type Stream } from '@dxos/codec-protobuf';
 import { Context } from '@dxos/context';
 import { invariant } from '@dxos/invariant';
@@ -14,6 +14,8 @@ import { schema, trace } from '@dxos/protocols';
 import { type Message as SignalMessage, type Signal } from '@dxos/protocols/proto/dxos/mesh/signal';
 import { createProtoRpcPeer, type ProtoRpcPeer } from '@dxos/rpc';
 
+import { SignalRpcClientMonitor } from './signal-rpc-client-monitor';
+
 const SIGNAL_KEEPALIVE_INTERVAL = 10000;
 
 interface Services {
@@ -39,9 +41,10 @@ export type SignalRPCClientParams = {
 };
 
 export class SignalRPCClient {
-  private _socket?: WebSocket;
-  private _rpc?: ProtoRpcPeer<Services>;
+  private readonly _socket: WebSocket;
+  private readonly _rpc: ProtoRpcPeer<Services>;
   private readonly _connectTrigger = new Trigger();
+
   private _keepaliveCtx?: Context;
 
   private _closed = false;
@@ -50,6 +53,8 @@ export class SignalRPCClient {
   private readonly _callbacks: SignalCallbacks;
   private readonly _closeComplete = new Trigger();
 
+  private readonly _monitor = new SignalRpcClientMonitor();
+
   constructor({ url, callbacks = {} }: SignalRPCClientParams) {
     const traceId = PublicKey.random().toHex();
     log.trace('dxos.mesh.signal-rpc-client.constructor', trace.begin({ id: traceId }));
@@ -92,6 +97,10 @@ export class SignalRPCClient {
     this._socket.onopen = async () => {
       try {
         await this._rpc!.open();
+        if (this._closed) {
+          await this._safeCloseRpc();
+          return;
+        }
         log(`RPC open ${this._url}`);
         this._callbacks.onConnected?.();
         this._connectTrigger.wake();
@@ -109,6 +118,8 @@ export class SignalRPCClient {
         );
       } catch (err: any) {
         this._callbacks.onError?.(err);
+        this._socket.close();
+        this._closed = true;
       }
     };
 
@@ -121,47 +132,46 @@ export class SignalRPCClient {
 
     this._socket.onerror = async (event: WebSocket.ErrorEvent) => {
       if (this._closed) {
-        // Ignore errors after close.
+        this._socket.close();
         return;
       }
+      this._closed = true;
 
       this._callbacks.onError?.(event.error ?? new Error(event.message));
-      this._connectTrigger.reset();
-
-      try {
-        await this._rpc?.close();
-      } catch (err) {
-        log.catch(err);
-      }
-      this._closed = true;
+      await this._safeCloseRpc();
 
-      log.warn(event.message ?? 'Socket error', { url: this._url });
+      log.warn(`Socket ${event.type ?? 'unknown'} error`, { message: event.message, url: this._url });
     };
 
     log.trace('dxos.mesh.signal-rpc-client.constructor', trace.end({ id: traceId }));
   }
 
   async close() {
-    await this._keepaliveCtx?.dispose();
+    if (this._closed) {
+      return;
+    }
     this._closed = true;
+
+    await this._keepaliveCtx?.dispose();
     try {
-      await this._rpc?.close();
+      await this._safeCloseRpc();
 
-      if (this._socket?.readyState === WebSocket.OPEN || this._socket?.readyState === WebSocket.CONNECTING) {
+      if (this._socket.readyState === WebSocket.OPEN || this._socket.readyState === WebSocket.CONNECTING) {
         // close() only starts the closing handshake.
         this._socket.close();
       }
+
       await this._closeComplete.wait({ timeout: 1_000 });
     } catch (err) {
-      log.warn('close error', err);
+      const failureReason = err instanceof TimeoutError ? 'timeout' : err?.constructor?.name ?? 'unknown';
+      this._monitor.recordClientCloseFailure({ failureReason });
     }
   }
 
   async join({ topic, peerId }: { topic: PublicKey; peerId: PublicKey }) {
     log('join', { topic, peerId, metadata: this._callbacks?.getMetadata?.() });
-    await this._connectTrigger.wait();
     invariant(!this._closed, 'SignalRPCClient is closed');
-    invariant(this._rpc, 'Rpc is not initialized');
+    await this._connectTrigger.wait();
     const swarmStream = this._rpc.rpc.Signal.join({
       swarm: topic.asUint8Array(),
       peer: peerId.asUint8Array(),
@@ -175,7 +185,6 @@ export class SignalRPCClient {
     log('receiveMessages', { peerId });
     invariant(!this._closed, 'SignalRPCClient is closed');
     await this._connectTrigger.wait();
-    invariant(this._rpc, 'Rpc is not initialized');
     const messageStream = this._rpc.rpc.Signal.receiveMessages({
       peer: peerId.asUint8Array(),
     });
@@ -187,7 +196,6 @@ export class SignalRPCClient {
     log('sendMessage', { author, recipient, payload, metadata: this._callbacks?.getMetadata?.() });
     invariant(!this._closed, 'SignalRPCClient is closed');
     await this._connectTrigger.wait();
-    invariant(this._rpc, 'Rpc is not initialized');
     await this._rpc.rpc.Signal.sendMessage({
       author: author.asUint8Array(),
       recipient: recipient.asUint8Array(),
@@ -195,4 +203,13 @@ export class SignalRPCClient {
       metadata: this._callbacks?.getMetadata?.(),
     });
   }
+
+  private async _safeCloseRpc() {
+    try {
+      this._connectTrigger.reset();
+      await this._rpc.close();
+    } catch (err) {
+      log.catch(err);
+    }
+  }
 }

package/src/signal-manager/memory-signal-manager.ts

@@ -13,7 +13,6 @@ import { type SwarmEvent } from '@dxos/protocols/proto/dxos/mesh/signal';
 import { ComplexMap, ComplexSet } from '@dxos/util';
 
 import { type SignalManager } from './signal-manager';
-import { type CommandTrace } from '../signal-client';
 import { type SignalStatus } from '../signal-methods';
 
 /**
@@ -38,7 +37,6 @@ export class MemorySignalManagerContext {
  */
 export class MemorySignalManager implements SignalManager {
   readonly statusChanged = new Event<SignalStatus[]>();
-  readonly commandTrace = new Event<CommandTrace>();
   readonly swarmEvent = new Event<{
     topic: PublicKey;
     swarmEvent: SwarmEvent;

package/src/signal-manager/signal-manager.ts

@@ -6,7 +6,6 @@ import { type Event } from '@dxos/async';
 import { type PublicKey } from '@dxos/keys';
 import { type SwarmEvent } from '@dxos/protocols/proto/dxos/mesh/signal';
 
-import { type CommandTrace } from '../signal-client';
 import { type Message, type SignalMethods, type SignalStatus } from '../signal-methods';
 
 /**
@@ -15,11 +14,9 @@ import { type Message, type SignalMethods, type SignalStatus } from '../signal-m
 export interface SignalManager extends SignalMethods {
   open(): Promise<void>;
   close(): Promise<void>;
-
   getStatus(): SignalStatus[];
 
   statusChanged: Event<SignalStatus[]>;
-  commandTrace: Event<CommandTrace>;
   swarmEvent: Event<{ topic: PublicKey; swarmEvent: SwarmEvent }>;
   onMessage: Event<Message>;
 }

package/src/signal-manager/websocket-signal-manager-monitor.ts

@@ -0,0 +1,20 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { trace } from '@dxos/tracing';
+
+export class WebsocketSignalManagerMonitor {
+  public recordRateLimitExceeded() {
+    trace.metrics.increment('mesh.signal.signal-manager.rate-limit-hit', 1);
+  }
+
+  public recordServerFailure(params: { serverName: string; willRestart: boolean }) {
+    trace.metrics.increment('mesh.signal.signal-manager.server-failure', 1, {
+      tags: {
+        server: params.serverName,
+        restarted: params.willRestart,
+      },
+    });
+  }
+}

package/src/signal-manager/websocket-signal-manager.ts

@@ -11,9 +11,11 @@ import { log } from '@dxos/log';
 import { RateLimitExceededError, TimeoutError, trace } from '@dxos/protocols';
 import { type Runtime } from '@dxos/protocols/proto/dxos/config';
 import { type SwarmEvent } from '@dxos/protocols/proto/dxos/mesh/signal';
+import { BitField, safeAwaitAll } from '@dxos/util';
 
 import { type SignalManager } from './signal-manager';
-import { type CommandTrace, SignalClient } from '../signal-client';
+import { WebsocketSignalManagerMonitor } from './websocket-signal-manager-monitor';
+import { SignalClient } from '../signal-client';
 import { type SignalClientMethods, type SignalMethods, type SignalStatus } from '../signal-methods';
 
 const MAX_SERVER_FAILURES = 5;
@@ -24,13 +26,18 @@ const WSS_SIGNAL_SERVER_REBOOT_DELAY = 3_000;
  */
 export class WebsocketSignalManager implements SignalManager {
   private readonly _servers = new Map<string, SignalClientMethods>();
+  private readonly _monitor = new WebsocketSignalManagerMonitor();
+
+  /**
+   * Used to avoid logging failed server restarts more than once until the server actually recovers.
+   */
+  private readonly _failedServersBitfield: Uint8Array;
 
   private _ctx!: Context;
   private _opened = false;
 
   readonly failureCount = new Map<string, number>();
   readonly statusChanged = new Event<SignalStatus[]>();
-  readonly commandTrace = new Event<CommandTrace>();
   readonly swarmEvent = new Event<{
     topic: PublicKey;
     swarmEvent: SwarmEvent;
@@ -66,8 +73,8 @@ export class WebsocketSignalManager implements SignalManager {
 
       this._servers.set(host.server, server);
       this.failureCount.set(host.server, 0);
-      server.commandTrace.on((trace) => this.commandTrace.emit(trace));
     }
+    this._failedServersBitfield = BitField.zeros(this._hosts.length);
   }
 
   @synchronized
@@ -80,8 +87,7 @@ export class WebsocketSignalManager implements SignalManager {
 
     this._initContext();
 
-    // TODO(burdon): Await.
-    [...this._servers.values()].forEach((server) => server.open());
+    await safeAwaitAll(this._servers.values(), (server) => server.open());
 
     this._opened = true;
     log.trace('dxos.mesh.websocket-signal-manager.open', trace.end({ id: this._instanceId }));
@@ -93,10 +99,8 @@ export class WebsocketSignalManager implements SignalManager {
       return;
     }
     this._opened = false;
-
     await this._ctx.dispose();
-
-    await Promise.all(Array.from(this._servers.values()).map((server) => server.close()));
+    await safeAwaitAll(this._servers.values(), (server) => server.close());
   }
 
   async restartServer(serverName: string) {
@@ -126,7 +130,6 @@ export class WebsocketSignalManager implements SignalManager {
   async leave({ topic, peerId }: { topic: PublicKey; peerId: PublicKey }) {
     log('leaving', { topic, peerId });
     invariant(this._opened, 'Closed');
-
     await this._forEachServer((server) => server.leave({ topic, peerId }));
   }
 
@@ -142,26 +145,35 @@ export class WebsocketSignalManager implements SignalManager {
     log('signal', { recipient });
     invariant(this._opened, 'Closed');
 
-    void this._forEachServer(async (server, serverName) => {
-      void server.sendMessage({ author, recipient, payload }).catch((err) => {
-        if (err instanceof RateLimitExceededError) {
-          log.info('WSS rate limit exceeded', { err });
-        } else if (err instanceof TimeoutError || err.constructor.name === 'TimeoutError') {
-          log.info('WSS sendMessage timeout', { err });
-          void this.checkServerFailure(serverName);
-        } else {
-          log.info(`error sending to ${serverName}`, { err });
-          void this.checkServerFailure(serverName);
-        }
-      });
+    void this._forEachServer(async (server, serverName, index) => {
+      void server
+        .sendMessage({ author, recipient, payload })
+        .then(() => this._clearServerFailedFlag(serverName, index))
+        .catch((err) => {
+          if (err instanceof RateLimitExceededError) {
+            log.info('WSS rate limit exceeded', { err });
+            this._monitor.recordRateLimitExceeded();
+          } else if (err instanceof TimeoutError || err.constructor.name === 'TimeoutError') {
+            log.info('WSS sendMessage timeout', { err });
+            void this.checkServerFailure(serverName, index);
+          } else {
+            log.warn(`error sending to ${serverName}`, { err });
+            void this.checkServerFailure(serverName, index);
+          }
+        });
     });
   }
 
   @synchronized
-  async checkServerFailure(serverName: string) {
+  async checkServerFailure(serverName: string, index: number) {
     const failureCount = this.failureCount.get(serverName!) ?? 0;
-    if (failureCount > MAX_SERVER_FAILURES) {
-      log.warn(`too many failures sending to ${serverName} (${failureCount} > ${MAX_SERVER_FAILURES}), restarting`);
+    const isRestartRequired = failureCount > MAX_SERVER_FAILURES;
+    this._monitor.recordServerFailure({ serverName, willRestart: isRestartRequired });
+    if (isRestartRequired) {
+      if (!BitField.get(this._failedServersBitfield, index)) {
+        log.warn('too many failures for ws-server, restarting', { serverName, failureCount });
+        BitField.set(this._failedServersBitfield, index, true);
+      }
       await this.restartServer(serverName!);
       this.failureCount.set(serverName!, 0);
       return;
@@ -170,6 +182,14 @@ export class WebsocketSignalManager implements SignalManager {
     this.failureCount.set(serverName!, (this.failureCount.get(serverName!) ?? 0) + 1);
   }
 
+  private _clearServerFailedFlag(serverName: string, index: number) {
+    if (BitField.get(this._failedServersBitfield, index)) {
+      log.info('server connection restored', { serverName });
+      BitField.set(this._failedServersBitfield, index, false);
+      this.failureCount.set(serverName!, 0);
+    }
+  }
+
   async subscribeMessages(peerId: PublicKey) {
     log('subscribed for message stream', { peerId });
     invariant(this._opened, 'Closed');
@@ -191,8 +211,10 @@ export class WebsocketSignalManager implements SignalManager {
   }
 
   private async _forEachServer<ReturnType>(
-    fn: (server: SignalMethods, serverName: string) => Promise<ReturnType>,
+    fn: (server: SignalMethods, serverName: string, index: number) => Promise<ReturnType>,
   ): Promise<ReturnType[]> {
-    return Promise.all(Array.from(this._servers.entries()).map(([serverName, server]) => fn(server, serverName)));
+    return Promise.all(
+      Array.from(this._servers.entries()).map(([serverName, server], idx) => fn(server, serverName, idx)),
+    );
   }
 }

package/src/signal-methods.ts

@@ -2,14 +2,16 @@
 // Copyright 2022 DXOS.org
 //
 
-import { type Any } from '@dxos/codec-protobuf';
 import { type PublicKey } from '@dxos/keys';
 import { type SignalState } from '@dxos/protocols/proto/dxos/mesh/signal';
 
 export interface Message {
   author: PublicKey;
   recipient: PublicKey;
-  payload: Any;
+  payload: {
+    type_url: string;
+    value: Uint8Array;
+  };
 }
 
 export type SignalStatus = {
@@ -43,6 +45,7 @@ export interface SignalMethods {
   /**
    * Start receiving messages from peer.
    */
+  // TODO(burdon): Return unsubscribe function. Encapsulate callback/routing here.
   subscribeMessages: (peerId: PublicKey) => Promise<void>;
 
   /**
@@ -55,7 +58,7 @@ export interface SignalMethods {
  * Signaling client.
  */
 export interface SignalClientMethods extends SignalMethods {
-  open(): Promise<void>;
-  close(): Promise<void>;
+  open(): Promise<this>;
+  close(): Promise<this>;
   getStatus(): SignalStatus;
 }