1import type { RelayManagedEndpoint, RelayManagedEndpointOrigin } from "@t3tools/contracts/relay";
2import { and, eq, exists, inArray, isNull, sql } from "drizzle-orm";
3import { QueryBuilder } from "drizzle-orm/pg-core";
4import * as Context from "effect/Context";
5import * as DateTime from "effect/DateTime";
6import * as Effect from "effect/Effect";
7import * as Layer from "effect/Layer";
8import * as Option from "effect/Option";
9import * as Schema from "effect/Schema";
10import * as SqlError from "effect/sql/SqlError";
11
12import * as RelayDb from "../db.ts";
13import { isManagedEndpointHostname, managedEndpointForHostname } from "../deploymentConfig.ts";
14import { relayEnvironmentLinks, relayManagedEndpointAllocations } from "../persistence/schema.ts";
15
16export interface ManagedEndpointAllocation {
17 readonly userId: string;
18 readonly environmentId: string;
19 readonly hostname: string;
20 readonly tunnelId: string | null;
21 readonly tunnelName: string;
22 readonly dnsRecordId: string | null;
23 readonly readyAt: string | null;
24 readonly origin: RelayManagedEndpointOrigin | null;
25 readonly updatedAt: string;
26 readonly generation: number;
27 readonly tunnelReleasedAt: string | null;
28}
29
30export interface ManagedEndpointTunnelAllocation extends ManagedEndpointAllocation {
31 readonly recoveryEnabled: boolean;
32}
33
34const MANAGED_ENDPOINT_ALLOCATION_LOOKUP_BATCH_SIZE = 500;
35
36export function resolveReadyManagedEndpoint(input: {
37 readonly allocation: ManagedEndpointAllocation;
38 readonly baseDomain: string | undefined;
39}): RelayManagedEndpoint | null {
40 if (
41 !input.baseDomain ||
42 input.allocation.readyAt === null ||
43 input.allocation.tunnelId === null ||
44 input.allocation.dnsRecordId === null ||
45 !isManagedEndpointHostname(input.allocation.hostname, input.baseDomain)
46 ) {
47 return null;
48 }
49 return managedEndpointForHostname(input.allocation.hostname);
50}
51
52export class ManagedEndpointAllocationPersistenceError extends Schema.TaggedError<ManagedEndpointAllocationPersistenceError>()(
53 "ManagedEndpointAllocationPersistenceError",
54 {
55 operation: Schema.Literals([
56 "get",
57 "reserve",
58 "record-tunnel",
59 "record-dns",
60 "mark-ready",
61 "enable-recovery",
62 "list-tunnels",
63 "lock-tunnel",
64 "claim-release",
65 "claim-deprovision",
66 "remove-claimed",
67 "get-by-tunnel-name",
68 ]),
69 stage: Schema.Literals(["database-request", "resolve-reservation"]),
70 userId: Schema.String,
71 environmentId: Schema.String,
72 hostname: Schema.optionalKey(Schema.String),
73 tunnelName: Schema.optionalKey(Schema.String),
74 tunnelId: Schema.optionalKey(Schema.String),
75 dnsRecordId: Schema.optionalKey(Schema.String),
76 cause: Schema.optional(Schema.Defect()),
77 },
78) {
79 override get message(): string {
80 return `Managed endpoint allocation '${this.operation}' failed during '${this.stage}' for user '${this.userId}', environment '${this.environmentId}'`;
81 }
82}
83
84interface ManagedEndpointAllocationKey {
85 readonly userId: string;
86 readonly environmentId: string;
87}
88
89interface ReserveManagedEndpointAllocationInput extends ManagedEndpointAllocationKey {
90 readonly hostname: string;
91 readonly tunnelName: string;
92}
93
94interface RecordManagedEndpointTunnelInput extends ManagedEndpointAllocationKey {
95 readonly tunnelId: string;
96 readonly generation: number;
97}
98
99interface RecordManagedEndpointDnsInput extends ManagedEndpointAllocationKey {
100 readonly dnsRecordId: string;
101 readonly tunnelId: string;
102 readonly generation: number;
103}
104
105interface MarkManagedEndpointReadyInput extends ManagedEndpointAllocationKey {
106 readonly tunnelId: string;
107 readonly generation: number;
108 readonly origin: RelayManagedEndpointOrigin;
109}
110
111interface ClaimManagedEndpointReleaseInput extends ManagedEndpointAllocationKey {
112 readonly tunnelId: string;
113 readonly generation: number;
114 readonly markReleased?: boolean;
115}
116
117interface EnableManagedEndpointRecoveryInput extends ManagedEndpointAllocationKey {
118 readonly tunnelId: string;
119 readonly environmentPublicKey: string;
120 readonly origin: RelayManagedEndpointOrigin;
121}
122
123interface ClaimManagedEndpointDeprovisionInput extends ManagedEndpointAllocationKey {
124 readonly generation: number;
125}
126
127interface RemoveClaimedManagedEndpointAllocationInput extends ManagedEndpointAllocationKey {
128 readonly generation: number;
129}
130
131export class ManagedEndpointAllocations extends Context.Service<
132 ManagedEndpointAllocations,
133 {
134 readonly get: (
135 input: ManagedEndpointAllocationKey,
136 ) => Effect.Effect<ManagedEndpointAllocation | null, ManagedEndpointAllocationPersistenceError>;
137 readonly getByTunnelName: (
138 tunnelName: string,
139 ) => Effect.Effect<ManagedEndpointAllocation | null, ManagedEndpointAllocationPersistenceError>;
140 readonly reserve: (
141 input: ReserveManagedEndpointAllocationInput,
142 ) => Effect.Effect<ManagedEndpointAllocation, ManagedEndpointAllocationPersistenceError>;
143 readonly recordTunnel: (
144 input: RecordManagedEndpointTunnelInput,
145 ) => Effect.Effect<number | null, ManagedEndpointAllocationPersistenceError>;
146 readonly recordDns: (
147 input: RecordManagedEndpointDnsInput,
148 ) => Effect.Effect<number | null, ManagedEndpointAllocationPersistenceError>;
149 readonly markReady: (
150 input: MarkManagedEndpointReadyInput,
151 ) => Effect.Effect<boolean, ManagedEndpointAllocationPersistenceError>;
152 readonly enableRecovery: (
153 input: EnableManagedEndpointRecoveryInput,
154 ) => Effect.Effect<boolean, ManagedEndpointAllocationPersistenceError>;
155 readonly listByTunnelNames: (
156 tunnelNames: ReadonlyArray<string>,
157 ) => Effect.Effect<
158 ReadonlyArray<ManagedEndpointTunnelAllocation>,
159 ManagedEndpointAllocationPersistenceError
160 >;
161 /** Claim deletion only while tunnel and generation match, so concurrent provisioning keeps its new tunnel. */
162 readonly claimRelease: (
163 input: ClaimManagedEndpointReleaseInput,
164 ) => Effect.Effect<number | null, ManagedEndpointAllocationPersistenceError>;
165 readonly withClaimedTunnel: <A, E, R>(
166 input: ClaimManagedEndpointReleaseInput,
167 effect: Effect.Effect<A, E, R>,
168 ) => Effect.Effect<Option.Option<A>, E | ManagedEndpointAllocationPersistenceError, R>;
169 /** Return a teardown claim generation, or null if concurrent provisioning superseded the captured snapshot. */
170 readonly claimDeprovision: (
171 input: ClaimManagedEndpointDeprovisionInput,
172 ) => Effect.Effect<number | null, ManagedEndpointAllocationPersistenceError>;
173 readonly removeClaimed: (
174 input: RemoveClaimedManagedEndpointAllocationInput,
175 ) => Effect.Effect<boolean, ManagedEndpointAllocationPersistenceError>;
176 }
177>()("t3code-relay/environments/ManagedEndpointAllocations") {}
178
179const allocationSelection = {
180 userId: relayManagedEndpointAllocations.userId,
181 environmentId: relayManagedEndpointAllocations.environmentId,
182 hostname: relayManagedEndpointAllocations.hostname,
183 tunnelId: relayManagedEndpointAllocations.tunnelId,
184 tunnelName: relayManagedEndpointAllocations.tunnelName,
185 dnsRecordId: relayManagedEndpointAllocations.dnsRecordId,
186 readyAt: relayManagedEndpointAllocations.readyAt,
187 origin: relayManagedEndpointAllocations.origin,
188 updatedAt: relayManagedEndpointAllocations.updatedAt,
189 generation: relayManagedEndpointAllocations.generation,
190 tunnelReleasedAt: relayManagedEndpointAllocations.tunnelReleasedAt,
191};
192
193const whereAllocation = (input: ManagedEndpointAllocationKey) =>
194 and(
195 eq(relayManagedEndpointAllocations.userId, input.userId),
196 eq(relayManagedEndpointAllocations.environmentId, input.environmentId),
197 );
198
199const make = Effect.gen(function* () {
200 const db = yield* RelayDb.RelayDb;
201
202 return ManagedEndpointAllocations.of({
203 get: Effect.fn("relay.managed_endpoint_allocations.get")(function* (
204 input: ManagedEndpointAllocationKey,
205 ) {
206 return yield* db
207 .select(allocationSelection)
208 .from(relayManagedEndpointAllocations)
209 .where(whereAllocation(input))
210 .limit(1)
211 .pipe(
212 Effect.map((rows) => rows[0] ?? null),
213 Effect.mapError(
214 (cause) =>
215 new ManagedEndpointAllocationPersistenceError({
216 operation: "get",
217 stage: "database-request",
218 ...input,
219 cause,
220 }),
221 ),
222 );
223 }),
224 getByTunnelName: Effect.fn("relay.managed_endpoint_allocations.get_by_tunnel_name")(function* (
225 tunnelName: string,
226 ) {
227 return yield* db
228 .select(allocationSelection)
229 .from(relayManagedEndpointAllocations)
230 .where(eq(relayManagedEndpointAllocations.tunnelName, tunnelName))
231 .limit(1)
232 .pipe(
233 Effect.map((rows) => rows[0] ?? null),
234 Effect.mapError(
235 (cause) =>
236 new ManagedEndpointAllocationPersistenceError({
237 operation: "get-by-tunnel-name",
238 stage: "database-request",
239 userId: "",
240 environmentId: "",
241 tunnelName,
242 cause,
243 }),
244 ),
245 );
246 }),
247 reserve: Effect.fn("relay.managed_endpoint_allocations.reserve")(function* (
248 input: ReserveManagedEndpointAllocationInput,
249 ) {
250 const now = DateTime.formatIso(yield* DateTime.now);
251 const inserted = yield* db
252 .insert(relayManagedEndpointAllocations)
253 .values({
254 ...input,
255 createdAt: now,
256 updatedAt: now,
257 })
258 // Refreshing updatedAt makes a re-reserved row count toward the tunnel limit again.
259 .onConflictDoUpdate({
260 target: [
261 relayManagedEndpointAllocations.userId,
262 relayManagedEndpointAllocations.environmentId,
263 ],
264 set: { updatedAt: now },
265 })
266 .returning(allocationSelection)
267 .pipe(
268 Effect.mapError(
269 (cause) =>
270 new ManagedEndpointAllocationPersistenceError({
271 operation: "reserve",
272 stage: "database-request",
273 ...input,
274 cause,
275 }),
276 ),
277 );
278
279 const allocation = inserted[0];
280 if (allocation === undefined) {
281 return yield* new ManagedEndpointAllocationPersistenceError({
282 operation: "reserve",
283 stage: "resolve-reservation",
284 ...input,
285 });
286 }
287
288 return allocation;
289 }),
290 recordTunnel: Effect.fn("relay.managed_endpoint_allocations.record_tunnel")(function* (
291 input: RecordManagedEndpointTunnelInput,
292 ) {
293 return yield* db
294 .update(relayManagedEndpointAllocations)
295 .set({
296 tunnelId: input.tunnelId,
297 readyAt: sql`case when ${relayManagedEndpointAllocations.tunnelId} = ${input.tunnelId} then ${relayManagedEndpointAllocations.readyAt} else null end`,
298 origin: sql`case when ${relayManagedEndpointAllocations.tunnelId} = ${input.tunnelId} then ${relayManagedEndpointAllocations.origin} else null end`,
299 // Replacement tunnels must re-register recovery before becoming reapable.
300 recoveryEnabledAt: sql`case when ${relayManagedEndpointAllocations.tunnelId} = ${input.tunnelId} then ${relayManagedEndpointAllocations.recoveryEnabledAt} else null end`,
301 recoveryEnvironmentPublicKey: sql`case when ${relayManagedEndpointAllocations.tunnelId} = ${input.tunnelId} then ${relayManagedEndpointAllocations.recoveryEnvironmentPublicKey} else null end`,
302 tunnelReleasedAt: null,
303 updatedAt: DateTime.formatIso(yield* DateTime.now),
304 generation: sql`${relayManagedEndpointAllocations.generation} + 1`,
305 })
306 .where(
307 and(
308 whereAllocation(input),
309 eq(relayManagedEndpointAllocations.generation, input.generation),
310 ),
311 )
312 .returning({ generation: relayManagedEndpointAllocations.generation })
313 .pipe(
314 Effect.map((rows) => rows[0]?.generation ?? null),
315 Effect.mapError(
316 (cause) =>
317 new ManagedEndpointAllocationPersistenceError({
318 operation: "record-tunnel",
319 stage: "database-request",
320 ...input,
321 cause,
322 }),
323 ),
324 );
325 }),
326 recordDns: Effect.fn("relay.managed_endpoint_allocations.record_dns")(function* (
327 input: RecordManagedEndpointDnsInput,
328 ) {
329 return yield* db
330 .update(relayManagedEndpointAllocations)
331 .set({
332 dnsRecordId: input.dnsRecordId,
333 updatedAt: DateTime.formatIso(yield* DateTime.now),
334 generation: sql`${relayManagedEndpointAllocations.generation} + 1`,
335 })
336 .where(
337 and(
338 whereAllocation(input),
339 eq(relayManagedEndpointAllocations.tunnelId, input.tunnelId),
340 eq(relayManagedEndpointAllocations.generation, input.generation),
341 ),
342 )
343 .returning({ generation: relayManagedEndpointAllocations.generation })
344 .pipe(
345 Effect.map((rows) => rows[0]?.generation ?? null),
346 Effect.mapError(
347 (cause) =>
348 new ManagedEndpointAllocationPersistenceError({
349 operation: "record-dns",
350 stage: "database-request",
351 ...input,
352 cause,
353 }),
354 ),
355 );
356 }),
357 markReady: Effect.fn("relay.managed_endpoint_allocations.mark_ready")(function* (
358 input: MarkManagedEndpointReadyInput,
359 ) {
360 const now = DateTime.formatIso(yield* DateTime.now);
361 return yield* db
362 .update(relayManagedEndpointAllocations)
363 .set({
364 readyAt: now,
365 origin: input.origin,
366 updatedAt: now,
367 generation: sql`${relayManagedEndpointAllocations.generation} + 1`,
368 })
369 .where(
370 and(
371 whereAllocation(input),
372 eq(relayManagedEndpointAllocations.tunnelId, input.tunnelId),
373 eq(relayManagedEndpointAllocations.generation, input.generation),
374 ),
375 )
376 .returning({ environmentId: relayManagedEndpointAllocations.environmentId })
377 .pipe(
378 Effect.map((rows) => rows.length > 0),
379 Effect.mapError(
380 (cause) =>
381 new ManagedEndpointAllocationPersistenceError({
382 operation: "mark-ready",
383 stage: "database-request",
384 ...input,
385 cause,
386 }),
387 ),
388 );
389 }),
390 enableRecovery: Effect.fn("relay.managed_endpoint_allocations.enable_recovery")(function* (
391 input: EnableManagedEndpointRecoveryInput,
392 ) {
393 const now = DateTime.formatIso(yield* DateTime.now);
394 return yield* db
395 .update(relayManagedEndpointAllocations)
396 .set({
397 recoveryEnabledAt: now,
398 recoveryEnvironmentPublicKey: input.environmentPublicKey,
399 updatedAt: now,
400 generation: sql`${relayManagedEndpointAllocations.generation} + 1`,
401 })
402 .where(
403 and(
404 whereAllocation(input),
405 eq(relayManagedEndpointAllocations.tunnelId, input.tunnelId),
406 eq(relayManagedEndpointAllocations.origin, input.origin),
407 exists(
408 new QueryBuilder()
409 .select({ userId: relayEnvironmentLinks.userId })
410 .from(relayEnvironmentLinks)
411 .where(
412 and(
413 eq(relayEnvironmentLinks.userId, input.userId),
414 eq(relayEnvironmentLinks.environmentId, input.environmentId),
415 eq(relayEnvironmentLinks.environmentPublicKey, input.environmentPublicKey),
416 eq(relayEnvironmentLinks.endpointProviderKind, "cloudflare_tunnel"),
417 isNull(relayEnvironmentLinks.revokedAt),
418 ),
419 )
420 .for("update"),
421 ),
422 ),
423 )
424 .returning({ environmentId: relayManagedEndpointAllocations.environmentId })
425 .pipe(
426 Effect.map((rows) => rows.length > 0),
427 Effect.mapError(
428 (cause) =>
429 new ManagedEndpointAllocationPersistenceError({
430 operation: "enable-recovery",
431 stage: "database-request",
432 ...input,
433 cause,
434 }),
435 ),
436 );
437 }),
438 listByTunnelNames: Effect.fn("relay.managed_endpoint_allocations.list_by_tunnel_names")(
439 function* (tunnelNames: ReadonlyArray<string>) {
440 if (tunnelNames.length === 0) {
441 return [];
442 }
443 const batches = Array.from(
444 { length: Math.ceil(tunnelNames.length / MANAGED_ENDPOINT_ALLOCATION_LOOKUP_BATCH_SIZE) },
445 (_, index) =>
446 tunnelNames.slice(
447 index * MANAGED_ENDPOINT_ALLOCATION_LOOKUP_BATCH_SIZE,
448 (index + 1) * MANAGED_ENDPOINT_ALLOCATION_LOOKUP_BATCH_SIZE,
449 ),
450 );
451 const results = yield* Effect.forEach(
452 batches,
453 (batch) =>
454 db
455 .select({
456 ...allocationSelection,
457 recoveryEnabledAt: relayManagedEndpointAllocations.recoveryEnabledAt,
458 recoveryEnvironmentPublicKey:
459 relayManagedEndpointAllocations.recoveryEnvironmentPublicKey,
460 linkedEnvironmentPublicKey: relayEnvironmentLinks.environmentPublicKey,
461 })
462 .from(relayManagedEndpointAllocations)
463 .leftJoin(
464 relayEnvironmentLinks,
465 and(
466 eq(relayEnvironmentLinks.userId, relayManagedEndpointAllocations.userId),
467 eq(
468 relayEnvironmentLinks.environmentId,
469 relayManagedEndpointAllocations.environmentId,
470 ),
471 isNull(relayEnvironmentLinks.revokedAt),
472 ),
473 )
474 .where(inArray(relayManagedEndpointAllocations.tunnelName, batch))
475 .pipe(
476 Effect.map((rows) =>
477 rows.map(
478 ({
479 recoveryEnabledAt,
480 recoveryEnvironmentPublicKey,
481 linkedEnvironmentPublicKey,
482 ...allocation
483 }) => ({
484 ...allocation,
485 recoveryEnabled:
486 recoveryEnabledAt !== null &&
487 recoveryEnvironmentPublicKey !== null &&
488 recoveryEnvironmentPublicKey === linkedEnvironmentPublicKey,
489 }),
490 ),
491 ),
492 Effect.mapError(
493 (cause) =>
494 new ManagedEndpointAllocationPersistenceError({
495 operation: "list-tunnels",
496 stage: "database-request",
497 userId: "*",
498 environmentId: "*",
499 cause,
500 }),
501 ),
502 ),
503 { concurrency: 1 },
504 );
505 return results.flat();
506 },
507 ),
508 claimRelease: Effect.fn("relay.managed_endpoint_allocations.claim_release")(function* (
509 input: ClaimManagedEndpointReleaseInput,
510 ) {
511 const now = DateTime.formatIso(yield* DateTime.now);
512 const claimed = yield* db
513 .update(relayManagedEndpointAllocations)
514 .set({
515 generation: sql`${relayManagedEndpointAllocations.generation} + 1`,
516 ...(input.markReleased === true ? { tunnelReleasedAt: now, updatedAt: now } : {}),
517 })
518 .where(
519 and(
520 whereAllocation(input),
521 eq(relayManagedEndpointAllocations.tunnelId, input.tunnelId),
522 eq(relayManagedEndpointAllocations.generation, input.generation),
523 ),
524 )
525 .returning({ generation: relayManagedEndpointAllocations.generation })
526 .pipe(
527 Effect.map((rows) => rows[0]?.generation ?? null),
528 Effect.mapError(
529 (cause) =>
530 new ManagedEndpointAllocationPersistenceError({
531 operation: "claim-release",
532 stage: "database-request",
533 userId: input.userId,
534 environmentId: input.environmentId,
535 tunnelId: input.tunnelId,
536 cause,
537 }),
538 ),
539 );
540 return claimed;
541 }),
542 withClaimedTunnel: Effect.fn("relay.managed_endpoint_allocations.with_claimed_tunnel")(
543 function* <A, E, R>(
544 input: ClaimManagedEndpointReleaseInput,
545 effect: Effect.Effect<A, E, R>,
546 ): Effect.fn.Return<Option.Option<A>, E | ManagedEndpointAllocationPersistenceError, R> {
547 const lockError = (cause: unknown) =>
548 new ManagedEndpointAllocationPersistenceError({
549 operation: "lock-tunnel",
550 stage: "database-request",
551 userId: input.userId,
552 environmentId: input.environmentId,
553 tunnelId: input.tunnelId,
554 cause,
555 });
556 return yield* db.$client
557 .withTransaction(
558 db
559 .select({ generation: relayManagedEndpointAllocations.generation })
560 .from(relayManagedEndpointAllocations)
561 .where(
562 and(
563 whereAllocation(input),
564 eq(relayManagedEndpointAllocations.tunnelId, input.tunnelId),
565 eq(relayManagedEndpointAllocations.generation, input.generation),
566 ),
567 )
568 .limit(1)
569 .for("update")
570 .pipe(
571 Effect.mapError(lockError),
572 Effect.flatMap((rows) =>
573 rows.length === 0 ? Effect.succeedNone : Effect.asSome(effect),
574 ),
575 ),
576 )
577 .pipe(
578 Effect.mapError((cause) => (SqlError.isSqlError(cause) ? lockError(cause) : cause)),
579 );
580 },
581 ),
582 claimDeprovision: Effect.fn("relay.managed_endpoint_allocations.claim_deprovision")(function* (
583 input: ClaimManagedEndpointDeprovisionInput,
584 ) {
585 const claimed = yield* db
586 .update(relayManagedEndpointAllocations)
587 // Claims move only the generation: the tunnel limit reads updatedAt as the last provisioning write.
588 .set({ generation: sql`${relayManagedEndpointAllocations.generation} + 1` })
589 .where(
590 and(
591 whereAllocation(input),
592 eq(relayManagedEndpointAllocations.generation, input.generation),
593 ),
594 )
595 .returning({ generation: relayManagedEndpointAllocations.generation })
596 .pipe(
597 Effect.map((rows) => rows[0]?.generation ?? null),
598 Effect.mapError(
599 (cause) =>
600 new ManagedEndpointAllocationPersistenceError({
601 operation: "claim-deprovision",
602 stage: "database-request",
603 userId: input.userId,
604 environmentId: input.environmentId,
605 cause,
606 }),
607 ),
608 );
609 return claimed;
610 }),
611 removeClaimed: Effect.fn("relay.managed_endpoint_allocations.remove_claimed")(function* (
612 input: RemoveClaimedManagedEndpointAllocationInput,
613 ) {
614 return yield* db
615 .delete(relayManagedEndpointAllocations)
616 .where(
617 and(
618 whereAllocation(input),
619 eq(relayManagedEndpointAllocations.generation, input.generation),
620 ),
621 )
622 .returning({ userId: relayManagedEndpointAllocations.userId })
623 .pipe(
624 Effect.map((rows) => rows.length > 0),
625 Effect.mapError(
626 (cause) =>
627 new ManagedEndpointAllocationPersistenceError({
628 operation: "remove-claimed",
629 stage: "database-request",
630 userId: input.userId,
631 environmentId: input.environmentId,
632 cause,
633 }),
634 ),
635 );
636 }),
637 });
638});
639
640export const layer = Layer.effect(ManagedEndpointAllocations, make);