All files / server/routes/simulation ws-session-manager.ts

100% Statements 50/50
100% Branches 29/29
100% Functions 11/11
100% Lines 49/49

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144                                                            21x   21x     54x 54x       189x       52x 52x 51x   52x       1x       1x       187x 187x 605x   187x       143x 143x 439x 127x                         49x 1x     48x 48x 48x 48x 48x   48x 47x 47x     48x 47x 47x   1x         47x 47x   1x           48x 46x 46x         58x 2x 2x         52x 52x 51x 51x 13x       52x 52x      
import { WebSocket } from "ws";
import { WSMessageType } from "@shared/schema";
import type { Logger } from "@shared/logger";
import type { SandboxRunner } from "../../services/sandbox-runner";
import type { SandboxRunnerPool } from "../../services/sandbox-runner-pool";
import { WsSessionLifecycle } from "../../services/ws-session-lifecycle";
import { sendMessageToClient } from "./ws-output-buffer";
import { webSocketMetricsTracker } from "../../services/server-metrics";
import type {
  SimulationAdmissionController,
  SimulationReservation,
} from "../../services/simulation-admission-controller";
 
export type ClientState = {
  subject: string;
  runner: SandboxRunner | null;
  isRunning: boolean;
  isPaused: boolean;
  testRunId?: string;
  queueAbortController: AbortController | null;
  reservation: SimulationReservation | null;
};
 
interface WsSessionManagerParams {
  pool: SandboxRunnerPool;
  logger: Logger;
  admissionController?: Pick<SimulationAdmissionController, "release">;
}
 
export class WsSessionManager {
  private readonly clientRunners = new WsSessionLifecycle<WebSocket, ClientState>();
 
  constructor(private readonly params: WsSessionManagerParams) {}
 
  register(ws: WebSocket, state: ClientState): void {
    webSocketMetricsTracker.onConnection();
    this.clientRunners.register(ws, state);
  }
 
  get(ws: WebSocket): ClientState | undefined {
    return this.clientRunners.get(ws);
  }
 
  remove(ws: WebSocket): ClientState | undefined {
    const state = this.clientRunners.remove(ws);
    if (state) {
      webSocketMetricsTracker.onDisconnection();
    }
    return state;
  }
 
  entries(): IterableIterator<[WebSocket, ClientState]> {
    return this.clientRunners.entries();
  }
 
  get size(): number {
    return this.clientRunners.size;
  }
 
  countRunningClients(): number {
    let count = 0;
    for (const state of this.clientRunners.values()) {
      if (state.isRunning) count++;
    }
    return count;
  }
 
  broadcastWorkerTotal(excludeWs?: WebSocket): void {
    const newTotal = this.countRunningClients();
    for (const [otherWs, otherState] of this.clientRunners.entries()) {
      if (otherWs !== excludeWs && otherState.isRunning) {
        sendMessageToClient(otherWs, {
          type: WSMessageType.COMPILATION_STATUS,
          workerTotal: newTotal,
        });
      }
    }
  }
 
  async safeReleaseRunner(
    state: ClientState,
    reason: string,
    expectedReservation: SimulationReservation | null = state.reservation,
  ): Promise<void> {
    if (expectedReservation && state.reservation !== expectedReservation) {
      return;
    }
 
    const runner = state.runner;
    state.runner = null;
    const wasRunning = state.isRunning;
    state.isRunning = false;
    state.isPaused = false;
 
    if (wasRunning) {
      webSocketMetricsTracker.onSessionStop();
      this.broadcastWorkerTotal();
    }
 
    if (runner) {
      try {
        await runner.stop();
      } catch (error) {
        this.params.logger.debug(
          `[SandboxRunnerPool] runner.stop() failed during ${reason}: ${error}`,
        );
      }
 
      try {
        await this.params.pool.releaseRunner(runner);
      } catch (error) {
        this.params.logger.warn(
          `[SandboxRunnerPool] releaseRunner failed during ${reason}: ${error}`,
        );
      }
    }
 
    if (expectedReservation && state.reservation === expectedReservation) {
      this.params.admissionController?.release(expectedReservation);
      state.reservation = null;
    }
  }
 
  abortQueuedAcquire(state: ClientState): void {
    if (state.queueAbortController) {
      state.queueAbortController.abort();
      state.queueAbortController = null;
    }
  }
 
  async cleanupClient(ws: WebSocket, reason: string): Promise<void> {
    const clientState = this.get(ws);
    if (clientState) {
      this.abortQueuedAcquire(clientState);
      if (clientState.runner || clientState.reservation) {
        await this.safeReleaseRunner(clientState, reason);
      }
    }
 
    this.remove(ws);
    this.broadcastWorkerTotal();
  }
}