Merge branch 'dev' into feat/geoip-country-is-not-rule

This commit is contained in:
Fred KISSIE
2026-07-03 22:50:56 +02:00
183 changed files with 11173 additions and 1880 deletions
+36 -26
View File
@@ -1,28 +1,39 @@
export enum FeatureId {
export enum LimitId {
USERS = "users",
SITES = "sites",
EGRESS_DATA_MB = "egressDataMb",
DOMAINS = "domains",
REMOTE_EXIT_NODES = "remoteExitNodes",
ORGINIZATIONS = "organizations",
ORGANIZATIONS = "organizations",
PUBLIC_RESOURCES = "publicResources",
PRIVATE_RESOURCES = "privateResources",
MACHINE_CLIENTS = "machineClients",
TIER1 = "tier1"
}
export async function getFeatureDisplayName(featureId: FeatureId): Promise<string> {
export async function getFeatureDisplayName(
featureId: LimitId
): Promise<string> {
switch (featureId) {
case FeatureId.USERS:
case LimitId.USERS:
return "Users";
case FeatureId.SITES:
case LimitId.SITES:
return "Sites";
case FeatureId.EGRESS_DATA_MB:
case LimitId.EGRESS_DATA_MB:
return "Egress Data (MB)";
case FeatureId.DOMAINS:
case LimitId.DOMAINS:
return "Domains";
case FeatureId.REMOTE_EXIT_NODES:
case LimitId.REMOTE_EXIT_NODES:
return "Remote Exit Nodes";
case FeatureId.ORGINIZATIONS:
case LimitId.ORGANIZATIONS:
return "Organizations";
case FeatureId.TIER1:
case LimitId.PUBLIC_RESOURCES:
return "Public Resources";
case LimitId.PRIVATE_RESOURCES:
return "Private Resources";
case LimitId.MACHINE_CLIENTS:
return "Machine Clients";
case LimitId.TIER1:
return "Home Lab";
default:
return featureId;
@@ -30,15 +41,16 @@ export async function getFeatureDisplayName(featureId: FeatureId): Promise<strin
}
// this is from the old system
export const FeatureMeterIds: Partial<Record<FeatureId, string>> = { // right now we are not charging for any data
export const FeatureMeterIds: Partial<Record<LimitId, string>> = {
// right now we are not charging for any data
// [FeatureId.EGRESS_DATA_MB]: "mtr_61Srreh9eWrExDSCe41D3Ee2Ir7Wm5YW"
};
export const FeatureMeterIdsSandbox: Partial<Record<FeatureId, string>> = {
export const FeatureMeterIdsSandbox: Partial<Record<LimitId, string>> = {
// [FeatureId.EGRESS_DATA_MB]: "mtr_test_61Snh2a2m6qome5Kv41DCpkOb237B3dQ"
};
export function getFeatureMeterId(featureId: FeatureId): string | undefined {
export function getFeatureMeterId(featureId: LimitId): string | undefined {
if (
process.env.ENVIRONMENT == "prod" &&
process.env.SANDBOX_MODE !== "true"
@@ -49,22 +61,20 @@ export function getFeatureMeterId(featureId: FeatureId): string | undefined {
}
}
export function getFeatureIdByMetricId(
metricId: string
): FeatureId | undefined {
return (Object.entries(FeatureMeterIds) as [FeatureId, string][]).find(
export function getFeatureIdByMetricId(metricId: string): LimitId | undefined {
return (Object.entries(FeatureMeterIds) as [LimitId, string][]).find(
([_, v]) => v === metricId
)?.[0];
}
export type FeaturePriceSet = Partial<Record<FeatureId, string>>;
export type FeaturePriceSet = Partial<Record<LimitId, string>>;
export const tier1FeaturePriceSet: FeaturePriceSet = {
[FeatureId.TIER1]: "price_1SzVE3D3Ee2Ir7Wm6wT5Dl3G"
[LimitId.TIER1]: "price_1SzVE3D3Ee2Ir7Wm6wT5Dl3G"
};
export const tier1FeaturePriceSetSandbox: FeaturePriceSet = {
[FeatureId.TIER1]: "price_1SxgpPDCpkOb237Bfo4rIsoT"
[LimitId.TIER1]: "price_1SxgpPDCpkOb237Bfo4rIsoT"
};
export function getTier1FeaturePriceSet(): FeaturePriceSet {
@@ -79,11 +89,11 @@ export function getTier1FeaturePriceSet(): FeaturePriceSet {
}
export const tier2FeaturePriceSet: FeaturePriceSet = {
[FeatureId.USERS]: "price_1SzVCcD3Ee2Ir7Wmn6U3KvPN"
[LimitId.USERS]: "price_1SzVCcD3Ee2Ir7Wmn6U3KvPN"
};
export const tier2FeaturePriceSetSandbox: FeaturePriceSet = {
[FeatureId.USERS]: "price_1SxaEHDCpkOb237BD9lBkPiR"
[LimitId.USERS]: "price_1SxaEHDCpkOb237BD9lBkPiR"
};
export function getTier2FeaturePriceSet(): FeaturePriceSet {
@@ -98,11 +108,11 @@ export function getTier2FeaturePriceSet(): FeaturePriceSet {
}
export const tier3FeaturePriceSet: FeaturePriceSet = {
[FeatureId.USERS]: "price_1SzVDKD3Ee2Ir7WmPtOKNusv"
[LimitId.USERS]: "price_1SzVDKD3Ee2Ir7WmPtOKNusv"
};
export const tier3FeaturePriceSetSandbox: FeaturePriceSet = {
[FeatureId.USERS]: "price_1SxaEODCpkOb237BiXdCBSfs"
[LimitId.USERS]: "price_1SxaEODCpkOb237BiXdCBSfs"
};
export function getTier3FeaturePriceSet(): FeaturePriceSet {
@@ -116,7 +126,7 @@ export function getTier3FeaturePriceSet(): FeaturePriceSet {
}
}
export function getFeatureIdByPriceId(priceId: string): FeatureId | undefined {
export function getFeatureIdByPriceId(priceId: string): LimitId | undefined {
// Check all feature price sets
const allPriceSets = [
getTier1FeaturePriceSet(),
@@ -125,7 +135,7 @@ export function getFeatureIdByPriceId(priceId: string): FeatureId | undefined {
];
for (const priceSet of allPriceSets) {
const entry = (Object.entries(priceSet) as [FeatureId, string][]).find(
const entry = (Object.entries(priceSet) as [LimitId, string][]).find(
([_, price]) => price === priceId
);
if (entry) {
+5 -5
View File
@@ -1,19 +1,19 @@
import Stripe from "stripe";
import { FeatureId, FeaturePriceSet } from "./features";
import { LimitId, FeaturePriceSet } from "./features";
import { usageService } from "./usageService";
export async function getLineItems(
featurePriceSet: FeaturePriceSet,
orgId: string,
orgId: string
): Promise<Stripe.Checkout.SessionCreateParams.LineItem[]> {
const users = await usageService.getUsage(orgId, FeatureId.USERS);
const users = await usageService.getUsage(orgId, LimitId.USERS);
return Object.entries(featurePriceSet).map(([featureId, priceId]) => {
let quantity: number | undefined;
if (featureId === FeatureId.USERS) {
if (featureId === LimitId.USERS) {
quantity = users?.instantaneousValue || 1;
} else if (featureId === FeatureId.TIER1) {
} else if (featureId === LimitId.TIER1) {
quantity = 1;
}
+35 -23
View File
@@ -1,70 +1,82 @@
import { FeatureId } from "./features";
import { LimitId } from "./features";
export type LimitSet = Partial<{
[key in FeatureId]: {
[key in LimitId]: {
value: number | null; // null indicates no limit
description?: string;
};
}>;
export const freeLimitSet: LimitSet = {
[FeatureId.SITES]: { value: 5, description: "Basic limit" },
[FeatureId.USERS]: { value: 5, description: "Basic limit" },
[FeatureId.DOMAINS]: { value: 5, description: "Basic limit" },
[FeatureId.REMOTE_EXIT_NODES]: { value: 1, description: "Basic limit" },
[FeatureId.ORGINIZATIONS]: { value: 1, description: "Basic limit" },
[LimitId.SITES]: { value: 5, description: "Basic limit" },
[LimitId.USERS]: { value: 5, description: "Basic limit" },
[LimitId.DOMAINS]: { value: 5, description: "Basic limit" },
[LimitId.REMOTE_EXIT_NODES]: { value: 1, description: "Basic limit" },
[LimitId.ORGANIZATIONS]: { value: 1, description: "Basic limit" },
[LimitId.PUBLIC_RESOURCES]: { value: 15, description: "Basic limit" },
[LimitId.PRIVATE_RESOURCES]: { value: 15, description: "Basic limit" },
[LimitId.MACHINE_CLIENTS]: { value: 5, description: "Basic limit" }
};
export const tier1LimitSet: LimitSet = {
[FeatureId.USERS]: { value: 7, description: "Home limit" },
[FeatureId.SITES]: { value: 10, description: "Home limit" },
[FeatureId.DOMAINS]: { value: 10, description: "Home limit" },
[FeatureId.REMOTE_EXIT_NODES]: { value: 1, description: "Home limit" },
[FeatureId.ORGINIZATIONS]: { value: 1, description: "Home limit" },
[LimitId.USERS]: { value: 7, description: "Home limit" },
[LimitId.SITES]: { value: 10, description: "Home limit" },
[LimitId.DOMAINS]: { value: 10, description: "Home limit" },
[LimitId.REMOTE_EXIT_NODES]: { value: 1, description: "Home limit" },
[LimitId.ORGANIZATIONS]: { value: 1, description: "Home limit" },
[LimitId.PUBLIC_RESOURCES]: { value: 30, description: "Home limit" },
[LimitId.PRIVATE_RESOURCES]: { value: 30, description: "Home limit" },
[LimitId.MACHINE_CLIENTS]: { value: 10, description: "Home limit" }
};
export const tier2LimitSet: LimitSet = {
[FeatureId.USERS]: {
[LimitId.USERS]: {
value: 50,
description: "Team limit"
},
[FeatureId.SITES]: {
[LimitId.SITES]: {
value: 50,
description: "Team limit"
},
[FeatureId.DOMAINS]: {
[LimitId.DOMAINS]: {
value: 50,
description: "Team limit"
},
[FeatureId.REMOTE_EXIT_NODES]: {
[LimitId.REMOTE_EXIT_NODES]: {
value: 3,
description: "Team limit"
},
[FeatureId.ORGINIZATIONS]: {
[LimitId.ORGANIZATIONS]: {
value: 1,
description: "Team limit"
}
},
[LimitId.PUBLIC_RESOURCES]: { value: 150, description: "Team limit" },
[LimitId.PRIVATE_RESOURCES]: { value: 150, description: "Team limit" },
[LimitId.MACHINE_CLIENTS]: { value: 25, description: "Team limit" }
};
export const tier3LimitSet: LimitSet = {
[FeatureId.USERS]: {
[LimitId.USERS]: {
value: 250,
description: "Business limit"
},
[FeatureId.SITES]: {
[LimitId.SITES]: {
value: 250,
description: "Business limit"
},
[FeatureId.DOMAINS]: {
[LimitId.DOMAINS]: {
value: 100,
description: "Business limit"
},
[FeatureId.REMOTE_EXIT_NODES]: {
[LimitId.REMOTE_EXIT_NODES]: {
value: 20,
description: "Business limit"
},
[FeatureId.ORGINIZATIONS]: {
[LimitId.ORGANIZATIONS]: {
value: 5,
description: "Business limit"
},
[LimitId.PUBLIC_RESOURCES]: { value: 750, description: "Business limit" },
[LimitId.PRIVATE_RESOURCES]: { value: 750, description: "Business limit" },
[LimitId.MACHINE_CLIENTS]: { value: 100, description: "Business limit" }
};
+2 -2
View File
@@ -1,7 +1,7 @@
import { db, limits } from "@server/db";
import { and, eq } from "drizzle-orm";
import { LimitSet } from "./limitSet";
import { FeatureId } from "./features";
import { LimitId } from "./features";
import logger from "@server/logger";
class LimitService {
@@ -38,7 +38,7 @@ class LimitService {
async getOrgLimit(
orgId: string,
featureId: FeatureId
featureId: LimitId
): Promise<number | null> {
const limitId = `${orgId}-${featureId}`;
const [limit] = await db
+21 -16
View File
@@ -9,7 +9,7 @@ import {
Transaction,
orgs
} from "@server/db";
import { FeatureId, getFeatureMeterId } from "./features";
import { LimitId, getFeatureMeterId } from "./features";
import logger from "@server/logger";
import { build } from "@server/build";
import { regionalCache as cache } from "#dynamic/lib/cache";
@@ -37,7 +37,7 @@ export class UsageService {
public async add(
orgId: string,
featureId: FeatureId,
featureId: LimitId,
value: number,
transaction: any = null
): Promise<Usage | null> {
@@ -114,7 +114,7 @@ export class UsageService {
private async internalAddUsage(
orgId: string, // here the orgId is the billing org already resolved by getBillingOrg in updateCount
featureId: FeatureId,
featureId: LimitId,
value: number,
trx: Transaction
): Promise<Usage> {
@@ -163,7 +163,7 @@ export class UsageService {
async updateCount(
orgId: string,
featureId: FeatureId,
featureId: LimitId,
value?: number,
customerId?: string
): Promise<void> {
@@ -227,7 +227,7 @@ export class UsageService {
private async getCustomerId(
orgId: string,
featureId: FeatureId
featureId: LimitId
): Promise<string | null> {
const orgIdToUse = await this.getBillingOrg(orgId);
@@ -269,18 +269,19 @@ export class UsageService {
public async getUsage(
orgId: string,
featureId: FeatureId,
featureId: LimitId,
trx: Transaction | typeof db = db
): Promise<Usage | null> {
if (noop()) {
return null;
}
const orgIdToUse = await this.getBillingOrg(orgId, trx);
const usageId = `${orgIdToUse}-${featureId}`;
let orgIdToUse = orgId;
try {
orgIdToUse = await this.getBillingOrg(orgId, trx);
const usageId = `${orgIdToUse}-${featureId}`;
const [result] = await trx
.select()
.from(usage)
@@ -340,8 +341,12 @@ export class UsageService {
`Failed to get usage for ${orgIdToUse}/${featureId}:`,
error
);
throw error;
if (process.env.NODE_ENV !== "development") {
throw error;
}
}
return null;
}
public async getBillingOrg(
@@ -376,7 +381,7 @@ export class UsageService {
public async checkLimitSet(
orgId: string,
featureId?: FeatureId,
featureId?: LimitId,
usage?: Usage,
trx: Transaction | typeof db = db
): Promise<boolean> {
@@ -384,13 +389,13 @@ export class UsageService {
return false;
}
const orgIdToUse = await this.getBillingOrg(orgId, trx);
// This method should check the current usage against the limits set for the organization
// and kick out all of the sites on the org
let hasExceededLimits = false;
let orgIdToUse = orgId;
try {
orgIdToUse = await this.getBillingOrg(orgId, trx);
let orgLimits: Limit[] = [];
if (featureId) {
// Get all limits set for this organization
@@ -424,7 +429,7 @@ export class UsageService {
} else {
currentUsage = await this.getUsage(
orgIdToUse,
limit.featureId as FeatureId,
limit.featureId as LimitId,
trx
);
}
+8 -1
View File
@@ -8,7 +8,7 @@ import {
userSiteResources,
clientSiteResources
} from "@server/db";
import { Config, ConfigSchema } from "./types";
import { Config, ConfigSchema, isTargetsOnlyResource } from "./types";
import {
PublicResourcesResults,
updatePublicResources
@@ -34,6 +34,12 @@ import {
rebuildClientAssociationsFromSiteResource,
waitForSiteResourceRebuildIdle
} from "../rebuildClientAssociations";
import { build } from "@server/build";
import HttpCode from "@server/types/HttpCode";
import createHttpError from "http-errors";
import next from "next";
import { LimitId } from "../billing";
import { usageService } from "../billing/usageService";
type ApplyBlueprintArgs = {
orgId: string;
@@ -64,6 +70,7 @@ export async function applyBlueprint({
let publicResourcesResults: PublicResourcesResults = [];
let privateResourcesResults: ClientResourcesResults = [];
await db.transaction(async (trx) => {
await updateResourcePolicies(orgId, config, trx);
+36
View File
@@ -25,6 +25,12 @@ import { getNextAvailableAliasAddress } from "../ip";
import { createCertificate } from "#dynamic/routers/certificates/createCertificate";
import { isLicensedOrSubscribed } from "#dynamic/lib/isLicencedOrSubscribed";
import { tierMatrix } from "../billing/tierMatrix";
import { build } from "@server/build";
import HttpCode from "@server/types/HttpCode";
import createHttpError from "http-errors";
import next from "next";
import { LimitId } from "../billing";
import { usageService } from "../billing/usageService";
async function getDomainForSiteResource(
siteResourceId: number | undefined,
@@ -413,6 +419,34 @@ export async function updatePrivateResources(
oldSites: existingSiteIds
});
} else {
// create a brand new resource
if (build == "saas") {
const usage = await usageService.getUsage(
orgId,
LimitId.PRIVATE_RESOURCES
);
if (!usage) {
throw new Error(
`Usage data not found for org ${orgId} and limit ${LimitId.PRIVATE_RESOURCES}`
);
}
const rejectResource = await usageService.checkLimitSet(
orgId,
LimitId.PRIVATE_RESOURCES,
{
...usage,
instantaneousValue: (usage.instantaneousValue || 0) + 1
} // We need to add one to know if we are violating the limit
);
if (rejectResource) {
throw new Error(
"Private resource limit exceeded. Please upgrade your plan."
);
}
}
let aliasAddress: string | null = null;
let releaseAliasLock: (() => Promise<void>) | null = null;
if (
@@ -609,6 +643,8 @@ export async function updatePrivateResources(
`Created new client resource ${newResource.name} (${newResource.siteResourceId}) for org ${orgId}`
);
await usageService.add(orgId, LimitId.PRIVATE_RESOURCES, 1, trx);
results.push({
newSiteResource: newResource,
newSites: allSites,
+34
View File
@@ -49,6 +49,11 @@ import { and, asc, eq, isNotNull, ne, or } from "drizzle-orm";
import { tierMatrix } from "../billing/tierMatrix";
import { isValidCIDR, isValidIP, isValidUrlGlobPattern } from "../validators";
import { Config, isTargetsOnlyResource, TargetData } from "./types";
import HttpCode from "@server/types/HttpCode";
import createHttpError from "http-errors";
import next from "next";
import { LimitId } from "../billing";
import { usageService } from "../billing/usageService";
export type PublicResourcesResults = {
proxyResource: Resource;
@@ -1003,6 +1008,33 @@ export async function updatePublicResources(
logger.debug(`Updated resource ${existingResource.resourceId}`);
} else {
// create a brand new resource
if (build === "saas") {
const usage = await usageService.getUsage(
orgId,
LimitId.PUBLIC_RESOURCES
);
if (!usage) {
throw new Error(
`Usage data not found for org ${orgId} and limit ${LimitId.PUBLIC_RESOURCES}`
);
}
const rejectResource = await usageService.checkLimitSet(
orgId,
LimitId.PUBLIC_RESOURCES,
{
...usage,
instantaneousValue: (usage.instantaneousValue || 0) + 1
} // We need to add one to know if we are violating the limit
);
if (rejectResource) {
throw new Error(
"Public resource limit exceeded. Please upgrade your plan."
);
}
}
let domain;
if (
["http", "ssh", "rdp", "vnc"].includes(resourceData.mode || "")
@@ -1292,6 +1324,8 @@ export async function updatePublicResources(
await createTarget(newResource.resourceId, targetData);
}
await usageService.add(orgId, LimitId.PUBLIC_RESOURCES, 1, trx);
logger.debug(`Created resource ${newResource.resourceId}`);
}
+9 -7
View File
@@ -24,7 +24,7 @@ import { deletePeer } from "@server/routers/gerbil/peers";
import { OlmErrorCodes } from "@server/routers/olm/error";
import { sendTerminateClient } from "@server/routers/client/terminate";
import { usageService } from "./billing/usageService";
import { FeatureId } from "./billing";
import { LimitId } from "./billing";
export type DeleteOrgByIdResult = {
deletedNewtIds: string[];
@@ -140,7 +140,9 @@ export async function deleteOrgById(
.select({ count: count() })
.from(orgDomains)
.where(eq(orgDomains.domainId, domainId));
logger.info(`Found ${orgCount.count} orgs using domain ${domainId}`);
logger.info(
`Found ${orgCount.count} orgs using domain ${domainId}`
);
if (orgCount.count === 1) {
domainIdsToDelete.push(domainId);
}
@@ -152,7 +154,7 @@ export async function deleteOrgById(
.where(inArray(domains.domainId, domainIdsToDelete));
}
await usageService.add(orgId, FeatureId.ORGINIZATIONS, -1, trx); // here we are decreasing the org count BEFORE deleting the org because we need to still be able to get the org to get the billing org inside of here
await usageService.add(orgId, LimitId.ORGANIZATIONS, -1, trx); // here we are decreasing the org count BEFORE deleting the org because we need to still be able to get the org to get the billing org inside of here
await trx.delete(orgs).where(eq(orgs.orgId, orgId));
@@ -199,22 +201,22 @@ export async function deleteOrgById(
if (org.billingOrgId) {
usageService.updateCount(
org.billingOrgId,
FeatureId.DOMAINS,
LimitId.DOMAINS,
domainCount ?? 0
);
usageService.updateCount(
org.billingOrgId,
FeatureId.SITES,
LimitId.SITES,
siteCount ?? 0
);
usageService.updateCount(
org.billingOrgId,
FeatureId.USERS,
LimitId.USERS,
userCount ?? 0
);
usageService.updateCount(
org.billingOrgId,
FeatureId.REMOTE_EXIT_NODES,
LimitId.REMOTE_EXIT_NODES,
remoteExitNodeCount ?? 0
);
}
+5 -1
View File
@@ -19,7 +19,11 @@ export async function verifyExitNodeOrgAccess(
export async function listExitNodes(
orgId: string,
filterOnline = false,
noCloud = false
noCloud = false,
// Accepted for parity with the enterprise implementation (used there for
// site-label filtering of remote exit nodes). The OSS build has no remote
// exit nodes, so it is unused here.
siteId?: number
) {
// TODO: pick which nodes to send and ping better than just all of them that are not remote
const allExitNodes = await db
+3 -3
View File
@@ -333,7 +333,7 @@ export async function getNextAvailableClientSubnet(
if (!acquired) {
throw new Error(`Failed to acquire lock: ${lockKey}`);
}
const release = () => lockManager.releaseLock(lockKey);
const release = () => lockManager.releaseLock(lockKey, acquired);
try {
const [org] = await transaction
@@ -395,7 +395,7 @@ export async function getNextAvailableAliasAddress(
if (!acquired) {
throw new Error(`Failed to acquire lock: ${lockKey}`);
}
const release = () => lockManager.releaseLock(lockKey);
const release = () => lockManager.releaseLock(lockKey, acquired);
try {
const [org] = await trx
@@ -463,7 +463,7 @@ export async function getNextAvailableOrgSubnet(): Promise<{
if (!acquired) {
throw new Error(`Failed to acquire lock: ${lockKey}`);
}
const release = () => lockManager.releaseLock(lockKey);
const release = () => lockManager.releaseLock(lockKey, acquired);
try {
const existingAddresses = await db
+41 -31
View File
@@ -1,3 +1,5 @@
import { randomUUID } from "crypto";
const instanceId = `local-${Math.random().toString(36).slice(2)}-${Date.now()}`;
type LocalLockRecord = {
@@ -15,58 +17,60 @@ export class LockManager {
}
}
private getLocalOwnerToken(): string {
return `${instanceId}:`;
}
/**
* Acquire a distributed lock using Redis SET with NX and PX options
* Acquire a local in-process lock using an optimistic Map-based check.
* @param lockKey - Unique identifier for the lock
* @param ttlMs - Time to live in milliseconds
* @returns Promise<boolean> - true if lock acquired, false otherwise
* @returns Promise<string | null> - a token identifying this specific acquisition
* (truthy) on success, or null if the lock could not be acquired.
*/
async acquireLock(
lockKey: string,
ttlMs: number = 30000,
maxRetries: number = 3,
retryDelayMs: number = 100
): Promise<boolean> {
): Promise<string | null> {
for (let attempt = 0; attempt < maxRetries; attempt++) {
this.clearExpiredLocalLock(lockKey);
const existing = localLocks.get(lockKey);
if (!existing) {
const token = `${instanceId}:${randomUUID()}`;
localLocks.set(lockKey, {
owner: this.getLocalOwnerToken(),
owner: token,
expiresAt: Date.now() + ttlMs
});
return true;
}
if (existing.owner === this.getLocalOwnerToken()) {
existing.expiresAt = Date.now() + ttlMs;
localLocks.set(lockKey, existing);
return true;
return token;
}
// The lock is currently held -- possibly by a different, unrelated
// caller in this same process. We intentionally do NOT treat
// same-process holders as automatically reentrant here: two
// independent logical operations (e.g. two different API requests)
// running concurrently in the same process must not both believe
// they hold the lock, or their writes under it can interleave
// unguarded. Just retry with backoff like any other contended lock.
if (attempt < maxRetries - 1) {
const delay = retryDelayMs * Math.pow(2, attempt);
await new Promise((resolve) => setTimeout(resolve, delay));
}
}
return false;
return null;
}
/**
* Release a lock using Lua script to ensure atomicity
* Release a lock previously acquired via acquireLock/acquireLockWithRetry.
* @param lockKey - Unique identifier for the lock
* @param token - the exact token returned by the acquisition being released.
* Required so a caller whose TTL already expired can't delete a
* different, currently-active holder's lock.
*/
async releaseLock(lockKey: string): Promise<void> {
async releaseLock(lockKey: string, token: string): Promise<void> {
this.clearExpiredLocalLock(lockKey);
const existing = localLocks.get(lockKey);
if (existing && existing.owner === this.getLocalOwnerToken()) {
if (existing && existing.owner === token) {
localLocks.delete(lockKey);
}
}
@@ -100,23 +104,29 @@ export class LockManager {
const ttl = Math.max(0, existing.expiresAt - Date.now());
return {
exists: true,
ownedByMe: existing.owner === this.getLocalOwnerToken(),
ownedByMe: existing.owner.startsWith(`${instanceId}:`),
ttl,
owner: existing.owner.split(":")[0]
};
}
/**
* Extend the TTL of an existing lock owned by this worker
* Extend the TTL of an existing lock, provided the token matches the
* acquisition currently holding it.
* @param lockKey - Unique identifier for the lock
* @param ttlMs - New TTL in milliseconds
* @param token - the token returned by the acquisition being extended
* @returns Promise<boolean> - true if extended successfully
*/
async extendLock(lockKey: string, ttlMs: number): Promise<boolean> {
async extendLock(
lockKey: string,
ttlMs: number,
token: string
): Promise<boolean> {
this.clearExpiredLocalLock(lockKey);
const existing = localLocks.get(lockKey);
if (!existing || existing.owner !== this.getLocalOwnerToken()) {
if (!existing || existing.owner !== token) {
return false;
}
@@ -131,14 +141,14 @@ export class LockManager {
* @param ttlMs - Time to live in milliseconds
* @param maxRetries - Maximum number of retry attempts
* @param baseDelayMs - Base delay between retries in milliseconds
* @returns Promise<boolean> - true if lock acquired
* @returns Promise<string | null> - token if acquired, null otherwise
*/
async acquireLockWithRetry(
lockKey: string,
ttlMs: number = 30000,
maxRetries: number = 5,
baseDelayMs: number = 100
): Promise<boolean> {
): Promise<string | null> {
for (let attempt = 0; attempt <= maxRetries; attempt++) {
const acquired = await this.acquireLock(
lockKey,
@@ -148,7 +158,7 @@ export class LockManager {
);
if (acquired) {
return true;
return acquired;
}
if (attempt < maxRetries) {
@@ -158,7 +168,7 @@ export class LockManager {
}
}
return false;
return null;
}
/**
@@ -173,16 +183,16 @@ export class LockManager {
fn: () => Promise<T>,
ttlMs: number = 30000
): Promise<T> {
const acquired = await this.acquireLock(lockKey, ttlMs);
const token = await this.acquireLock(lockKey, ttlMs);
if (!acquired) {
if (!token) {
throw new Error(`Failed to acquire lock: ${lockKey}`);
}
try {
return await fn();
} finally {
await this.releaseLock(lockKey);
await this.releaseLock(lockKey, token);
}
}
@@ -204,7 +214,7 @@ export class LockManager {
let locksOwnedByMe = 0;
for (const value of localLocks.values()) {
if (value.owner === this.getLocalOwnerToken()) {
if (value.owner.startsWith(`${instanceId}:`)) {
locksOwnedByMe++;
}
}
+24
View File
@@ -0,0 +1,24 @@
export const ORG_REBUILD_CONCURRENCY_LIMIT = 10;
const orgActiveRebuilds = new Map<string, number>();
export async function incrementOrgRebuildCount(orgId: string): Promise<void> {
orgActiveRebuilds.set(orgId, (orgActiveRebuilds.get(orgId) ?? 0) + 1);
}
export async function decrementOrgRebuildCount(orgId: string): Promise<void> {
const current = orgActiveRebuilds.get(orgId) ?? 0;
if (current <= 1) {
orgActiveRebuilds.delete(orgId);
} else {
orgActiveRebuilds.set(orgId, current - 1);
}
}
export async function getOrgActiveRebuildCount(orgId: string): Promise<number> {
return orgActiveRebuilds.get(orgId) ?? 0;
}
export async function checkOrgRebuildRateLimit(orgId: string): Promise<boolean> {
return (orgActiveRebuilds.get(orgId) ?? 0) >= ORG_REBUILD_CONCURRENCY_LIMIT;
}
+19 -1
View File
@@ -45,11 +45,23 @@ import {
} from "@server/routers/client/targets";
import { lockManager } from "#dynamic/lib/lock";
import { rebuildQueue } from "#dynamic/lib/rebuildQueue";
import {
checkOrgRebuildRateLimit,
decrementOrgRebuildCount,
incrementOrgRebuildCount,
ORG_REBUILD_CONCURRENCY_LIMIT
} from "#dynamic/lib/orgRebuildCounter";
export { ORG_REBUILD_CONCURRENCY_LIMIT };
// TTL for rebuild-association locks. These functions can fan out into many
// peer/proxy updates, so give them a generous window.
const REBUILD_ASSOCIATIONS_LOCK_TTL_MS = 120000;
export async function isOrgRebuildRateLimited(orgId: string): Promise<boolean> {
return checkOrgRebuildRateLimit(orgId);
}
const REBUILD_IDLE_POLL_INTERVAL_MS = 300;
const REBUILD_IDLE_DEFAULT_TIMEOUT_MS = 130_000; // slightly longer than lock TTL
const REBUILD_IDLE_HANDLER_TIMEOUT_MS = 5_000;
@@ -271,6 +283,7 @@ export async function getClientSiteResourceAccess(
export async function rebuildClientAssociationsFromSiteResource(
siteResource: SiteResource
) {
await incrementOrgRebuildCount(siteResource.orgId);
try {
return await lockManager.withLock(
`rebuild-client-associations:site-resource:${siteResource.siteResourceId}`,
@@ -292,6 +305,8 @@ export async function rebuildClientAssociationsFromSiteResource(
return { mergedAllClients: [] };
}
throw err;
} finally {
await decrementOrgRebuildCount(siteResource.orgId);
}
}
@@ -1638,8 +1653,9 @@ export async function handleMessagingForUpdatedSiteResource(
export async function rebuildClientAssociationsFromClient(
client: Client
): Promise<void> {
const trx = primaryDb;
await incrementOrgRebuildCount(client.orgId);
try {
const trx = primaryDb;
return await lockManager.withLock(
`rebuild-client-associations:client:${client.clientId}`,
() => rebuildClientAssociationsFromClientImpl(client, trx),
@@ -1660,6 +1676,8 @@ export async function rebuildClientAssociationsFromClient(
return;
}
throw err;
} finally {
await decrementOrgRebuildCount(client.orgId);
}
}
+100 -6
View File
@@ -1,3 +1,5 @@
import logger from "@server/logger";
export type RebuildJobType = "site-resource" | "client";
export interface RebuildJob {
@@ -16,12 +18,104 @@ export interface RebuildQueueManager {
isQueued(job: RebuildJob): Promise<boolean>;
}
class NoopRebuildQueue implements RebuildQueueManager {
async enqueue(_job: RebuildJob): Promise<void> {}
startProcessing(_handlers: RebuildJobHandlers): void {}
async isQueued(_job: RebuildJob): Promise<boolean> {
return false;
// In-process FIFO used when there is no Redis to back a distributed queue
// (OSS build, or Redis unavailable). A job that loses the per-resource
// rebuild lock race lands here instead of being silently dropped, and gets
// retried shortly after against fresh DB state.
const POLL_INTERVAL_MS = 500;
const BATCH_SIZE = 5;
function dedupeKey(job: RebuildJob): string {
return `${job.type}:${job.id}`;
}
class InMemoryRebuildQueue implements RebuildQueueManager {
private queue: RebuildJob[] = [];
private queuedSet = new Set<string>();
private processing = false;
private processingStarted = false;
private handlers: RebuildJobHandlers | null = null;
async isQueued(job: RebuildJob): Promise<boolean> {
return this.queuedSet.has(dedupeKey(job));
}
async enqueue(job: RebuildJob): Promise<void> {
const key = dedupeKey(job);
if (this.queuedSet.has(key)) {
logger.debug(
`Rebuild queue: skipped duplicate queued job ${job.type}:${job.id}`
);
return;
}
this.queuedSet.add(key);
this.queue.push(job);
logger.debug(
`Rebuild queue: enqueued ${job.type}:${job.id} (queue position: tail)`
);
}
startProcessing(handlers: RebuildJobHandlers): void {
if (this.processingStarted) return;
this.processingStarted = true;
this.handlers = handlers;
setInterval(() => {
this.tryProcessBatch().catch((err) => {
logger.error(
"Rebuild queue: unhandled error in process loop:",
err
);
});
}, POLL_INTERVAL_MS);
logger.info("Rebuild queue processor started (in-memory)");
}
private async tryProcessBatch(): Promise<void> {
if (this.processing || !this.handlers || this.queue.length === 0) {
return;
}
this.processing = true;
try {
for (let i = 0; i < BATCH_SIZE; i++) {
const job = this.queue.shift();
if (!job) break; // queue drained
// Remove from the dedupe set once dequeued so the same job
// can be re-queued while this one is in progress.
this.queuedSet.delete(dedupeKey(job));
logger.debug(
`Rebuild queue: processing ${job.type}:${job.id}`
);
try {
if (job.type === "site-resource") {
await this.handlers.onSiteResource(job.id);
} else if (job.type === "client") {
await this.handlers.onClient(job.id);
} else {
logger.warn(
`Rebuild queue: unknown job type "${(job as any).type}", discarding`
);
}
logger.debug(
`Rebuild queue: completed ${job.type}:${job.id}`
);
} catch (err) {
logger.error(
`Rebuild queue: job ${job.type}:${job.id} threw an error:`,
err
);
}
}
} finally {
this.processing = false;
}
}
}
export const rebuildQueue: RebuildQueueManager = new NoopRebuildQueue();
export const rebuildQueue: RebuildQueueManager = new InMemoryRebuildQueue();
+3 -3
View File
@@ -14,7 +14,7 @@ import {
} from "@server/db";
import { eq, and, inArray, ne, exists } from "drizzle-orm";
import { usageService } from "@server/lib/billing/usageService";
import { FeatureId } from "@server/lib/billing";
import { LimitId } from "@server/lib/billing";
export async function assignUserToOrg(
org: Org,
@@ -61,7 +61,7 @@ export async function assignUserToOrg(
);
if (orgsInBillingDomainThatTheUserIsStillIn.length === 0) {
await usageService.add(org.orgId, FeatureId.USERS, 1, trx);
await usageService.add(org.orgId, LimitId.USERS, 1, trx);
}
}
}
@@ -157,7 +157,7 @@ export async function removeUserFromOrg(
);
if (orgsInBillingDomainThatTheUserIsStillIn.length === 0) {
await usageService.add(org.orgId, FeatureId.USERS, -1, trx);
await usageService.add(org.orgId, LimitId.USERS, -1, trx);
}
}
}