Wesley Israel

Posted on Jul 18

Implementing a Leaky Bucket Rate Limiting System in Node.js:

#node #typescript #graphql #k6

📋 Introduction

Recently, I faced a fascinating technical challenge: implementing a complete rate limiting system using the Leaky Bucket strategy in Node.js. The project involved creating an HTTP server with authentication, GraphQL for PIX queries, and a multi-tenant system with granular token control.

In this article, I'll share my complete implementation journey, from initial architecture to final testing, including all technical decisions and challenges faced.

🎯 The Challenge

The goal was to build a system that:

✅ Node.js HTTP server with Koa.js and TypeScript
✅ Multi-tenancy strategy (each user with their own bucket)
✅ Bearer Token authentication (JWT)
✅ GraphQL mutation for PIX query
✅ Leaky Bucket strategy for token control
✅ Complete tests with Jest
✅ Postman documentation
✅ Load testing with k6

🏗️ System Architecture

Project Structure

leaky-bucket/
├── src/
│   ├── middleware/          # Authentication and rate limiting middlewares
│   ├── services/           # Business logic (auth, leaky bucket, PIX)
│   ├── graphql/            # GraphQL schema and resolvers
│   ├── models/             # Data models
│   ├── utils/              # Utilities (JWT, logger, metrics)
│   └── server.ts           # Main server
├── tests/                  # Jest and load tests
├── docs/                   # Documentation and Postman
└── scripts/                # Automation scripts

Data Flow

Client → Auth Middleware → Leaky Bucket → GraphQL Resolver → PIX Service

🔧 Detailed Implementation

1. Initial Configuration

I started by configuring the environment with TypeScript and necessary dependencies:

{
  "dependencies": {
    "koa": "^2.14.2",
    "koa-router": "^12.0.0",
    "graphql": "^16.8.1",
    "apollo-server-koa": "^3.12.1",
    "jsonwebtoken": "^9.0.2",
    "bcryptjs": "^2.4.3"
  }
}

2. Multi-Tenancy Strategy

I implemented a system where each user has their own token bucket:

// src/models/user.ts
interface User {
  id: string;
  email: string;
  name: string;
  password: string;
  tokens: number;
  maxTokens: number;
  lastRefill: Date;
}

class UserRepository {
  private users: Map<string, User> = new Map();

  createUser(email: string, password: string, name: string): User {
    const user: User = {
      id: generateId(),
      email,
      name,
      password: hashPassword(password),
      tokens: 10, // Initial tokens
      maxTokens: 10,
      lastRefill: new Date(),
    };

    this.users.set(user.id, user);
    return user;
  }
}

3. JWT Authentication

The authentication middleware validates tokens and extracts user information:

// src/middleware/auth.ts
export const authMiddleware = async (ctx: Context, next: Next) => {
  const authHeader = ctx.headers.authorization;

  if (!authHeader || !authHeader.startsWith('Bearer ')) {
    throw new UnauthorizedError('Authentication token required');
  }

  const token = authHeader.substring(7);
  const decoded = jwt.verify(token, JWT_SECRET) as JWTPayload;

  ctx.state.user = decoded;
  await next();
};

4. Leaky Bucket Strategy

The Leaky Bucket implementation was the heart of the system:

// src/services/leakyBucketService.ts
export class LeakyBucketService {
  private readonly LEAK_RATE = 1; // tokens per second
  private readonly REFILL_INTERVAL = 1000; // 1 second

  async consumeToken(userId: string): Promise<boolean> {
    const user = await this.userRepository.findById(userId);
    if (!user) throw new Error('User not found');

    // Calculate leaked tokens since last refill
    const now = new Date();
    const timeDiff = now.getTime() - user.lastRefill.getTime();
    const leakedTokens =
      Math.floor(timeDiff / this.REFILL_INTERVAL) * this.LEAK_RATE;

    // Update available tokens
    const newTokens = Math.min(user.maxTokens, user.tokens + leakedTokens);
    const newLastRefill = new Date(
      now.getTime() - (timeDiff % this.REFILL_INTERVAL)
    );

    if (newTokens < 1) {
      return false; // No tokens available
    }

    // Consume one token
    await this.userRepository.updateTokens(
      userId,
      newTokens - 1,
      newLastRefill
    );
    return true;
  }
}

5. GraphQL for PIX Query

I implemented a complete GraphQL schema:

// src/graphql/types.ts
const typeDefs = gql`
  type User {
    id: ID!
    email: String!
    name: String!
  }

  type PixQueryResult {
    success: Boolean!
    pixKey: String!
    value: Float!
    accountHolder: String!
    accountType: String!
    bankName: String!
    message: String!
  }

  type TokenStatus {
    tokens: Int!
    maxTokens: Int!
    lastRefill: String!
  }

  type Query {
    tokenStatus: TokenStatus!
  }

  type Mutation {
    register(email: String!, password: String!, name: String!): AuthResult!
    login(email: String!, password: String!): AuthResult!
    queryPixKey(pixKey: String!, value: Float!): PixQueryResult!
  }
`;

6. Rate Limiting Middleware

The middleware integrates Leaky Bucket with requests:

// src/middleware/leakyBucket.ts
export const leakyBucketMiddleware = async (ctx: Context, next: Next) => {
  const user = ctx.state.user;

  const hasToken = await leakyBucketService.consumeToken(user.userId);

  if (!hasToken) {
    ctx.status = 429; // Too Many Requests
    ctx.body = {
      error: 'Rate limit exceeded',
      message:
        'You have exceeded the request limit. Try again in a few seconds.',
    };
    return;
  }

  await next();
};

