fix(relay): guard unlink teardown generation

Co-authored-by: codex <codex@users.noreply.github.com>

Author: juliusmarminge

infra/relay/src/environments/EnvironmentConnector.test.ts

@@ -198,7 +198,9 @@ function makeAllocations(
     recordDns: () => Effect.die("unused"),
     markReady: () => Effect.die("unused"),
     claimRelease: () => Effect.die("unused"),
+    claimDeprovision: () => Effect.die("unused"),
     remove: () => Effect.die("unused"),
+    removeClaimed: () => Effect.die("unused"),
   };
 }
 

infra/relay/src/environments/EnvironmentLinker.test.ts

@@ -136,6 +136,7 @@ function testLayer(input?: {
           revokeForEnvironmentPublicKey: () => Effect.succeed(false),
         }),
         Layer.succeed(ManagedEndpointProvider.ManagedEndpointProvider, {
+          prepareDeprovision: () => Effect.succeed(null),
           deprovision: input?.deprovision ?? (() => Effect.void),
           release: () => Effect.succeed(true),
           provision: () =>

infra/relay/src/environments/ManagedEndpointAllocations.test.ts

@@ -10,6 +10,61 @@ const layerWithDb = (db: RelayDb.RelayDb["Service"]) =>
   ManagedEndpointAllocations.layer.pipe(Layer.provide(Layer.succeed(RelayDb.RelayDb, db)));
 
 describe("ManagedEndpointAllocations", () => {
+  it.effect("returns a claim generation only when deprovision wins the allocation CAS", () => {
+    let claimedAt: string | undefined;
+    const fakeDb = {
+      update: (table: unknown) => {
+        expect(table).toBe(relayManagedEndpointAllocations);
+        return {
+          set: (values: { readonly updatedAt: string }) => {
+            claimedAt = values.updatedAt;
+            return {
+              where: () => ({
+                returning: () => Effect.succeed([{ userId: "user-1" }]),
+              }),
+            };
+          },
+        };
+      },
+    } as unknown as RelayDb.RelayDb["Service"];
+
+    return Effect.gen(function* () {
+      const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
+      const generation = yield* allocations.claimDeprovision({
+        userId: "user-1",
+        environmentId: "environment-1",
+        updatedAt: "captured-generation",
+      });
+
+      expect(generation).toBe(claimedAt);
+      expect(generation).not.toBeNull();
+    }).pipe(Effect.provide(layerWithDb(fakeDb)));
+  });
+
+  it.effect("does not remove an allocation superseded after a deprovision claim", () => {
+    const fakeDb = {
+      delete: (table: unknown) => {
+        expect(table).toBe(relayManagedEndpointAllocations);
+        return {
+          where: () => ({
+            returning: () => Effect.succeed([]),
+          }),
+        };
+      },
+    } as unknown as RelayDb.RelayDb["Service"];
+
+    return Effect.gen(function* () {
+      const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
+      expect(
+        yield* allocations.removeClaimed({
+          userId: "user-1",
+          environmentId: "environment-1",
+          updatedAt: "outdated-claim-generation",
+        }),
+      ).toBe(false);
+    }).pipe(Effect.provide(layerWithDb(fakeDb)));
+  });
+
   it.effect("retains database failures with allocation operation and identity", () => {
     const cause = new Error("database unavailable");
     const fakeDb = {

infra/relay/src/environments/ManagedEndpointAllocations.ts

@@ -51,7 +51,9 @@ export class ManagedEndpointAllocationPersistenceError extends Schema.TaggedErro
       "record-dns",
       "mark-ready",
       "claim-release",
+      "claim-deprovision",
       "remove",
+      "remove-claimed",
     ]),
     stage: Schema.Literals(["database-request", "resolve-reservation"]),
     userId: Schema.String,
@@ -91,6 +93,14 @@ interface ClaimManagedEndpointReleaseInput extends ManagedEndpointAllocationKey
   readonly updatedAt: string;
 }
 
+interface ClaimManagedEndpointDeprovisionInput extends ManagedEndpointAllocationKey {
+  readonly updatedAt: string;
+}
+
+interface RemoveClaimedManagedEndpointAllocationInput extends ManagedEndpointAllocationKey {
+  readonly updatedAt: string;
+}
+
 export class ManagedEndpointAllocations extends Context.Service<
   ManagedEndpointAllocations,
   {
@@ -119,9 +129,22 @@ export class ManagedEndpointAllocations extends Context.Service<
     readonly claimRelease: (
       input: ClaimManagedEndpointReleaseInput,
     ) => Effect.Effect<boolean, ManagedEndpointAllocationPersistenceError>;
+    /**
+     * Claims the complete allocation for teardown only if its generation still
+     * matches the snapshot captured by the unlink operation.
+     *
+     * Returns the claim generation used by `removeClaimed`, or null when a
+     * concurrent provision has already superseded the snapshot.
+     */
+    readonly claimDeprovision: (
+      input: ClaimManagedEndpointDeprovisionInput,
+    ) => Effect.Effect<string | null, ManagedEndpointAllocationPersistenceError>;
     readonly remove: (
       input: ManagedEndpointAllocationKey,
     ) => Effect.Effect<void, ManagedEndpointAllocationPersistenceError>;
+    readonly removeClaimed: (
+      input: RemoveClaimedManagedEndpointAllocationInput,
+    ) => Effect.Effect<boolean, ManagedEndpointAllocationPersistenceError>;
   }
 >()("t3code-relay/environments/ManagedEndpointAllocations") {}
 
@@ -321,6 +344,35 @@ export const make = Effect.gen(function* () {
         );
       return claimed;
     }),
+    claimDeprovision: Effect.fn("relay.managed_endpoint_allocations.claim_deprovision")(function* (
+      input: ClaimManagedEndpointDeprovisionInput,
+    ) {
+      const claimedAt = DateTime.formatIso(yield* DateTime.now);
+      const claimed = yield* db
+        .update(relayManagedEndpointAllocations)
+        .set({ updatedAt: claimedAt })
+        .where(
+          and(
+            whereAllocation(input),
+            eq(relayManagedEndpointAllocations.updatedAt, input.updatedAt),
+          ),
+        )
+        .returning({ userId: relayManagedEndpointAllocations.userId })
+        .pipe(
+          Effect.map((rows) => rows.length > 0),
+          Effect.mapError(
+            (cause) =>
+              new ManagedEndpointAllocationPersistenceError({
+                operation: "claim-deprovision",
+                stage: "database-request",
+                userId: input.userId,
+                environmentId: input.environmentId,
+                cause,
+              }),
+          ),
+        );
+      return claimed ? claimedAt : null;
+    }),
     remove: Effect.fn("relay.managed_endpoint_allocations.remove")(function* (
       input: ManagedEndpointAllocationKey,
     ) {
@@ -339,6 +391,32 @@ export const make = Effect.gen(function* () {
           ),
         );
     }),
+    removeClaimed: Effect.fn("relay.managed_endpoint_allocations.remove_claimed")(function* (
+      input: RemoveClaimedManagedEndpointAllocationInput,
+    ) {
+      return yield* db
+        .delete(relayManagedEndpointAllocations)
+        .where(
+          and(
+            whereAllocation(input),
+            eq(relayManagedEndpointAllocations.updatedAt, input.updatedAt),
+          ),
+        )
+        .returning({ userId: relayManagedEndpointAllocations.userId })
+        .pipe(
+          Effect.map((rows) => rows.length > 0),
+          Effect.mapError(
+            (cause) =>
+              new ManagedEndpointAllocationPersistenceError({
+                operation: "remove-claimed",
+                stage: "database-request",
+                userId: input.userId,
+                environmentId: input.environmentId,
+                cause,
+              }),
+          ),
+        );
+    }),
   });
 });
 

