@nmtjs/gateway 0.15.0-beta.2 → 0.15.0-beta.4

package/src/injectables.ts

@@ -0,0 +1,88 @@
+import type { Readable } from 'node:stream'
+
+import type {
+  ProtocolBlobInterface,
+  ProtocolBlobMetadata,
+} from '@nmtjs/protocol'
+import { anyAbortSignal } from '@nmtjs/common'
+import {
+  createFactoryInjectable,
+  createLazyInjectable,
+  createOptionalInjectable,
+  Scope,
+} from '@nmtjs/core'
+
+import type { GatewayConnection } from './connections.ts'
+
+export const connection = createLazyInjectable<
+  GatewayConnection,
+  Scope.Connection
+>(Scope.Connection, 'Gateway connection')
+
+export const connectionId = createLazyInjectable<
+  GatewayConnection['id'],
+  Scope.Connection
+>(Scope.Connection, 'Gateway connection id')
+
+export const connectionData = createLazyInjectable<unknown, Scope.Connection>(
+  Scope.Connection,
+  "Gateway connection's data",
+)
+
+export const connectionAbortSignal = createLazyInjectable<
+  AbortSignal,
+  Scope.Connection
+>(Scope.Connection, 'Connection abort signal')
+
+export const rpcClientAbortSignal = createLazyInjectable<
+  AbortSignal,
+  Scope.Call
+>(Scope.Call, 'RPC client abort signal')
+
+export const rpcStreamAbortSignal = createLazyInjectable<
+  AbortSignal,
+  Scope.Call
+>(Scope.Call, 'RPC stream abort signal')
+
+export const rpcAbortSignal = createFactoryInjectable(
+  {
+    dependencies: {
+      rpcClientAbortSignal,
+      connectionAbortSignal,
+      rpcStreamAbortSignal: createOptionalInjectable(rpcStreamAbortSignal),
+    },
+    factory: (ctx) =>
+      anyAbortSignal(
+        ctx.rpcClientAbortSignal,
+        ctx.connectionAbortSignal,
+        ctx.rpcStreamAbortSignal,
+      ),
+  },
+  'Any RPC abort signal',
+)
+
+export const createBlob = createLazyInjectable<
+  (
+    source:
+      | Readable
+      | globalThis.ReadableStream
+      | File
+      | Blob
+      | string
+      | ArrayBuffer
+      | Uint8Array,
+    metadata?: ProtocolBlobMetadata,
+  ) => ProtocolBlobInterface,
+  Scope.Call
+>(Scope.Call, 'Create RPC blob')
+
+export const GatewayInjectables = {
+  connection,
+  connectionId,
+  connectionData,
+  connectionAbortSignal,
+  rpcClientAbortSignal,
+  rpcStreamAbortSignal,
+  rpcAbortSignal,
+  createBlob,
+}

package/src/rpcs.ts

@@ -0,0 +1,90 @@
+import type { Future } from '@nmtjs/common'
+import { createFuture } from '@nmtjs/common'
+
+export class RpcManager {
+  // connectionId:callId -> AbortController
+  readonly rpcs = new Map<string, AbortController>()
+  // connectionId:callId -> Future<void>
+  readonly streams = new Map<string, Future<void>>()
+
+  set(connectionId: string, callId: number, controller: AbortController) {
+    const key = this.getKey(connectionId, callId)
+    this.rpcs.set(key, controller)
+  }
+
+  get(connectionId: string, callId: number) {
+    const key = this.getKey(connectionId, callId)
+    return this.rpcs.get(key)
+  }
+
+  delete(connectionId: string, callId: number) {
+    const key = this.getKey(connectionId, callId)
+    this.rpcs.delete(key)
+  }
+
+  abort(connectionId: string, callId: number) {
+    const key = this.getKey(connectionId, callId)
+    const controller = this.rpcs.get(key)
+    if (controller) {
+      controller.abort()
+      this.rpcs.delete(key)
+      this.releasePull(connectionId, callId)
+    }
+  }
+
+  awaitPull(
+    connectionId: string,
+    callId: number,
+    signal?: AbortSignal,
+  ): Promise<void> {
+    const key = this.getKey(connectionId, callId)
+    const rpc = this.rpcs.get(key)
+    if (!rpc) throw new Error(`RPC not found`)
+    const future = this.streams.get(key)
+    if (future) {
+      return future.promise
+    } else {
+      const newFuture = createFuture<void>()
+      if (signal)
+        signal.addEventListener('abort', () => newFuture.resolve(), {
+          once: true,
+        })
+      this.streams.set(key, newFuture)
+      return newFuture.promise
+    }
+  }
+
+  releasePull(connectionId: string, callId: number) {
+    const key = this.getKey(connectionId, callId)
+    const future = this.streams.get(key)
+    if (future) {
+      future.resolve()
+      this.streams.delete(key)
+    }
+  }
+
+  close(connectionId: string) {
+    // Iterate all RPCs and abort those belonging to this connection
+    // Optimization: Maintain a Set<callId> per connectionId
+    for (const [key, controller] of this.rpcs) {
+      if (key.startsWith(`${connectionId}:`)) {
+        controller.abort()
+        this.rpcs.delete(key)
+      }
+    }
+    // Also release any pending pulls for this connection
+    for (const key of this.streams.keys()) {
+      if (key.startsWith(`${connectionId}:`)) {
+        const future = this.streams.get(key)
+        if (future) {
+          future.resolve()
+          this.streams.delete(key)
+        }
+      }
+    }
+  }
+
+  private getKey(connectionId: string, callId: number) {
+    return `${connectionId}:${callId}`
+  }
+}

