theater-client/src/connection/TheaterConnection.ts

TypeScript client library for Theater actor system TCP protocol

/**
 * Low-level TCP connection to Theater server
 * Each connection is dedicated to a specific operation to avoid response multiplexing
 */

import { EventEmitter } from 'node:events';
import net from 'node:net';
import type { ManagementCommand, ManagementResponse, FrameMessage, FragmentData } from '../types/protocol.js';
import type { TheaterClientConfig } from '../types/client.js';
import { TheaterConnectionError, TheaterTimeoutError, TheaterProtocolError } from '../types/client.js';
import { createLogger } from '../utils/logger.js';
import { serializeMessage, createFrame, encodeJson } from '../utils/serialization.js';

const log = createLogger('TheaterConnection');

/**
 * Fragment reassembly state for a single message
 */
interface PartialMessage {
  messageId: number;
  totalFragments: number;
  fragments: Map<number, Buffer>;
  createdAt: number;
}

/**
 * A single TCP connection to the Theater server
 * Implements the hygienic connection pattern - each operation gets its own connection
 */
export class TheaterConnection extends EventEmitter {
  private socket: net.Socket | null = null;
  private dataBuffer: Buffer = Buffer.alloc(0);
  private messageQueue: ManagementResponse[] = [];
  private messageWaiters: Array<(msg: ManagementResponse) => void> = [];
  private messageListenerSetup = false;
  private _connected = false;
  
  // Fragment reassembly state
  private partialMessages: Map<number, PartialMessage> = new Map();
  private lastCleanup: number = Date.now();
  
  private readonly timeout: number;
  private readonly FRAGMENT_TIMEOUT = 30000; // 30 seconds
  private readonly CLEANUP_INTERVAL = 10000; // 10 seconds

  constructor(
    private readonly host: string,
    private readonly port: number,
    config: Partial<TheaterClientConfig> = {}
  ) {
    super();
    
    // Set timeout: 0 means no timeout, undefined defaults to 30 seconds
    this.timeout = config.timeout ?? 30000; // 30 second default timeout, 0 = no timeout
    
    // Prevent EventTarget memory leak warnings
    this.setMaxListeners(20);
    
    // Set up persistent message listener
    this.setupMessageListener();
  }

  get connected(): boolean {
    return this._connected && this.socket !== null;
  }

  private setupMessageListener(): void {
    if (this.messageListenerSetup) return;
    
    this.messageListenerSetup = true;
    
    // Persistent listener that handles the queue
    this.on('message', (message: ManagementResponse) => {
      if (this.messageWaiters.length > 0) {
        // Someone is waiting for a message, deliver immediately
        log.debug(`Delivering message immediately to waiter, ${this.messageWaiters.length - 1} waiters remaining`);
        const waiter = this.messageWaiters.shift()!;
        waiter(message);
      } else {
        // No one waiting, add to queue
        log.debug(`No waiters, queuing message. Queue size: ${this.messageQueue.length + 1}`);
        this.messageQueue.push(message);
      }
    });
  }

  /**
   * Clean up expired partial messages
   */
  private cleanupExpiredFragments(): void {
    const now = Date.now();
    
    // Only cleanup periodically to avoid overhead
    if (now - this.lastCleanup < this.CLEANUP_INTERVAL) {
      return;
    }

    const beforeCount = this.partialMessages.size;
    
    for (const [messageId, partial] of this.partialMessages.entries()) {
      if (now - partial.createdAt > this.FRAGMENT_TIMEOUT) {
        log.warn(`Cleaning up expired partial message ${messageId}`);
        this.partialMessages.delete(messageId);
      }
    }

    const cleaned = beforeCount - this.partialMessages.size;
    if (cleaned > 0) {
      log.debug(`Cleaned up ${cleaned} expired partial messages`);
    }

    this.lastCleanup = now;
  }

