pangolin_mirror/pangolin
1
0
Fork 0
mirror of https://github.com/fosrl/pangolin.git synced 2026-10-23 21:19:07 +02:00
Code Issues Packages Projects Releases Wiki Activity

Improve latency on register messages for olm

Browse Source
This commit is contained in:
Owen
2026-10-22 16:04:38 -04:00
parent 290f1f2ff0
commit 5ff8ecdf78
9 changed files with 224 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/lib/ip.ts
+9 -4
View File
@@ -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
+15
View File
@@ -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
+26 -16
View File
@@ -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
+5
View File
@@ -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
+36
View File
@@ -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
+18
View File
@@ -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
+12 -4
View File
@@ -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})`
server/routers/newt/peers.ts
+31
View File
@@ -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
+72 -56
View File
@@ -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<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 exitNodesById = new Map<number, typeof exitNodes.$inferSelect>();
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>[] = [];
const peersToDelete: Parameters<typeof deletePeersBatch>[0] = [];
const peersToAdd: Parameters<typeof addPeersBatch>[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);
});
}