diff --git a/infra/relay/src/environments/ManagedEndpointAllocations.ts b/infra/relay/src/environments/ManagedEndpointAllocations.ts index 09433cc872520..d432f45c61763 100644 --- a/infra/relay/src/environments/ManagedEndpointAllocations.ts +++ b/infra/relay/src/environments/ManagedEndpointAllocations.ts @@ -255,7 +255,14 @@ const make = Effect.gen(function* () { createdAt: now, updatedAt: now, }) - .onConflictDoNothing() + // Refreshing updatedAt makes a re-reserved row count toward the tunnel limit again. + .onConflictDoUpdate({ + target: [ + relayManagedEndpointAllocations.userId, + relayManagedEndpointAllocations.environmentId, + ], + set: { updatedAt: now }, + }) .returning(allocationSelection) .pipe( Effect.mapError( @@ -269,26 +276,7 @@ const make = Effect.gen(function* () { ), ); - const allocation = - inserted[0] ?? - (yield* db - .select(allocationSelection) - .from(relayManagedEndpointAllocations) - .where(whereAllocation(input)) - .limit(1) - .pipe( - Effect.map((rows) => rows[0]), - Effect.mapError( - (cause) => - new ManagedEndpointAllocationPersistenceError({ - operation: "reserve", - stage: "database-request", - ...input, - cause, - }), - ), - )); - + const allocation = inserted[0]; if (allocation === undefined) { return yield* new ManagedEndpointAllocationPersistenceError({ operation: "reserve", diff --git a/infra/relay/src/environments/ManagedEndpointProvider.test.ts b/infra/relay/src/environments/ManagedEndpointProvider.test.ts index 7fc21ee959564..36ea69d5a2694 100644 --- a/infra/relay/src/environments/ManagedEndpointProvider.test.ts +++ b/infra/relay/src/environments/ManagedEndpointProvider.test.ts @@ -335,10 +335,13 @@ function makeTunnelLimits( result: ManagedTunnelLimits.ManagedTunnelLimitExceeded | null = null, ) { return ManagedTunnelLimits.ManagedTunnelLimits.of({ - ensureCapacity: (input) => - Effect.suspend(() => { + withCapacity: (input, reserve) => + Effect.gen(function* () { calls.push(input); - return result === null ? Effect.void : Effect.fail(result); + if (result !== null) { + return yield* result; + } + return yield* reserve; }), }); } diff --git a/infra/relay/src/environments/ManagedEndpointProvider.ts b/infra/relay/src/environments/ManagedEndpointProvider.ts index d63286467b893..5524b5117faaf 100644 --- a/infra/relay/src/environments/ManagedEndpointProvider.ts +++ b/infra/relay/src/environments/ManagedEndpointProvider.ts @@ -1169,27 +1169,7 @@ const make = Effect.gen(function* () { environmentHash, ); const requestedTunnelName = managedEndpointTunnelName(cf.namespace, environmentHash); - yield* tunnelLimits - .ensureCapacity({ - userId: input.userId, - environmentId: input.environmentId, - }) - .pipe( - Effect.catchTags({ - ManagedTunnelLimitPersistenceError: (cause) => - Effect.fail( - new ManagedEndpointProvisioningFailed({ - userId: input.userId, - environmentId: input.environmentId, - stage: "check-tunnel-limit", - hostname: requestedHostname, - tunnelName: requestedTunnelName, - cause, - }), - ), - }), - ); - const allocation = yield* allocations + const reserveAllocation = allocations .reserve({ userId: input.userId, environmentId: input.environmentId, @@ -1209,6 +1189,26 @@ const make = Effect.gen(function* () { }), ), ); + const allocation = yield* tunnelLimits + .withCapacity( + { userId: input.userId, environmentId: input.environmentId }, + reserveAllocation, + ) + .pipe( + Effect.catchTags({ + ManagedTunnelLimitPersistenceError: (cause) => + Effect.fail( + new ManagedEndpointProvisioningFailed({ + userId: input.userId, + environmentId: input.environmentId, + stage: "check-tunnel-limit", + hostname: requestedHostname, + tunnelName: requestedTunnelName, + cause, + }), + ), + }), + ); const { hostname, tunnelName } = allocation; if (input.reuseReadyTunnel === true) { diff --git a/infra/relay/src/environments/ManagedTunnelLimits.test.ts b/infra/relay/src/environments/ManagedTunnelLimits.test.ts new file mode 100644 index 0000000000000..9879e38f18a37 --- /dev/null +++ b/infra/relay/src/environments/ManagedTunnelLimits.test.ts @@ -0,0 +1,115 @@ +import { describe, expect, it } from "@effect/vitest"; +import * as Context from "effect/Context"; +import * as Effect from "effect/Effect"; +import * as Layer from "effect/Layer"; +import * as Semaphore from "effect/Semaphore"; +import { PgDialect } from "drizzle-orm/pg-core"; + +import * as RelayDb from "../db.ts"; +import { relayManagedEndpointAllocations } from "../persistence/schema.ts"; +import * as ManagedTunnelLimits from "./ManagedTunnelLimits.ts"; + +class FakeTransaction extends Context.Service()( + "t3code-relay/environments/ManagedTunnelLimits.test/FakeTransaction", +) {} + +// Models Postgres just enough: the advisory lock only works inside a transaction and is held until it ends. +function makeFakeDb(reserved: ReadonlyArray, countConditions: Array = []) { + const userLock = Semaphore.makeUnsafe(1); + return { + $client: { + withTransaction: (effect: Effect.Effect) => { + const transaction = { holdsLock: false }; + return effect.pipe( + Effect.provideService(FakeTransaction, transaction), + Effect.ensuring( + Effect.suspend(() => (transaction.holdsLock ? userLock.release(1) : Effect.void)), + ), + ); + }, + }, + execute: () => + Effect.gen(function* () { + const transaction = yield* Effect.serviceOption(FakeTransaction); + if (transaction._tag === "None") { + return yield* Effect.die("advisory lock taken outside a transaction"); + } + yield* userLock.take(1); + transaction.value.holdsLock = true; + return []; + }), + select: () => ({ + from: (table: unknown) => ({ + where: (condition: unknown) => { + if (table !== relayManagedEndpointAllocations) { + return { limit: () => Effect.succeed([]) }; + } + countConditions.push(condition); + return Effect.yieldNow.pipe(Effect.as([{ activeTunnels: reserved.length }])); + }, + }), + }), + } as unknown as RelayDb.RelayDb["Service"]; +} + +describe("ManagedTunnelLimits", () => { + it.effect("rejects concurrent reservations past the limit", () => { + const reserved: Array = []; + const attempts = ManagedTunnelLimits.DEFAULT_MANAGED_TUNNEL_LIMIT + 2; + + return Effect.gen(function* () { + const limits = yield* ManagedTunnelLimits.ManagedTunnelLimits; + const results = yield* Effect.forEach( + Array.from({ length: attempts }, (_, index) => `environment-${index}`), + (environmentId) => + limits + .withCapacity( + { userId: "user-1", environmentId }, + Effect.yieldNow.pipe(Effect.andThen(Effect.sync(() => reserved.push(environmentId)))), + ) + .pipe(Effect.result), + { concurrency: "unbounded" }, + ); + + expect(reserved).toHaveLength(ManagedTunnelLimits.DEFAULT_MANAGED_TUNNEL_LIMIT); + expect( + results.filter( + (result) => + result._tag === "Failure" && result.failure._tag === "ManagedTunnelLimitExceeded", + ), + ).toHaveLength(2); + }).pipe( + Effect.provide( + ManagedTunnelLimits.layer.pipe( + Layer.provide(Layer.succeed(RelayDb.RelayDb, makeFakeDb(reserved))), + ), + ), + ); + }); + + it.effect("pending reservations stop counting once released or withdrawn", () => { + const countConditions: Array = []; + + return Effect.gen(function* () { + const limits = yield* ManagedTunnelLimits.ManagedTunnelLimits; + yield* limits.withCapacity({ userId: "user-1", environmentId: "environment-1" }, Effect.void); + + expect(countConditions).toHaveLength(1); + const query = new PgDialect().sqlToQuery(countConditions[0] as never); + // claimRelease writes updatedAt = tunnelReleasedAt; a later reserve moves updatedAt past it. + expect(query.sql).toContain( + '"relay_managed_endpoint_allocations"."updated_at" > "relay_managed_endpoint_allocations"."tunnel_released_at"', + ); + // A link revoked or made publish-only after the last reservation no longer holds a slot. + expect(query.sql).toContain( + '"relay_environment_links"."updated_at" >= "relay_managed_endpoint_allocations"."updated_at"', + ); + }).pipe( + Effect.provide( + ManagedTunnelLimits.layer.pipe( + Layer.provide(Layer.succeed(RelayDb.RelayDb, makeFakeDb([], countConditions))), + ), + ), + ); + }); +}); diff --git a/infra/relay/src/environments/ManagedTunnelLimits.ts b/infra/relay/src/environments/ManagedTunnelLimits.ts index f6a271c3a92c8..15c50ae7efb1f 100644 --- a/infra/relay/src/environments/ManagedTunnelLimits.ts +++ b/infra/relay/src/environments/ManagedTunnelLimits.ts @@ -1,8 +1,11 @@ -import { and, count, eq, isNull, ne } from "drizzle-orm"; +import { and, count, eq, exists, gt, gte, isNotNull, isNull, ne, not, or, sql } from "drizzle-orm"; +import { QueryBuilder } from "drizzle-orm/pg-core"; import * as Context from "effect/Context"; +import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; import * as Schema from "effect/Schema"; +import * as SqlError from "effect/sql/SqlError"; import * as RelayDb from "../db.ts"; import { @@ -11,12 +14,14 @@ import { relayManagedTunnelLimits, } from "../persistence/schema.ts"; -const DEFAULT_MANAGED_TUNNEL_LIMIT = 3; +export const DEFAULT_MANAGED_TUNNEL_LIMIT = 3; +// Reserve-to-tunnel takes two Cloudflare calls capped at 8 seconds each; a reservation older than this was abandoned. +const MANAGED_TUNNEL_PENDING_WINDOW_MINUTES = 15; export class ManagedTunnelLimitPersistenceError extends Schema.TaggedError()( "ManagedTunnelLimitPersistenceError", { - operation: Schema.Literals(["load-limit", "count-tunnels"]), + operation: Schema.Literals(["transaction", "lock-user", "load-limit", "count-tunnels"]), userId: Schema.String, cause: Schema.Defect(), }, @@ -43,10 +48,11 @@ export class ManagedTunnelLimitExceeded extends Schema.TaggedError Effect.Effect; + /** Runs `reserve` only while the user is under their limit, serialized per user in one transaction. */ + readonly withCapacity: ( + input: { readonly userId: string; readonly environmentId: string }, + reserve: Effect.Effect, + ) => Effect.Effect; } >()("t3code-relay/environments/ManagedTunnelLimits") {} @@ -54,63 +60,130 @@ const make = Effect.gen(function* () { const db = yield* RelayDb.RelayDb; return ManagedTunnelLimits.of({ - ensureCapacity: Effect.fn("relay.managed_tunnel_limits.ensure_capacity")(function* (input) { - const overrides = yield* db - .select({ maxTunnels: relayManagedTunnelLimits.maxTunnels }) - .from(relayManagedTunnelLimits) - .where(eq(relayManagedTunnelLimits.userId, input.userId)) - .limit(1) - .pipe( - Effect.mapError( - (cause) => - new ManagedTunnelLimitPersistenceError({ - operation: "load-limit", - userId: input.userId, - cause, - }), - ), + withCapacity: Effect.fn("relay.managed_tunnel_limits.with_capacity")(function* ( + input: { readonly userId: string; readonly environmentId: string }, + reserve: Effect.Effect, + ): Effect.fn.Return { + const checkAndReserve = Effect.gen(function* () { + // Held until commit, so the count below sees every reservation that took the lock first. + yield* db + .execute( + sql`select pg_advisory_xact_lock(hashtextextended(${`managed-tunnel-limit:${input.userId}`}, 0))`, + ) + .pipe( + Effect.mapError( + (cause) => + new ManagedTunnelLimitPersistenceError({ + operation: "lock-user", + userId: input.userId, + cause, + }), + ), + ); + + const overrides = yield* db + .select({ maxTunnels: relayManagedTunnelLimits.maxTunnels }) + .from(relayManagedTunnelLimits) + .where(eq(relayManagedTunnelLimits.userId, input.userId)) + .limit(1) + .pipe( + Effect.mapError( + (cause) => + new ManagedTunnelLimitPersistenceError({ + operation: "load-limit", + userId: input.userId, + cause, + }), + ), + ); + const maxTunnels = overrides[0]?.maxTunnels ?? DEFAULT_MANAGED_TUNNEL_LIMIT; + + const pendingCutoff = DateTime.formatIso( + DateTime.subtract(yield* DateTime.now, { + minutes: MANAGED_TUNNEL_PENDING_WINDOW_MINUTES, + }), ); - const maxTunnels = overrides[0]?.maxTunnels ?? DEFAULT_MANAGED_TUNNEL_LIMIT; + // Counts allocations, not links: a link is written only after its tunnel exists. + // Excludes this environment so relinking at the limit stays idempotent. + const allocation = relayManagedEndpointAllocations; + const withdrawnLink = (since?: typeof allocation.updatedAt) => + exists( + new QueryBuilder() + .select({ userId: relayEnvironmentLinks.userId }) + .from(relayEnvironmentLinks) + .where( + and( + eq(relayEnvironmentLinks.userId, allocation.userId), + eq(relayEnvironmentLinks.environmentId, allocation.environmentId), + or( + isNotNull(relayEnvironmentLinks.revokedAt), + eq(relayEnvironmentLinks.managedTunnelsEnabled, false), + ), + since === undefined ? undefined : gte(relayEnvironmentLinks.updatedAt, since), + ), + ), + ); + const counted = yield* db + .select({ activeTunnels: count() }) + .from(allocation) + .where( + and( + eq(allocation.userId, input.userId), + ne(allocation.environmentId, input.environmentId), + or( + // Reserved but not yet holding a tunnel. Release writes updatedAt = tunnelReleasedAt and + // a relink reserves before rewriting its link, so only a newer reservation counts here. + and( + gt(allocation.updatedAt, pendingCutoff), + or( + isNull(allocation.tunnelReleasedAt), + gt(allocation.updatedAt, allocation.tunnelReleasedAt), + ), + not(withdrawnLink(allocation.updatedAt)), + ), + and( + isNotNull(allocation.tunnelId), + isNull(allocation.tunnelReleasedAt), + not(withdrawnLink()), + ), + ), + ), + ) + .pipe( + Effect.mapError( + (cause) => + new ManagedTunnelLimitPersistenceError({ + operation: "count-tunnels", + userId: input.userId, + cause, + }), + ), + ); + const activeTunnels = counted[0]?.activeTunnels ?? 0; - // Exclude this environment so relinking remains idempotent at the account limit. - const counted = yield* db - .select({ activeTunnels: count() }) - .from(relayManagedEndpointAllocations) - .innerJoin( - relayEnvironmentLinks, - and( - eq(relayEnvironmentLinks.userId, relayManagedEndpointAllocations.userId), - eq(relayEnvironmentLinks.environmentId, relayManagedEndpointAllocations.environmentId), - ), - ) - .where( - and( - eq(relayManagedEndpointAllocations.userId, input.userId), - ne(relayManagedEndpointAllocations.environmentId, input.environmentId), - isNull(relayEnvironmentLinks.revokedAt), - eq(relayEnvironmentLinks.managedTunnelsEnabled, true), - ), - ) - .pipe( - Effect.mapError( - (cause) => - new ManagedTunnelLimitPersistenceError({ - operation: "count-tunnels", + if (activeTunnels >= maxTunnels) { + return yield* new ManagedTunnelLimitExceeded({ + userId: input.userId, + environmentId: input.environmentId, + maxTunnels, + activeTunnels, + }); + } + return yield* reserve; + }); + + // Statements inside map their own SqlErrors, so any left here come from opening or committing the transaction. + return yield* db.$client.withTransaction(checkAndReserve).pipe( + Effect.mapError((cause) => + SqlError.isSqlError(cause) + ? new ManagedTunnelLimitPersistenceError({ + operation: "transaction", userId: input.userId, cause, - }), - ), - ); - const activeTunnels = counted[0]?.activeTunnels ?? 0; - - if (activeTunnels >= maxTunnels) { - return yield* new ManagedTunnelLimitExceeded({ - userId: input.userId, - environmentId: input.environmentId, - maxTunnels, - activeTunnels, - }); - } + }) + : cause, + ), + ); }), }); });