1
import type { CacheContext } from "@/worker/cache";
2
import type { IdxView } from "./types";
3
 
4
import { bytesToHex, createLogger, hexToBytes } from "@/worker/common";
5
import { packIndexKey } from "@/worker/keys";
6
import { countSubrequest, getLimiter } from "@/worker/git/operations/limits";
7
 
8
const IDX_VIEW_CACHE_MAX_BYTES = 16 * 1024 * 1024;
9
const IDX_TRAILER_BYTES = 40;
10
const UINT32_SPAN = 0x1_0000_0000;
11
type CachedIdxView = {
12
  view: IdxView;
13
  bytes: number;
14
};
15
 
16
const idxViewCache = new Map<string, CachedIdxView>();
17
let idxViewCacheBytes = 0;
18
 
19
function getIdxViewCacheKey(packKey: string, packSize: number): string {
20
  // Hinted loads can be wrong if a stale catalog row slips through. Key the
21
  // isolate-shared cache by both pack key and pack size so a bad hint can only
22
  // waste cache space, not poison the correct entry for later requests.
23
  return `${packKey}\0${packSize}`;
24
}
25
 
26
function estimateIdxViewBytes(view: IdxView): number {
27
  return (
28
    view.fanout.byteLength +
29
    view.rawNames.byteLength +
30
    view.offsets.byteLength +
31
    view.nextOffsetByIndex.byteLength +
32
    view.sortedOffsets.byteLength +
33
    view.sortedOffsetIndices.byteLength +
34
    view.packChecksum.byteLength +
35
    view.idxChecksum.byteLength
36
  );
37
}
38
 
39
function touchIdxViewCache(packKey: string, packSize: number, value: IdxView) {
40
  const key = getIdxViewCacheKey(packKey, packSize);
41
  const existing = idxViewCache.get(key);
42
  if (existing) {
43
    idxViewCache.delete(key);
44
    idxViewCacheBytes -= existing.bytes;
45
  }
46
 
47
  const bytes = estimateIdxViewBytes(value);
48
  // The request-local memo already pins hot idx views for the lifetime of one
49
  // request. The global cache is only a cross-request accelerator, so cap it by
50
  // bytes instead of entry count to avoid a few large idx files crowding out the
51
  // worker heap.
52
  if (bytes > IDX_VIEW_CACHE_MAX_BYTES) return;
53
 
54
  idxViewCache.set(key, { view: value, bytes });
55
  idxViewCacheBytes += bytes;
56
 
57
  while (idxViewCacheBytes > IDX_VIEW_CACHE_MAX_BYTES) {
58
    const firstKey = idxViewCache.keys().next().value;
59
    if (!firstKey) break;
60
    const first = idxViewCache.get(firstKey);
61
    idxViewCache.delete(firstKey);
62
    idxViewCacheBytes -= first?.bytes ?? 0;
63
  }
64
}
65
 
66
function logOnce(cacheCtx: CacheContext | undefined, flag: string, fn: () => void) {
67
  if (!cacheCtx) {
68
    fn();
69
    return;
70
  }
71
  cacheCtx.memo = cacheCtx.memo || {};
72
  cacheCtx.memo.flags = cacheCtx.memo.flags || new Set<string>();
73
  if (cacheCtx.memo.flags.has(flag)) return;
74
  fn();
75
  cacheCtx.memo.flags.add(flag);
76
}
77
 
78
export function getOidHexAt(view: IdxView, index: number): string {
79
  const start = index * 20;
80
  return bytesToHex(view.rawNames.subarray(start, start + 20));
81
}
82
 
83
function compareOidAt(view: IdxView, index: number, needle: Uint8Array, needleStart = 0): number {
84
  const start = index * 20;
85
  for (let i = 0; i < 20; i++) {
86
    const diff = view.rawNames[start + i] - needle[needleStart + i];
87
    if (diff !== 0) return diff;
88
  }
89
  return 0;
90
}
91
 
92
function readUint64AsNumber(dv: DataView, pos: number): number {
93
  const hi = dv.getUint32(pos, false);
94
  const lo = dv.getUint32(pos + 4, false);
95
  const value = hi * UINT32_SPAN + lo;
96
  if (!Number.isSafeInteger(value)) {
97
    throw new Error(
98
      `idx: 64-bit offset 0x${hi.toString(16)}${lo.toString(16).padStart(8, "0")} exceeds safe integer support`
99
    );
100
  }
101
  return value;
102
}
103
 
104
/**
105
 * Binary search for the entry index at a given pack byte offset.
106
 * Returns the entry index, or undefined if not found.
107
 */
108
export function findOffsetIndex(view: IdxView, offset: number): number | undefined {
109
  const arr = view.sortedOffsets;
110
  let lo = 0;
111
  let hi = arr.length - 1;
112
  while (lo <= hi) {
113
    const mid = (lo + hi) >>> 1;
114
    const v = arr[mid];
115
    if (v === offset) return view.sortedOffsetIndices[mid];
116
    if (v < offset) lo = mid + 1;
117
    else hi = mid - 1;
118
  }
119
  return undefined;
120
}
121
 
