Files
watchparty/server/vm/base.ts
T
2026-06-21 03:53:06 +03:00

626 lines
21 KiB
TypeScript

import config from "../config.ts";
import axios from "axios";
import { redis, redisCount } from "../utils/redis.ts";
import { postgres as pg } from "../utils/postgres.ts";
import type { PoolConfig, PoolRegion } from "./utils.ts";
import type { Client } from "pg";
const incrInterval = 5 * 1000;
const decrInterval = 15 * 1000;
const cleanupInterval = 5 * 60 * 1000;
// If postgres isn't configured we can still run in stateless mode
// Only start/get/terminate can be used, otherwise exception will be thrown
const postgres = !pg ? (null as unknown as Client) : pg;
export abstract class VMManager {
protected isLarge = false;
protected region: PoolRegion = "US";
private limitSize = 0;
private minSize = 0;
protected hostname: string | undefined;
constructor({ isLarge, region, limitSize, minSize, hostname }: PoolConfig) {
this.isLarge = isLarge;
this.region = region;
this.limitSize = Number(limitSize) || 0;
this.minSize = Number(minSize) || 0;
this.hostname = hostname;
}
public getIsLarge = () => {
return this.isLarge;
};
public getRegion = () => {
return this.region;
};
public getMinSize = () => {
return this.minSize;
};
public getLimitSize = () => {
return this.limitSize;
};
public getTargetBuffer = () => {
let buffer = this.limitSize * 0.02;
// If ramping config, adjust buffer based on the hour
// During ramp down hours, keep a smaller buffer
// During ramp up hours, keep a larger buffer
// const rampDownHours = config.VM_POOL_RAMP_DOWN_HOURS.split(",").map(Number);
// const rampUpHours = config.VM_POOL_RAMP_UP_HOURS.split(",").map(Number);
// const nowHour = new Date().getUTCHours();
// const isRampDown =
// rampDownHours.length &&
// pointInInterval24(nowHour, rampDownHours[0], rampDownHours[1]);
// const isRampUp =
// rampUpHours.length &&
// pointInInterval24(nowHour, rampUpHours[0], rampUpHours[1]);
// if (isRampDown) {
// buffer *= 0.5;
// } else if (isRampUp) {
// buffer *= 1.5;
// }
return Math.ceil(buffer);
};
public getCurrentSize = async () => {
const { rows } = await postgres.query(
`SELECT count(1) FROM vbrowser WHERE pool = $1`,
[this.getPoolName()],
);
return Number(rows[0]?.count);
};
public getPoolName = () => {
return this.id + (this.isLarge ? "Large" : "") + this.region;
};
public getAvailableCount = async (): Promise<number> => {
const { rows } = await postgres.query(
`SELECT count(1) FROM vbrowser WHERE pool = $1 and state = 'available'`,
[this.getPoolName()],
);
return Number(rows[0]?.count);
};
public getStagingCount = async (): Promise<number> => {
const { rows } = await postgres.query(
`SELECT count(1) FROM vbrowser WHERE pool = $1 and state = 'staging'`,
[this.getPoolName()],
);
return Number(rows[0]?.count);
};
public getAvailableVBrowsers = async (): Promise<string[]> => {
const { rows } = await postgres.query(
`SELECT vmid from vbrowser WHERE pool = $1 and state = 'available'`,
[this.getPoolName()],
);
return rows.map((row: any) => row.vmid);
};
public getStagingVBrowsers = async (): Promise<string[]> => {
const { rows } = await postgres.query(
`SELECT vmid from vbrowser WHERE pool = $1 and state = 'staging'`,
[this.getPoolName()],
);
return rows.map((row: any) => row.vmid);
};
public getTag = () => {
return (
(config.VBROWSER_TAG || "vbrowser") +
this.region +
(this.isLarge ? "Large" : "")
);
};
public assignVM = async (
roomId: string,
uid: string,
): Promise<AssignedVM | undefined> => {
if (!roomId || !uid) {
return undefined;
}
// Update and use SKIP LOCKED to ensure each consumer gets a different one
const { rows } = await postgres.query(
`
UPDATE vbrowser
SET "roomId" = $1, uid = $2, "heartbeatTime" = NOW(), "assignTime" = NOW(), state = 'used'
WHERE id = (
SELECT id
FROM vbrowser
WHERE state = 'available'
AND pool = $3
ORDER BY "creationTime" DESC
FOR UPDATE SKIP LOCKED
LIMIT 1
)
RETURNING data, pass`,
[roomId, uid, this.getPoolName()],
);
let vm: VM | undefined = rows[0]?.data;
let pass: string = rows[0]?.pass;
if (!vm) {
return;
}
return { ...vm, pass, assignTime: Date.now() };
};
public resetVM = async (vmid: string, roomId?: string): Promise<void> => {
if (roomId !== undefined) {
// verify the roomId matches if user initiated
const { rows } = await postgres.query(
`SELECT "roomId" FROM vbrowser WHERE pool = $1 AND vmid = $2`,
[this.getPoolName(), vmid],
);
if (rows[0]?.roomId && rows[0]?.roomId !== roomId) {
console.log(
"[RESET] %s: roomId mismatch on %s, expected %s, got %s",
this.getPoolName(),
vmid,
rows[0]?.roomId,
roomId,
);
return;
}
}
console.log("[RESET]", this.getPoolName(), vmid, roomId);
// We generally want to reuse if the provider has per-hour billing
// Since most user sessions are less than an hour
// Otherwise if it's per-second or Docker, it's easier to just terminate it on reboot
if (this.reuseVMs) {
// To reset without API calls/reimaging (logic can be reused on any cloud provider so we can avoid individual reboot implementations):
// Update vbrowser script to generate uuid password
// add admin endpoints to vbrowser to get password and reboot
// to reset, hit reboot endpoint with admin key OR call cloud reboot method
// in checkstaging, check password endpoint with admin key, update DB once available
const { rows } = await postgres.query(
`SELECT image FROM vbrowser WHERE pool = $1 AND vmid = $2`,
[this.getPoolName(), vmid],
);
const vmImageId = rows[0]?.image;
// Check if VM needs to be reimaged
if (this.imageId !== vmImageId) {
await this.reimageVM(vmid);
redisCount("vBrowserReimage");
// Update the vmImageId
await postgres.query(
`UPDATE vbrowser SET image = $3 WHERE pool = $1 AND vmid = $2`,
[this.getPoolName(), vmid, this.imageId],
);
} else {
await this.rebootVM(vmid);
}
// we could crash here and then row will remain in used state
// Once the heartbeat becomes stale cleanup will reset it again
const result = await postgres.query(
`
INSERT INTO vbrowser(pool, vmid, "creationTime", state)
VALUES($1, $2, NOW(), 'staging')
ON CONFLICT(pool, vmid) DO
UPDATE SET state = 'staging',
"roomId" = NULL, uid = NULL, retries = 0, "heartbeatTime" = NULL, "assignTime" = NULL, data = NULL
`,
[this.getPoolName(), vmid],
);
console.log("UPSERT", result.rowCount);
// Normally this should be an update, but we could insert if:
// if cleaning up a VM we didn't record in db on create
// if we resized down and deleted db row but didn't complete the termination
} else {
this.terminateVMWrapper(vmid);
}
};
public startVMWrapper = async () => {
// generate credentials and boot a VM
const password = crypto.randomUUID();
const id = await this.startVM(password);
// We might fail to record it if crashing here but cleanup will reset it
await postgres.query(
`
INSERT INTO vbrowser(pool, vmid, "creationTime", state, image)
VALUES($1, $2, NOW(), 'staging', $3)`,
[this.getPoolName(), id, this.imageId],
);
redisCount("vBrowserLaunches");
return id;
};
protected terminateVMWrapper = async (vmid: string) => {
console.log("[TERMINATE]", this.getPoolName(), vmid);
// Update the DB before calling terminate
// If we don't actually complete the termination, cleanup will reset it
const { command, rowCount } = await postgres.query(
`DELETE FROM vbrowser WHERE pool = $1 AND vmid = $2 RETURNING id`,
[this.getPoolName(), vmid],
);
console.log(command, rowCount);
// We can log the VM lifetime by returning the creationTime and diffing
await this.terminateVM(vmid);
};
checkVMReady = async (host: string) => {
// NOTE: This URL doesn't work for docker since it doesn't have /
// But since we don't reuse docker VMs it's ok
const url = "https://" + host.replace("/", "/health");
try {
// const out = execSync(`curl -i -L -v --ipv4 '${host}'`);
// if (!out.toString().startsWith('OK') && !out.toString().startsWith('404 page not found')) {
// throw new Error('mismatched response from health');
// }
const resp = await axios({
method: "GET",
url,
timeout: 1000,
});
// Check to make sure the VM was recently rebooted (we could also check on the password to ensure reset)
const timeSinceBoot = Date.now() / 1000 - Number(resp.data);
// console.log(timeSinceBoot);
return this.reuseVMs ? timeSinceBoot < 150 * 1000 : true;
} catch (e) {
// console.log(url, e.message, e.response?.status);
return false;
}
};
public runBackgroundJobs = async () => {
const resizeVMGroupIncr = async () => {
const availableCount = await this.getAvailableCount();
const stagingCount = await this.getStagingCount();
const currentSize = await this.getCurrentSize();
let launch = false;
launch =
availableCount + stagingCount < this.getTargetBuffer() &&
currentSize < (this.getLimitSize() || Infinity);
if (launch) {
console.log(
"[RESIZE-INCR]",
this.getPoolName(),
"target:",
this.getTargetBuffer(),
"available:",
availableCount,
"staging:",
stagingCount,
"currentSize:",
currentSize,
"limit:",
this.getLimitSize(),
);
try {
await this.startVMWrapper();
} catch (e: any) {
console.log(
e.response?.status,
JSON.stringify(e.response?.data),
e.config?.url,
e.config?.data,
);
}
}
};
const resizeVMGroupDecr = async () => {
const availableCount = await this.getAvailableCount();
const stagingCount = await this.getStagingCount();
let unlaunch = false;
unlaunch = availableCount + stagingCount > this.getTargetBuffer();
if (unlaunch) {
// use SKIP LOCKED to delete to avoid deleting VM that might be assigning
// filter to only VMs eligible for deletion
// they must be up for long enough
// keep the oldest min pool size number of VMs
// Hetzner/DO/Scaleway rounds up to nearest hour
let modulo = 3600;
const { rows } = await postgres.query(
`
DELETE FROM vbrowser
WHERE id = (
SELECT id
FROM vbrowser
WHERE pool = $1
AND state = 'available'
AND id >= (SELECT id from vbrowser WHERE pool = $1 ORDER BY id ASC LIMIT 1 OFFSET $2)
AND CAST(extract(epoch from now() - "creationTime") as INT) % $3 > $4
FOR UPDATE SKIP LOCKED
LIMIT 1
) RETURNING vmid`,
[
this.getPoolName(),
this.getMinSize(),
modulo,
config.VM_MIN_UPTIME_MINUTES * 60, // to seconds
],
);
const first = rows[0];
if (first) {
console.log("[RESIZE-DECR] %s: %s", this.getPoolName(), first.vmid);
await this.terminateVMWrapper(first.vmid);
}
}
};
const cleanupVMGroup = async () => {
// Reset hanging VMs
// It's possible we created a VM but lost track of it
// Take the list of VMs from API
// subtract VMs that have a heartbeat or available or staging
let allVMs = [];
try {
allVMs = await this.listVMs(this.getTag());
} catch (e) {
console.log(
"[CLEANUP] %s: failed to fetch VM list",
this.getPoolName(),
);
return;
}
const { rows } = await postgres.query(
`
SELECT vmid from vbrowser
WHERE pool = $1
AND
("heartbeatTime" > (NOW() - INTERVAL '5 minutes')
OR state = 'staging'
OR state = 'available')
`,
[this.getPoolName()],
);
const inUse = new Set(rows.map((row: any) => row.vmid));
console.log(
"[CLEANUP] %s: found %s VMs, %s to keep",
this.getPoolName(),
allVMs.length,
inUse.size,
);
for (let server of allVMs) {
if (!inUse.has(server.id)) {
redisCount("vBrowserCleanup");
console.log("[CLEANUP]", this.getPoolName(), server.id);
try {
await this.resetVM(server.id);
//this.terminateVMWrapper(server.id);
} catch (e: any) {
console.warn("[CLEANUP]", this.getPoolName(), e.response?.data);
}
await new Promise((resolve) => setTimeout(resolve, 2000));
}
}
};
const checkStaging = async () => {
// Increment retry count and return data
const { rows } = await postgres.query(
`
UPDATE vbrowser
SET retries = retries + 1
WHERE pool = $1 and state = 'staging'
RETURNING id, vmid, data, retries
`,
[this.getPoolName()],
);
const stagingPromises = rows.map(async (row: any): Promise<string> => {
const rowid = row.id;
const vmid: string = row.vmid;
const retryCount: number = row.retries;
let vm: VM | null = row.data;
if (retryCount < this.minRetries) {
if (config.NODE_ENV === "development") {
console.log(
"[CHECKSTAGING] %s: [vmid: %s] [attempt: %s] waiting for minRetries",
this.getPoolName(),
vmid,
retryCount,
);
}
// Do a minimum # of retries to give reboot time
return [vmid, retryCount, false].join(",");
}
// Fetch data on first attempt
// Try again only every once in a while to reduce load on API
const shouldFetchVM =
retryCount === this.minRetries + 1 || retryCount % 30 === 0;
// Refetch the VM every once in a while even if cached to check if it should be deleted
if (shouldFetchVM) {
try {
vm = await this.getVM(vmid);
} catch (e: any) {
console.warn(e.response?.data);
if (e.response?.status === 404) {
// Remove the VM because the provider says it doesn't exist
await postgres.query("DELETE FROM vbrowser WHERE id = $1", [
rowid,
]);
throw new Error("failed to find vm " + vmid);
}
}
if (vm?.host) {
// Save the VM data
await postgres.query(
`UPDATE vbrowser SET data = $1 WHERE id = $2`,
[vm, rowid],
);
}
}
if (!vm?.host) {
console.log(
"[CHECKSTAGING] %s: no host for vm %s",
this.getPoolName(),
vmid,
);
throw new Error("no host for vm " + vmid);
}
if (retryCount % 150 === 0) {
console.log(
"[CHECKSTAGING] %s: %s poweron, attach to network",
this.getPoolName(),
vmid,
);
this.powerOn(vmid);
//this.attachToNetwork(vmid);
}
if (retryCount % 180 === 0) {
console.log("[CHECKSTAGING]", this.getPoolName(), "giving up:", vmid);
redisCount("vBrowserStagingFails");
await redis?.lpush("vBrowserStageFails", vmid);
await redis?.ltrim("vBrowserStageFails", 0, 19);
// VM didn't come up. set image to null so we reimage
await postgres.query(
`UPDATE vbrowser SET image = NULL WHERE pool = $1 AND vmid = $2`,
[this.getPoolName(), vmid],
);
await this.resetVM(vmid);
}
if (retryCount >= 180) {
throw new Error("too many attempts on vm " + vmid);
}
const ready = await this.checkVMReady(vm.host);
if (
ready ||
retryCount % (config.NODE_ENV === "development" ? 1 : 30) === 0
) {
console.log(
"[CHECKSTAGING] %s: [ready: %s] [vmid: %s] [retries: %s] [host: %s]",
this.getPoolName(),
ready,
vmid,
retryCount,
vm?.host,
);
}
if (ready) {
const passUrl =
"https://" +
(vm.host.includes("/")
? vm.host.replace("/", "/password")
: vm.host + "/password") +
(vm.host.includes("?") ? "&" : "?") +
"key=" +
config.VBROWSER_ADMIN_KEY;
// console.log(passUrl);
const resp2 = await axios({
method: "GET",
url: passUrl,
timeout: 1000,
});
const pass = resp2.data;
// password is a uuid so it should never match the previous value, if it does we probably didn't reset properly
const { rows } = await postgres.query(
`UPDATE vbrowser SET state = 'available', pass = $2 WHERE id = $1 AND pass IS DISTINCT FROM $2`,
[rowid, pass],
);
// console.log(rows);
await redis?.lpush("vBrowserStageRetries", retryCount);
await redis?.ltrim("vBrowserStageRetries", 0, 19);
}
return [vmid, retryCount, ready].join(",");
});
// TODO log something if we timeout
const result = await Promise.race([
Promise.allSettled(stagingPromises),
new Promise((resolve) => setTimeout(resolve, 30000)),
]);
return result;
};
console.log("[VMWORKER] %s: starting background jobs", this.getPoolName());
setInterval(resizeVMGroupIncr, incrInterval);
setInterval(resizeVMGroupDecr, decrInterval);
setInterval(async () => {
console.log(
"[STATS] %s: currentSize %s, available %s, staging %s, target %s",
this.getPoolName(),
await this.getCurrentSize(),
await this.getAvailableCount(),
await this.getStagingCount(),
this.getTargetBuffer(),
);
}, 10000);
// The following may take a while per iteration
// Use while loop and delay between iterations rather than setInterval to avoid stacking requests
setImmediate(async () => {
while (true) {
console.time(this.getPoolName() + ":cleanup");
try {
await cleanupVMGroup();
} catch (e: any) {
console.warn(
"[CLEANUPVMGROUP-ERROR]",
this.getPoolName(),
e.response?.data,
);
}
console.timeEnd(this.getPoolName() + ":cleanup");
await new Promise((resolve) => setTimeout(resolve, cleanupInterval));
}
});
setImmediate(async () => {
while (true) {
// console.time(this.getPoolName() + ':checkstaging');
try {
await checkStaging();
} catch (e) {
console.warn("[CHECKSTAGING-ERROR]", this.getPoolName(), e);
}
// console.timeEnd(this.getPoolName() + ':checkstaging');
await new Promise((resolve) => setTimeout(resolve, 1000));
}
});
};
public abstract id: string;
protected abstract size: string;
protected abstract largeSize: string;
protected abstract minRetries: number;
protected abstract reuseVMs: boolean;
protected abstract imageId: string;
protected abstract startVM: (name: string) => Promise<string>;
protected abstract rebootVM: (id: string) => Promise<void>;
protected abstract reimageVM: (id: string) => Promise<void>;
protected abstract terminateVM: (id: string) => Promise<void>;
public abstract getVM: (id: string) => Promise<VM>;
protected abstract listVMs: (filter: string) => Promise<VM[]>;
protected abstract powerOn: (id: string) => Promise<void>;
protected abstract attachToNetwork: (id: string) => Promise<void>;
protected abstract mapServerObject: (server: any) => VM;
public abstract updateSnapshot: () => Promise<string>;
}
function pointInInterval24(x: number, a: number, b: number) {
return nonNegativeMod(x - a, 24) <= nonNegativeMod(b - a, 24);
}
function nonNegativeMod(n: number, m: number) {
return ((n % m) + m) % m;
}
export interface VM {
id: string;
host: string;
provider: string;
large: boolean;
region: string;
}
export interface AssignedVM extends VM {
pass: string;
assignTime: number;
controllerClient?: string;
creatorUID?: string;
creatorClientID?: string;
}
export interface VMManagers {
standard: VMManager | null;
large: VMManager | null;
US: VMManager | null;
}