@milaboratories/pl-client 2.16.11 → 2.16.12

package/src/core/websocket_stream.ts

@@ -0,0 +1,412 @@
+import { WebSocket } from 'undici';
+import {
+  TxAPI_ClientMessage as ClientMessageType,
+  TxAPI_ServerMessage as ServerMessageType,
+} from '../proto-grpc/github.com/milaboratory/pl/plapi/plapiproto/api';
+import type { BiDiStream } from './abstract_stream';
+import Denque from 'denque';
+import type { RetryConfig } from '../helpers/retry_strategy';
+import { RetryStrategy } from '../helpers/retry_strategy';
+
+type ConnectionState = 'disconnected' | 'connecting' | 'connected' | 'closing' | 'closed';
+
+interface QueuedMessage {
+  message: ClientMessageType;
+  resolve: () => void;
+  reject: (error: Error) => void;
+}
+
+interface ResponseResolver {
+  resolve: (value: IteratorResult<ServerMessageType>) => void;
+  reject: (error: Error) => void;
+}
+
+/**
+ * WebSocket-based bidirectional stream implementation for LLTransaction.
+ * Implements BiDiStream interface which is compatible with DuplexStreamingCall.
+ */
+export class WebSocketBiDiStream implements BiDiStream<ClientMessageType, ServerMessageType> {
+  // Connection
+  private ws: WebSocket | null = null;
+  private connectionState: ConnectionState = 'disconnected';
+  private readonly url: string;
+  private readonly jwtToken?: string;
+  private readonly abortSignal: AbortSignal;
+  private readonly reconnection: RetryStrategy;
+
+  // Send management
+  private readonly sendQueue = new Denque<QueuedMessage>();
+  private sendCompleted = false;
+
+  // Response management
+  private readonly responseQueue = new Denque<ServerMessageType>();
+  private responseResolvers: ResponseResolver[] = [];
+
+  // Error tracking
+  private connectionError: Error | null = null;
+
+  // === Public API ===
+
+  public readonly requests = {
+    send: async (message: ClientMessageType): Promise<void> => {
+      this.validateSendState();
+      return this.enqueueSend(message);
+    },
+
+    complete: async (): Promise<void> => {
+      if (this.sendCompleted) return;
+
+      this.sendCompleted = true;
+      await this.drainSendQueue();
+      this.closeConnection();
+    },
+  };
+
+  public readonly responses: AsyncIterable<ServerMessageType> = {
+    [Symbol.asyncIterator]: () => this.createResponseIterator(),
+  };
+
+  constructor(
+    url: string,
+    abortSignal: AbortSignal,
+    jwtToken?: string,
+    retryConfig: Partial<RetryConfig> = {},
+  ) {
+    this.url = url;
+    this.jwtToken = jwtToken;
+    this.abortSignal = abortSignal;
+
+    this.reconnection = new RetryStrategy(retryConfig, {
+      onRetry: () => { void this.connect(); },
+      onMaxAttemptsReached: (error) => this.handleError(error),
+    });
+
+    if (abortSignal.aborted) {
+      this.connectionState = 'closed';
+      return;
+    }
+
+    this.attachAbortSignalHandler();
+    void this.connect();
+  }
+
+  // === Connection Lifecycle ===
+
+  private connect(): void {
+    if (this.isConnectingOrConnected() || this.abortSignal.aborted) return;
+
+    this.connectionState = 'connecting';
+    this.connectionError = null;
+
+    try {
+      this.ws = this.createWebSocket();
+      this.attachWebSocketHandlers();
+    } catch (error) {
+      this.connectionError = this.toError(error);
+      this.connectionState = 'disconnected';
+      this.reconnection.schedule();
+    }
+  }
+
+  private createWebSocket(): WebSocket {
+    const options = this.jwtToken
+      ? { headers: { authorization: `Bearer ${this.jwtToken}` } }
+      : undefined;
+
+    const ws = new (WebSocket as any)(this.url, options);
+    if (ws) {
+      ws.binaryType = 'arraybuffer';
+    }
+    return ws;
+  }
+
+  private attachWebSocketHandlers(): void {
+    if (!this.ws) return;
+
+    this.ws.addEventListener('open', () => this.onOpen());
+    this.ws.addEventListener('message', (event) => this.onMessage(event.data));
+    this.ws.addEventListener('error', (error) => this.onError(error));
+    this.ws.addEventListener('close', () => this.onClose());
+  }
+
+  private attachAbortSignalHandler(): void {
+    this.abortSignal.addEventListener('abort', () => this.close());
+  }
+
+  private onOpen(): void {
+    this.connectionState = 'connected';
+    this.reconnection.reset();
+    void this.processSendQueue();
+  }
+
+  private onClose(): void {
+    this.ws = null;
+
+    if (this.isClosed() || this.abortSignal.aborted) return;
+
+    if (this.sendCompleted) {
+      this.finalizeStream();
+    } else {
+      this.connectionState = 'disconnected';
+      this.reconnection.schedule();
+    }
+  }
+
+  private onError(error: unknown): void {
+    this.handleError(this.toError(error));
+  }
+
+  private onMessage(data: unknown): void {
+    try {
+      const message = this.parseMessage(data);
+      this.deliverResponse(message);
+    } catch (error) {
+      this.handleError(this.toError(error));
+    }
+  }
+
+  private closeConnection(): void {
+    if (this.ws?.readyState === WebSocket.OPEN) {
+      this.ws.close();
+    }
+  }
+
+  private close(): void {
+    if (this.isClosed()) return;
+
+    this.connectionState = 'closed';
+    this.reconnection.cancel();
+    this.closeWebSocket();
+    this.rejectAllPendingOperations();
+  }
+
+  private closeWebSocket(): void {
+    if (!this.ws) return;
+
+    try {
+      this.ws.close();
+    } catch {
+      // Suppress close errors
+    }
+
+    this.ws = null;
+  }
+
+  private finalizeStream(): void {
+    this.connectionState = 'closed';
+    this.resolveAllPendingResponses();
+  }
+
+  private resolveAllPendingResponses(): void {
+    while (this.responseResolvers.length > 0) {
+      const resolver = this.responseResolvers.shift()!;
+      resolver.resolve({ value: undefined as any, done: true });
+    }
+  }
+
+  private parseMessage(data: unknown): ServerMessageType {
+    if (data instanceof ArrayBuffer) {
+      return ServerMessageType.fromBinary(new Uint8Array(data));
+    }
+
+    throw new Error(`Unsupported message format: ${typeof data}`);
+  }
+
+  // === Send Queue Management ===
+
+  private validateSendState(): void {
+    if (this.sendCompleted) {
+      throw new Error('Cannot send: stream already completed');
+    }
+
+    if (this.abortSignal.aborted) {
+      throw new Error('Cannot send: stream aborted');
+    }
+  }
+
+  private enqueueSend(message: ClientMessageType): Promise<void> {
+    return new Promise<void>((resolve, reject) => {
+      this.sendQueue.push({ message, resolve, reject });
+      void this.processSendQueue();
+    });
+  }
+
+  private processSendQueue(): void {
+    if (!this.canSendMessages()) return;
+
+    while (this.sendQueue.length > 0) {
+      const queued = this.sendQueue.shift()!;
+      this.sendQueuedMessage(queued);
+    }
+  }
+
+  private canSendMessages(): boolean {
+    return this.connectionState === 'connected' && this.ws !== null;
+  }
+
+  private sendQueuedMessage(queued: QueuedMessage): void {
+    try {
+      const ws = this.ws;
+      if (!ws) {
+        throw new Error('WebSocket is not connected');
+      }
+
+      // Check if WebSocket is in a valid state for sending
+      if (ws.readyState !== WebSocket.OPEN) {
+        throw new Error(`WebSocket is not open (readyState: ${ws.readyState})`);
+      }
+
+      const binary = ClientMessageType.toBinary(queued.message);
+      ws.send(binary);
+      queued.resolve();
+    } catch (error) {
+      queued.reject(this.toError(error));
+    }
+  }
+
+  private async drainSendQueue(): Promise<void> {
+    const POLL_INTERVAL_MS = 10;
+
+    while (this.sendQueue.length > 0) {
+      await this.waitForCondition(
+        () => this.sendQueue.length === 0,
+        POLL_INTERVAL_MS,
+      );
+    }
+  }
+
+  private waitForCondition(
+    condition: () => boolean,
+    intervalMs: number,
+  ): Promise<void> {
+    return new Promise<void>((resolve, reject) => {
+      if (this.abortSignal.aborted) {
+        return reject(this.toError(this.abortSignal.reason) ?? new Error('Stream aborted'));
+      }
+
+      let timeoutId: ReturnType<typeof setTimeout>;
+      const onAbort = () => {
+        clearTimeout(timeoutId);
+        reject(this.toError(this.abortSignal.reason) ?? new Error('Stream aborted'));
+      };
+
+      this.abortSignal.addEventListener('abort', onAbort, { once: true });
+
+      const check = () => {
+        if (condition() || this.isStreamEnded()) {
+          this.abortSignal.removeEventListener('abort', onAbort);
+          resolve();
+        } else {
+          timeoutId = setTimeout(check, intervalMs);
+        }
+      };
+
+      check();
+    });
+  }
+
+  // === Response Delivery ===
+
+  private deliverResponse(message: ServerMessageType): void {
+    if (this.responseResolvers.length > 0) {
+      const resolver = this.responseResolvers.shift()!;
+      resolver.resolve({ value: message, done: false });
+    } else {
+      this.responseQueue.push(message);
+    }
+  }
+
+  private async *createResponseIterator(): AsyncIterator<ServerMessageType> {
+    while (true) {
+      const result = await this.nextResponse();
+
+      if (result.done) break;
+
+      yield result.value;
+    }
+  }
+
+  private nextResponse(): Promise<IteratorResult<ServerMessageType>> {
+    return new Promise<IteratorResult<ServerMessageType>>((resolve, reject) => {
+      // Fast path: message already available
+      if (this.responseQueue.length > 0) {
+        const message = this.responseQueue.shift()!;
+        resolve({ value: message, done: false });
+        return;
+      }
+
+      // Stream ended
+      if (this.isStreamEnded()) {
+        if (this.connectionError) {
+          reject(this.connectionError);
+        } else {
+          resolve({ value: undefined as any, done: true });
+        }
+        return;
+      }
+
+      // Wait for next message
+      this.responseResolvers.push({ resolve, reject });
+    });
+  }
+
+  // === Error Handling ===
+  private handleError(error: Error): void {
+    if (this.isClosed()) return;
+
+    this.connectionState = 'closed';
+    this.connectionError = error;
+    this.reconnection.cancel();
+    this.closeWebSocket();
+    this.rejectAllPendingOperations(error);
+  }
+
+  private rejectAllPendingOperations(error?: Error): void {
+    const err = error ?? this.createStreamClosedError();
+    this.rejectAllSendOperations(err);
+    this.rejectAllResponseResolvers(err);
+  }
+
+  private rejectAllSendOperations(error: Error): void {
+    while (this.sendQueue.length > 0) {
+      const queued = this.sendQueue.shift()!;
+      queued.reject(error);
+    }
+  }
+
+  private rejectAllResponseResolvers(error: Error): void {
+    while (this.responseResolvers.length > 0) {
+      const resolver = this.responseResolvers.shift()!;
+      resolver.reject(error);
+    }
+  }
+
+  private createStreamClosedError(): Error {
+    if (this.abortSignal.aborted) {
+      const reason = this.abortSignal.reason;
+      if (reason instanceof Error) {
+        return reason;
+      }
+      return new Error('Stream aborted', { cause: reason });
+    }
+    return new Error('Stream closed');
+  }
+  // === State Checks ===
+
+  private isConnectingOrConnected(): boolean {
+    return this.connectionState === 'connecting'
+      || this.connectionState === 'connected';
+  }
+
+  private isClosed(): boolean {
+    return this.connectionState === 'closed';
+  }
+
+  private isStreamEnded(): boolean {
+    return this.isClosed() || this.abortSignal.aborted;
+  }
+
+  private toError(error: unknown): Error {
+    return error instanceof Error ? error : new Error(String(error));
+  }
+}

