Skip to content

Commit 50383a6

Browse files
fforbeckfforbeck
committed
reviewer suggestions implemented
1 parent b6fffc0 commit 50383a6

File tree

3 files changed

+46
-71
lines changed

3 files changed

+46
-71
lines changed

‎src/middleware/withEgressTracker.js‎

Lines changed: 22 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,7 @@ export function withEgressTracker (handler) {
2121
return handler(req, env, ctx)
2222
}
2323

24-
let response
25-
try {
26-
response = await handler(req, env, ctx)
27-
} catch (error) {
28-
console.error('Error in egress tracker handler:', error)
29-
throw error
30-
}
31-
24+
const response = await handler(req, env, ctx)
3225
if (!response.ok || !response.body) {
3326
return response
3427
}
@@ -38,17 +31,18 @@ export function withEgressTracker (handler) {
3831
serviceURL: env.ACCOUNTING_SERVICE_URL
3932
})
4033

41-
const { readable, writable } = createEgressPassThroughStream(ctx, accounting, dataCid)
42-
43-
try {
44-
ctx.waitUntil(response.body.pipeTo(writable))
45-
} catch (error) {
46-
console.error('Error in egress tracker handler:', error)
47-
// Original response in case of an error to avoid breaking the chain and serve the content
48-
return response
49-
}
34+
const responseBody = response.body.pipeThrough(
35+
createByteCountStream((totalBytesServed) => {
36+
// Non-blocking call to the accounting service to record egress
37+
if (totalBytesServed > 0) {
38+
ctx.waitUntil(
39+
accounting.record(dataCid, totalBytesServed, new Date().toISOString())
40+
);
41+
}
42+
})
43+
);
5044