122
/**
123
 * Get the next pack offset after the given offset (the start of the next entry),
124
 * or `packSize - 20` for the last entry. Returns undefined if offset is not found.
125
 */
126
export function getNextOffset(view: IdxView, offset: number): number | undefined {
127
  const index = findOffsetIndex(view, offset);
128
  return index === undefined ? undefined : view.nextOffsetByIndex[index];
129
}
130
 
131
export function getNextOffsetByIndex(view: IdxView, index: number): number | undefined {
132
  if (index < 0 || index >= view.count) return undefined;
133
  return view.nextOffsetByIndex[index];
134
}
135
 
136
export function findOidIndexFromBytes(view: IdxView, needle: Uint8Array, needleStart = 0): number {
137
  if (needleStart < 0 || needleStart + 20 > needle.byteLength) return -1;
138
 
139
  const first = needle[needleStart];
140
  const lo = first === 0 ? 0 : view.fanout[first - 1] || 0;
141
  const hi = (view.fanout[first] || 0) - 1;
142
  if (hi < lo) return -1;
143
 
144
  let left = lo;
145
  let right = hi;
146
  while (left <= right) {
147
    const mid = (left + right) >> 1;
148
    const cmp = compareOidAt(view, mid, needle, needleStart);
149
    if (cmp === 0) return mid;
150
    if (cmp < 0) left = mid + 1;
151
    else right = mid - 1;
152
  }
153
  return -1;
154
}
155
 
156
export function findOidIndex(view: IdxView, oid: string | Uint8Array): number {
157
  const needle = typeof oid === "string" ? hexToBytes(oid.toLowerCase()) : oid;
158
  if (needle.length !== 20) return -1;
159
  return findOidIndexFromBytes(view, needle);
160
}
161
 
162
export function parseIdxView(
163
  packKey: string,
164
  idxBuf: Uint8Array,
165
  packSize: number
166
): IdxView | undefined {
167
  if (!Number.isSafeInteger(packSize) || packSize < 20) {
168
    throw new Error(`idx: unsupported pack size ${packSize}`);
169
  }
170
  if (idxBuf.byteLength < 8 + 256 * 4 + IDX_TRAILER_BYTES) return undefined;
171
  if (!(idxBuf[0] === 0xff && idxBuf[1] === 0x74 && idxBuf[2] === 0x4f && idxBuf[3] === 0x63)) {
172
    return undefined;
173
  }
174
  const dv = new DataView(idxBuf.buffer, idxBuf.byteOffset, idxBuf.byteLength);
175
  const version = dv.getUint32(4, false);
176
  if (version !== 2 && version !== 3) return undefined;
177
 
178
  const fanout = new Uint32Array(256);
179
  let pos = 8;
180
  for (let i = 0; i < 256; i++) {
181
    fanout[i] = dv.getUint32(pos, false);
182
    pos += 4;
183
  }
184
 
185
  const count = fanout[255] || 0;
186
  const namesStart = pos;
187
  const namesEnd = namesStart + count * 20;
188
  const crcsStart = namesEnd;
189
  const crcsEnd = crcsStart + count * 4;
190
  const offsetsStart = crcsEnd;
191
  const offsetsEnd = offsetsStart + count * 4;
192
  const largeOffsetsStart = offsetsEnd;
193
  const largeOffsetsLimit = idxBuf.byteLength - IDX_TRAILER_BYTES;
194
  if (namesEnd > idxBuf.byteLength || offsetsEnd > largeOffsetsLimit) return undefined;
195
 
196
  // Copy the name table into its own buffer so the cache does not keep the
197
  // full idx object alive just to service OID binary searches.
198
  const rawNames = idxBuf.slice(namesStart, namesEnd);
199
  const offsets = new Float64Array(count);
200
  for (let i = 0; i < count; i++) {
201
    const u32 = dv.getUint32(offsetsStart + i * 4, false);
202
    if (u32 & 0x80000000) {
203
      const li = u32 & 0x7fffffff;
204
      const largeOffsetPos = largeOffsetsStart + li * 8;
205
      if (largeOffsetPos + 8 > largeOffsetsLimit) return undefined;
206
      offsets[i] = readUint64AsNumber(dv, largeOffsetPos);
207
    } else {
208
      offsets[i] = u32 >>> 0;
209
    }
210
  }
211
 
212
  // Build sorted offset arrays for binary-search lookups (replaces Maps,
213
  // saving ~24 MB of overhead for 97k entries).
214
  const sortedOffsetIndices = new Uint32Array(count);
215
  for (let i = 0; i < count; i++) sortedOffsetIndices[i] = i;
216
  sortedOffsetIndices.sort((a, b) => offsets[a] - offsets[b]);
217
 
218
  const sortedOffsets = new Float64Array(count);
219
  const nextOffsetByIndex = new Float64Array(count);
220
  for (let i = 0; i < count; i++) {
221
    sortedOffsets[i] = offsets[sortedOffsetIndices[i]];
222
    const nextOffset = i + 1 < count ? offsets[sortedOffsetIndices[i + 1]] : packSize - 20;
223
    nextOffsetByIndex[sortedOffsetIndices[i]] = nextOffset;
224
  }
225
 
226
  return {
227
    packKey,
228
    count,
229
    fanout,
230
    rawNames,
231
    offsets,
232
    nextOffsetByIndex,
233
    sortedOffsets,
234
    sortedOffsetIndices,
235
    packSize,
236
    packChecksum: idxBuf.slice(idxBuf.byteLength - IDX_TRAILER_BYTES, idxBuf.byteLength - 20),
237
    idxChecksum: idxBuf.slice(idxBuf.byteLength - 20),
238
  };
239
}
240
 
