Pull jit limit from org

This commit is contained in:
Owen
2026-09-23 10:40:21 -04:00
parent 537649af18
commit d406cc61b5
7 changed files with 113 additions and 19 deletions
+55 -9
View File
@@ -17,7 +17,8 @@ import {
sites,
Transaction,
userOrgRoles,
userSiteResources
userSiteResources,
orgs
} from "@server/db";
import { and, count, eq, inArray, isNotNull, ne } from "drizzle-orm";
@@ -234,7 +235,8 @@ export async function getClientSiteResourceAccess(
.select({
clientId: clients.clientId,
pubKey: clients.pubKey,
subnet: clients.subnet
subnet: clients.subnet,
orgId: clients.orgId
})
.from(clients)
.where(
@@ -260,7 +262,8 @@ export async function getClientSiteResourceAccess(
.select({
clientId: clients.clientId,
pubKey: clients.pubKey,
subnet: clients.subnet
subnet: clients.subnet,
orgId: clients.orgId
})
.from(clients)
.where(
@@ -463,7 +466,8 @@ async function rebuildClientAssociationsFromSiteResourceImpl(
.select({
clientId: clients.clientId,
pubKey: clients.pubKey,
subnet: clients.subnet
subnet: clients.subnet,
orgId: clients.orgId
})
.from(clients)
.where(
@@ -568,7 +572,8 @@ async function rebuildClientAssociationsFromSiteResourceImpl(
.select({
clientId: clients.clientId,
pubKey: clients.pubKey,
subnet: clients.subnet
subnet: clients.subnet,
orgId: clients.orgId
})
.from(clients)
.where(
@@ -720,11 +725,13 @@ async function handleMessagesForSiteClients(
clientId: number;
pubKey: string | null;
subnet: string | null;
orgId: string;
}[],
existingClients: {
clientId: number;
pubKey: string | null;
subnet: string | null;
orgId: string;
}[],
clientSitesToAdd: number[],
clientSitesToRemove: number[],
@@ -805,6 +812,7 @@ async function handleMessagesForSiteClients(
clientId: number;
pubKey: string | null;
subnet: string | null;
orgId: string;
}
>();
@@ -860,6 +868,22 @@ async function handleMessagesForSiteClients(
.map((r) => [r.clientId as number, r.olmId])
);
// Batch-fetch the orgs for all clients we need to process so we don't
// issue a redundant query per client in the loop below
const orgIdsToProcess = Array.from(
new Set(
Array.from(clientsToProcess.values()).map((client) => client.orgId)
)
);
const orgRows =
orgIdsToProcess.length > 0
? await trx
.select()
.from(orgs)
.where(inArray(orgs.orgId, orgIdsToProcess))
: [];
const orgByOrgId = new Map(orgRows.map((org) => [org.orgId, org]));
for (const client of clientsToProcess.values()) {
// UPDATE THE NEWT
if (!client.subnet || !client.pubKey) {
@@ -899,7 +923,14 @@ async function handleMessagesForSiteClients(
}
if (isAdd) {
if (clientSiteCounts[client.clientId] > 250) {
const org = orgByOrgId.get(client.orgId);
if (!org) {
logger.warn(`Client ${client.clientId} org not found`);
continue;
}
if (clientSiteCounts[client.clientId] > org.settingsJitModeLimit) {
// skip adding the peer if we have more than 250 sites because we are in jit mode anyway
logger.info(
`rebuildClientAssociations: Client ${client.clientId} has ${clientSiteCounts[client.clientId]} sites so skipping adding peer to newt and olm because it is likely in jit mode`
@@ -1570,7 +1601,8 @@ export async function handleMessagingForUpdatedSiteResource(
.select({
clientId: clientSiteResourcesAssociationsCache.clientId,
pubKey: clients.pubKey,
subnet: clients.subnet
subnet: clients.subnet,
orgId: clients.orgId
})
.from(clientSiteResourcesAssociationsCache)
.innerJoin(
@@ -2391,6 +2423,19 @@ async function handleMessagesForClientSites(
.where(eq(clientSitesAssociationsCache.clientId, client.clientId))
.then((rows) => Number(rows[0].count));
// client.orgId is constant for this call, so fetch the org once
// instead of re-querying it for every site in the loop below
const [org] = await trx
.select()
.from(orgs)
.where(eq(orgs.orgId, client.orgId))
.limit(1);
if (!org) {
logger.warn(`Client ${client.clientId} org not found`);
return;
}
for (const siteData of sitesData) {
const site = siteData.sites;
const exitNode = siteData.exitNodes;
@@ -2451,7 +2496,7 @@ async function handleMessagesForClientSites(
continue;
}
if (totalSitesOnClient > 250) {
if (totalSitesOnClient > org.settingsJitModeLimit) {
// skip adding the site if we have more than 250 because we are in jit mode anyway
logger.info(
`rebuildClientAssociations: Client ${client.clientId} has ${totalSitesOnClient} sites so skipping adding peer to newt and olm because it is likely in jit mode`
@@ -3061,7 +3106,8 @@ export async function cleanupSiteAssociations(
.select({
clientId: clients.clientId,
pubKey: clients.pubKey,
subnet: clients.subnet
subnet: clients.subnet,
orgId: clients.orgId
})
.from(clients)
.where(inArray(clients.clientId, cachedClientIds))