51-
return new Response(readable, {
45+
return new Response(responseBody, {
5246
status: response.status,
5347
statusText: response.statusText,
5448
headers: response.headers
@@ -60,58 +54,39 @@ export function withEgressTracker (handler) {
6054
* Creates a TransformStream to count bytes served to the client.
6155
* It records egress when the stream is finalized without an error.
6256
*
63-
* @param {import('@web3-storage/gateway-lib/middleware').Context} ctx - The context object.
64-
* @param {AccountingService} accounting - The accounting service instance to record egress.
65-
* @param {import('@web3-storage/gateway-lib/handlers').CID} dataCid - The CID of the served content.
66-
* @returns {TransformStream} - The created TransformStream.
57+
* @param {(totalBytesServed: number) => void} onClose
58+
* @template {Uint8Array} T
59+
* @returns {TransformStream<T, T>} - The created TransformStream.
6760
*/
68-
function createEgressPassThroughStream (ctx, accounting, dataCid) {
61+
function createByteCountStream(onClose) {
6962
let totalBytesServed = 0
7063

7164
return new TransformStream({
72-
/**
73-
* The start function is called when the stream is being initialized.
74-
* It resets the total bytes served to 0.
75-
*/
76-
start () {
77-
totalBytesServed = 0
78-
},
7965
/**
8066
* The transform function is called for each chunk of the response body.
8167
* It enqueues the chunk and updates the total bytes served.
8268
* If an error occurs, it signals an error to the controller and logs it.
8369
* The bytes are not counted in case of enqueuing an error.
84-
* @param {Uint8Array} chunk
85-
* @param {TransformStreamDefaultController} controller
8670
*/
87-
async transform (chunk, controller) {
71+
async transform(chunk, controller) {
8872
try {
8973
controller.enqueue(chunk)
9074
totalBytesServed += chunk.byteLength
9175
} catch (error) {
92-
console.error('Error while counting egress bytes:', error)
76+
console.error('Error while counting bytes:', error)
9377
controller.error(error)
9478
}
9579
},
9680

9781
/**
9882
* The flush function is called when the stream is being finalized,
9983
* which is when the response is being sent to the client.
100-
* So before the response is sent, we record the egress.
101-
* It is called only once and it triggers a non-blocking call to the accounting service.
84+
* So before the response is sent, we record the egress using the callback.
10285
* If an error occurs, the egress is not recorded.
103-
* NOTE: The flush function is NOT called in case of an stream error.
86+
* NOTE: The flush function is NOT called in case of a stream error.
10487
*/
105-
async flush (controller) {
106-
try {
107-
// Non-blocking call to the accounting service to record egress
108-
if (totalBytesServed > 0) {
109-
ctx.waitUntil(accounting.record(dataCid, totalBytesServed, new Date().toISOString()))
110-
}
111-
} catch (error) {
112-
console.error('Error while recording egress:', error)
113-
controller.error(error)
114-
}
88+
async flush() {
89+
onClose(totalBytesServed)
11590
}
11691
})
11792
}

‎test/fixtures/worker-fixture.js‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ const __dirname = path.dirname(__filename)
1212
*/
1313
const wranglerEnv = process.env.WRANGLER_ENV || 'integration'
1414

15-
const DEBUG = process.env.DEBUG === 'true' || false
15+
const DEBUG = process.env.DEBUG === 'true'
1616

1717
/**
1818
* Worker information object

‎test/unit/middleware/withEgressTracker.spec.js‎

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import { describe, it, afterEach, before } from 'mocha'
88
import { assert, expect } from 'chai'
99
import sinon from 'sinon'
1010
import { CID } from 'multiformats'
11-
import { withEgressHandler } from '../../../src/handlers/egress-tracker.js'
11+
import { withEgressTracker } from '../../../src/middleware/withEgressTracker.js'
1212
import { Builder, toBlobKey } from '../../helpers/builder.js'
1313
import { CARReaderStream } from 'carstream'
1414

@@ -27,19 +27,19 @@ const createRequest = async ({ authorization } = {}) =>
2727
})
2828

2929
const env =
30-
/** @satisfies {import('../../../src/handlers/egress-tracker.types.js').Environment} */
30+
/** @satisfies {import('../../../src/middleware/withEgressTracker.types.js').Environment} */
3131
({
3232
DEBUG: 'true',
3333
ACCOUNTING_SERVICE_URL: 'http://example.com',
3434
FF_EGRESS_TRACKER_ENABLED: 'true'
3535
})
3636

3737
const accountingRecordMethodStub = sinon.stub()
38-
// @ts-expect-error
39-
.returns(async (cid, bytes, servedAt) => {
40-
console.log(`[mock] record called with cid: ${cid}, bytes: ${bytes}, servedAt: ${servedAt}`)
41-
return Promise.resolve()
42-
})
38+
.returns(
39+
/** @type {import('../../../src/bindings.js').AccountingService['record']} */
40+
async (cid, bytes, servedAt) => {
41+
console.log(`[mock] record called with cid: ${cid}, bytes: ${bytes}, servedAt: ${servedAt}`)
42+
})
4343

4444
/**
4545
* Mock implementation of the AccountingService.
@@ -58,7 +58,7 @@ const AccountingService = ({ serviceURL }) => {
5858
}
5959

6060
const ctx =
61-
/** @satisfies {import('../../../src/handlers/egress-tracker.js').EgressTrackerContext} */
61+
/** @satisfies {import('../../../src/middleware/withEgressTracker.js').EgressTrackerContext} */
6262
({
6363
dataCid: CID.parse('bafybeibv7vzycdcnydl5n5lbws6lul2omkm6a6b5wmqt77sicrwnhesy7y'),
6464
waitUntil: sinon.stub().returns(undefined),
@@ -110,7 +110,7 @@ describe('withEgressTracker', async () => {
110110

111111
const innerHandler = sinon.stub().resolves(mockResponse)
112112

113-
const handler = withEgressHandler(innerHandler)
113+
const handler = withEgressTracker(innerHandler)
114114
const request = await createRequest()
115115
const response = await handler(request, env, ctx)
116116
// Ensure the response body is fully consumed
@@ -134,7 +134,7 @@ describe('withEgressTracker', async () => {
134134
}), { status: 200 })
135135

136136
const innerHandler = sinon.stub().resolves(mockResponse)
137-
const handler = withEgressHandler(innerHandler)
137+
const handler = withEgressTracker(innerHandler)
138138
const request = await createRequest()
139139

140140
const response = await handler(request, env, ctx)
@@ -160,7 +160,7 @@ describe('withEgressTracker', async () => {
160160
}), { status: 200 })
161161

162162
const innerHandler = sinon.stub().resolves(mockResponse)
163-
const handler = withEgressHandler(innerHandler)
163+
const handler = withEgressTracker(innerHandler)
164164
const request = await createRequest()
165165

166166
const response = await handler(request, env, ctx)
@@ -197,7 +197,7 @@ describe('withEgressTracker', async () => {
197197
})
198198

199199
const innerHandler = sinon.stub().resolves(mockResponse)
200-
const handler = withEgressHandler(innerHandler)
200+
const handler = withEgressTracker(innerHandler)
201201
const request = await createRequest()
202202

203203
const response = await handler(request, env, ctx)
@@ -233,7 +233,7 @@ describe('withEgressTracker', async () => {
233233
}), { status: 200 })
234234

235235
const innerHandler = sinon.stub().resolves(mockResponse)
236-
const handler = withEgressHandler(innerHandler)
236+
const handler = withEgressTracker(innerHandler)
237237
const request = await createRequest()
238238

239239
const response = await handler(request, env, ctx)
@@ -249,7 +249,7 @@ describe('withEgressTracker', async () => {
249249
describe('withEgressTracker -> Feature Flag', () => {
250250
it('should not track egress if the feature flag is disabled', async () => {
251251
const innerHandler = sinon.stub().resolves(new Response(null, { status: 200 }))
252-
const handler = withEgressHandler(innerHandler)
252+
const handler = withEgressTracker(innerHandler)
253253
const request = await createRequest()
254254
const envDisabled = { ...env, FF_EGRESS_TRACKER_ENABLED: 'false' }
255255

@@ -264,7 +264,7 @@ describe('withEgressTracker', async () => {
264264
it('should not track egress for non-OK responses', async () => {
265265
const mockResponse = new Response(null, { status: 404 })
266266
const innerHandler = sinon.stub().resolves(mockResponse)
267-
const handler = withEgressHandler(innerHandler)
267+
const handler = withEgressTracker(innerHandler)
268268
const request = await createRequest()
269269

270270
const response = await handler(request, env, ctx)
@@ -276,7 +276,7 @@ describe('withEgressTracker', async () => {
276276
it('should not track egress if the response has no body', async () => {
277277
const mockResponse = new Response(null, { status: 200 })
278278
const innerHandler = sinon.stub().resolves(mockResponse)
279-
const handler = withEgressHandler(innerHandler)
279+
const handler = withEgressTracker(innerHandler)
280280
const request = await createRequest()
281281

282282
const response = await handler(request, env, ctx)
@@ -310,8 +310,8 @@ describe('withEgressTracker', async () => {
310310
const innerHandler1 = sinon.stub().resolves(mockResponse1)
311311
const innerHandler2 = sinon.stub().resolves(mockResponse2)
312312

313-
const handler1 = withEgressHandler(innerHandler1)
314-
const handler2 = withEgressHandler(innerHandler2)
313+
const handler1 = withEgressTracker(innerHandler1)
314+
const handler2 = withEgressTracker(innerHandler2)
315315

316316
const request1 = await createRequest()
317317
const request2 = await createRequest()
@@ -347,7 +347,7 @@ describe('withEgressTracker', async () => {
347347
}), { status: 200, headers: { 'Content-Type': 'application/json' } })
348348

349349
const innerHandler = sinon.stub().resolves(mockResponse)
350-
const handler = withEgressHandler(innerHandler)
350+
const handler = withEgressTracker(innerHandler)
351351
const request = await createRequest()
352352

353353
const response = await handler(request, env, ctx)
@@ -370,7 +370,7 @@ describe('withEgressTracker', async () => {
370370
}), { status: 200 })
371371

372372
const innerHandler = sinon.stub().resolves(mockResponse)
373-
const handler = withEgressHandler(innerHandler)
373+
const handler = withEgressTracker(innerHandler)
374374
const request = await createRequest()
375375

376376
const response = await handler(request, env, ctx)
@@ -391,7 +391,7 @@ describe('withEgressTracker', async () => {
391391
}), { status: 200 })
392392

393393
const innerHandler = sinon.stub().resolves(mockResponse)
394-
const handler = withEgressHandler(innerHandler)
394+
const handler = withEgressTracker(innerHandler)
395395
const request = await createRequest()
396396

397397
const response = await handler(request, env, ctx)
@@ -418,7 +418,7 @@ describe('withEgressTracker', async () => {
418418
}), { status: 200 })
419419

420420
const innerHandler = sinon.stub().resolves(mockResponse)
421-
const handler = withEgressHandler(innerHandler)
421+
const handler = withEgressTracker(innerHandler)
422422
const request = await createRequest()
423423
const response = await handler(request, env, ctx)
424424

@@ -445,7 +445,7 @@ describe('withEgressTracker', async () => {
445445
}), { status: 200 })
446446

447447
const innerHandler = sinon.stub().resolves(mockResponse)
448-
const handler = withEgressHandler(innerHandler)
448+
const handler = withEgressTracker(innerHandler)
449449
const request = await createRequest()
450450

451451
// Simulate an error in the accounting service record method

0 commit comments

Comments
 (0)