huge optimizations

This commit is contained in:
KingKDot
2026-06-08 22:33:34 -05:00
parent 9a1fc96e2b
commit d6ad2cadea
21 changed files with 649 additions and 72 deletions
+16 -9
View File
@@ -1,6 +1,10 @@
import { logger } from "../logger";
import { timingSafeEqual } from "crypto";
let warnedDisableAuthIgnored = false;
let loggedAuthDisabledByEnv = false;
let loggedAuthDisabledNoToken = false;
function safeCompare(a: string, b: string): boolean {
const bufA = Buffer.from(a);
const bufB = Buffer.from(b);
@@ -19,16 +23,25 @@ export function isAuthorizedAgentRequest(
if (disableAuth) {
const nodeEnv = String(process.env.NODE_ENV || "development").toLowerCase();
if (nodeEnv === "production") {
logger.warn("[auth] OVERLORD_DISABLE_AGENT_AUTH is ignored in production mode");
if (!warnedDisableAuthIgnored) {
warnedDisableAuthIgnored = true;
logger.warn("[auth] OVERLORD_DISABLE_AGENT_AUTH is ignored in production mode");
}
} else {
logger.info("[auth] Agent auth explicitly disabled by OVERLORD_DISABLE_AGENT_AUTH=true (non-production mode)");
if (!loggedAuthDisabledByEnv) {
loggedAuthDisabledByEnv = true;
logger.info("[auth] Agent auth explicitly disabled by OVERLORD_DISABLE_AGENT_AUTH=true (non-production mode)");
}
return true;
}
}
const token = agentToken?.trim();
if (!token) {
logger.info("[auth] Agent auth disabled");
if (!loggedAuthDisabledNoToken) {
loggedAuthDisabledNoToken = true;
logger.info("[auth] Agent auth disabled");
}
return true;
}
@@ -38,11 +51,5 @@ export function isAuthorizedAgentRequest(
(headerToken !== null && safeCompare(headerToken, token)) ||
(queryToken !== null && safeCompare(queryToken, token));
if (!isAuthed) {
logger.info("[auth] Agent auth failed");
} else {
logger.info("[auth] Agent authenticated successfully");
}
return isAuthed;
}
@@ -0,0 +1,16 @@
import { describe, expect, test } from "bun:test";
import { getHeartbeatBatchSize } from "./maintenance-loops";
describe("maintenance heartbeat batching", () => {
test("spreads large fleets across the heartbeat interval", () => {
expect(getHeartbeatBatchSize(1000, 30_000)).toBe(34);
});
test("keeps small fleets responsive", () => {
expect(getHeartbeatBatchSize(10, 30_000)).toBe(1);
});
test("handles empty fleets", () => {
expect(getHeartbeatBatchSize(0, 30_000)).toBe(0);
});
});
@@ -1,4 +1,5 @@
import { logger } from "../logger";
import { metrics } from "../metrics";
import { sendPingRequest } from "../wsHandlers";
import type { ClientInfo } from "../types";
import { pruneStaleClients } from "./stale-prune";
@@ -13,8 +14,30 @@ type StartMaintenanceParams = {
disconnectTimeoutMs: number;
};
function nonNegativeIntEnv(name: string, fallback: number): number {
const raw = String(process.env[name] || "").trim();
if (!raw) return fallback;
const parsed = Number(raw);
if (!Number.isFinite(parsed) || parsed < 0) return fallback;
return Math.floor(parsed);
}
const HEARTBEAT_SWEEP_TICK_MS = Math.max(
250,
nonNegativeIntEnv("OVERLORD_HEARTBEAT_SWEEP_TICK_MS", 1_000),
);
const HEARTBEAT_BATCH_SIZE = nonNegativeIntEnv("OVERLORD_HEARTBEAT_BATCH_SIZE", 0);
export function getHeartbeatBatchSize(totalClients: number, heartbeatIntervalMs: number): number {
if (totalClients <= 0) return 0;
if (HEARTBEAT_BATCH_SIZE > 0) return Math.min(totalClients, HEARTBEAT_BATCH_SIZE);
const ticksPerSweep = Math.max(1, Math.floor(heartbeatIntervalMs / HEARTBEAT_SWEEP_TICK_MS));
return Math.max(1, Math.ceil(totalClients / ticksPerSweep));
}
export function startMaintenanceLoops(params: StartMaintenanceParams): void {
setInterval(() => {
const startedAt = Date.now();
pruneStaleClients({
clients: params.getClients(),
staleMs: params.staleMs,
@@ -22,6 +45,7 @@ export function startMaintenanceLoops(params: StartMaintenanceParams): void {
setOnlineState: params.setOnlineState,
deleteClient: params.deleteClient,
});
metrics.recordInternalTask("stale-prune", Date.now() - startedAt);
}, 5000);
const livenessTimeoutMs = Math.max(
@@ -29,10 +53,23 @@ export function startMaintenanceLoops(params: StartMaintenanceParams): void {
60_000,
);
let heartbeatCursor = 0;
setInterval(() => {
const startedAt = Date.now();
const now = Date.now();
for (const [id, info] of params.getClients().entries()) {
if (info.role !== "client") continue;
const clients = Array.from(params.getClients().entries()).filter(
([, info]) => info.role === "client",
);
if (clients.length === 0) {
heartbeatCursor = 0;
metrics.recordInternalTask("heartbeat-sweep", Date.now() - startedAt);
return;
}
if (heartbeatCursor >= clients.length) heartbeatCursor = 0;
const batchSize = getHeartbeatBatchSize(clients.length, params.heartbeatIntervalMs);
for (let i = 0; i < batchSize; i++) {
const [id, info] = clients[(heartbeatCursor + i) % clients.length];
const lastActivity = info.lastSeen || 0;
if (lastActivity && now - lastActivity > livenessTimeoutMs) {
logger.warn(
@@ -51,5 +88,7 @@ export function startMaintenanceLoops(params: StartMaintenanceParams): void {
logger.debug(`[ping] heartbeat failed for ${id}`, err);
}
}
}, params.heartbeatIntervalMs);
heartbeatCursor = (heartbeatCursor + batchSize) % clients.length;
metrics.recordInternalTask("heartbeat-sweep", Date.now() - startedAt);
}, HEARTBEAT_SWEEP_TICK_MS);
}
@@ -23,6 +23,7 @@ import {
} from "../../db";
import { metrics } from "../../metrics";
import { encodeMessage } from "../../protocol";
import { getConfig } from "../../config";
import { requireClientAccess, requireFeatureAccess, requirePermission } from "../../rbac";
import {
canUserAccessClient,
@@ -176,6 +177,12 @@ export async function handleClientRoutes(
if (error instanceof Response) return error;
return new Response("Forbidden", { status: 403 });
}
if (
req.headers.get("x-overlord-thumbnail-source") === "dashboard" &&
getConfig().thumbnails.dashboardEnabled === false
) {
return Response.json({ ok: false, disabled: true }, { headers: deps.CORS_HEADERS });
}
const { requestThumbnailRegen, markThumbnailRequested, getThumbnailVersion, waitForThumbnail, clearThumbnailRequest } = await import("../../thumbnails");
markThumbnailRequested(clientId);
const target = clientManager.getClient(clientId);
@@ -9,7 +9,7 @@ async function getGeoip() {
}
import { logAudit, AuditAction } from "../../auditLog";
import * as clientManager from "../../clientManager";
import { clientExists, setOnlineState, upsertClientRow, getClientEnrollmentStatus, setClientEnrollmentStatus, lookupClientByPublicKey, getClientPublicKeyById, getBuildByTag, computeClientSuspiciousFlags } from "../../db";
import { clientExists, setOnlineState, setOfflineStates, upsertClientRow, getClientEnrollmentStatus, setClientEnrollmentStatus, lookupClientByPublicKey, getClientPublicKeyById, getBuildByTag, computeClientSuspiciousFlags, type OfflineStateUpdate } from "../../db";
import { getConfig } from "../../config";
import { logger } from "../../logger";
import { metrics } from "../../metrics";
@@ -18,6 +18,7 @@ import * as sessionManager from "../../sessions/sessionManager";
import type { SocketData } from "../../sessions/types";
import type { ClientInfo } from "../../types";
import { clearClientSyncState, handleFrame, handleHello, handlePing, handlePong } from "../../wsHandlers";
import { queueClientDbUpdate } from "../../client-db-sync";
import { getMaxPayloadLimit, getMessageByteLength, isAllowedClientMessageType } from "../../wsValidation";
import { stopAllProxiesForClient } from "../socks5-proxy-manager";
import { verifyBuildToken, isBuildBanned } from "../build-signing";
@@ -42,19 +43,96 @@ type ClientLifecyclePayload = {
};
type PendingOffline = {
timer: ReturnType<typeof setTimeout>;
dueAt: number;
payload: ClientLifecyclePayload;
disconnectReason: string | undefined;
disconnectDetail: string | undefined;
deps: WsLifecycleDeps;
};
const pendingOffline = new Map<string, PendingOffline>();
let pendingOfflineFlushTimer: ReturnType<typeof setTimeout> | null = null;
const OFFLINE_BATCH_SIZE = Math.max(
1,
Number(process.env.OVERLORD_OFFLINE_BATCH_SIZE || 250),
);
const OFFLINE_BATCH_DELAY_MS = Math.max(
1,
Number(process.env.OVERLORD_OFFLINE_BATCH_DELAY_MS || 25),
);
function cancelPendingOffline(clientId: string): boolean {
const pending = pendingOffline.get(clientId);
if (!pending) return false;
clearTimeout(pending.timer);
pendingOffline.delete(clientId);
return true;
}
function applyPendingOffline(
clientId: string,
pending: PendingOffline,
offlineUpdates: OfflineStateUpdate[],
): void {
if (clientManager.hasClient(clientId)) return;
pending.deps.notifyDashboardClientEvent("client_offline", pending.payload);
pending.deps.broadcastClientEvent("client_offline", pending.payload);
pending.deps.notifyRemoteDesktopStatus(clientId, "offline", "Client disconnected");
offlineUpdates.push({
id: clientId,
disconnectReason: pending.disconnectReason,
disconnectDetail: pending.disconnectDetail,
});
pending.deps.notifyDashboard();
}
function schedulePendingOfflineFlush(delayMs: number): void {
if (pendingOfflineFlushTimer) return;
pendingOfflineFlushTimer = setTimeout(() => {
pendingOfflineFlushTimer = null;
flushPendingOffline();
}, delayMs);
}
function flushPendingOffline(): void {
if (pendingOffline.size === 0) return;
const startedAt = Date.now();
const now = Date.now();
let processed = 0;
let nextDueAt = Number.POSITIVE_INFINITY;
const offlineUpdates: OfflineStateUpdate[] = [];
for (const [clientId, pending] of pendingOffline) {
if (pending.dueAt > now) {
if (pending.dueAt < nextDueAt) nextDueAt = pending.dueAt;
continue;
}
pendingOffline.delete(clientId);
applyPendingOffline(clientId, pending, offlineUpdates);
processed += 1;
if (processed >= OFFLINE_BATCH_SIZE) {
break;
}
}
if (offlineUpdates.length > 0) {
const dbStartedAt = Date.now();
setOfflineStates(offlineUpdates);
metrics.recordInternalTask("offline-db-flush", Date.now() - dbStartedAt);
}
metrics.recordInternalTask("offline-flush", Date.now() - startedAt);
if (pendingOffline.size === 0) return;
if (processed >= OFFLINE_BATCH_SIZE) {
schedulePendingOfflineFlush(OFFLINE_BATCH_DELAY_MS);
return;
}
if (Number.isFinite(nextDueAt)) {
schedulePendingOfflineFlush(Math.max(1, nextDueAt - Date.now()));
}
}
function schedulePendingOffline(
clientId: string,
payload: ClientLifecyclePayload,
@@ -65,24 +143,25 @@ function schedulePendingOffline(
cancelPendingOffline(clientId);
if (OFFLINE_GRACE_MS <= 0) {
deps.notifyDashboardClientEvent("client_offline", payload);
deps.broadcastClientEvent("client_offline", payload);
setOnlineState(clientId, false, disconnectReason, disconnectDetail);
deps.notifyRemoteDesktopStatus(clientId, "offline", "Client disconnected");
deps.notifyDashboard();
pendingOffline.set(clientId, {
dueAt: Date.now(),
payload,
disconnectReason,
disconnectDetail,
deps,
});
schedulePendingOfflineFlush(1);
return;
}
const timer = setTimeout(() => {
pendingOffline.delete(clientId);
deps.notifyDashboardClientEvent("client_offline", payload);
deps.broadcastClientEvent("client_offline", payload);
setOnlineState(clientId, false, disconnectReason, disconnectDetail);
deps.notifyRemoteDesktopStatus(clientId, "offline", "Client disconnected");
deps.notifyDashboard();
}, OFFLINE_GRACE_MS);
pendingOffline.set(clientId, { timer });
pendingOffline.set(clientId, {
dueAt: Date.now() + OFFLINE_GRACE_MS,
payload,
disconnectReason,
disconnectDetail,
deps,
});
schedulePendingOfflineFlush(OFFLINE_GRACE_MS);
}
type PendingScript = {
@@ -263,8 +342,6 @@ export function handleWebSocketOpen(ws: ServerWebSocket<SocketData>, deps: WsLif
} catch {}
}, ENROLLMENT_TIMEOUT_MS);
enrollmentTimeouts.set(id, timeout);
logger.info(`[open] ${id} role=${role} — purgatory challenge sent`);
}
export async function handleWebSocketMessage(
@@ -560,7 +637,7 @@ export async function handleWebSocketMessage(
};
clientManager.addClient(resolvedId, infoObj);
upsertClientRow({
queueClientDbUpdate({
id: resolvedId,
publicKey,
keyFingerprint,
@@ -644,7 +721,6 @@ export async function handleWebSocketMessage(
break;
case "pong":
handlePong(client, payload);
deps.notifyDashboard();
break;
case "frame":
if ((payload as any)?.header?.fps === 0) {
@@ -963,7 +1039,6 @@ export function handleWebSocketClose(
const currentClient = clientManager.getClient(clientId);
if (currentClient && currentClient.ws !== ws) {
logger.info(`[close] ${clientId} code=${code} (superseded socket, skipping cleanup)`);
return;
}
@@ -994,7 +1069,6 @@ export function handleWebSocketClose(
deps.hvncStreamingState.delete(clientId);
deps.webcamStreamingState.delete(clientId);
stopRemoteDesktopRecording(clientId, "client disconnected");
logger.info(`[close] ${clientId} code=${code} reason=${reason} disconnect_reason=${storedDisconnectReason || "unknown"}`);
if (role === "client" && currentClient) {
schedulePendingOffline(
@@ -14,8 +14,16 @@ type WsUpgradeDeps = {
isAuthorizedAgentRequest: (req: Request, url: URL) => boolean;
};
const WS_RATE_WINDOW_MS = 10_000;
const WS_RATE_MAX = 30;
function positiveIntEnv(name: string, fallback: number): number {
const raw = String(process.env[name] || "").trim();
if (!raw) return fallback;
const parsed = Number(raw);
if (!Number.isFinite(parsed) || parsed <= 0) return fallback;
return Math.floor(parsed);
}
const WS_RATE_WINDOW_MS = positiveIntEnv("OVERLORD_WS_UPGRADE_RATE_WINDOW_MS", 10_000);
const WS_RATE_MAX = positiveIntEnv("OVERLORD_WS_UPGRADE_RATE_MAX", 30);
const wsRateMap = new Map<string, { count: number; windowStart: number }>();
setInterval(() => {
@@ -238,7 +246,6 @@ export async function handleWsUpgradeRoutes(
const wsMatch = url.pathname.match(/^\/api\/clients\/(.+)\/stream\/ws$/);
if (wsMatch) {
logger.info(`[auth] Checking agent authorization for client connection`);
if (!deps.isAuthorizedAgentRequest(req, url)) {
return new Response("Unauthorized", { status: 401 });
}
@@ -1,6 +1,14 @@
export const HEARTBEAT_INTERVAL_MS = 15_000;
export const STALE_MS = 90_000;
export const DISCONNECT_TIMEOUT_MS = 10_000;
function positiveIntEnv(name: string, fallback: number): number {
const raw = String(process.env[name] || "").trim();
if (!raw) return fallback;
const parsed = Number(raw);
if (!Number.isFinite(parsed) || parsed <= 0) return fallback;
return Math.floor(parsed);
}
export const HEARTBEAT_INTERVAL_MS = positiveIntEnv("OVERLORD_HEARTBEAT_INTERVAL_MS", 15_000);
export const STALE_MS = positiveIntEnv("OVERLORD_STALE_MS", 90_000);
export const DISCONNECT_TIMEOUT_MS = positiveIntEnv("OVERLORD_DISCONNECT_TIMEOUT_MS", 10_000);
export const PRUNE_BATCH = Number(process.env.PRUNE_BATCH || 500);
export const MAX_WS_MESSAGE_BYTES_VIEWER = Number(