241
export async function loadIdxView(
242
  env: Env,
243
  packKey: string,
244
  cacheCtx?: CacheContext,
245
  packSizeHint?: number
246
): Promise<IdxView | undefined> {
247
  if (cacheCtx && !cacheCtx.memo) {
248
    cacheCtx.memo = {};
249
  }
250
  const cached =
251
    packSizeHint === undefined
252
      ? undefined
253
      : idxViewCache.get(getIdxViewCacheKey(packKey, packSizeHint));
254
  const memoView = cacheCtx?.memo?.idxViews?.get(packKey);
255
  if (memoView) {
256
    if (packSizeHint === undefined || memoView.packSize === packSizeHint) {
257
      return memoView;
258
    }
259
    if (cached) {
260
      // When the request-local memo was populated from a stale size hint, prefer
261
      // the size-matched isolate cache entry instead of re-fetching the idx.
262
      touchIdxViewCache(packKey, cached.view.packSize, cached.view);
263
      if (cacheCtx?.memo) {
264
        cacheCtx.memo.idxViews = cacheCtx.memo.idxViews || new Map();
265
        cacheCtx.memo.idxViews.set(packKey, cached.view);
266
      }
267
      return cached.view;
268
    }
269
    cacheCtx?.memo?.idxViews?.delete(packKey);
270
  }
271
  const promiseKey =
272
    packSizeHint === undefined ? packKey : getIdxViewCacheKey(packKey, packSizeHint);
273
  if (cacheCtx?.memo?.idxViewPromises?.has(promiseKey)) {
274
    return await cacheCtx.memo.idxViewPromises.get(promiseKey);
275
  }
276
 
277
  if (cached) {
278
    touchIdxViewCache(packKey, cached.view.packSize, cached.view);
279
    if (cacheCtx?.memo) {
280
      cacheCtx.memo.idxViews = cacheCtx.memo.idxViews || new Map();
281
      cacheCtx.memo.idxViews.set(packKey, cached.view);
282
    }
283
    return cached.view;
284
  }
285
 
286
  const log = createLogger(env.LOG_LEVEL, { service: "PackedIdxView" });
287
  const limiter = getLimiter(cacheCtx);
288
  const inflight = (async () => {
289
    // Coalesce concurrent cold loads for the same pack inside one request so
290
    // membership probes do not multiply the idx/head R2 traffic.
291
    const idxObj = await limiter.run("r2:get-pack-idx", async () => {
292
      if (!countSubrequest(cacheCtx)) {
293
        logOnce(cacheCtx, "packed-idx-soft-budget-warned", () => {
294
          log.warn("soft-budget-exhausted", {
295
            op: "r2:get-pack-idx",
296
          });
297
        });
298
      }
299
      return await env.REPO_BUCKET.get(packIndexKey(packKey));
300
    });
301
    if (!idxObj) return undefined;
302
 
303
    let packSize = packSizeHint;
304
    if (packSize === undefined) {
305
      const packHead = await limiter.run("r2:head-pack", async () => {
306
        if (!countSubrequest(cacheCtx)) {
307
          logOnce(cacheCtx, "packed-head-soft-budget-warned", () => {
308
            log.warn("soft-budget-exhausted", {
309
              op: "r2:head-pack",
310
            });
311
          });
312
        }
313
        return await env.REPO_BUCKET.head(packKey);
314
      });
315
      if (!packHead) return undefined;
316
      packSize = packHead.size;
317
    }
318
 
319
    const view = parseIdxView(packKey, new Uint8Array(await idxObj.arrayBuffer()), packSize);
320
    if (!view) return undefined;
321
 
322
    // `nextOffsetByIndex` for the final entry depends on the pack size, so the
323
    // isolate cache is keyed by both pack key and pack size. That lets normal
324
    // hinted reads reuse idx views across requests without allowing a stale hint
325
    // to overwrite the correct entry for the same pack key.
326
    touchIdxViewCache(packKey, packSize, view);
327
    if (cacheCtx?.memo) {
328
      cacheCtx.memo.idxViews = cacheCtx.memo.idxViews || new Map();
329
      cacheCtx.memo.idxViews.set(packKey, view);
330
    }
331
    return view;
332
  })();
333
 
334
  if (cacheCtx?.memo) {
335
    cacheCtx.memo.idxViewPromises = cacheCtx.memo.idxViewPromises || new Map();
336
    cacheCtx.memo.idxViewPromises.set(promiseKey, inflight);
337
  }
338
  try {
339
    return await inflight;
340
  } finally {
341
    cacheCtx?.memo?.idxViewPromises?.delete(promiseKey);
342
  }
343
}