fix(relay): concurrent link requests can no longer exceed the managed tunnel limit

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",

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;
       }),
   });
 }

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) {

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<FakeTransaction, { holdsLock: boolean }>()(
+  "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<unknown>, countConditions: Array<unknown> = []) {
+  const userLock = Semaphore.makeUnsafe(1);
+  return {
+    $client: {
+      withTransaction: <A, E, R>(effect: Effect.Effect<A, E, R>) => {
+        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<string> = [];
+    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<unknown> = [];
+
+    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))),
+        ),
+      ),
+    );
+  });
+});