infra/relay/src/environments/ManagedEndpointProvider.test.ts

@@ -50,7 +50,9 @@ interface AllocationCall {
     | "recordDns"
     | "markReady"
     | "claimRelease"
-    | "remove";
+    | "claimDeprovision"
+    | "remove"
+    | "removeClaimed";
   readonly input: unknown;
 }
 
@@ -226,11 +228,31 @@ function makeAllocations(calls: AllocationCall[] = []) {
         mutate(allocationKey(input), (current) => current);
         return true;
       }),
+    claimDeprovision: (input) =>
+      Effect.sync(() => {
+        calls.push({ operation: "claimDeprovision", input });
+        const allocation = allocations.get(allocationKey(input));
+        if (allocation === undefined || allocation.updatedAt !== input.updatedAt) {
+          return null;
+        }
+        mutate(allocationKey(input), (current) => current);
+        return allocations.get(allocationKey(input))?.updatedAt ?? null;
+      }),
     remove: (input) =>
       Effect.sync(() => {
         calls.push({ operation: "remove", input });
         allocations.delete(allocationKey(input));
       }),
+    removeClaimed: (input) =>
+      Effect.sync(() => {
+        calls.push({ operation: "removeClaimed", input });
+        const allocation = allocations.get(allocationKey(input));
+        if (allocation === undefined || allocation.updatedAt !== input.updatedAt) {
+          return false;
+        }
+        allocations.delete(allocationKey(input));
+        return true;
+      }),
   });
 }
 
