Skip to content

Commit fc4690b

Browse files
alexey-yarmoshalexey-yarmosh
committed
feat: add credits-required headers
1 parent 4ca340c commit fc4690b

File tree

4 files changed

+81
-52
lines changed

4 files changed

+81
-52
lines changed

Diff for: ‎src/lib/credits.ts

+25-17
Original file line numberDiff line numberDiff line change
@@ -7,30 +7,38 @@ const ER_CONSTRAINT_FAILED_CODE = 4025;
77
export class Credits {
88
constructor (private readonly sql: Knex) {}
99

10-
async consume (userId: string, credits: number): Promise<{ isConsumed: boolean, remainingCredits?: number }> {
11-
try {
12-
const result = await this.sql.raw<[[{amount: number | null}]]>(`
13-
INSERT INTO ?? (user_id, amount)
14-
VALUES (?, ?)
15-
ON DUPLICATE KEY UPDATE
16-
amount = amount - ?
17-
RETURNING amount
18-
`, [ CREDITS_TABLE, userId, null, credits ]);
19-
20-
const remainingCredits = result[0]?.[0]?.amount;
21-
22-
if (remainingCredits || remainingCredits === 0) {
23-
return { isConsumed: true, remainingCredits };
24-
}
10+
async consume (userId: string, credits: number): Promise<{ isConsumed: boolean, remainingCredits: number }> {
11+
let numberOfUpdates = null;
2512

26-
return { isConsumed: false };
13+
try {
14+
numberOfUpdates = await this.sql(CREDITS_TABLE).where({ user_id: userId }).update({ amount: this.sql.raw('amount - ?', [ credits ]) });
2715
} catch (error) {
2816
if (error && (error as Error & {errno?: number}).errno === ER_CONSTRAINT_FAILED_CODE) {
29-
return { isConsumed: false };
17+
const remainingCredits = await this.getRemainingCredits(userId);
18+
return { isConsumed: false, remainingCredits };
3019
}
3120

3221
throw error;
3322
}
23+
24+
if (numberOfUpdates === 0) {
25+
return { isConsumed: false, remainingCredits: 0 };
26+
}
27+
28+
const remainingCredits = await this.getRemainingCredits(userId);
29+
return { isConsumed: true, remainingCredits };
30+
}
31+
32+
async getRemainingCredits (userId: string): Promise<number> {
33+
const result = await this.sql(CREDITS_TABLE).where({ user_id: userId }).select<[{ amount: number }]>('amount');
34+
35+
const remainingCredits = result[0]?.amount;
36+
37+
if (remainingCredits || remainingCredits === 0) {
38+
return remainingCredits;
39+
}
40+
41+
throw new Error('Credits data for the user not found.');
3442
}
3543
}
3644

Diff for: ‎src/lib/rate-limiter.ts

+18-17
Original file line numberDiff line numberDiff line change
@@ -44,14 +44,16 @@ export const rateLimit = async (ctx: ExtendedContext, numberOfProbes: number) =>
4444
} catch (error) {
4545
if (error instanceof RateLimiterRes) {
4646
if (ctx.state.userId) {
47-
const { isConsumed, consumedCredits, remainingCredits } = await consumeCredits(ctx.state.userId, error, numberOfProbes);
47+
const { isConsumed, requiredCredits, remainingCredits } = await consumeCredits(ctx.state.userId, error, numberOfProbes);
4848

4949
if (isConsumed) {
50-
const result = await rateLimiter.reward(id, consumedCredits);
51-
setCreditsHeaders(ctx, consumedCredits!, remainingCredits!);
50+
const result = await rateLimiter.reward(id, requiredCredits);
51+
setConsumedHeaders(ctx, requiredCredits, remainingCredits);
5252
setRateLimitHeaders(ctx, result, rateLimiter);
5353
return;
5454
}
55+
56+
setRequiredHeaders(ctx, requiredCredits, remainingCredits);
5557
}
5658

5759
const result = await rateLimiter.reward(id, numberOfProbes);
@@ -64,20 +66,14 @@ export const rateLimit = async (ctx: ExtendedContext, numberOfProbes: number) =>
6466
};
6567

6668
const consumeCredits = async (userId: string, rateLimiterRes: RateLimiterRes, numberOfProbes: number) => {
67-
const freePoints = config.get<number>('measurement.authenticatedRateLimit');
68-
const requiredPoints = Math.min(rateLimiterRes.consumedPoints - freePoints, numberOfProbes);
69-
const { isConsumed, remainingCredits } = await credits.consume(userId, requiredPoints);
70-
71-
if (isConsumed) {
72-
return {
73-
isConsumed,
74-
consumedCredits: requiredPoints,
75-
remainingCredits,
76-
};
77-
}
69+
const freeCredits = config.get<number>('measurement.authenticatedRateLimit');
70+
const requiredCredits = Math.min(rateLimiterRes.consumedPoints - freeCredits, numberOfProbes);
71+
const { isConsumed, remainingCredits } = await credits.consume(userId, requiredCredits);
7872

7973
return {
80-
isConsumed: false,
74+
isConsumed,
75+
requiredCredits,
76+
remainingCredits,
8177
};
8278
};
8379

@@ -87,7 +83,12 @@ const setRateLimitHeaders = (ctx: ExtendedContext, result: RateLimiterRes, rateL
8783
ctx.set('X-RateLimit-Remaining', `${result.remainingPoints}`);
8884
};
8985

90-
const setCreditsHeaders = (ctx: ExtendedContext, consumedCredits: number, remainingCredits: number) => {
91-
ctx.set('X-Credits-Cost', `${consumedCredits}`);
86+
const setConsumedHeaders = (ctx: ExtendedContext, consumedCredits: number, remainingCredits: number) => {
87+
ctx.set('X-Credits-Consumed', `${consumedCredits}`);
88+
ctx.set('X-Credits-Remaining', `${remainingCredits}`);
89+
};
90+
91+
const setRequiredHeaders = (ctx: ExtendedContext, requiredCredits: number, remainingCredits: number) => {
92+
ctx.set('X-Credits-Required', `${requiredCredits}`);
9293
ctx.set('X-Credits-Remaining', `${remainingCredits}`);
9394
};

Diff for: ‎test/tests/integration/ratelimit.test.ts

+19-12
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,8 @@ describe('rate limiter', () => {
236236
}).expect(202) as Response;
237237

238238
expect(response.headers['x-ratelimit-remaining']).to.equal('248');
239-
expect(response.headers['x-credits-cost']).to.not.exist;
239+
expect(response.headers['x-credits-consumed']).to.not.exist;
240+
expect(response.headers['x-credits-required']).to.not.exist;
240241
expect(response.headers['x-credits-remaining']).to.not.exist;
241242
const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
242243
expect(amount).to.equal(10);
@@ -254,7 +255,8 @@ describe('rate limiter', () => {
254255
}).expect(202) as Response;
255256

256257
expect(response.headers['x-ratelimit-remaining']).to.equal('0');
257-
expect(response.headers['x-credits-cost']).to.equal('2');
258+
expect(response.headers['x-credits-consumed']).to.equal('2');
259+
expect(response.headers['x-credits-required']).to.not.exist;
258260
expect((response.headers['x-credits-remaining'])).to.equal('8');
259261
const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
260262
expect(amount).to.equal(8);
@@ -272,7 +274,8 @@ describe('rate limiter', () => {
272274
}).expect(202) as Response;
273275

274276
expect(response.headers['x-ratelimit-remaining']).to.equal('0');
275-
expect(response.headers['x-credits-cost']).to.equal('1');
277+
expect(response.headers['x-credits-consumed']).to.equal('1');
278+
expect(response.headers['x-credits-required']).to.not.exist;
276279
expect((response.headers['x-credits-remaining'])).to.equal('9');
277280
const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
278281
expect(amount).to.equal(9);
@@ -296,8 +299,9 @@ describe('rate limiter', () => {
296299
}).expect(429) as Response;
297300

298301
expect(response.headers['x-ratelimit-remaining']).to.equal('0');
299-
expect(response.headers['x-credits-cost']).to.not.exist;
300-
expect(response.headers['x-credits-remaining']).to.not.exist;
302+
expect(response.headers['x-credits-consumed']).to.not.exist;
303+
expect(response.headers['x-credits-required']).to.equal('2');
304+
expect(response.headers['x-credits-remaining']).to.equal('1');
301305
const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
302306
expect(amount).to.equal(1);
303307
});
@@ -314,7 +318,8 @@ describe('rate limiter', () => {
314318
}).expect(202) as Response;
315319

316320
expect(response.headers['x-ratelimit-remaining']).to.equal('0');
317-
expect(response.headers['x-credits-cost']).to.equal('2');
321+
expect(response.headers['x-credits-consumed']).to.equal('2');
322+
expect(response.headers['x-credits-required']).to.not.exist;
318323
expect((response.headers['x-credits-remaining'])).to.equal('8');
319324
const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
320325
expect(amount).to.equal(8);
@@ -338,8 +343,9 @@ describe('rate limiter', () => {
338343
}).expect(429) as Response;
339344

340345
expect(response.headers['x-ratelimit-remaining']).to.equal('1');
341-
expect(response.headers['x-credits-cost']).to.not.exist;
342-
expect(response.headers['x-credits-remaining']).to.not.exist;
346+
expect(response.headers['x-credits-consumed']).to.not.exist;
347+
expect(response.headers['x-credits-required']).to.equal('1');
348+
expect(response.headers['x-credits-remaining']).to.equal('0');
343349
const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
344350
expect(amount).to.equal(0);
345351
});
@@ -360,10 +366,11 @@ describe('rate limiter', () => {
360366
}).expect(429) as Response;
361367