package/src/streams.ts

@@ -0,0 +1,408 @@
+import type Stream from 'node:stream'
+
+import type {
+  EncodeRPCStreams,
+  ProtocolBlob,
+  ProtocolBlobMetadata,
+} from '@nmtjs/protocol'
+import type { ProtocolRPCEncode } from '@nmtjs/protocol/client'
+import { noopFn } from '@nmtjs/common'
+import {
+  ProtocolClientStream,
+  ProtocolServerStream,
+} from '@nmtjs/protocol/server'
+
+import { StreamTimeout } from './enums.ts'
+
+export type StreamConfig = {
+  timeouts: {
+    [StreamTimeout.Pull]: number
+    [StreamTimeout.Consume]: number
+    [StreamTimeout.Finish]: number
+  }
+}
+
+type StreamTimeouts = Record<StreamTimeout, any>
+
+type ClientStreamState = {
+  connectionId: string
+  callId: number
+  stream: ProtocolClientStream
+  timeouts: StreamTimeouts
+}
+
+type ServerStreamState = {
+  connectionId: string
+  callId: number
+  stream: ProtocolServerStream
+  timeouts: StreamTimeouts
+}
+
+type StreamState = ClientStreamState | ServerStreamState
+
+/**
+ * @todo Clarify Pull/Consume timeout semantics - currently ambiguous whether
+ *       Pull timeout means "client not pulling" or "server not producing" for server streams
+ */
+export class BlobStreamsManager {
+  readonly clientStreams = new Map<string, ClientStreamState>()
+  readonly serverStreams = new Map<string, ServerStreamState>()
+
+  // Index for quick lookup by callId (connectionId:callId -> Set<streamId>)
+  readonly connectionClientStreams = new Map<string, Set<number>>()
+  readonly connectionServerStreams = new Map<string, Set<number>>()
+  readonly clientCallStreams = new Map<string, Set<number>>()
+  readonly serverCallStreams = new Map<string, Set<number>>()
+
+  readonly timeoutDurations: Record<StreamTimeout, number>
+
+  constructor(config: StreamConfig) {
+    this.timeoutDurations = {
+      [StreamTimeout.Pull]: config.timeouts[StreamTimeout.Pull],
+      [StreamTimeout.Consume]: config.timeouts[StreamTimeout.Consume],
+      [StreamTimeout.Finish]: config.timeouts[StreamTimeout.Finish],
+    }
+  }
+
+  // --- Client Streams (Upload) ---
+
+  createClientStream(
+    connectionId: string,
+    callId: number,
+    streamId: number,
+    metadata: ProtocolBlobMetadata,
+    options: Stream.ReadableOptions,
+  ) {
+    const stream = new ProtocolClientStream(streamId, metadata, options)
+    stream.on('error', noopFn)
+
+    const key = this.getKey(connectionId, streamId)
+    const state: ClientStreamState = {
+      connectionId,
+      callId,
+      stream,
+      timeouts: {
+        [StreamTimeout.Pull]: undefined,
+        [StreamTimeout.Consume]: undefined,
+        [StreamTimeout.Finish]: undefined,
+      },
+    }
+    this.clientStreams.set(key, state)
+    this.trackClientCall(connectionId, callId, streamId)
+    this.trackConnectionClientStream(connectionId, streamId)
+
+    this.startTimeout(state, StreamTimeout.Consume)
+
+    return stream
+  }
+
+  pushToClientStream(
+    connectionId: string,
+    streamId: number,
+    chunk: ArrayBufferView,
+  ) {
+    const key = this.getKey(connectionId, streamId)
+    const state = this.clientStreams.get(key)
+    if (state) {
+      state.stream.write(chunk)
+      this.clearTimeout(state, StreamTimeout.Consume)
+      this.startTimeout(state, StreamTimeout.Pull)
+    }
+  }
+
+  endClientStream(connectionId: string, streamId: number) {
+    const key = this.getKey(connectionId, streamId)
+    const state = this.clientStreams.get(key)
+    if (state) {
+      state.stream.end(null)
+      this.removeClientStream(connectionId, streamId)
+    }
+  }
+
+  abortClientStream(connectionId: string, streamId: number, error = 'Aborted') {
+    const key = this.getKey(connectionId, streamId)
+    const state = this.clientStreams.get(key)
+    if (state) {
+      state.stream.destroy(new Error(error))
+      this.removeClientStream(connectionId, streamId)
+    }
+  }
+
+  consumeClientStream(connectionId: string, callId: number, streamId: number) {
+    this.untrackClientCall(connectionId, callId, streamId)
+  }
+
+  private removeClientStream(connectionId: string, streamId: number) {
+    const key = this.getKey(connectionId, streamId)
+    const state = this.clientStreams.get(key)
+    if (state) {
+      this.clientStreams.delete(key)
+      this.clearTimeout(state, StreamTimeout.Finish)
+      this.clearTimeout(state, StreamTimeout.Pull)
+      this.clearTimeout(state, StreamTimeout.Consume)
+      this.untrackClientCall(connectionId, state.callId, streamId)
+      this.untrackConnectionClientStream(connectionId, streamId)
+    }
+  }
+
+  // --- Server Streams (Download) ---
+
+  getServerStreamsMetadata(connectionId: string, callId: number) {
+    const key = this.getCallKey(connectionId, callId)
+    const streamIds = this.serverCallStreams.get(key)
+    const streams: EncodeRPCStreams = {}
+
+    if (streamIds) {
+      for (const streamId of streamIds) {
+        const streamKey = this.getKey(connectionId, streamId)
+        const state = this.serverStreams.get(streamKey)
+        if (state) {
+          streams[streamId] = state.stream.metadata
+        }
+      }
+    }
+
+    return streams
+  }
+
+  createServerStream(
+    connectionId: string,
+    callId: number,
+    streamId: number,
+    blob: ProtocolBlob,
+  ) {
+    const stream = new ProtocolServerStream(streamId, blob)
+    const key = this.getKey(connectionId, streamId)
+
+    const state: ServerStreamState = {
+      connectionId,
+      callId,
+      stream,
+      timeouts: {
+        [StreamTimeout.Pull]: undefined,
+        [StreamTimeout.Consume]: undefined,
+        [StreamTimeout.Finish]: undefined,
+      },
+    }
+
+    // Prevent unhandled 'error' events, in case the user did not subscribe to them
+    stream.on('error', noopFn)
+
+    this.serverStreams.set(key, state)
+    this.trackServerCall(connectionId, callId, streamId)
+    this.trackConnectionServerStream(connectionId, streamId)
+
+    this.startTimeout(state, StreamTimeout.Finish)
+    this.startTimeout(state, StreamTimeout.Consume)
+
+    return stream
+  }
+
+  pullServerStream(connectionId: string, streamId: number) {
+    const key = this.getKey(connectionId, streamId)
+    const state = this.serverStreams.get(key)
+    if (state) {
+      state.stream.resume()
+      this.clearTimeout(state, StreamTimeout.Consume)
+      this.startTimeout(state, StreamTimeout.Pull)
+    }
+  }
+
+  abortServerStream(connectionId: string, streamId: number, error = 'Aborted') {
+    const key = this.getKey(connectionId, streamId)
+    const state = this.serverStreams.get(key)
+    if (state) {
+      state.stream.destroy(new Error(error))
+      this.removeServerStream(connectionId, streamId)
+    }
+  }
+
+  removeServerStream(connectionId: string, streamId: number) {
+    const key = this.getKey(connectionId, streamId)
+    const state = this.serverStreams.get(key)
+    if (state) {
+      this.serverStreams.delete(key)
+      this.clearTimeout(state, StreamTimeout.Pull)
+      this.clearTimeout(state, StreamTimeout.Consume)
+      this.clearTimeout(state, StreamTimeout.Finish)
+      this.untrackServerCall(connectionId, state.callId, streamId)
+      this.untrackConnectionServerStream(connectionId, streamId)
+    }
+  }
+
+  // --- Timeouts ---
+
+  private startTimeout(state: StreamState, type: StreamTimeout) {
+    this.clearTimeout(state, type)
+    const duration = this.timeoutDurations[type]
+    const timeout = setTimeout(() => {
+      if (state.stream instanceof ProtocolClientStream) {
+        this.abortClientStream(
+          state.connectionId,
+          state.stream.id,
+          `${type} timeout`,
+        )
+      } else {
+        this.abortServerStream(
+          state.connectionId,
+          state.stream.id,
+          `${type} timeout`,
+        )
+      }
+      state.timeouts[type] = undefined
+    }, duration)
+    state.timeouts[type] = timeout
+  }
+
+  private clearTimeout(state: StreamState, type: StreamTimeout) {
+    const timeout = state.timeouts[type]
+    if (timeout) {
+      clearTimeout(timeout)
+      state.timeouts[type] = undefined
+    }
+  }
+
+  // --- Helpers ---
+
+  private getKey(connectionId: string, streamId: number) {
+    return `${connectionId}:${streamId}`
+  }
+
+  private getCallKey(connectionId: string, callId: number) {
+    return `${connectionId}:${callId}`
+  }
+
+  private trackClientCall(
+    connectionId: string,
+    callId: number,
+    streamId: number,
+  ) {
+    const key = this.getCallKey(connectionId, callId)
+    let set = this.clientCallStreams.get(key)
+    if (!set) {
+      set = new Set()
+      this.clientCallStreams.set(key, set)
+    }
+    set.add(streamId)
+  }
+
+  private untrackClientCall(
+    connectionId: string,
+    callId: number,
+    streamId: number,
+  ) {
+    const key = this.getCallKey(connectionId, callId)
+    const set = this.clientCallStreams.get(key)
+    if (set) {
+      set.delete(streamId)
+      if (set.size === 0) {
+        this.clientCallStreams.delete(key)
+      }
+    }
+  }
+
+  private trackServerCall(
+    connectionId: string,
+    callId: number,
+    streamId: number,
+  ) {
+    const key = this.getCallKey(connectionId, callId)
+    let set = this.serverCallStreams.get(key)
+    if (!set) {
+      set = new Set()
+      this.serverCallStreams.set(key, set)
+    }
+    set.add(streamId)
+  }
+
+  private untrackServerCall(
+    connectionId: string,
+    callId: number,
+    streamId: number,
+  ) {
+    const key = this.getCallKey(connectionId, callId)
+    const set = this.serverCallStreams.get(key)
+    if (set) {
+      set.delete(streamId)
+      if (set.size === 0) {
+        this.serverCallStreams.delete(key)
+      }
+    }
+  }
+
+  private trackConnectionClientStream(connectionId: string, streamId: number) {
+    let set = this.connectionClientStreams.get(connectionId)
+    if (!set) {
+      set = new Set()
+      this.connectionClientStreams.set(connectionId, set)
+    }
+    set.add(streamId)
+  }
+
+  private untrackConnectionClientStream(
+    connectionId: string,
+    streamId: number,
+  ) {
+    const set = this.connectionClientStreams.get(connectionId)
+    if (set) {
+      set.delete(streamId)
+      if (set.size === 0) {
+        this.connectionClientStreams.delete(connectionId)
+      }
+    }
+  }
+
+  private trackConnectionServerStream(connectionId: string, streamId: number) {
+    let set = this.connectionServerStreams.get(connectionId)
+    if (!set) {
+      set = new Set()
+      this.connectionServerStreams.set(connectionId, set)
+    }
+    set.add(streamId)
+  }
+
+  private untrackConnectionServerStream(
+    connectionId: string,
+    streamId: number,
+  ) {
+    const set = this.connectionServerStreams.get(connectionId)
+    if (set) {
+      set.delete(streamId)
+      if (set.size === 0) {
+        this.connectionServerStreams.delete(connectionId)
+      }
+    }
+  }
+
+  // --- Cleanup ---
+
+  abortClientCallStreams(
+    connectionId: string,
+    callId: number,
+    reason = 'Call aborted',
+  ) {
+    const key = this.getCallKey(connectionId, callId)
+    const clientStreamIds = this.clientCallStreams.get(key)
+    if (clientStreamIds) {
+      for (const streamId of [...clientStreamIds]) {
+        this.abortClientStream(connectionId, streamId, reason)
+      }
+    }
+  }
+
+  cleanupConnection(connectionId: string) {
+    const clientStreamIds = this.connectionClientStreams.get(connectionId)
+    if (clientStreamIds) {
+      for (const streamId of [...clientStreamIds]) {
+        this.abortClientStream(connectionId, streamId, 'Connection closed')
+      }
+    }
+
+    const serverStreamIds = this.connectionServerStreams.get(connectionId)
+    if (serverStreamIds) {
+      for (const streamId of [...serverStreamIds]) {
+        this.abortServerStream(connectionId, streamId, 'Connection closed')
+      }
+    }
+  }
+}

