1
import type { CacheContext } from "@/worker/cache";
2
 
3
// Global caps
4
export const MAX_SIMULTANEOUS_CONNECTIONS = 6; // Cloudflare per-request connection limit
5
export const DEFAULT_SUBREQUEST_BUDGET = 900; // soft budget before hard cap (~1000)
6
 
7
// Structural limiter type compatible with RequestMemo.limiter
8
export type Limiter = { run<T>(label: string, fn: () => Promise<T>): Promise<T> };
9
 
10
// Lightweight semaphore to cap concurrent upstream calls per request
11
export class SubrequestLimiter {
12
  private max: number;
13
  private cur = 0;
14
  private queue: Array<() => void> = [];
15
 
16
  constructor(max: number) {
17
    this.max = Math.max(1, max | 0);
18
  }
19
 
20
  private acquire(): Promise<void> {
21
    if (this.cur < this.max) {
22
      this.cur++;
23
      return Promise.resolve();
24
    }
25
    return new Promise((resolve) => {
26
      this.queue.push(() => {
27
        this.cur++;
28
        resolve();
29
      });
30
    });
31
  }
32
 
33
  private release() {
34
    this.cur--;
35
    if (this.cur < 0) this.cur = 0;
36
    const next = this.queue.shift();
37
    if (next) next();
38
  }
39
 
40
  async run<T>(_label: string, fn: () => Promise<T>): Promise<T> {
41
    await this.acquire();
42
    try {
43
      return await fn();
44
    } finally {
45
      this.release();
46
    }
47
  }
48
}
49
 
50
export function getLimiter(cacheCtx?: CacheContext): Limiter {
51
  const fallback: Limiter = new SubrequestLimiter(MAX_SIMULTANEOUS_CONNECTIONS);
52
  if (!cacheCtx) return fallback as Limiter;
53
  cacheCtx.memo = cacheCtx.memo || {};
54
  if (!cacheCtx.memo.limiter) {
55
    cacheCtx.memo.limiter = new SubrequestLimiter(MAX_SIMULTANEOUS_CONNECTIONS);
56
  }
57
  return cacheCtx.memo.limiter as Limiter;
58
}
59
 
60
// Decrement subrequest soft budget; returns false when budget exhausted
61
export function countSubrequest(cacheCtx?: CacheContext, n = 1): boolean {
62
  if (!cacheCtx) return true; // nothing to track
63
  cacheCtx.memo = cacheCtx.memo || {};
64
  const cur = cacheCtx.memo.subreqBudget ?? DEFAULT_SUBREQUEST_BUDGET;
65
  const next = cur - Math.max(1, n);
66
  cacheCtx.memo.subreqBudget = next;
67
  return next >= 0;
68
}