diff --git a/.server-changes/billing-limits-page-timeout.md b/.server-changes/billing-limits-page-timeout.md new file mode 100644 index 0000000000..6edbccf48e --- /dev/null +++ b/.server-changes/billing-limits-page-timeout.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: fix +--- + +Fixed the billing limits page timing out for organizations with many preview branches, especially while a spend limit was being enforced. The page now loads quickly, so you can raise or resolve your limit without delay. diff --git a/apps/webapp/app/v3/services/billingLimit/billingLimitConstants.ts b/apps/webapp/app/v3/services/billingLimit/billingLimitConstants.ts index fbaeea845c..5aea0f1afd 100644 --- a/apps/webapp/app/v3/services/billingLimit/billingLimitConstants.ts +++ b/apps/webapp/app/v3/services/billingLimit/billingLimitConstants.ts @@ -16,6 +16,9 @@ export const BILLING_LIMIT_RECONCILE_LOOKUP_CONCURRENCY = 10; /** Inline bulk-cancel budget for billing limit resolve (worker visibility is 10 min). */ export const BILLING_LIMIT_RESOLVE_BULK_CANCEL_BUDGET_MS = 8 * 60_000; +/** Server-side cap (seconds) on the org-level queued-run count; it renders in a loader. */ +export const BILLING_LIMIT_QUEUED_COUNT_MAX_EXECUTION_S = 10; + export type BillingLimitConvergeTargetState = "grace" | "rejected" | "ok"; export function isBillableEnvironmentType(type: RuntimeEnvironmentType): boolean { diff --git a/apps/webapp/app/v3/services/billingLimit/billingLimitQueuedRuns.server.ts b/apps/webapp/app/v3/services/billingLimit/billingLimitQueuedRuns.server.ts index 7cf3cf7cd9..066af4be10 100644 --- a/apps/webapp/app/v3/services/billingLimit/billingLimitQueuedRuns.server.ts +++ b/apps/webapp/app/v3/services/billingLimit/billingLimitQueuedRuns.server.ts @@ -1,9 +1,13 @@ +import type { ClickHouse } from "@internal/clickhouse"; import type { PrismaClient, TaskRunStatus } from "@trigger.dev/database"; import { QUEUED_STATUSES, RUNNING_STATUSES } from "~/components/runs/v3/TaskRunStatus"; import { prisma } from "~/db.server"; import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstance.server"; import { RunsRepository } from "~/services/runsRepository/runsRepository.server"; -import { BILLABLE_ENVIRONMENT_TYPES } from "./billingLimitConstants"; +import { + BILLABLE_ENVIRONMENT_TYPES, + BILLING_LIMIT_QUEUED_COUNT_MAX_EXECUTION_S, +} from "./billingLimitConstants"; import { boundedIn } from "@trigger.dev/database"; export type BillableEnvironmentRef = { @@ -11,6 +15,11 @@ export type BillableEnvironmentRef = { projectId: string; }; +/** + * Environments whose runs a billing limit must act on. Archived branches stay included: + * archiving is a soft update that cancels nothing, so an archived branch can still hold + * executing runs that enforcement has to cancel. + */ export async function getBillableEnvironmentsForBillingLimit( organizationId: string, prismaClient: PrismaClient = prisma @@ -73,27 +82,37 @@ async function countRunsForBillableEnvironment( }); } -/** Same source as BillingLimitBulkCancelService — ClickHouse countRuns(QUEUED_STATUSES). */ +/** + * Same table and statuses as BillingLimitBulkCancelService's per-environment counts, but a + * single org-level ClickHouse query filtered on environment_type. Deliberately NOT a + * per-environment loop: an org can have thousands of (mostly archived) preview environments, + * and sequential per-env counts hold the billing-limits loader open past the edge timeout. + * The count is display-only, so a server-side execution cap beats an unbounded query. + */ export async function countBillableQueuedRunsForOrganization( - organizationId: string + organizationId: string, + clickhouse?: ClickHouse ): Promise { - const environments = await getBillableEnvironmentsForBillingLimit(organizationId); + const client = + clickhouse ?? + (await clickhouseFactory.getClickhouseForOrganization(organizationId, "standard")); - if (environments.length === 0) { - return 0; - } + const queryBuilder = client.taskRuns.countQueryBuilder({ + settings: { max_execution_time: BILLING_LIMIT_QUEUED_COUNT_MAX_EXECUTION_S }, + }); - const runsRepository = await createBillingLimitRunsRepository(organizationId); + queryBuilder + .where("organization_id = {organizationId: String}", { organizationId }) + .where("environment_type IN {environmentTypes: Array(String)}", { + environmentTypes: [...BILLABLE_ENVIRONMENT_TYPES], + }) + .where("status IN {statuses: Array(String)}", { statuses: [...QUEUED_STATUSES] }); - let total = 0; + const [queryError, result] = await queryBuilder.execute(); - for (const environment of environments) { - total += await countQueuedRunsForBillableEnvironment( - runsRepository, - organizationId, - environment - ); + if (queryError) { + throw queryError; } - return total; + return result[0]?.count ?? 0; } diff --git a/apps/webapp/app/v3/services/billingLimit/getBillingLimitQueuedRunCount.server.ts b/apps/webapp/app/v3/services/billingLimit/getBillingLimitQueuedRunCount.server.ts index 2689a3dc9b..af0e20a565 100644 --- a/apps/webapp/app/v3/services/billingLimit/getBillingLimitQueuedRunCount.server.ts +++ b/apps/webapp/app/v3/services/billingLimit/getBillingLimitQueuedRunCount.server.ts @@ -1,9 +1,26 @@ +import { tryCatch } from "@trigger.dev/core/utils"; import { EnvironmentPauseSource } from "@trigger.dev/database"; import { prisma } from "~/db.server"; +import { logger } from "~/services/logger.server"; import { countBillableQueuedRunsForOrganization } from "./billingLimitQueuedRuns.server"; +/** + * Display-only count for the billing-limits page. Falls back to 0 on failure (the page hides + * the count label at 0) — the recovery panel must stay reachable even when the count errors, + * because it is the customer's only self-serve path out of an enforced limit. + */ export async function getBillingLimitQueuedRunCount(organizationId: string): Promise { - return countBillableQueuedRunsForOrganization(organizationId); + const [error, count] = await tryCatch(countBillableQueuedRunsForOrganization(organizationId)); + + if (error) { + logger.error("getBillingLimitQueuedRunCount failed, returning 0", { + organizationId, + error, + }); + return 0; + } + + return count ?? 0; } export async function countBillingLimitPausedEnvironments(organizationId: string): Promise { diff --git a/apps/webapp/test/billingLimitQueuedRuns.test.ts b/apps/webapp/test/billingLimitQueuedRuns.test.ts index 9c923890eb..14dc0c0db8 100644 --- a/apps/webapp/test/billingLimitQueuedRuns.test.ts +++ b/apps/webapp/test/billingLimitQueuedRuns.test.ts @@ -2,7 +2,11 @@ import { describe, expect, vi } from "vitest"; import { setTimeout } from "node:timers/promises"; import { replicationContainerTest } from "@internal/testcontainers"; import { RunsRepository } from "~/services/runsRepository/runsRepository.server"; -import { countQueuedRunsForBillableEnvironment } from "~/v3/services/billingLimit/billingLimitQueuedRuns.server"; +import { + countBillableQueuedRunsForOrganization, + countQueuedRunsForBillableEnvironment, + getBillableEnvironmentsForBillingLimit, +} from "~/v3/services/billingLimit/billingLimitQueuedRuns.server"; import { setupClickhouseReplication } from "./utils/replicationUtils"; vi.setConfig({ testTimeout: 60_000 }); @@ -109,4 +113,188 @@ describe("billingLimitQueuedRuns", () => { expect(developmentCount).toBe(1); } ); + + replicationContainerTest( + "counts queued runs org-wide with a single query, spanning billable environment types only", + async ({ clickhouseContainer, redisOptions, postgresContainer, prisma }) => { + const { clickhouse } = await setupClickhouseReplication({ + prisma, + databaseUrl: postgresContainer.getConnectionUri(), + clickhouseUrl: clickhouseContainer.getConnectionUrl(), + redisOptions, + }); + + const organization = await prisma.organization.create({ + data: { title: "billing-limit-org-count", slug: "billing-limit-org-count" }, + }); + + const project = await prisma.project.create({ + data: { + name: "billing-limit-org-count", + slug: "billing-limit-org-count", + organizationId: organization.id, + externalRef: "billing-limit-org-count", + }, + }); + + const productionEnv = await prisma.runtimeEnvironment.create({ + data: { + slug: "prod", + type: "PRODUCTION", + projectId: project.id, + organizationId: organization.id, + apiKey: "prod-org-count", + pkApiKey: "prod-org-count", + shortcode: "prod-org-count", + }, + }); + + const developmentEnv = await prisma.runtimeEnvironment.create({ + data: { + slug: "dev", + type: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + apiKey: "dev-org-count", + pkApiKey: "dev-org-count", + shortcode: "dev-org-count", + }, + }); + + const previewEnv = await prisma.runtimeEnvironment.create({ + data: { + slug: "preview", + type: "PREVIEW", + projectId: project.id, + organizationId: organization.id, + apiKey: "preview-org-count", + pkApiKey: "preview-org-count", + shortcode: "preview-org-count", + branchName: "feature-branch", + archivedAt: new Date(), + }, + }); + + const runRows = [ + { + friendlyId: "run_org_prod_pending", + env: productionEnv, + type: "PRODUCTION", + status: "PENDING", + }, + { + friendlyId: "run_org_prod_delayed", + env: productionEnv, + type: "PRODUCTION", + status: "DELAYED", + }, + { + friendlyId: "run_org_prod_done", + env: productionEnv, + type: "PRODUCTION", + status: "COMPLETED_SUCCESSFULLY", + }, + { + friendlyId: "run_org_dev_pending", + env: developmentEnv, + type: "DEVELOPMENT", + status: "PENDING", + }, + { + friendlyId: "run_org_preview_pending", + env: previewEnv, + type: "PREVIEW", + status: "PENDING", + }, + ] as const; + + for (const row of runRows) { + await prisma.taskRun.create({ + data: { + friendlyId: row.friendlyId, + taskIdentifier: "queued-task", + status: row.status, + payload: JSON.stringify({}), + traceId: "trace", + spanId: "span", + queue: "main", + runtimeEnvironmentId: row.env.id, + projectId: project.id, + organizationId: organization.id, + environmentType: row.type, + engine: "V2", + }, + }); + } + + await setTimeout(1000); + + const count = await countBillableQueuedRunsForOrganization(organization.id, clickhouse); + + expect(count).toBe(3); + } + ); + + replicationContainerTest( + "keeps archived environments in the billable environment list so enforcement can cancel their runs", + async ({ prisma }) => { + const organization = await prisma.organization.create({ + data: { title: "billing-limit-archived", slug: "billing-limit-archived" }, + }); + + const project = await prisma.project.create({ + data: { + name: "billing-limit-archived", + slug: "billing-limit-archived", + organizationId: organization.id, + externalRef: "billing-limit-archived", + }, + }); + + const activeEnv = await prisma.runtimeEnvironment.create({ + data: { + slug: "preview-active", + type: "PREVIEW", + projectId: project.id, + organizationId: organization.id, + apiKey: "preview-active-test", + pkApiKey: "preview-active-test", + shortcode: "preview-active-test", + branchName: "live-branch", + }, + }); + + const archivedEnv = await prisma.runtimeEnvironment.create({ + data: { + slug: "preview-archived", + type: "PREVIEW", + projectId: project.id, + organizationId: organization.id, + apiKey: "preview-archived-test", + pkApiKey: "preview-archived-test", + shortcode: "preview-archived-test", + branchName: "old-branch", + archivedAt: new Date(), + }, + }); + + const developmentEnv = await prisma.runtimeEnvironment.create({ + data: { + slug: "dev", + type: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + apiKey: "dev-archived-test", + pkApiKey: "dev-archived-test", + shortcode: "dev-archived-test", + }, + }); + + const environments = await getBillableEnvironmentsForBillingLimit(organization.id, prisma); + const environmentIds = environments.map((environment) => environment.id).sort(); + + expect(environmentIds).toEqual([activeEnv.id, archivedEnv.id].sort()); + expect(environmentIds).not.toContain(developmentEnv.id); + } + ); });