@@ -774,12 +796,54 @@ describe("ManagedEndpointProvider", () => {
           "recordDns",
           "markReady",
           "get",
-          "remove",
+          "claimDeprovision",
+          "removeClaimed",
         ]);
       }).pipe(Effect.provide(layer));
     },
   );
 
+  it.effect("does not deprovision an allocation superseded by a concurrent relink", () => {
+    const tunnelCalls: TunnelCall[] = [];
+    const dnsCalls: DnsCall[] = [];
+    const allocationCalls: AllocationCall[] = [];
+    const layer = providerLayer(
+      makePersistentTunnelClient(tunnelCalls),
+      makeDnsClient(dnsCalls),
+      makeAllocations(allocationCalls),
+    );
+
+    return Effect.gen(function* () {
+      const provider = yield* ManagedEndpointProvider.ManagedEndpointProvider;
+      const key = { userId: "user_ABC", environmentId: "env_ABC" } as const;
+      const request = {
+        ...key,
+        origin: { localHttpHost: "127.0.0.1", localHttpPort: 3773 },
+      } as const;
+      yield* provider.provision(request);
+      const unlinkTarget = yield* provider.prepareDeprovision(key);
+      expect(unlinkTarget).not.toBeNull();
+      if (unlinkTarget === null) {
+        return;
+      }
+
+      // A relink refreshes the allocation generation after unlink captured its
+      // target but before unlink begins external teardown.
+      yield* provider.provision(request);
+      const tunnelCallCount = tunnelCalls.length;
+      const dnsCallCount = dnsCalls.length;
+      const allocationCallCount = allocationCalls.length;
+
+      yield* provider.deprovision({ ...key, target: unlinkTarget });
+
+      expect(tunnelCalls).toHaveLength(tunnelCallCount);
+      expect(dnsCalls).toHaveLength(dnsCallCount);
+      expect(allocationCalls.slice(allocationCallCount).map((call) => call.operation)).toEqual([
+        "claimDeprovision",
+      ]);
+    }).pipe(Effect.provide(layer));
+  });
+
   it.effect("releases the tunnel while keeping the allocation, DNS record, and hostname", () => {
     const tunnelCalls: TunnelCall[] = [];
     const dnsCalls: DnsCall[] = [];
@@ -1004,8 +1068,10 @@ describe("ManagedEndpointProvider", () => {
         "recordDns",
         "markReady",
         "get",
+        "claimDeprovision",
         "get",
-        "remove",
+        "claimDeprovision",
+        "removeClaimed",
       ]);
     }).pipe(Effect.provide(layer));
   });
@@ -1046,7 +1112,7 @@ describe("ManagedEndpointProvider", () => {
       });
       yield* provider.deprovision(key);
 
-      expect(allocationCalls.map((call) => call.operation)).toContain("remove");
+      expect(allocationCalls.map((call) => call.operation)).toContain("removeClaimed");
     }).pipe(Effect.provide(layer));
   });