diff --git a/server/lib/ip.ts b/server/lib/ip.ts index c3567b3ef..f2a646114 100644 --- a/server/lib/ip.ts +++ b/server/lib/ip.ts @@ -492,6 +492,9 @@ export async function getNextAvailableOrgSubnet(): Promise<{ } } +const remoteSubnetCidrSchema = z.union([z.cidrv4(), z.cidrv6()]); +const remoteSubnetIpSchema = z.union([z.ipv4(), z.ipv6()]); + export function generateRemoteSubnets( allSiteResources: SiteResource[] ): string[] { @@ -503,14 +506,16 @@ export function generateRemoteSubnets( // we purposely filter out the gateway resource here because we add it manually on the client if (sr.mode === "cidr") { // check if its a valid CIDR using zod - const cidrSchema = z.union([z.cidrv4(), z.cidrv6()]); - const parseResult = cidrSchema.safeParse(sr.destination); + const parseResult = remoteSubnetCidrSchema.safeParse( + sr.destination + ); return parseResult.success; } if (sr.mode === "host" || sr.mode === "ssh") { // check if its a valid IP using zod - const ipSchema = z.union([z.ipv4(), z.ipv6()]); - const parseResult = ipSchema.safeParse(sr.destination); + const parseResult = remoteSubnetIpSchema.safeParse( + sr.destination + ); return parseResult.success; } return false; diff --git a/server/lib/lock.ts b/server/lib/lock.ts index 2f6fad673..a565aaffa 100644 --- a/server/lib/lock.ts +++ b/server/lib/lock.ts @@ -110,6 +110,21 @@ export class LockManager { }; } + /** + * Check whether any of the given locks currently exist + * @param lockKeys - Unique identifiers for the locks + * @returns Promise - true if at least one lock is held + */ + async anyLockExists(lockKeys: string[]): Promise { + for (const lockKey of lockKeys) { + this.clearExpiredLocalLock(lockKey); + if (localLocks.has(lockKey)) { + return true; + } + } + return false; + } + /** * Extend the TTL of an existing lock, provided the token matches the * acquisition currently holding it. diff --git a/server/lib/rebuildClientAssociations.ts b/server/lib/rebuildClientAssociations.ts index 323a8fe65..9cbf192ae 100644 --- a/server/lib/rebuildClientAssociations.ts +++ b/server/lib/rebuildClientAssociations.ts @@ -92,6 +92,24 @@ export async function hasActiveSiteResourceRebuild( return rebuildQueue.isQueued({ type: "site-resource", id: siteResourceId }); } +/** + * Returns true if a rebuild for any of the given site resources is currently + * active or pending in the rebuild queue. Checks all of them in a couple of + * round trips instead of one per resource. + */ +export async function hasAnyActiveSiteResourceRebuild( + siteResourceIds: number[] +): Promise { + if (siteResourceIds.length === 0) return false; + const lockKeys = siteResourceIds.map( + (id) => `rebuild-client-associations:site-resource:${id}` + ); + if (await lockManager.anyLockExists(lockKeys)) return true; + return rebuildQueue.anyQueued( + siteResourceIds.map((id) => ({ type: "site-resource", id })) + ); +} + /** * Resolves once there is no active or queued rebuild for the given site resource. * Logs a warning and resolves early if the timeout is reached. @@ -130,14 +148,10 @@ export async function waitForSiteRebuildIdle( eq(siteNetworks.networkId, siteResources.networkId) ) .where(eq(siteNetworks.siteId, siteId)); - let allIdle = true; - for (const { siteResourceId } of resourceRows) { - if (await hasActiveSiteResourceRebuild(siteResourceId)) { - allIdle = false; - break; - } - } - if (allIdle) return; + const busy = await hasAnyActiveSiteResourceRebuild( + resourceRows.map((row) => row.siteResourceId) + ); + if (!busy) return; await new Promise((r) => setTimeout(r, REBUILD_IDLE_POLL_INTERVAL_MS) ); @@ -164,14 +178,10 @@ export async function waitForClientRebuildIdle( }) .from(clientSiteResourcesAssociationsCache) .where(eq(clientSiteResourcesAssociationsCache.clientId, clientId)); - let allIdle = true; - for (const { siteResourceId } of resourceRows) { - if (await hasActiveSiteResourceRebuild(siteResourceId)) { - allIdle = false; - break; - } - } - if (allIdle) return; + const busy = await hasAnyActiveSiteResourceRebuild( + resourceRows.map((row) => row.siteResourceId) + ); + if (!busy) return; await new Promise((r) => setTimeout(r, REBUILD_IDLE_POLL_INTERVAL_MS) ); diff --git a/server/lib/rebuildQueue.ts b/server/lib/rebuildQueue.ts index 7491f5f52..f1e2f9c15 100644 --- a/server/lib/rebuildQueue.ts +++ b/server/lib/rebuildQueue.ts @@ -20,6 +20,7 @@ export interface RebuildQueueManager { enqueue(job: RebuildJob): Promise; startProcessing(handlers: RebuildJobHandlers): void; isQueued(job: RebuildJob): Promise; + anyQueued(jobs: RebuildJob[]): Promise; } // In-process FIFO used when there is no Redis to back a distributed queue @@ -48,6 +49,10 @@ class InMemoryRebuildQueue implements RebuildQueueManager { return this.queuedSet.has(dedupeKey(job)); } + async anyQueued(jobs: RebuildJob[]): Promise { + return jobs.some((job) => this.queuedSet.has(dedupeKey(job))); + } + async enqueue(job: RebuildJob): Promise { const key = dedupeKey(job); if (this.queuedSet.has(key)) { diff --git a/server/private/lib/lock.ts b/server/private/lib/lock.ts index 94b27a337..17f4cd817 100644 --- a/server/private/lib/lock.ts +++ b/server/private/lib/lock.ts @@ -17,6 +17,9 @@ import { v4 as uuidv4 } from "uuid"; const instanceId = uuidv4(); +// Max keys passed to a single EXISTS call in anyLockExists +const LOCK_EXISTS_CHUNK_SIZE = 1000; + type LocalLockRecord = { owner: string; expiresAt: number; @@ -254,6 +257,39 @@ export class LockManager { } } + /** + * Check whether any of the given locks currently exist, using one EXISTS + * call per chunk of keys instead of a round trip per key + * @param lockKeys - Unique identifiers for the locks + * @returns Promise - true if at least one lock is held + */ + async anyLockExists(lockKeys: string[]): Promise { + if (!redis || !redis.status || redis.status !== "ready") { + for (const lockKey of lockKeys) { + this.clearExpiredLocalLock(lockKey); + if (localLocks.has(lockKey)) { + return true; + } + } + return false; + } + + try { + for (let i = 0; i < lockKeys.length; i += LOCK_EXISTS_CHUNK_SIZE) { + const redisKeys = lockKeys + .slice(i, i + LOCK_EXISTS_CHUNK_SIZE) + .map((lockKey) => `lock:${lockKey}`); + if ((await redis.exists(...redisKeys)) > 0) { + return true; + } + } + return false; + } catch (error) { + logger.error("Failed to check lock existence:", error); + return false; + } + } + /** * Extend the TTL of an existing lock, provided the token matches the * acquisition currently holding it. diff --git a/server/private/lib/rebuildQueue.ts b/server/private/lib/rebuildQueue.ts index 87da95f60..4fe3c77af 100644 --- a/server/private/lib/rebuildQueue.ts +++ b/server/private/lib/rebuildQueue.ts @@ -66,6 +66,24 @@ class RedisRebuildQueue { } } + // Pipelines the membership checks so any number of jobs costs one round trip + async anyQueued(jobs: RebuildJob[]): Promise { + if (jobs.length === 0) return false; + if (!redis || redis.status !== "ready") return false; + try { + const pipeline = redis.pipeline(); + for (const job of jobs) { + pipeline.sismember(QUEUED_SET_KEY, `${job.type}:${job.id}`); + } + const results = await pipeline.exec(); + return (results ?? []).some( + ([err, member]) => !err && member === 1 + ); + } catch { + return false; + } + } + async enqueue(job: RebuildJob): Promise { if (!redis || redis.status !== "ready") { logger.warn( diff --git a/server/private/routers/ws/ws.ts b/server/private/routers/ws/ws.ts index 7327c7067..db27d591f 100644 --- a/server/private/routers/ws/ws.ts +++ b/server/private/routers/ws/ws.ts @@ -792,13 +792,21 @@ const sendToClientsBatch = async ( ) ); - for (const entry of entries) { + // Issue all increments at once so Redis receives them back to back on one + // connection (in entry order) instead of waiting a round trip per entry + const configVersions = await Promise.all( + entries.map((entry) => + entry.options?.incrementConfigVersion + ? incrementClientConfigVersion(entry.clientId) + : stableConfigVersionByClient.get(entry.clientId) + ) + ); + + for (const [index, entry] of entries.entries()) { const options = entry.options || {}; const { clientId, message } = entry; - const configVersion = options.incrementConfigVersion - ? await incrementClientConfigVersion(clientId) - : stableConfigVersionByClient.get(clientId); + const configVersion = configVersions[index]; logger.debug( `sendToClientsBatch: Message type ${message.type} queued for clientId ${clientId} (new configVersion: ${configVersion})` diff --git a/server/routers/newt/peers.ts b/server/routers/newt/peers.ts index cb69d7480..6831229ff 100644 --- a/server/routers/newt/peers.ts +++ b/server/routers/newt/peers.ts @@ -98,6 +98,37 @@ export async function deletePeer( return site; } +export async function addPeersBatch( + peers: { + siteId: number; + newtId: string; + peer: { + publicKey: string; + allowedIps: string[]; + endpoint: string; + }; + }[] +) { + if (peers.length === 0) { + return; + } + + await sendToClientsBatch( + peers.map((peer) => ({ + clientId: peer.newtId, + message: { + type: "newt/wg/peer/add", + data: peer.peer + }, + options: { incrementConfigVersion: true } + })) + ).catch((error) => { + logger.warn(`Error sending batched newt peer additions:`, error); + }); + + logger.info(`Added ${peers.length} peer(s) to newts (batch)`); +} + export async function deletePeersBatch( peers: { siteId: number; diff --git a/server/routers/olm/buildConfiguration.ts b/server/routers/olm/buildConfiguration.ts index 79f2fc248..043d997e8 100644 --- a/server/routers/olm/buildConfiguration.ts +++ b/server/routers/olm/buildConfiguration.ts @@ -3,8 +3,8 @@ import { clientSiteResourcesAssociationsCache, clientSitesAssociationsCache, db, - exitNodes, networks, + newts, SiteResource, siteNetworks, siteResources, @@ -16,9 +16,8 @@ import { generateRemoteSubnets } from "@server/lib/ip"; import logger from "@server/logger"; -import { and, eq, inArray } from "drizzle-orm"; -import { addPeer, deletePeer } from "../newt/peers"; -import { SiR } from "react-icons/si"; +import { and, eq } from "drizzle-orm"; +import { addPeersBatch, deletePeersBatch } from "../newt/peers"; export async function buildSiteConfigurationForOlmClient( client: Client, @@ -41,7 +40,7 @@ export async function buildSiteConfigurationForOlmClient( let exitNodeAliases: string[] = []; // Get all sites data - const sitesData = await db + const sitesDataQuery = db .select() .from(sites) .innerJoin( @@ -50,10 +49,26 @@ export async function buildSiteConfigurationForOlmClient( ) .where(eq(clientSitesAssociationsCache.clientId, client.clientId)); + // Batch-fetch the newt for every site so peer updates can be sent as one + // batch instead of looking up the newt per site. Not needed in JIT mode + // because no peers are pushed to the sites there. + const newtRowsQuery = jitMode + ? Promise.resolve([]) + : db + .select({ siteId: newts.siteId, newtId: newts.newtId }) + .from(newts) + .innerJoin( + clientSitesAssociationsCache, + eq(newts.siteId, clientSitesAssociationsCache.siteId) + ) + .where( + eq(clientSitesAssociationsCache.clientId, client.clientId) + ); + // Batch-fetch every site resource this client has access to across ALL sites // in a single query, then group by siteId in memory. This avoids issuing one // query per site (which would be N round-trips for N sites). - const allClientSiteResources = await db + const allClientSiteResourcesQuery = db .select({ siteResource: siteResources, siteId: siteNetworks.siteId @@ -78,6 +93,19 @@ export async function buildSiteConfigurationForOlmClient( ) ); + const [sitesData, newtRows, allClientSiteResources] = await Promise.all([ + sitesDataQuery, + newtRowsQuery, + allClientSiteResourcesQuery + ]); + + const newtIdBySiteId = new Map(); + for (const row of newtRows) { + if (row.siteId != null && !newtIdBySiteId.has(row.siteId)) { + newtIdBySiteId.set(row.siteId, row.newtId); + } + } + const siteResourcesBySiteId = new Map(); const siteResourcesForExitNode = []; for (const row of allClientSiteResources) { @@ -107,28 +135,8 @@ export async function buildSiteConfigurationForOlmClient( }; } - // Batch-fetch exit nodes for all sites in one query (only needed in relay mode). - const exitNodesById = new Map(); - if (!jitMode && relay) { - const exitNodeIds = Array.from( - new Set( - sitesData - .map(({ sites: s }) => s.exitNodeId) - .filter((id): id is number => id != null) - ) - ); - if (exitNodeIds.length > 0) { - const nodes = await db - .select() - .from(exitNodes) - .where(inArray(exitNodes.exitNodeId, exitNodeIds)); - for (const n of nodes) { - exitNodesById.set(n.exitNodeId, n); - } - } - } - - const peerOps: Promise[] = []; + const peersToDelete: Parameters[0] = []; + const peersToAdd: Parameters[0] = []; // Process each site for (const { @@ -175,12 +183,23 @@ export async function buildSiteConfigurationForOlmClient( // continue; // } + const newtId = newtIdBySiteId.get(site.siteId); + if (!newtId) { + logger.warn( + `Newt not found for site ${site.siteId}, skipping peer updates` + ); + } + // If public key changed, delete old peer from this site - if (client.pubKey && client.pubKey != publicKey) { - logger.info( + if (newtId && client.pubKey && client.pubKey != publicKey) { + logger.debug( `Public key mismatch. Deleting old peer from site ${site.siteId}...` ); - peerOps.push(deletePeer(site.siteId, client.pubKey!)); + peersToDelete.push({ + siteId: site.siteId, + publicKey: client.pubKey, + newtId + }); } if (!site.exitNodeSubnet) { @@ -191,37 +210,30 @@ export async function buildSiteConfigurationForOlmClient( // Add the peer to the exit node for this site. The endpoint comes from // the already-joined association row above, so no extra query needed. if (association.endpoint && publicKey) { - logger.info( - `Adding peer ${publicKey} to site ${site.siteId} with endpoint ${association.endpoint}` - ); - peerOps.push( - addPeer(site.siteId, { - publicKey: publicKey, - allowedIps: [`${client.subnet.split("/")[0]}/32`], // we want to only allow from that client - endpoint: relay ? "" : association.endpoint - }) - ); + if (newtId) { + logger.debug( + `Adding peer ${publicKey} to site ${site.siteId} with endpoint ${association.endpoint}` + ); + peersToAdd.push({ + siteId: site.siteId, + newtId, + peer: { + publicKey: publicKey, + allowedIps: [`${client.subnet.split("/")[0]}/32`], // we want to only allow from that client + endpoint: relay ? "" : association.endpoint + } + }); + } } else { logger.warn( `Client ${client.clientId} has no endpoint, skipping peer addition` ); } - let relayEndpoint: string | undefined = undefined; - if (relay) { - const exitNode = exitNodesById.get(site.exitNodeId); - if (!exitNode) { - logger.warn(`Exit node not found for site ${site.siteId}`); - continue; - } - relayEndpoint = `${exitNode.endpoint}:${exitNode.relayPort}`; - } - // Add site configuration to the array siteConfigurations.push({ siteId: site.siteId, name: site.name, - // relayEndpoint: relayEndpoint, // this can be undefined now if not relayed // lets not do this for now because it would conflict with the hole punch testing endpoint: site.endpoint, localEndpoints: site.localEndpoints ? JSON.parse(site.localEndpoints) @@ -234,10 +246,14 @@ export async function buildSiteConfigurationForOlmClient( }); } - // Run all peer add/delete operations concurrently rather than serially per - // site, so total time is bounded by the slowest call instead of the sum. - if (peerOps.length > 0) { - Promise.allSettled(peerOps).catch((err) => { + // Send the peer removals and additions to the newts as one batch each + // without blocking the response on them. Removals go first so the old key + // is gone before the new one is added. + if (peersToDelete.length > 0 || peersToAdd.length > 0) { + (async () => { + await deletePeersBatch(peersToDelete); + await addPeersBatch(peersToAdd); + })().catch((err) => { logger.error("Error processing peer operations: ", err); }); }