362368
expect(response.headers['x-ratelimit-remaining']).to.equal('0');
363-
expect(response.headers['x-credits-cost']).to.not.exist;
364-
expect(response.headers['x-credits-remaining']).to.not.exist;
365-
const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
366-
expect(amount).to.equal(null);
369+
expect(response.headers['x-credits-consumed']).to.not.exist;
370+
expect(response.headers['x-credits-required']).to.equal('2');
371+
expect(response.headers['x-credits-remaining']).to.equal('0');
372+
const credits = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
373+
expect(credits).to.deep.equal([]);
367374
});
368375
});
369376
});

Diff for: ‎test/tests/unit/credits.test.ts

+19-6
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,10 @@ import * as sinon from 'sinon';
44
import { Credits } from '../../../src/lib/credits.js';
55

66
describe('Credits', () => {
7+
const updateStub = sinon.stub();
78
const selectStub = sinon.stub();
89
const whereStub = sinon.stub().returns({
10+
update: updateStub,
911
select: selectStub,
1012
});
1113
const sqlStub = sinon.stub().returns({
@@ -18,31 +20,42 @@ describe('Credits', () => {
1820
});
1921

2022
it('should return true if row was updated', async () => {
21-
sqlStub.raw.resolves([ [{ amount: 0 }] ]);
23+
updateStub.resolves(1);
24+
selectStub.resolves([{ amount: 5 }]);
25+
const credits = new Credits(sqlStub as unknown as Knex);
26+
const result = await credits.consume('userId', 10);
27+
expect(result).to.deep.equal({ isConsumed: true, remainingCredits: 5 });
28+
});
29+
30+
it('should return true if row was updated to 0', async () => {
31+
updateStub.resolves(1);
32+
selectStub.resolves([{ amount: 0 }]);
2233
const credits = new Credits(sqlStub as unknown as Knex);
2334
const result = await credits.consume('userId', 10);
2435
expect(result).to.deep.equal({ isConsumed: true, remainingCredits: 0 });
2536
});
2637

2738
it(`should return false if row wasn't updated`, async () => {
28-
sqlStub.raw.resolves([ [{ amount: null }] ]);
39+
updateStub.resolves(0);
2940
const credits = new Credits(sqlStub as unknown as Knex);
3041
const result = await credits.consume('userId', 10);
31-
expect(result).to.deep.equal({ isConsumed: false });
42+
expect(selectStub.callCount).to.equal(0);
43+
expect(result).to.deep.equal({ isConsumed: false, remainingCredits: 0 });
3244
});
3345

3446
it(`should return false if update throws ER_CONSTRAINT_FAILED_CODE`, async () => {
3547
const error: Error & {errno?: number} = new Error('constraint');
3648
error.errno = 4025;
37-
sqlStub.raw.rejects(error);
49+
updateStub.rejects(error);
50+
selectStub.resolves([{ amount: 5 }]);
3851
const credits = new Credits(sqlStub as unknown as Knex);
3952
const result = await credits.consume('userId', 10);
40-
expect(result).to.deep.equal({ isConsumed: false });
53+
expect(result).to.deep.equal({ isConsumed: false, remainingCredits: 5 });
4154
});
4255

4356
it(`should throw if update throws other error`, async () => {
4457
const error = new Error('other error');
45-
sqlStub.raw.rejects(error);
58+
updateStub.rejects(error);
4659
const credits = new Credits(sqlStub as unknown as Knex);
4760
const result = await credits.consume('userId', 10).catch(err => err);
4861
expect(result).to.equal(error);

0 commit comments

Comments
 (0)