Add batch messaging functions to rebuild function

server/routers/olm/peers.ts

@@ -1,9 +1,9 @@
-import { sendToClient } from "#dynamic/routers/ws";
+import { sendToClient, sendToClientsBatch } from "#dynamic/routers/ws";
 import { clientSitesAssociationsCache, db, olms } from "@server/db";
 import { canCompress } from "@server/lib/clientVersionChecks";
 import config from "@server/lib/config";
 import logger from "@server/logger";
-import { and, eq } from "drizzle-orm";
+import { and, eq, inArray } from "drizzle-orm";
 import { Alias } from "yaml";
 
 export async function addPeer(
@@ -205,3 +205,150 @@ export async function initPeerAddHandshake(
         `Initiated peer add handshake for site ${peer.siteId} to olm ${olmId}`
     );
 }
+
+export async function deletePeersBatch(
+    peers: {
+        clientId: number;
+        siteId: number;
+        publicKey: string;
+        olmId?: string;
+        version?: string | null;
+    }[]
+) {
+    if (peers.length === 0) {
+        return;
+    }
+
+    const unresolvedClientIds = peers
+        .filter((peer) => !peer.olmId)
+        .map((peer) => peer.clientId);
+
+    const olmByClientId = new Map<
+        number,
+        { olmId: string; version: string | null }
+    >();
+
+    if (unresolvedClientIds.length > 0) {
+        const olmRows = await db
+            .select({
+                clientId: olms.clientId,
+                olmId: olms.olmId,
+                version: olms.version
+            })
+            .from(olms)
+            .where(inArray(olms.clientId, unresolvedClientIds));
+
+        for (const row of olmRows) {
+            if (row.clientId !== null) {
+                olmByClientId.set(row.clientId, {
+                    olmId: row.olmId,
+                    version: row.version
+                });
+            }
+        }
+    }
+
+    const batchPayloads = peers
+        .map((peer) => {
+            const resolved = peer.olmId
+                ? { olmId: peer.olmId, version: peer.version ?? null }
+                : olmByClientId.get(peer.clientId);
+            if (!resolved) {
+                return null;
+            }
+
+            return {
+                clientId: resolved.olmId,
+                message: {
+                    type: "olm/wg/peer/remove",
+                    data: {
+                        publicKey: peer.publicKey,
+                        siteId: peer.siteId
+                    }
+                },
+                options: {
+                    incrementConfigVersion: true,
+                    compress: canCompress(
+                        peer.version ?? resolved.version,
+                        "olm"
+                    )
+                }
+            };
+        })
+        .filter((payload) => payload !== null);
+
+    if (batchPayloads.length === 0) {
+        return;
+    }
+
+    await sendToClientsBatch(batchPayloads).catch((error) => {
+        logger.warn(`Error sending batched olm peer removals:`, error);
+    });
+
+    logger.info(`Deleted ${batchPayloads.length} peer(s) from olms (batch)`);
+}
+
+export async function initPeerAddHandshakeBatch(
+    handshakes: {
+        clientId: number;
+        peer: {
+            siteId: number;
+            exitNode: {
+                publicKey: string;
+                endpoint: string;
+            };
+        };
+        olmId: string;
+        chainId?: string;
+    }[]
+) {
+    if (handshakes.length === 0) {
+        return;
+    }
+
+    await sendToClientsBatch(
+        handshakes.map((item) => ({
+            clientId: item.olmId,
+            message: {
+                type: "olm/wg/peer/holepunch/site/add",
+                data: {
+                    siteId: item.peer.siteId,
+                    exitNode: {
+                        publicKey: item.peer.exitNode.publicKey,
+                        relayPort:
+                            config.getRawConfig().gerbil.clients_start_port,
+                        endpoint: item.peer.exitNode.endpoint
+                    },
+                    chainId: item.chainId
+                }
+            },
+            options: { incrementConfigVersion: true }
+        }))
+    ).catch((error) => {
+        logger.warn(`Error sending batched olm handshakes:`, error);
+    });
+
+    await Promise.all(
+        handshakes.map((item) =>
+            db
+                .update(clientSitesAssociationsCache)
+                .set({ isJitMode: false })
+                .where(
+                    and(
+                        eq(
+                            clientSitesAssociationsCache.clientId,
+                            item.clientId
+                        ),
+                        eq(
+                            clientSitesAssociationsCache.siteId,
+                            item.peer.siteId
+                        )
+                    )
+                )
+        )
+    );
+
+    logger.info(
+        `Initiated ${handshakes.length} peer add handshake(s) to olms (batch)`
+    );
+}