feat: respect context in server

This commit is contained in:
Aarnav Tale
2025-02-19 18:09:42 -05:00
parent 06049169a2
commit f5436f5ee3
32 changed files with 359 additions and 204 deletions
+39 -32
View File
@@ -2,8 +2,8 @@
import { readFile, writeFile } from 'node:fs/promises';
import { setTimeout as pSetTimeout } from 'node:timers/promises';
import type { LoaderFunctionArgs } from 'react-router';
import type { HostInfo } from '~/types';
import { WebSocket } from 'ws';
import type { HostInfo } from '~/types';
import log from './log';
// Essentially a HashMap which invalidates entries after a certain time.
@@ -68,7 +68,7 @@ class TimedCache<K, V> {
this.writeLock = true;
const data = Array.from(this._cache.entries()).map(([key, value]) => {
return { key, value, expires: this._timeCache.get(key) }
return { key, value, expires: this._timeCache.get(key) };
});
await writeFile(this.filepath, JSON.stringify(data), 'utf-8');
@@ -85,10 +85,11 @@ export async function initAgentCache(defaultTTL: number, filepath: string) {
}
let agentSocket: WebSocket | undefined;
// TODO: Actually type this?
export function initAgentSocket(context: LoaderFunctionArgs['context']) {
if (!context.ws) {
return;
};
}
const client = context.ws.clients.values().next().value;
agentSocket = client;
@@ -97,21 +98,24 @@ export function initAgentSocket(context: LoaderFunctionArgs['context']) {
// Check the cache and then attempt the websocket query
// If we aren't connected to an agent, then debug log and return the cache
export async function queryAgent(nodes: string[]) {
return;
return;
// biome-ignore lint: bruh
if (!cache) {
log.error('CACH', 'Cache not initialized');
return;
}
const cached: Record<string, HostInfo> = {};
await Promise.all(nodes.map(async node => {
const cachedData = await cache?.get(node);
if (cachedData) {
cached[node] = cachedData;
}
}))
await Promise.all(
nodes.map(async (node) => {
const cachedData = await cache?.get(node);
if (cachedData) {
cached[node] = cachedData;
}
}),
);
const uncached = nodes.filter(node => !cached[node]);
const uncached = nodes.filter((node) => !cached[node]);
// No need to query the agent if we have all the data cached
if (uncached.length === 0) {
@@ -124,29 +128,32 @@ export async function queryAgent(nodes: string[]) {
return cached;
}
agentSocket.send(JSON.stringify({ NodeIDs: uncached }));
const returnData = await new Promise<Record<string, HostInfo> | void>((resolve, reject) => {
const timeout = setTimeout(() => {
agentSocket?.removeAllListeners('message');
resolve();
}, 3000);
agentSocket?.on('message', async (message: string) => {
const data = JSON.parse(message.toString());
if (Object.keys(data).length === 0) {
agentSocket?.send(JSON.stringify({ NodeIDs: uncached }));
// biome-ignore lint: bruh
const returnData = await new Promise<Record<string, HostInfo> | void>(
(resolve, reject) => {
const timeout = setTimeout(() => {
agentSocket?.removeAllListeners('message');
resolve();
}
}, 3000);
agentSocket?.removeAllListeners('message');
resolve(data);
});
});
if (returnData) {
for await (const [node, info] of Object.entries(returnData)) {
await cache.set(node, info);
}
}
agentSocket?.on('message', async (message: string) => {
const data = JSON.parse(message.toString());
if (Object.keys(data).length === 0) {
resolve();
}
agentSocket?.removeAllListeners('message');
resolve(data);
});
},
);
// if (returnData) {
// for await (const [node, info] of Object.entries(returnData)) {
// await cache?.set(node, info);
// }
// }
return returnData ? { ...cached, ...returnData } : cached;
}