package/src/transport.ts

@@ -0,0 +1,95 @@
+import type { Async } from '@nmtjs/common'
+import type { Injection, LazyInjectable, Scope } from '@nmtjs/core'
+import type { ConnectionType, ProtocolVersion } from '@nmtjs/protocol'
+import type { ProtocolFormats } from '@nmtjs/protocol/server'
+
+import type { GatewayConnection } from './connections.ts'
+import type { ProxyableTransportType } from './enums.ts'
+import type { GatewayRpc } from './types.ts'
+
+export interface TransportConnection {
+  connectionId: string
+  type: ConnectionType
+}
+
+export interface TransportOnConnectOptions<
+  Type extends ConnectionType = ConnectionType,
+> {
+  type: Type extends ConnectionType.Bidirectional
+    ? Type
+    : ConnectionType.Unidirectional
+  protocolVersion: ProtocolVersion
+  accept: string | null
+  contentType: string | null
+  data: unknown
+}
+
+export interface TransportOnMessageOptions {
+  connectionId: string
+  data: ArrayBuffer
+}
+
+export type TransportWorkerParams<
+  Type extends ConnectionType = ConnectionType,
+> = {
+  formats: ProtocolFormats
+  onConnect: (
+    options: TransportOnConnectOptions<Type>,
+    ...injections: Injection[]
+  ) => Promise<GatewayConnection & AsyncDisposable>
+  onDisconnect: (connectionId: GatewayConnection['id']) => Promise<void>
+  onMessage: (
+    options: TransportOnMessageOptions,
+    ...injections: Injection[]
+  ) => Promise<void>
+  onRpc: (
+    connection: GatewayConnection,
+    rpc: GatewayRpc,
+    signal: AbortSignal,
+    ...injections: Injection[]
+  ) => Promise<unknown>
+}
+
+export interface TransportWorkerStartOptions<
+  Type extends ConnectionType = ConnectionType,
+> extends TransportWorkerParams<Type> {
+  // for extra props in the future
+}
+
+export interface TransportWorker<Type extends ConnectionType = ConnectionType> {
+  start: (params: TransportWorkerParams<Type>) => Async<string>
+  stop: (params: Pick<TransportWorkerParams<Type>, 'formats'>) => Async<void>
+  send?: Type extends 'unidirectional'
+    ? never
+    : (connectionId: string, buffer: ArrayBufferView) => boolean | null
+}
+
+export interface Transport<
+  Type extends ConnectionType = ConnectionType,
+  TransportOptions = any,
+  Injections extends {
+    [key: string]: LazyInjectable<any, Scope.Connection | Scope.Call>
+  } = { [key: string]: LazyInjectable<any, Scope.Connection | Scope.Call> },
+  Proxyable extends ProxyableTransportType | undefined =
+    | ProxyableTransportType
+    | undefined,
+> {
+  proxyable: Proxyable
+  injectables?: Injections
+  factory: (options: TransportOptions) => Async<TransportWorker<Type>>
+}
+
+export function createTransport<
+  Type extends ConnectionType = ConnectionType,
+  TransportOptions = any,
+  Injections extends {
+    [key: string]: LazyInjectable<any, Scope.Connection | Scope.Call>
+  } = { [key: string]: LazyInjectable<any, Scope.Connection | Scope.Call> },
+  Proxyable extends ProxyableTransportType | undefined =
+    | ProxyableTransportType
+    | undefined,
+>(
+  config: Transport<Type, TransportOptions, Injections, Proxyable>,
+): Transport<Type, TransportOptions, Injections, Proxyable> {
+  return config
+}

package/src/types.ts

@@ -0,0 +1,24 @@
+import type { AnyInjectable, Container, Logger, Scope } from '@nmtjs/core'
+import type { MessageContext as ProtocolMessageContext } from '@nmtjs/protocol/server'
+
+import type { GatewayApiCallOptions } from './api.ts'
+
+export type ConnectionIdentityType = string
+export type ConnectionIdentity = AnyInjectable<
+  ConnectionIdentityType,
+  Scope.Global
+>
+
+export interface GatewayRpc {
+  callId: number
+  procedure: string
+  payload: unknown
+  metadata?: GatewayApiCallOptions['metadata']
+}
+
+export interface GatewayRpcContext extends ProtocolMessageContext, GatewayRpc {
+  container: Container
+  signal: AbortSignal
+  logger: Logger
+  [Symbol.asyncDispose](): Promise<void>
+}