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
+2 -1
View File
@@ -78,7 +78,8 @@ export const orgs = pgTable("orgs", {
"settingsEnableGlobalNewtAutoUpdate"
)
.notNull()
.default(false)
.default(false),
settingsJitModeLimit: integer("settingsJitModeLimit").notNull().default(250)
});
export const orgDomains = pgTable("orgDomains", {
+2 -1
View File
@@ -78,7 +78,8 @@ export const orgs = sqliteTable("orgs", {
{ mode: "boolean" }
)
.notNull()
.default(false)
.default(false),
settingsJitModeLimit: integer("settingsJitModeLimit").notNull().default(250)
});
export const userDomains = sqliteTable("userDomains", {
+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))
+20 -3
View File
@@ -8,7 +8,8 @@ import {
sites,
clientSitesAssociationsCache,
exitNodes,
ExitNode
ExitNode,
orgs
} from "@server/db";
import { db } from "@server/db";
import { eq, and, inArray } from "drizzle-orm";
@@ -112,7 +113,12 @@ export async function updateHolePunch(
destinations: destinations
});
} catch (error) {
if (!(error instanceof Error && error.message === "Exit node not allowed")) {
if (
!(
error instanceof Error &&
error.message === "Exit node not allowed"
)
) {
logger.error(error);
}
return next(
@@ -460,7 +466,18 @@ async function handleClientEndpointChange(
return;
}
if (sitesWithNewtsToUpdate.length > 250) {
const [org] = await db
.select()
.from(orgs)
.where(eq(orgs.orgId, client.orgId))
.limit(1);
if (!org) {
logger.warn(`Client ${clientId} org not found`);
return;
}
if (sitesWithNewtsToUpdate.length > org.settingsJitModeLimit) {
logger.warn(
`Client ${clientId} has ${sitesWithNewtsToUpdate.length} connected sites so the client will be in jit mode anyway, skipping endpoint updates`
);
+19 -2
View File
@@ -8,7 +8,8 @@ import {
ExitNode,
exitNodes,
sites,
clientSitesAssociationsCache
clientSitesAssociationsCache,
orgs
} from "@server/db";
import { olms } from "@server/db";
import HttpCode from "@server/types/HttpCode";
@@ -225,7 +226,23 @@ export async function getOlmToken(
)
.where(eq(clientSitesAssociationsCache.clientId, clientIdToUse!));
if (clientSites.length > 250 && build == "saas") {
const [org] = await db
.select()
.from(orgs)
.where(eq(orgs.orgId, orgIdToUse))
.limit(1);
if (!org) {
logger.warn(`Client ${clientIdToUse} org not found`);
return next(
createHttpError(
HttpCode.INTERNAL_SERVER_ERROR,
"Client's org not found"
)
);
}
if (clientSites.length > org.settingsJitModeLimit && build == "saas") {
// set all of the cache rows isJitMode to true
await db
.update(clientSitesAssociationsCache)
@@ -277,7 +277,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
);
let jitMode = false;
if (sitesCount > 250 && build == "saas") {
if (sitesCount > org.settingsJitModeLimit && build == "saas") {
// THIS IS THE MAX ON THE BUSINESS TIER
// we have too many sites
// If we have too many sites we need to drop into fully JIT mode by not sending any of the sites
+14 -2
View File
@@ -5,7 +5,8 @@ import {
exitNodes,
Olm,
sites,
clientSitesAssociationsCache
clientSitesAssociationsCache,
orgs
} from "@server/db";
import { buildSiteConfigurationForOlmClient } from "./buildConfiguration";
import { sendToClient } from "#dynamic/routers/ws";
@@ -36,8 +37,19 @@ export async function sendOlmSyncMessage(olm: Olm, client: Client) {
{ orgId: client.orgId }
);
const [org] = await db
.select()
.from(orgs)
.where(eq(orgs.orgId, client.orgId))
.limit(1);
if (!org) {
logger.warn(`Client ${client.clientId} org not found`);
return;
}
let jitMode = false;
if (sitesCount > 250 && build == "saas") {
if (sitesCount > org.settingsJitModeLimit && build == "saas") {
// THIS IS THE MAX ON THE BUSINESS TIER
// we have too many sites
// If we have too many sites we need to drop into fully JIT mode by not sending any of the sites