Improve latency on register messages for olm

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 = remoteSubnetCidrSchema.safeParse(
-                const parseResult = cidrSchema.safeParse(sr.destination);
+                    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 = remoteSubnetIpSchema.safeParse(
-                const parseResult = ipSchema.safeParse(sr.destination);
+                    sr.destination
+                );
                 return parseResult.success;
             }
             return false;

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<boolean> - true if at least one lock is held
+     */
+    async anyLockExists(lockKeys: string[]): Promise<boolean> {
+        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.

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<boolean> {
+    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;
+        const busy = await hasAnyActiveSiteResourceRebuild(
-        for (const { siteResourceId } of resourceRows) {
+            resourceRows.map((row) => row.siteResourceId)
-            if (await hasActiveSiteResourceRebuild(siteResourceId)) {
+        );
-                allIdle = false;
+        if (!busy) return;
-                break;
-            }
-        }
-        if (allIdle) return;
         await new Promise<void>((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;
+        const busy = await hasAnyActiveSiteResourceRebuild(
-        for (const { siteResourceId } of resourceRows) {
+            resourceRows.map((row) => row.siteResourceId)
-            if (await hasActiveSiteResourceRebuild(siteResourceId)) {
+        );
-                allIdle = false;
+        if (!busy) return;
-                break;
-            }
-        }
-        if (allIdle) return;
         await new Promise<void>((r) =>
             setTimeout(r, REBUILD_IDLE_POLL_INTERVAL_MS)
         );

server/lib/rebuildQueue.ts

@@ -20,6 +20,7 @@ export interface RebuildQueueManager {
     enqueue(job: RebuildJob): Promise<void>;
     startProcessing(handlers: RebuildJobHandlers): void;
     isQueued(job: RebuildJob): Promise<boolean>;
+    anyQueued(jobs: RebuildJob[]): Promise<boolean>;
 }
 
 // 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<boolean> {
+        return jobs.some((job) => this.queuedSet.has(dedupeKey(job)));
+    }
+
     async enqueue(job: RebuildJob): Promise<void> {
         const key = dedupeKey(job);
         if (this.queuedSet.has(key)) {

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<boolean> - true if at least one lock is held
+     */
+    async anyLockExists(lockKeys: string[]): Promise<boolean> {
+        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.

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<boolean> {
+        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<void> {
         if (!redis || redis.status !== "ready") {
             logger.warn(

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
+        const configVersion = configVersions[index];
-            ? await incrementClientConfigVersion(clientId)
-            : stableConfigVersionByClient.get(clientId);
 
         logger.debug(
             `sendToClientsBatch: Message type ${message.type} queued for clientId ${clientId} (new configVersion: ${configVersion})`

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;

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 { and, eq } from "drizzle-orm";
-import { addPeer, deletePeer } from "../newt/peers";
+import { addPeersBatch, deletePeersBatch } from "../newt/peers";
-import { SiR } from "react-icons/si";
 
 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<number, string>();
+    for (const row of newtRows) {
+        if (row.siteId != null && !newtIdBySiteId.has(row.siteId)) {
+            newtIdBySiteId.set(row.siteId, row.newtId);
+        }
+    }
+
     const siteResourcesBySiteId = new Map<number, SiteResource[]>();
     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 peersToDelete: Parameters<typeof deletePeersBatch>[0] = [];
-    const exitNodesById = new Map<number, typeof exitNodes.$inferSelect>();
+    const peersToAdd: Parameters<typeof addPeersBatch>[0] = [];
-    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<unknown>[] = [];
 
     // 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) {
+        if (newtId && client.pubKey && client.pubKey != publicKey) {
-            logger.info(
+            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(
+            if (newtId) {
-                `Adding peer ${publicKey} to site ${site.siteId} with endpoint ${association.endpoint}`
+                logger.debug(
-            );
+                    `Adding peer ${publicKey} to site ${site.siteId} with endpoint ${association.endpoint}`
-            peerOps.push(
+                );
-                addPeer(site.siteId, {
+                peersToAdd.push({
-                    publicKey: publicKey,
+                    siteId: site.siteId,
-                    allowedIps: [`${client.subnet.split("/")[0]}/32`], // we want to only allow from that client
+                    newtId,
-                    endpoint: relay ? "" : association.endpoint
+                    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
+    // Send the peer removals and additions to the newts as one batch each
-    // site, so total time is bounded by the slowest call instead of the sum.
+    // without blocking the response on them. Removals go first so the old key
-    if (peerOps.length > 0) {
+    // is gone before the new one is added.
-        Promise.allSettled(peerOps).catch((err) => {
+    if (peersToDelete.length > 0 || peersToAdd.length > 0) {
+        (async () => {
+            await deletePeersBatch(peersToDelete);
+            await addPeersBatch(peersToAdd);
+        })().catch((err) => {
             logger.error("Error processing peer operations: ", err);
         });
     }