🧪 Comprehensive Testing

Unit Tests with Jest

I implemented 68 tests covering all scenarios:

// src/__tests__/leakyBucket.test.ts
describe('LeakyBucketService', () => {
  test('should consume token when available', async () => {
    const service = new LeakyBucketService();
    const result = await service.consumeToken('user1');
    expect(result).toBe(true);
  });

  test('should reject when no tokens', async () => {
    // Consume all tokens
    for (let i = 0; i < 10; i++) {
      await service.consumeToken('user1');
    }

    const result = await service.consumeToken('user1');
    expect(result).toBe(false);
  });
});

Load Testing with k6

I created load testing scripts to validate behavior under stress:

// tests/load/stress-test.js
import http from 'k6/http';
import { check, sleep } from 'k6';

export const options = {
  stages: [
    { duration: '2m', target: 100 }, // Ramp up
    { duration: '5m', target: 100 }, // Constant load
    { duration: '2m', target: 0 }, // Ramp down
  ],
};

export default function () {
  const token = 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9...';

  const response = http.post(
    'http://localhost:3000/graphql',
    {
      query: `mutation QueryPixKey($pixKey: String!, $value: Float!) {
      queryPixKey(pixKey: $pixKey, value: $value) {
        success pixKey value accountHolder
      }
    }`,
      variables: { pixKey: 'test@example.com', value: 100.5 },
    },
    {
      headers: {
        'Content-Type': 'application/json',
        Authorization: `Bearer ${token}`,
      },
    }
  );

  check(response, {
    'status is 200': r => r.status === 200,
    'rate limited when needed': r => r.status === 429 || r.status === 200,
  });

  sleep(1);
}

📚 Complete Documentation

Postman Collection

I created a complete collection with all endpoints:

{
  "info": {
    "name": "Leaky Bucket API",
    "description": "Complete API with authentication and rate limiting"
  },
  "item": [
    {
      "name": "Auth",
      "item": [
        {
          "name": "Register",
          "request": {
            "method": "POST",
            "url": "http://localhost:3000/graphql",
            "body": {
              "mode": "raw",
              "raw": "{\"query\":\"mutation Register($email: String!, $password: String!, $name: String!) { register(email: $email, password: $password, name: $name) { success token user { id email name } } }\",\"variables\":{\"email\":\"test@example.com\",\"password\":\"password123\",\"name\":\"Test User\"}}",
              "options": {
                "raw": {
                  "language": "json"
                }
              }
            }
          }
        }
      ]
    }
  ]
}

Detailed README

I documented the entire project with:

📖 Installation and configuration guide
🏗️ Architecture and data flow
🔧 Environment configuration
🧪 How to run tests
📊 Metrics and monitoring
🔒 Security considerations

🚀 Deployment and Configuration

Automation Scripts

I created scripts to facilitate development:

#!/bin/bash
# scripts/setup-env.sh

echo "🚀 Setting up Leaky Bucket environment..."

# Create .env file if it doesn't exist
if [ ! -f .env ]; then
    cat > .env << EOF
# Server Configuration
PORT=3000
NODE_ENV=development

# JWT
JWT_SECRET=your-super-secret-jwt-key-change-in-production
JWT_EXPIRES_IN=24h

# Leaky Bucket
INITIAL_TOKENS=10
MAX_TOKENS=10
LEAK_RATE=1
REFILL_INTERVAL=1000

# Logging
LOG_LEVEL=info

# Tests
TEST_USER_EMAIL=test@example.com
TEST_USER_PASSWORD=password123
TEST_USER_NAME=Test User
EOF
    echo "✅ .env file created successfully!"
else
    echo "ℹ️  .env file already exists"
fi

echo "🎉 Environment configured! Run 'npm run dev' to start the server"

📊 Results and Metrics

Test Performance

68 Jest tests: ✅ All passing
Load tests: ✅ Rate limiting working
Response time: < 100ms for normal requests
Throughput: 1000+ req/s with active rate limiting

Feature Coverage

✅ 100% of mandatory requirements implemented
✅ 93% of bonus requirements implemented
✅ 100% of tests passing
✅ 100% of documentation complete

🔍 Challenges and Solutions

1. Token Synchronization

Challenge: Ensuring tokens are consumed atomically in multi-thread environment.

Solution: Implemented a user-based locking system using Map and atomic operations.

2. Environment Configuration

Challenge: Maintaining consistent configurations between development and production.

Solution: Created environment variable system with validation and default values.

3. Load Testing

Challenge: Simulating realistic load and validating rate limiting.

Solution: Implemented k6 tests with different scenarios (spike, stress, load).

🎯 Lessons Learned

Modular Architecture: Clear separation of responsibilities facilitated testing and maintenance
Comprehensive Testing: Unit + integration + load tests are essential
Documentation: Documenting from the beginning saves time in the future
Configuration: Automating environment setup improves DX
Rate Limiting: Leaky Bucket is more efficient than Token Bucket for specific use cases

🔮 Next Steps

To evolve the system, I would consider:

🔄 Implement Redis for token persistence
📈 Add metrics with Prometheus
🔐 Implement refresh tokens
🌐 Add IP-based rate limiting
📱 Create React + Relay frontend

📝 Conclusion

This project demonstrated the importance of well-thought architecture, comprehensive testing, and complete documentation. The resulting system is robust, scalable, and production-ready.

The Leaky Bucket strategy proved effective for rate limiting, and GraphQL integration provided a flexible and well-documented API.

Complete code: GitHub Repository

Liked the article? Leave a ❤️ and share with other developers!

DEV Community