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 = 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;

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;
-        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<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;
-        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<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)) {