fix(supply): harden shard budget guard (#7596)

* fix(supply): harden shard budget regression guard

* fix(review): release shard bodies between retries

Author: koala73

scripts/seed-supply-vulnerability.mjs

@@ -60,7 +60,7 @@ export const MIN_CHOKEPOINT_COVERAGE = 7;
 export const FLOW_STATUS_CHOKEPOINT_IDS = Object.freeze(
   new Set(CHOKEPOINT_MAP.map((entry) => entry.canonicalId)),
 );
-const MAX_SHARD_PIPELINE_BATCHES = 15;
+export const MAX_SHARD_PIPELINE_BATCHES = 15;
 const PIPELINE_TIMEOUT_MS = 15_000;
 const AUXILIARY_WRITE_TIMEOUT_MS = 5_000;
 // Bulk shard batches carry up to MAX_SHARD_PIPELINE_BYTES; the small
@@ -700,7 +700,11 @@ export async function readRedisSnapshots(keys, {
   return result;
 }
 
-const MAX_SHARD_PIPELINE_BYTES = 4.5 * 1024 * 1024;
+export const MAX_SHARD_PIPELINE_BYTES = 4.5 * 1024 * 1024;
+
+// This local packing target keeps shard writes within the 30-second write
+// timeout and the process heap budget. Upstash's per-command limit does not
+// bound the size of the multi-command pipeline body.
 
 export function buildShardPipelineBatches(payload) {
   const commands = buildShardEntries(payload).map(({ key, value }) => {
@@ -746,16 +750,19 @@ export async function publishShardCohort(payload, {
   if (batches.length > MAX_SHARD_PIPELINE_BATCHES) {
     throw new Error(`Supply vulnerability needs ${batches.length} shard batches, above the ${MAX_SHARD_PIPELINE_BATCHES}-batch runtime budget`);
   }
+  const requestBytes = [];
   const normalizedConcurrency = Math.max(1, Math.trunc(concurrency));
   const writeBatchOnce = async (commands) => {
+    const body = JSON.stringify(commands);
+    requestBytes.push(Buffer.byteLength(body, 'utf8'));
     const response = await fetchImpl(`${url}/pipeline`, {
       method: 'POST',
       headers: {
         Authorization: `Bearer ${token}`,
         'Content-Type': 'application/json',
         'User-Agent': CHROME_UA,
       },
-      body: JSON.stringify(commands),
+      body,
       signal: AbortSignal.timeout(SHARD_WRITE_TIMEOUT_MS),
     });
     if (!response.ok) {
@@ -780,7 +787,12 @@ export async function publishShardCohort(payload, {
   for (let start = 0; start < batches.length; start += normalizedConcurrency) {
     await Promise.all(batches.slice(start, start + normalizedConcurrency).map(writeBatch));
   }
-  return { batches: batches.length, shards: batches.reduce((count, batch) => count + batch.length, 0) };
+  return {
+    batches: batches.length,
+    shards: batches.reduce((count, batch) => count + batch.length, 0),
+    totalRequestBytes: requestBytes.reduce((total, bytes) => total + bytes, 0),
+    maxRequestBytes: Math.max(0, ...requestBytes),
+  };
 }
 
 export async function fetchSourceSnapshots() {
@@ -883,7 +895,18 @@ if (isMain) {
         expectedHs4: COMMODITY_REGISTRY.commodities.flatMap((mapping) => mapping.hs4 || []),
         expectedHs2: COMMODITY_REGISTRY.commodities.map((mapping) => mapping.transitHs2).filter(Boolean),
       })) throw new Error('supply vulnerability payload failed coverage or cross-index validation');
-      await publishShardCohort(payload);
+      const publication = await publishShardCohort(payload);
+      console.log(
+        `[SupplyVulnerability] wrote ${publication.shards} shards in `
+        + `${publication.batches}/${MAX_SHARD_PIPELINE_BATCHES} batches; `
+        + `largest request ${publication.maxRequestBytes} bytes`,
+      );
+      if (publication.batches >= MAX_SHARD_PIPELINE_BATCHES - 1) {
+        console.warn(
+          `[SupplyVulnerability] shard batch headroom is low: `
+          + `${publication.batches}/${MAX_SHARD_PIPELINE_BATCHES} batches`,
+        );
+      }
     },
     publishTransform: projectCountryIndex,
     extraKeys: [{

tests/supply-vulnerability-seeder.test.mjs

@@ -25,15 +25,28 @@ import {
   publishVulnerabilityCohort,
   readRedisSnapshots,
   validatePayload,
+  MAX_SHARD_PIPELINE_BATCHES,
+  MAX_SHARD_PIPELINE_BYTES,
 } from '../scripts/seed-supply-vulnerability.mjs';
+import { loadSharedConfig } from '../scripts/_seed-utils.mjs';
 import {
+  COMMODITY_REGISTRY,
   MIN_COUNTRY_COVERAGE,
   MIN_SCORED_COMMODITIES_PER_RANKABLE_COUNTRY,
   buildVulnerabilityCoverageRequirements,
 } from '../scripts/shared/supply-vulnerability-coverage.mjs';
 import { computeCountryLevelExposure } from '../scripts/seed-hs2-chokepoint-exposure.mjs';
+import { CHOKEPOINT_REGISTRY } from '../src/config/chokepoint-registry.ts';
 
 const evaluatedAt = '2026-08-30T00:00:00.000Z';
+const COUNTRY_PORT_CLUSTERS = loadSharedConfig('country-port-clusters.json');
+const countryCount = Object.keys(COUNTRY_PORT_CLUSTERS)
+  .filter((iso2) => /^[A-Z]{2}$/.test(iso2)).length;
+const commodityCount = COMMODITY_REGISTRY.commodities.length;
+const chokepointCount = CHOKEPOINT_REGISTRY.length;
+const dependencyRatios = [0.76, 0.76, 0.72, 0.82, 0.55, 0.69, 0.44, 0.56, 0.34];
+const PREVIOUS_MAX_SHARD_PIPELINE_BATCHES = 15;
+const PREVIOUS_SHARD_PIPELINE_BYTES = 4 * 1024 * 1024;
 
 function snapshot(data, fetchedAt = '2026-08-29T00:00:00.000Z') {
   return { data, fetchedAt };
@@ -402,7 +415,9 @@ describe('supply-vulnerability publication', () => {
     assert.equal(dependency.score, wheat.score);
     assert.equal(dependency.transitShare, wheat.components.transitExposure.chokepoints[0].transitShare);
     assert.ok(buildShardEntries(payload).every(({ value }) => Buffer.byteLength(JSON.stringify(value)) < 5 * 1024 * 1024));
-    assert.ok(buildShardPipelineBatches(payload).every((batch) => Buffer.byteLength(JSON.stringify(batch)) < 4 * 1024 * 1024));
+    assert.ok(buildShardPipelineBatches(payload).every((batch) => (
+      Buffer.byteLength(JSON.stringify(batch)) <= MAX_SHARD_PIPELINE_BYTES
+    )));
     assert.equal(nextShardSlot(0), 1);
     assert.equal(nextShardSlot(1), 0);
   });
@@ -485,12 +500,12 @@ describe('supply-vulnerability publication', () => {
     });
     const template = base.countries.AE.vulnerabilities[0];
     const countries = {};
-    for (let countryIndex = 0; countryIndex < 197; countryIndex += 1) {
+    for (let countryIndex = 0; countryIndex < countryCount; countryIndex += 1) {
       const iso2 = `X${String(countryIndex).padStart(3, '0')}`;
       countries[iso2] = {
         iso2,
         name: `Country ${countryIndex}`,
-        vulnerabilities: Array.from({ length: 23 }, (_, commodityIndex) => ({
+        vulnerabilities: Array.from({ length: commodityCount }, (_, commodityIndex) => ({
           ...structuredClone(template),
           countryIso2: iso2,
           countryName: `Country ${countryIndex}`,
@@ -500,7 +515,7 @@ describe('supply-vulnerability publication', () => {
             ...structuredClone(template.components),
             transitExposure: {
               ...structuredClone(template.components.transitExposure),
-              chokepoints: Array.from({ length: 13 }, (_, routeIndex) => ({
+              chokepoints: Array.from({ length: chokepointCount }, (_, routeIndex) => ({
                 ...structuredClone(template.components.transitExposure.chokepoints[0]),
                 id: `route-${routeIndex}`,
                 name: `Representative route ${routeIndex}`,
@@ -511,7 +526,9 @@ describe('supply-vulnerability publication', () => {
       };
     }
     const records = Object.values(countries).flatMap((country) => country.vulnerabilities);
-    const dependencyCounts = [3441, 3441, 3248, 3717, 2505, 3115, 1984, 2518, 1535];
+    const dependencyCounts = Array.from({ length: chokepointCount }, (_, routeIndex) => (
+      Math.max(1, Math.round(records.length * dependencyRatios[routeIndex % dependencyRatios.length]))
+    ));
     const chokepoints = Object.fromEntries(dependencyCounts.map((count, routeIndex) => {
       const id = `route-${routeIndex}`;
       return [id, {
@@ -537,23 +554,31 @@ describe('supply-vulnerability publication', () => {
       }];
     }));
     const payload = { ...base, countries, chokepoints };
-    const legacyBytes = Buffer.byteLength(JSON.stringify({
-      generatedAt: payload.generatedAt,
-      methodologyVersion: payload.methodologyVersion,
-      countries,
-    }));
-    assert.ok(legacyBytes > 5 * 1024 * 1024, `expected legacy projection above 5MB, got ${legacyBytes}`);
+    const shardEntries = buildShardEntries(payload);
+    const totalCommandBytes = shardEntries.reduce((total, { key, value }) => (
+      total + Buffer.byteLength(JSON.stringify([
+        'SET',
+        key,
+        JSON.stringify(value),
+        'EX',
+        SHARD_TTL_SECONDS,
+      ]), 'utf8')
+    ), 0);
+    assert.ok(
+      totalCommandBytes > PREVIOUS_MAX_SHARD_PIPELINE_BATCHES * PREVIOUS_SHARD_PIPELINE_BYTES,
+      `fixture must exceed the pre-fix 15-request capacity, got ${totalCommandBytes}`,
+    );
     assert.ok(Buffer.byteLength(JSON.stringify(projectCountryIndex(payload))) < 5 * 1024 * 1024);
-    assert.ok(buildShardEntries(payload).every(({ value }) => Buffer.byteLength(JSON.stringify(value)) < 5 * 1024 * 1024));
+    assert.ok(shardEntries.every(({ value }) => Buffer.byteLength(JSON.stringify(value)) < 5 * 1024 * 1024));
 
-    const batches = buildShardPipelineBatches(payload);
-    assert.ok(batches.length <= 15, `all shard batches must fit one write wave, got ${batches.length}`);
     let active = 0;
     let peak = 0;
+    const requestBytes = [];
     const result = await publishShardCohort(payload, {
       credentials: { url: 'https://redis.example', token: 'test-token' },
       retryDelayMs: 0,
       fetchImpl: async (_url, request) => {
+        requestBytes.push(Buffer.byteLength(request.body, 'utf8'));
         active += 1;
         peak = Math.max(peak, active);
         await new Promise((resolve) => setImmediate(resolve));
@@ -565,9 +590,12 @@ describe('supply-vulnerability publication', () => {
         };
       },
     });
-    assert.equal(result.batches, batches.length);
-    assert.equal(result.shards, buildShardEntries(payload).length);
-    assert.equal(peak, batches.length, 'all shard batches must fit in one write wave');
+    assert.ok(result.batches <= MAX_SHARD_PIPELINE_BATCHES);
+    assert.equal(result.shards, shardEntries.length);
+    assert.equal(result.totalRequestBytes, requestBytes.reduce((sum, bytes) => sum + bytes, 0));
+    assert.equal(result.maxRequestBytes, Math.max(...requestBytes));
+    assert.ok(result.maxRequestBytes <= MAX_SHARD_PIPELINE_BYTES);
+    assert.equal(peak, result.batches, 'all shard batches must fit in one write wave');
   });
 
   it('bounds source-read concurrency across Redis pipeline batches', async () => {
@@ -661,12 +689,14 @@ describe('supply-vulnerability publication', () => {
       countryNames: { AE: 'United Arab Emirates' },
     });
     let calls = 0;
-    await publishShardCohort(payload, {
+    const requestBytes = [];
+    const result = await publishShardCohort(payload, {
       credentials: { url: 'https://redis.example', token: 'test-token' },
       concurrency: 1,
       retryDelayMs: 0,
       fetchImpl: async (_url, request) => {
         calls += 1;
+        requestBytes.push(Buffer.byteLength(request.body, 'utf8'));
         if (calls === 1) return { ok: false, status: 503, headers: new Headers() };
         return {
           ok: true,
@@ -675,6 +705,9 @@ describe('supply-vulnerability publication', () => {
       },
     });
     assert.equal(calls, 2);
+    assert.equal(result.totalRequestBytes, requestBytes.reduce((sum, bytes) => sum + bytes, 0));
+    assert.equal(result.maxRequestBytes, Math.max(...requestBytes));
+    assert.equal(requestBytes[0], requestBytes[1]);
   });
 
   it('atomically switches the shared cohort and activation marker and propagates failure', async () => {