1import chalk from 'chalk';
2
3export default class TokenBucket {
4  constructor(capacity, refillAmount, refillTime) {
5    this.capacity = capacity;
6    this.refillTime = refillTime; // amount of time between refills (in sec)
7    this.refillAmount = refillAmount; // number of tokens to add per refill cycle
8    this.db = {};
9  }
10
11  refillBucket(key) {
12    if (this.db[key] === undefined) return null;
13
14      const { tokens, ts } = this.db[key];
15      const currentTime = Date.now();
16      const elapsedTime = Math.floor(
17        (currentTime - ts) / (this.refillTime * 1000) // convert to seconds
18      );
19
20      const newTokens = elapsedTime * this.refillAmount;
21
22      this.db[key] = {
23        tokens: Math.min(this.capacity, tokens + newTokens),
24        ts: currentTime,
25      };
26
27      return this.db[key];
28
29  }
30
31  createBucket(key) {
32    if (this.db[key] === undefined) {
33      this.db[key] = {
34        tokens: this.capacity,
35        ts: Date.now(),
36      };
37    }
38      return this.db[key];
39  }
40
41  handleRequest(key) {
42    let bucket = this.createBucket(key);
43    const currentTime = Date.now();
44
45    // check if the time elapsed since the (convert to seconds)
46    const elapsedTime = Math.floor((currentTime - bucket.ts) / 1000);
47
48    if (elapsedTime >= this.refillTime) {
49      bucket = this.refillBucket(key);
50    } else {
51      if (bucket?.tokens <= 0) {
52        console.log(
53          chalk.red(
54            `Request[REJECTED] for ${key} (tokens - ${
55              bucket.tokens
56            }) -- ${new Date().toLocaleTimeString()}
57`
58          )
59        );
60        return false;
61      }
62    }
63
64    if (!bucket) {
65      chalk.red(
66        `Request[REJECTED] for ${key} -- ${new Date().toLocaleTimeString()} -- BUCKET NOT FOUND
67`
68      );
69      return false;
70    }
71
72    console.log(
73      chalk.green(
74        `Request[ACCEPTED] for ${key} (tokens - ${
75          bucket.tokens
76        }) -- 3:06:53 AM
77`
78      )
79    );
80    bucket.tokens -= 1;
81    return true;
82
83  }
84}
85

1import TokenBucket from './TokenBucket.js';
2
3// capacity, refillAmount, refillTime (sec)
4// token bucket with 2 token every 1 min
5// const bucket = new TokenBucket(4, 2, 60);
6
7// token bucket with capacity 4 and ading 4 tokens every 5 sec
8// const bucket = new TokenBucket(4, 4, 5);
9
10// token bucket with capacity 4 and ading 4 tokens every 2 sec
11const bucket = new TokenBucket(4, 4, 2);
12
13bucket.handleRequest('user1');
14bucket.handleRequest('user1');
15bucket.handleRequest('user1');
16bucket.handleRequest('user1');
17bucket.handleRequest('user1');
18
19setTimeout(() => {
20  bucket.handleRequest('user1');
21  bucket.handleRequest('user1');
22  bucket.handleRequest('user1');
23  bucket.handleRequest('user1');
24  bucket.handleRequest('user1');
25  bucket.handleRequest('user1');
26
27  setTimeout(() => {
28    bucket.handleRequest('user1');
29  }, 3000);
30}, 3000);