  /**
   * Handle fragment reassembly
   */
  private handleFragment(fragment: FragmentData): ManagementResponse | null {
    this.cleanupExpiredFragments();

    const { message_id, fragment_index, total_fragments, data } = fragment;
    
    log.debug(`Received fragment ${fragment_index + 1}/${total_fragments} for message ${message_id}`);

    // Decode the base64 data
    let fragmentData: Buffer;
    try {
      fragmentData = Buffer.from(data, 'base64');
    } catch (error) {
      throw new TheaterProtocolError(`Failed to decode fragment data: ${error}`);
    }

    // Get or create partial message
    let partial = this.partialMessages.get(message_id);
    if (!partial) {
      partial = {
        messageId: message_id,
        totalFragments: total_fragments,
        fragments: new Map(),
        createdAt: Date.now()
      };
      this.partialMessages.set(message_id, partial);
    }

    // Validate fragment consistency
    if (partial.totalFragments !== total_fragments) {
      throw new TheaterProtocolError(
        `Fragment total mismatch for message ${message_id}: expected ${partial.totalFragments}, got ${total_fragments}`
      );
    }

    // Add this fragment
    partial.fragments.set(fragment_index, fragmentData);

    // Check if message is complete
    if (partial.fragments.size === total_fragments) {
      log.debug(`Message ${message_id} is complete, reassembling`);

      // Remove from partial messages
      this.partialMessages.delete(message_id);

      // Reassemble fragments in order
      const reassembledBuffers: Buffer[] = [];
      for (let i = 0; i < total_fragments; i++) {
        const fragmentBuffer = partial.fragments.get(i);
        if (!fragmentBuffer) {
          throw new TheaterProtocolError(`Missing fragment ${i} for message ${message_id}`);
        }
        reassembledBuffers.push(fragmentBuffer);
      }

      // Combine all fragments
      const completeMessage = Buffer.concat(reassembledBuffers);
      
      // Parse the complete message as JSON
      try {
        const messageStr = completeMessage.toString('utf8');
        return JSON.parse(messageStr) as ManagementResponse;
      } catch (error) {
        throw new TheaterProtocolError(`Failed to parse reassembled message: ${error}`);
      }
    } else {
      // Still waiting for more fragments
      log.debug(
        `Message ${message_id} still incomplete (${partial.fragments.size}/${total_fragments} fragments)`
      );
      return null;
    }
  }

  async connect(): Promise<void> {
    if (this.connected) {
      return;
    }

    return new Promise((resolve, reject) => {
      this.socket = new net.Socket();

      // Set up timeout (only if timeout > 0)
      let timeoutHandle: NodeJS.Timeout | null = null;
      if (this.timeout > 0) {
        timeoutHandle = setTimeout(() => {
          this.socket?.destroy();
          reject(new TheaterTimeoutError('connect', this.timeout));
        }, this.timeout);
      }

      this.socket.connect(this.port, this.host, () => {
        if (timeoutHandle) clearTimeout(timeoutHandle);
        this._connected = true;
        log.debug(`Connected to Theater server at ${this.host}:${this.port}`);
        resolve();
      });

      this.socket.on('error', (error) => {
        if (timeoutHandle) clearTimeout(timeoutHandle);
        this._connected = false;
        reject(new TheaterConnectionError(`TCP connection failed: ${error.message}`, error));
      });

      this.socket.on('data', (data) => {
        log.debug(`Received data: ${data.length} bytes`);
        this.handleData(data);
      });

      this.socket.on('close', () => {
        if (timeoutHandle) clearTimeout(timeoutHandle);
        this._connected = false;
        this.emit('disconnect');
      });
    });
  }

