diff --git a/apps/backend/prisma/migrations/20260623010000_add_plan_usage_range_indexes/migration.sql b/apps/backend/prisma/migrations/20260623010000_add_plan_usage_range_indexes/migration.sql new file mode 100644 index 000000000..cebc53f24 --- /dev/null +++ b/apps/backend/prisma/migrations/20260623010000_add_plan_usage_range_indexes/migration.sql @@ -0,0 +1,11 @@ +-- SPLIT_STATEMENT_SENTINEL +-- SINGLE_STATEMENT_SENTINEL +-- RUN_OUTSIDE_TRANSACTION_SENTINEL +CREATE INDEX CONCURRENTLY IF NOT EXISTS "EmailOutbox_tenancyId_startedSendingAt_idx" + ON /* SCHEMA_NAME_SENTINEL */."EmailOutbox"("tenancyId", "startedSendingAt"); + +-- SPLIT_STATEMENT_SENTINEL +-- SINGLE_STATEMENT_SENTINEL +-- RUN_OUTSIDE_TRANSACTION_SENTINEL +CREATE INDEX CONCURRENTLY IF NOT EXISTS "SessionReplay_tenancyId_startedAt_idx" + ON /* SCHEMA_NAME_SENTINEL */."SessionReplay"("tenancyId", "startedAt"); diff --git a/apps/backend/prisma/schema.prisma b/apps/backend/prisma/schema.prisma index 8eb12390e..685c77d26 100644 --- a/apps/backend/prisma/schema.prisma +++ b/apps/backend/prisma/schema.prisma @@ -405,6 +405,7 @@ model SessionReplay { @@id([tenancyId, id]) @@index([tenancyId, projectUserId, startedAt]) + @@index([tenancyId, startedAt], name: "SessionReplay_tenancyId_startedAt_idx") @@index([tenancyId, lastEventAt]) // index by updatedAt instead of lastEventAt because event timing can be spoofed @@index([tenancyId, refreshTokenId, updatedAt]) @@ -1067,6 +1068,7 @@ model EmailOutbox { @@id([tenancyId, id]) @@index([tenancyId, finishedSendingAt(sort: Desc), scheduledAtIfNotYetQueued(sort: Desc), priority, id], map: "EmailOutbox_ordering_idx") + @@index([tenancyId, startedSendingAt], name: "EmailOutbox_tenancyId_startedSendingAt_idx") @@index([tenancyId, simpleStatus], map: "EmailOutbox_simple_status_tenancy_idx") @@index([tenancyId, status], map: "EmailOutbox_status_tenancy_idx") @@index([isQueued], map: "EmailOutbox_isQueued_idx") diff --git a/apps/backend/src/lib/plan-entitlements.test.ts b/apps/backend/src/lib/plan-entitlements.test.ts index aec03993e..7d9f0d50c 100644 --- a/apps/backend/src/lib/plan-entitlements.test.ts +++ b/apps/backend/src/lib/plan-entitlements.test.ts @@ -4,6 +4,8 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { arePlanLimitsEnforced, getBillingTeamId, + getNonAnonymousUserCountForTenancies, + getOwnedProjectAndTenancyIdsForBillingTeam, getOwnedProjectIdsForBillingTeam, getOwnedTenancyIdsForBillingTeam, getTeamWideItemCapacityForTests, @@ -98,6 +100,28 @@ describe("team-wide ownership aggregation", () => { expect(tenancyIds).toEqual(["tenancy-a-main", "tenancy-a-dev", "tenancy-b-main"]); }); + it("lists owned project and tenancy ids from one ownership scope", async () => { + const scope = await getOwnedProjectAndTenancyIdsForBillingTeam("team-1", globalPrisma); + expect(scope).toMatchInlineSnapshot(` + { + "projectIds": [ + "project-a", + "project-b", + ], + "tenancyIds": [ + "tenancy-a-main", + "tenancy-a-dev", + "tenancy-b-main", + ], + } + `); + }); + + it("counts non-anonymous users from already-resolved tenancies", async () => { + const usage = await getNonAnonymousUserCountForTenancies(["tenancy-a-main", "tenancy-b-main"], globalPrisma); + expect(usage).toBe(2); + }); + it("counts only non-anonymous users across all owned tenancies", async () => { const usage = await getTeamWideNonAnonymousUserCount("team-1", globalPrisma); expect(usage).toBe(3); diff --git a/apps/backend/src/lib/plan-entitlements.ts b/apps/backend/src/lib/plan-entitlements.ts index 932ab977c..1b8527061 100644 --- a/apps/backend/src/lib/plan-entitlements.ts +++ b/apps/backend/src/lib/plan-entitlements.ts @@ -39,6 +39,11 @@ type GlobalPrismaLike = { }, }; +type OwnedBillingScope = { + projectIds: string[], + tenancyIds: string[], +}; + type ItemCapacityReaders = { getPrismaForTenancy: (tenancy: Tenancy) => Promise, getItemQuantityForCustomer: (options: { @@ -85,13 +90,16 @@ export async function getOwnedProjectIdsForBillingTeam( return projects.map((project) => project.id); } -export async function getOwnedTenancyIdsForBillingTeam( +export async function getOwnedProjectAndTenancyIdsForBillingTeam( billingTeamId: string, globalPrisma: GlobalPrismaLike = globalPrismaClient, -): Promise { +): Promise { const projectIds = await getOwnedProjectIdsForBillingTeam(billingTeamId, globalPrisma); if (projectIds.length === 0) { - return []; + return { + projectIds, + tenancyIds: [], + }; } const tenancies = await globalPrisma.tenancy.findMany({ where: { @@ -103,16 +111,23 @@ export async function getOwnedTenancyIdsForBillingTeam( id: true, }, }); - return tenancies.map((tenancy) => tenancy.id); + return { + projectIds, + tenancyIds: tenancies.map((tenancy) => tenancy.id), + }; } -export async function getTeamWideNonAnonymousUserCount( +export async function getOwnedTenancyIdsForBillingTeam( billingTeamId: string, globalPrisma: GlobalPrismaLike = globalPrismaClient, +): Promise { + return (await getOwnedProjectAndTenancyIdsForBillingTeam(billingTeamId, globalPrisma)).tenancyIds; +} + +export async function getNonAnonymousUserCountForTenancies( + tenancyIds: string[], + globalPrisma: GlobalPrismaLike = globalPrismaClient, ): Promise { - // Usage metric: how many non-anonymous users are currently consumed by this billing team. - // This is compared against auth user capacity to determine over-limit conditions. - const tenancyIds = await getOwnedTenancyIdsForBillingTeam(billingTeamId, globalPrisma); if (tenancyIds.length === 0) { return 0; } @@ -126,6 +141,16 @@ export async function getTeamWideNonAnonymousUserCount( }); } +export async function getTeamWideNonAnonymousUserCount( + billingTeamId: string, + globalPrisma: GlobalPrismaLike = globalPrismaClient, +): Promise { + // Usage metric: how many non-anonymous users are currently consumed by this billing team. + // This is compared against auth user capacity to determine over-limit conditions. + const tenancyIds = await getOwnedTenancyIdsForBillingTeam(billingTeamId, globalPrisma); + return await getNonAnonymousUserCountForTenancies(tenancyIds, globalPrisma); +} + async function getTeamWideItemCapacity( billingTeamId: string, itemId: string, diff --git a/apps/backend/src/lib/plan-usage.ts b/apps/backend/src/lib/plan-usage.ts index b53c01391..50b8be096 100644 --- a/apps/backend/src/lib/plan-usage.ts +++ b/apps/backend/src/lib/plan-usage.ts @@ -4,9 +4,8 @@ import { getSubscriptionMapForCustomer } from "@/lib/payments/customer-data"; import { isActiveSubscription } from "@/lib/payments"; import { getBillingTeamId, - getOwnedProjectIdsForBillingTeam, - getOwnedTenancyIdsForBillingTeam, - getTeamWideNonAnonymousUserCount, + getNonAnonymousUserCountForTenancies, + getOwnedProjectAndTenancyIdsForBillingTeam, } from "@/lib/plan-entitlements"; import { DEFAULT_BRANCH_ID, getSoleTenancyFromProjectBranch, getTenancy, type Tenancy } from "@/lib/tenancies"; import { getPrismaClientForTenancy, getPrismaSchemaForTenancy, globalPrismaClient, sqlQuoteIdent } from "@/prisma-client"; @@ -18,6 +17,10 @@ import type { SubscriptionRow } from "./payments/schema/types"; type PlanUsageKind = PlanUsageResponse["rows"][number]["kind"]; type PlanUsageRow = PlanUsageResponse["rows"][number]; type UsageLimit = number | null; +type TenancyMeteredUsage = { + emails: number, + sessionReplays: number, +}; type UsagePeriod = { start: Date, @@ -46,6 +49,8 @@ const PLAN_LABELS = new Map([ ["growth", "Growth"], ]); +const PLAN_USAGE_TENANCY_COUNTER_CONCURRENCY = 4; + export function getNextPlanId(planId: PlanId): "team" | "growth" | null { if (planId === "free") { return "team"; @@ -202,38 +207,57 @@ async function getOwnerTeamDisplayName(internalTenancy: Tenancy, ownerTeamId: st return team?.displayName ?? throwErr(`Owner team ${ownerTeamId} not found in the internal tenancy`); } -async function countEmailsForTenancy(tenancyId: string, period: UsagePeriod): Promise { - const tenancy = await getTenancy(tenancyId) ?? throwErr(`Tenancy ${tenancyId} not found while counting email usage`); +async function countMeteredUsageForTenancy(tenancyId: string, period: UsagePeriod): Promise { + const tenancy = await getTenancy(tenancyId) ?? throwErr(`Tenancy ${tenancyId} not found while counting plan usage`); const schema = await getPrismaSchemaForTenancy(tenancy); const prisma = await getPrismaClientForTenancy(tenancy); - const rows = await prisma.$replica().$queryRaw<[{ count: number }]>` - SELECT COUNT(*)::int AS count - FROM ${sqlQuoteIdent(schema)}."EmailOutbox" - WHERE "tenancyId" = ${tenancy.id}::uuid - AND "startedSendingAt" IS NOT NULL - AND "startedSendingAt" >= ${period.start} - AND "startedSendingAt" < ${period.end} + const rows = await prisma.$replica().$queryRaw>` + SELECT + ( + SELECT COUNT(*)::int + FROM ${sqlQuoteIdent(schema)}."EmailOutbox" + WHERE "tenancyId" = ${tenancy.id}::uuid + AND "startedSendingAt" IS NOT NULL + AND "startedSendingAt" >= ${period.start} + AND "startedSendingAt" < ${period.end} + ) AS "emails", + ( + SELECT COUNT(*)::int + FROM ${sqlQuoteIdent(schema)}."SessionReplay" + WHERE "tenancyId" = ${tenancy.id}::uuid + AND "startedAt" >= ${period.start} + AND "startedAt" < ${period.end} + ) AS "sessionReplays" `; - return Number(rows[0].count); + const row = rows[0] ?? throwErr(`Missing plan usage count row for tenancy ${tenancy.id}`); + return { + emails: Number(row.emails), + sessionReplays: Number(row.sessionReplays), + }; } -async function countSessionReplaysForTenancy(tenancyId: string, period: UsagePeriod): Promise { - const tenancy = await getTenancy(tenancyId) ?? throwErr(`Tenancy ${tenancyId} not found while counting session replay usage`); - const schema = await getPrismaSchemaForTenancy(tenancy); - const prisma = await getPrismaClientForTenancy(tenancy); - const rows = await prisma.$replica().$queryRaw<[{ count: number }]>` - SELECT COUNT(*)::int AS count - FROM ${sqlQuoteIdent(schema)}."SessionReplay" - WHERE "tenancyId" = ${tenancy.id}::uuid - AND "startedAt" >= ${period.start} - AND "startedAt" < ${period.end} - `; - return Number(rows[0].count); -} +async function sumTenancyMeteredUsage(tenancyIds: string[], period: UsagePeriod): Promise { + const totals: TenancyMeteredUsage = { + emails: 0, + sessionReplays: 0, + }; + let nextIndex = 0; -async function sumTenancyUsage(tenancyIds: string[], counter: (tenancyId: string) => Promise): Promise { - const counts = await Promise.all(tenancyIds.map(counter)); - return counts.reduce((sum, count) => sum + count, 0); + // Keep this page from turning a team with many tenancies into an unbounded burst of replica COUNTs. + async function worker(): Promise { + while (nextIndex < tenancyIds.length) { + const index = nextIndex; + nextIndex++; + const tenancyId = tenancyIds[index] ?? throwErr(`Missing tenancy ID at index ${index} while counting plan usage`); + const usage = await countMeteredUsageForTenancy(tenancyId, period); + totals.emails += usage.emails; + totals.sessionReplays += usage.sessionReplays; + } + } + + const workerCount = Math.min(PLAN_USAGE_TENANCY_COUNTER_CONCURRENCY, tenancyIds.length); + await Promise.all(Array.from({ length: workerCount }, async () => await worker())); + return totals; } async function countAnalyticsEventsForProjects(projectIds: string[], period: UsagePeriod): Promise { @@ -336,18 +360,16 @@ export async function getPlanUsageForProject(project: UsageSourceProject, now: D const planId = resolveActivePlanId(activePlanSubscription); const period = getPlanUsagePeriod(activePlanSubscription, now); - const [ownerTeamDisplayName, ownedProjectIds, ownedTenancyIds, dashboardAdmins, authUsers] = await Promise.all([ + const [ownerTeamDisplayName, ownedScope, dashboardAdmins] = await Promise.all([ getOwnerTeamDisplayName(internalTenancy, ownerTeamId), - getOwnedProjectIdsForBillingTeam(ownerTeamId), - getOwnedTenancyIdsForBillingTeam(ownerTeamId), + getOwnedProjectAndTenancyIdsForBillingTeam(ownerTeamId), countDashboardAdmins(internalTenancy, ownerTeamId, now), - getTeamWideNonAnonymousUserCount(ownerTeamId), ]); - const [emails, analyticsEvents, sessionReplays] = await Promise.all([ - sumTenancyUsage(ownedTenancyIds, async (tenancyId) => await countEmailsForTenancy(tenancyId, period)), - countAnalyticsEventsForProjects(ownedProjectIds, period), - sumTenancyUsage(ownedTenancyIds, async (tenancyId) => await countSessionReplaysForTenancy(tenancyId, period)), + const [authUsers, meteredUsage, analyticsEvents] = await Promise.all([ + getNonAnonymousUserCountForTenancies(ownedScope.tenancyIds), + sumTenancyMeteredUsage(ownedScope.tenancyIds, period), + countAnalyticsEventsForProjects(ownedScope.projectIds, period), ]); return { @@ -362,9 +384,9 @@ export async function getPlanUsageForProject(project: UsageSourceProject, now: D planId, dashboardAdmins, authUsers, - emails, + emails: meteredUsage.emails, analyticsEvents, - sessionReplays, + sessionReplays: meteredUsage.sessionReplays, }), }; }