package/src/helpers/retry_strategy.ts

@@ -0,0 +1,123 @@
+export interface RetryConfig {
+  maxAttempts: number;
+  initialDelay: number;
+  maxDelay: number;
+}
+
+export const DEFAULT_RETRY_CONFIG: RetryConfig = {
+  maxAttempts: 10,
+  initialDelay: 100,
+  maxDelay: 30000,
+};
+
+export interface RetryCallbacks {
+  onRetry: () => void;
+  onMaxAttemptsReached: (error: Error) => void;
+}
+
+export class RetryStrategy {
+  private attempts = 0;
+  private timer: ReturnType<typeof setTimeout> | null = null;
+  private readonly config: RetryConfig;
+  private readonly callbacks: RetryCallbacks;
+  private readonly backoff: ExponentialBackoff;
+
+  constructor(config: Partial<RetryConfig>, callbacks: RetryCallbacks) {
+    this.config = { ...DEFAULT_RETRY_CONFIG, ...config };
+    this.callbacks = callbacks;
+    this.backoff = new ExponentialBackoff({
+      initialDelay: this.config.initialDelay,
+      maxDelay: this.config.maxDelay,
+      factor: 2,
+      jitter: 0.1,
+    });
+  }
+
+  schedule(): void {
+    if (this.timer) return;
+    if (this.hasExceededLimit()) {
+      this.notifyMaxAttemptsReached();
+      return;
+    }
+
+    this.timer = setTimeout(() => {
+      this.timer = null;
+      this.attempts++;
+      this.callbacks.onRetry();
+    }, this.backoff.delay());
+  }
+
+  cancel(): void {
+    if (this.timer) {
+      clearTimeout(this.timer);
+      this.timer = null;
+    }
+  }
+
+  reset(): void {
+    this.attempts = 0;
+    this.backoff.reset();
+  }
+
+  private hasExceededLimit(): boolean {
+    return this.attempts >= this.config.maxAttempts;
+  }
+
+  private notifyMaxAttemptsReached(): void {
+    const error = new Error(
+      `Max retry attempts (${this.config.maxAttempts}) reached`,
+    );
+    this.callbacks.onMaxAttemptsReached(error);
+  }
+}
+
+interface ExponentialBackoffConfig {
+  initialDelay: number;
+  maxDelay: number;
+  factor: number;
+  jitter: number;
+}
+
+class ExponentialBackoff {
+  private readonly initialDelay: number;
+  private readonly maxDelay: number;
+
+  private currentDelay: number;
+
+  private readonly factor: number;
+  private readonly jitter: number;
+
+  constructor(config: ExponentialBackoffConfig) {
+    this.initialDelay = config.initialDelay;
+    this.maxDelay = config.maxDelay;
+    this.factor = config.factor;
+    this.jitter = config.jitter;
+    this.currentDelay = config.initialDelay;
+  }
+
+  delay(): number {
+    if (this.currentDelay >= this.maxDelay) {
+      return this.applyJitter(this.maxDelay);
+    }
+
+    this.currentDelay = this.currentDelay * this.factor;
+
+    if (this.currentDelay > this.maxDelay) {
+      this.currentDelay = this.maxDelay;
+    }
+
+    return this.applyJitter(this.currentDelay);
+  }
+
+  reset(): void {
+    this.currentDelay = this.initialDelay;
+  }
+
+  private applyJitter(delay: number): number {
+    if (delay === 0 || this.jitter === 0) {
+      return delay;
+    }
+    const delayFactor = 1 - (this.jitter / 2) + Math.random() * this.jitter;
+    return delay * delayFactor;
+  }
+}