  private handleData(data: Buffer): void {
    // Accumulate data in buffer
    this.dataBuffer = Buffer.concat([this.dataBuffer, data]);

    // Try to parse length-delimited messages
    while (this.dataBuffer.length >= 4) {
      // Read the length prefix (4 bytes, big-endian)
      const messageLength = this.dataBuffer.readUInt32BE(0);

      // Check if we have the complete message
      if (this.dataBuffer.length >= 4 + messageLength) {
        // Extract the message
        const messageBytes = this.dataBuffer.subarray(4, 4 + messageLength);

        // Remove processed data from buffer
        this.dataBuffer = this.dataBuffer.subarray(4 + messageLength);

        try {
          const messageStr = messageBytes.toString('utf8');
          const frameMessage: FrameMessage = JSON.parse(messageStr);

          // Unwrap FragmentingCodec format
          let actualMessage: ManagementResponse | null = null;
          
          if (frameMessage.Complete) {
            // FrameType::Complete - convert byte array back to JSON
            const messageBytes = Buffer.from(frameMessage.Complete);
            actualMessage = JSON.parse(messageBytes.toString('utf8'));
          } else if (frameMessage.Fragment) {
            // FrameType::Fragment - handle fragment reassembly
            try {
              actualMessage = this.handleFragment(frameMessage.Fragment);
            } catch (error) {
              log.error(`Fragment handling error: ${error}`);
              this.emit('error', error);
              return;
            }
          } else {
            // Unknown frame type
            log.error(`Unknown frame type: ${JSON.stringify(frameMessage)}`);
            this.emit('error', new TheaterProtocolError(`Unknown frame type: ${JSON.stringify(frameMessage)}`));
            return;
          }

          // Only emit if we have a complete message
          if (actualMessage) {
            this.emit('message', actualMessage);
            log.debug(`Emitted message: ${JSON.stringify(actualMessage)}`);
          }
        } catch (error) {
          const errorMessage = error instanceof Error ? error.message : String(error);
          log.error(`Failed to parse message: ${errorMessage}`);
          this.emit('error', new TheaterProtocolError(`Failed to parse message: ${errorMessage}`));
        }
      } else {
        // Not enough data for complete message, wait for more
        break;
      }
    }
  }

  async send(command: ManagementCommand): Promise<void> {
    if (!this.connected || !this.socket) {
      throw new TheaterConnectionError('Connection not established');
    }

    try {
      // Wrap command in FragmentingCodec format (FrameType::Complete for small messages)
      const commandBytes = encodeJson(command);
      const frameMessage = createFrame(commandBytes);

      // Serialize with length prefix
      const serializedMessage = serializeMessage(frameMessage);

      log.debug(`Sending command: ${JSON.stringify(command)}`);
      this.socket.write(serializedMessage);
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : String(error);
      throw new TheaterConnectionError(`Failed to send command: ${errorMessage}`, error as Error);
    }
  }

  async waitForMessage(): Promise<ManagementResponse> {
    // Check queue first
    if (this.messageQueue.length > 0) {
      const message = this.messageQueue.shift()!;
      log.debug(`Retrieved queued message, ${this.messageQueue.length} remaining`);
      return message;
    }
    
    // No queued messages, wait for next one
    log.debug('No queued messages, waiting for next...');
    return new Promise<ManagementResponse>((resolve) => {
      this.messageWaiters.push(resolve);
    });
  }

  async receive(): Promise<ManagementResponse> {
    return new Promise((resolve, reject) => {
      let timeoutHandle: NodeJS.Timeout | null = null;
      
      // Only set timeout if timeout > 0 (0 means no timeout)
      if (this.timeout > 0) {
        timeoutHandle = setTimeout(() => {
          cleanup();
          reject(new TheaterTimeoutError('receive', this.timeout));
        }, this.timeout);
      }

      const cleanup = () => {
        if (timeoutHandle) clearTimeout(timeoutHandle);
        this.removeListener('error', onError);
        this.removeListener('disconnect', onDisconnect);
      };

      const onError = (error: Error) => {
        cleanup();
        reject(error);
      };

      const onDisconnect = () => {
        cleanup();
        reject(new TheaterConnectionError('Connection closed'));
      };

      // Set up error and disconnect listeners
      this.once('error', onError);
      this.once('disconnect', onDisconnect);
      
      // Use the queue-based waiting
      this.waitForMessage().then((message) => {
        cleanup();
        resolve(message);
      }).catch((error) => {
        cleanup();
        reject(error);
      });
    });
  }

  /**
   * Register a handler for incoming messages (for event streaming)
   */
  onMessage(handler: (message: ManagementResponse) => void): () => void {
    this.on('message', handler);
    return () => this.removeListener('message', handler);
  }

  /**
   * Register a handler for connection errors (for event streaming)
   */
  onError(handler: (error: Error) => void): () => void {
    this.on('error', handler);
    return () => this.removeListener('error', handler);
  }

  close(): void {
    this._connected = false;
    
    // Clean up fragment reassembly state
    this.partialMessages.clear();
    
    // Reject any pending waiters
    // Note: We can't call waiters here since they expect ManagementResponse
    // The waiters will naturally timeout or error when the connection closes
    this.messageWaiters = [];
    this.messageQueue = [];
    
    if (this.socket) {
      this.socket.end();
      this.socket = null;
    }
    
    log.debug('Connection closed');
  }
}