dahjah/headplane
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

feat: keep track of agent ids

Browse Source
This commit is contained in:
Aarnav Tale 2025-03-06 17:37:41 -05:00
parent 5dd4c41291
commit d9314887fb
No known key found for this signature in database
3 changed files with 14 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
server/context/app.ts
View File

@ -1,19 +1,22 @@
import type { HostInfo } from '~/types'; import type { HostInfo } from '~/types';
import { TimedCache } from '~server/ws/cache'; import { TimedCache } from '~server/ws/cache';
import { hp_agentRequest, hp_getAgentCache } from '~server/ws/data'; import { hp_agentRequest, hp_getAgentCache } from '~server/ws/data';
import { hp_getAgents } from '~server/ws/socket';
import { hp_getConfig } from './loader'; import { hp_getConfig } from './loader';
import type { HeadplaneConfig } from './parser'; import type { HeadplaneConfig } from './parser';
export interface AppContext { export interface AppContext {
context: HeadplaneConfig; context: HeadplaneConfig;
agentData?: TimedCache<HostInfo>;
hp_agentRequest: typeof hp_agentRequest; hp_agentRequest: typeof hp_agentRequest;
agents: string[];
agentData?: TimedCache<HostInfo>;
} }
export default function appContext() { export default function appContext(): AppContext {
return { return {
context: hp_getConfig(), context: hp_getConfig(),
agentData: hp_getAgentCache(),
hp_agentRequest, hp_agentRequest,
agents: [...hp_getAgents().keys()],
agentData: hp_getAgentCache(),
}; };
} }

6
server/ws/data.ts
View File

@ -28,8 +28,7 @@ export async function hp_agentRequest(nodeList: string[]) {
// Request to all connected agents (we can have multiple) // Request to all connected agents (we can have multiple)
// Luckily we can parse all the data at once through message parsing // Luckily we can parse all the data at once through message parsing
// and then overlapping cache entries will be overwritten by time // and then overlapping cache entries will be overwritten by time
const agents = [...hp_getAgents()]; const agents = hp_getAgents();
console.log(agents);
// Deduplicate the list of nodes // Deduplicate the list of nodes
const NodeIDs = [...new Set(nodeList)]; const NodeIDs = [...new Set(nodeList)];
@ -40,7 +39,7 @@ export async function hp_agentRequest(nodeList: string[]) {
// Await so that data loads on first request without racing // Await so that data loads on first request without racing
// Since we do agent.once() we NEED to wait for it to finish // Since we do agent.once() we NEED to wait for it to finish
await Promise.allSettled( await Promise.allSettled(
agents.map(async (agent) => { [...agents].map(async ([id, agent]) => {
agent.send(JSON.stringify({ NodeIDs })); agent.send(JSON.stringify({ NodeIDs }));
await new Promise<void>((resolve) => { await new Promise<void>((resolve) => {
// Just as a safety measure, we set a maximum timeout of 3 seconds // Just as a safety measure, we set a maximum timeout of 3 seconds
@ -48,6 +47,7 @@ export async function hp_agentRequest(nodeList: string[]) {
agent.once('message', (data) => { agent.once('message', (data) => {
const parsed = JSON.parse(data.toString()); const parsed = JSON.parse(data.toString());
log.debug('CACH', 'Received agent data from %s', id);
for (const [node, info] of Object.entries<HostInfo>(parsed)) { for (const [node, info] of Object.entries<HostInfo>(parsed)) {
cache?.set(node, info); cache?.set(node, info);
log.debug('CACH', 'Cached %s', node); log.debug('CACH', 'Cached %s', node);

7
server/ws/socket.ts
View File

@ -10,7 +10,7 @@ export function initWebsocket(authKey: string) {
log.info('SRVX', 'Starting a WebSocket server for agent connections'); log.info('SRVX', 'Starting a WebSocket server for agent connections');
server.on('connection', (ws, req) => { server.on('connection', (ws, req) => {
const tailnetID = req.headers['x-headplane-tailnet-id']; const tailnetID = req.headers['x-headplane-tailnet-id'];
if (!tailnetID) { if (!tailnetID || typeof tailnetID !== 'string') {
log.warn( log.warn(
'SRVX', 'SRVX',
'Rejecting an agent WebSocket connection without a tailnet ID', 'Rejecting an agent WebSocket connection without a tailnet ID',
@ -29,6 +29,7 @@ export function initWebsocket(authKey: string) {
return; return;
} }
agents.set(tailnetID, ws);
const pinger = setInterval(() => { const pinger = setInterval(() => {
if (ws.readyState !== WebSocket.OPEN) { if (ws.readyState !== WebSocket.OPEN) {
clearInterval(pinger); clearInterval(pinger);
@ -40,6 +41,7 @@ export function initWebsocket(authKey: string) {
ws.on('close', () => { ws.on('close', () => {
clearInterval(pinger); clearInterval(pinger);
agents.delete(tailnetID);
}); });
ws.on('error', (error) => { ws.on('error', (error) => {
@ -54,6 +56,7 @@ export function initWebsocket(authKey: string) {
return server; return server;
} }
const agents = new Map<string, WebSocket>();
export function hp_getAgents() { export function hp_getAgents() {
return server.clients; return agents;
} }