feat: report remaining egress quotas

piece-retriever/bin/piece-retriever.js

@@ -103,6 +103,7 @@ export default {
       let retrievalCandidate
       let retrievalResult
       const retrievalAttempts = []
+      const allCandidates = [...retrievalCandidates]
 
       while (retrievalCandidates.length > 0) {
         const retrievalCandidateIndex = Math.floor(
@@ -200,7 +201,7 @@ export default {
           retrievalResult.response,
         )
         setContentSecurityPolicy(response)
-        response.headers.set('X-Data-Set-ID', retrievalCandidate.dataSetId)
+        setDiagnosticHeaders(response, retrievalCandidate, allCandidates)
         response.headers.set(
           'Cache-Control',
           `public, max-age=${env.CLIENT_CACHE_TTL}`,
@@ -327,7 +328,7 @@ export default {
         headers: retrievalResult.response.headers,
       })
       setContentSecurityPolicy(response)
-      response.headers.set('X-Data-Set-ID', retrievalCandidate.dataSetId)
+      setDiagnosticHeaders(response, retrievalCandidate, allCandidates)
       response.headers.set(
         'Cache-Control',
         `public, max-age=${env.CLIENT_CACHE_TTL}`,
@@ -394,3 +395,45 @@ function getErrorHttpStatusMessage(error) {
 
   return { status, message }
 }
+
+/**
+ * Sets diagnostic headers on the response.
+ *
+ * @param {Response} response - The response to set headers on.
+ * @param {any} selectedCandidate - The candidate that was used for retrieval.
+ * @param {any[]} allCandidates - All retrieval candidates.
+ */
+function setDiagnosticHeaders(response, selectedCandidate, allCandidates) {
+  response.headers.set('X-Data-Set-ID', selectedCandidate.dataSetId)
+
+  const totalCdnEgressQuota = allCandidates.reduce(
+    (acc, curr) => acc + curr.cdnEgressQuota,
+    0n,
+  )
+  const totalCacheMissEgressQuota = allCandidates.reduce(
+    (acc, curr) => acc + curr.cacheMissEgressQuota,
+    0n,
+  )
+
+  const datasetCdnEgressQuota = allCandidates
+    .filter((c) => c.dataSetId === selectedCandidate.dataSetId)
+    .reduce((acc, curr) => acc + curr.cdnEgressQuota, 0n)
+
+  const datasetCacheMissEgressQuota = allCandidates
+    .filter((c) => c.dataSetId === selectedCandidate.dataSetId)
+    .reduce((acc, curr) => acc + curr.cacheMissEgressQuota, 0n)
+
+  response.headers.set('X-Cdn-Egress-Remaining', String(datasetCdnEgressQuota))
+  response.headers.set(
+    'X-Cache-Miss-Egress-Remaining',
+    String(datasetCacheMissEgressQuota),
+  )
+  response.headers.set(
+    'X-Total-Cdn-Egress-Remaining',
+    String(totalCdnEgressQuota),
+  )
+  response.headers.set(
+    'X-Total-Cache-Miss-Egress-Remaining',
+    String(totalCacheMissEgressQuota),
+  )
+}

piece-retriever/test/retriever.test.js

@@ -1043,6 +1043,188 @@ describe('piece-retriever.fetch', () => {
     expect(result).toMatchObject({ egress_bytes: 0, bot_name: botName })
   })
 
+  it('sets egress remaining headers for empty response body', async () => {
+    const dataSetId = 'snapshot-empty'
+    const pieceCid = 'bagaSnapshotEmpty'
+    const cdnEgressQuota = 123
+    const cacheMissEgressQuota = 456
+
+    await withDataSetPieces(env, {
+      dataSetId,
+      serviceProviderId: 'svc-snap-empty',
+      payerAddress: defaultPayerAddress,
+      withCDN: true,
+      cdnEgressQuota,
+      cacheMissEgressQuota,
+      pieceCid,
+    })
+    await withApprovedProvider(env, {
+      id: 'svc-snap-empty',
+      serviceUrl: 'https://snap-empty.example/',
+    })
+
+    const mockRetrieveFile = vi.fn().mockResolvedValue({
+      response: new Response(null, { status: 200 }),
+      cacheMiss: true,
+    })
+
+    const ctx = createExecutionContext()
+    const req = withRequest(defaultPayerAddress, pieceCid)
+    const res = await worker.fetch(req, env, ctx, {
+      retrieveFile: mockRetrieveFile,
+    })
+    await waitOnExecutionContext(ctx)
+
+    expect(res.status).toBe(200)
+    expect(res.headers.get('X-Total-Cdn-Egress-Remaining')).toBe(
+      String(cdnEgressQuota),
+    )
+    expect(res.headers.get('X-Cdn-Egress-Remaining')).toBe(
+      String(cdnEgressQuota),
+    )
+    expect(res.headers.get('X-Total-Cache-Miss-Egress-Remaining')).toBe(
+      String(cacheMissEgressQuota),
+    )
+    expect(res.headers.get('X-Cache-Miss-Egress-Remaining')).toBe(
+      String(cacheMissEgressQuota),
+    )
+  })
+
+  it('sets egress remaining headers for streaming responses', async () => {
+    const dataSetId = 'snapshot-stream'
+    const pieceCid = 'bagaSnapshotStream'
+    const cdnEgressQuota = 789
+    const cacheMissEgressQuota = 1000
+
+    await withDataSetPieces(env, {
+      dataSetId,
+      serviceProviderId: 'svc-snap-stream',
+      payerAddress: defaultPayerAddress,
+      withCDN: true,
+      cdnEgressQuota,
+      cacheMissEgressQuota,
+      pieceId: 'p-snap-stream',
+      pieceCid,
+    })
+    await withApprovedProvider(env, {
+      id: 'svc-snap-stream',
+      serviceUrl: 'https://snap-stream.example/',
+    })
+
+    const fakeBody = new Uint8Array([1, 2, 3, 4])
+    const fakeResponse = new Response(fakeBody, {
+      status: 200,
+      headers: { 'Content-Length': String(fakeBody.length) },
+    })
+
+    const mockRetrieveFile = vi.fn().mockResolvedValue({
+      response: fakeResponse,
+      cacheMiss: true,
+      validate: () => true,
+    })
+
+    const ctx = createExecutionContext()
+    const req = withRequest(defaultPayerAddress, pieceCid)
+    const res = await worker.fetch(req, env, ctx, {
+      retrieveFile: mockRetrieveFile,
+    })
+
+    const body = await res.arrayBuffer()
+    expect(body.byteLength).toBe(fakeBody.length)
+    await waitOnExecutionContext(ctx)
+
+    expect(res.status).toBe(200)
+    expect(res.headers.get('X-Total-Cdn-Egress-Remaining')).toBe(
+      String(cdnEgressQuota),
+    )
+    expect(res.headers.get('X-Cdn-Egress-Remaining')).toBe(
+      String(cdnEgressQuota),
+    )
+    expect(res.headers.get('X-Total-Cache-Miss-Egress-Remaining')).toBe(
+      String(cacheMissEgressQuota),
+    )
+    expect(res.headers.get('X-Cache-Miss-Egress-Remaining')).toBe(
+      String(cacheMissEgressQuota),
+    )
+  })
+  it('sums egress remaining headers for multiple candidates', async () => {
+    const dataSetId1 = 'dataset-1'
+    const dataSetId2 = 'dataset-2'
+    const pieceCid = 'bagaMultipleCandidates'
+    const cdnEgressQuota1 = 100
+    const cacheMissEgressQuota1 = 50
+    const cdnEgressQuota2 = 200
+    const cacheMissEgressQuota2 = 150
+
+    // Candidate for dataset 1
+    await withDataSetPieces(env, {
+      dataSetId: dataSetId1,
+      serviceProviderId: 'svc-1',
+      payerAddress: defaultPayerAddress,
+      withCDN: true,
+      cdnEgressQuota: cdnEgressQuota1,
+      cacheMissEgressQuota: cacheMissEgressQuota1,
+      pieceCid,
+    })
+    await withApprovedProvider(env, {
+      id: 'svc-1',
+      serviceUrl: 'https://svc-1.example/',
+    })
+
+    // Candidate for dataset 2
+    await withDataSetPieces(env, {
+      dataSetId: dataSetId2,
+      serviceProviderId: 'svc-2',
+      payerAddress: defaultPayerAddress,
+      withCDN: true,
+      cdnEgressQuota: cdnEgressQuota2,
+      cacheMissEgressQuota: cacheMissEgressQuota2,
+      pieceCid,
+    })
+    await withApprovedProvider(env, {
+      id: 'svc-2',
+      serviceUrl: 'https://svc-2.example/',
+    })
+
+    const mockRetrieveFile = vi.fn().mockResolvedValue({
+      response: new Response('hello', { status: 402 }),
+      cacheMiss: true,
+    })
+
+    const ctx = createExecutionContext()
+    const req = withRequest(defaultPayerAddress, pieceCid)
+    const res = await worker.fetch(req, env, ctx, {
+      retrieveFile: mockRetrieveFile,
+    })
+    await res.text()
+    await waitOnExecutionContext(ctx)
+
+    expect(res.status).toBe(402)
+
+    const totalCdn = Number(res.headers.get('X-Total-Cdn-Egress-Remaining'))
+    const totalCacheMiss = Number(
+      res.headers.get('X-Total-Cache-Miss-Egress-Remaining'),
+    )
+    expect(totalCdn).toBe(cdnEgressQuota1 + cdnEgressQuota2)
+    expect(totalCacheMiss).toBe(cacheMissEgressQuota1 + cacheMissEgressQuota2)
+
+    const pickedDataSetId = res.headers.get('X-Data-Set-ID')
+    const datasetCdn = Number(res.headers.get('X-Cdn-Egress-Remaining'))
+    const datasetCacheMiss = Number(
+      res.headers.get('X-Cache-Miss-Egress-Remaining'),
+    )
+
+    if (pickedDataSetId === dataSetId1) {
+      expect(datasetCdn).toBe(cdnEgressQuota1)
+      expect(datasetCacheMiss).toBe(cacheMissEgressQuota1)
+    } else if (pickedDataSetId === dataSetId2) {
+      expect(datasetCdn).toBe(cdnEgressQuota2)
+      expect(datasetCacheMiss).toBe(cacheMissEgressQuota2)
+    } else {
+      throw new Error(`Unexpected pickedDataSetId: ${pickedDataSetId}`)
+    }
+  })
+
   it('stores bot name in retrieval logs when SP returns 502', async () => {
     const { pieceCid, dataSetId } = CONTENT_STORED_ON_CALIBRATION[0]
     const url = 'https://example.com/piece/502test'