chore: wip

Author: chrisbbreuer

examples/key-rate-limiting.ts

@@ -0,0 +1,81 @@
+import { Queue, RateLimiter } from '../src'
+
+interface EmailData {
+  userId: string
+  to: string
+  subject: string
+  content: string
+}
+
+async function main() {
+  console.log('🔒 Key-based Rate Limiting Example')
+
+  // Create an email queue with key-based rate limiting
+  const emailQueue = new Queue<EmailData>('emails', {
+    verbose: true,
+    logLevel: 'info',
+  })
+
+  // Create a rate limiter with key-based options
+  // This limits each user to 2 emails per 10 seconds
+  const emailRateLimiter = new RateLimiter(emailQueue, {
+    max: 2, // 2 emails
+    duration: 10000, // per 10 seconds
+    keyPrefix: 'userId' // Rate limit by userId field
+  })
+
+  // Demonstrate sending emails for different users
+  const users = ['user1', 'user2', 'user3']
+
+  // Send multiple emails from different users
+  for (let i = 0; i < 10; i++) {
+    const userId = users[i % users.length]
+
+    // Send email
+    console.log(`📧 Attempting to send email #${i+1} for user ${userId}`)
+
+    // Check if rate limited first
+    const rateLimitCheck = await emailQueue.isRateLimited(undefined, { userId, to: '', subject: '', content: '' })
+
+    if (rateLimitCheck.limited) {
+      console.log(`⛔ Rate limit reached for user ${userId}. Try again in ${rateLimitCheck.resetIn}ms`)
+      continue
+    }
+
+    // Add job to queue
+    const job = await emailQueue.add({
+      userId,
+      to: `${userId}@example.com`,
+      subject: `Test email #${i+1}`,
+      content: `This is test email #${i+1} for ${userId}`
+    })
+
+    console.log(`✅ Email #${i+1} queued for ${userId} with job ID: ${job.id}`)
+
+    // Wait a bit between requests to see the rate limiting in action
+    await new Promise(resolve => setTimeout(resolve, 1000))
+  }
+
+  // Process queue with a handler that logs emails
+  emailQueue.process(1, async (job) => {
+    const { userId, to, subject } = job.data
+    console.log(`🚀 Processing email for ${userId} to ${to}: "${subject}"`)
+
+    // Simulate sending email
+    await new Promise(resolve => setTimeout(resolve, 500))
+
+    return { success: true, sentAt: new Date() }
+  })
+
+  // Let the queue process for a while
+  await new Promise(resolve => setTimeout(resolve, 15000))
+
+  // Close the queue
+  await emailQueue.close()
+  console.log('👋 Queue closed')
+}
+
+main().catch(error => {
+  console.error('Error in example:', error)
+  process.exit(1)
+})

examples/priority-queue.ts

@@ -75,4 +75,4 @@ async function main() {
   console.log('\n👋 All tasks completed, queue closed')
 }
 
-main().catch(console.error)
+main().catch(console.error)

package.json

@@ -65,7 +65,8 @@
     "example:basic": "bun examples/basic.ts",
     "example:advanced": "bun examples/advanced.ts",
     "example:advanced-features": "bun examples/advanced-features.ts",
-    "example:priority-queue": "bun examples/priority-queue.ts"
+    "example:priority-queue": "bun examples/priority-queue.ts",
+    "example:key-rate-limiting": "bun examples/key-rate-limiting.ts"
   },
   "devDependencies": {
     "@stacksjs/docs": "^0.70.23",

src/priority-queue.ts

@@ -256,7 +256,7 @@ export class PriorityQueue<T = any> {
 
     for (let priority = this.priorityLevels - 1; priority >= 0; priority--) {
       const priorityKey = this.getPriorityKey(priority)
-      const jobIds = await this.redisClient.send('LRANGE', [priorityKey, start, end])
+      const jobIds = await this.redisClient.send('LRANGE', [priorityKey, start.toString(), end.toString()])
 
       if (jobIds && jobIds.length > 0) {
         for (const jobId of jobIds) {
@@ -323,7 +323,7 @@ export class PriorityQueue<T = any> {
 
       if (length && length > 0) {
         // Get all jobs from this priority level
-        const jobIds = await this.redisClient.send('LRANGE', [priorityKey, 0, -1])
+        const jobIds = await this.redisClient.send('LRANGE', [priorityKey, '0', '-1'])
 
         if (jobIds && jobIds.length > 0) {
           // Move jobs to the front of the waiting queue in reverse order
@@ -374,7 +374,7 @@ export class PriorityQueue<T = any> {
 
       for (let priority = 0; priority < this.priorityLevels; priority++) {
         const priorityKey = this.getPriorityKey(priority)
-        const jobIds = await this.redisClient.send('LRANGE', [priorityKey, 0, -1])
+        const jobIds = await this.redisClient.send('LRANGE', [priorityKey, '0', '-1'])
 
         if (jobIds && jobIds.length > 0) {
           for (const jobId of jobIds) {

src/queue.ts

@@ -10,6 +10,8 @@ import { Metrics } from './metrics'
 import { StalledJobChecker } from './stalled-checker'
 import { generateId, getRedisClient, mergeOptions } from './utils'
 import { Worker } from './worker'
+import { RateLimiter } from './rate-limiter'
+import type { RateLimitResult } from './rate-limiter'
 
 export class Queue<T = any> {
   name: string
@@ -22,13 +24,16 @@ export class Queue<T = any> {
   private cleanupService: CleanupService | null = null
   private stalledChecker: StalledJobChecker | null = null
   private readonly logger = createLogger()
+  private limiter: RateLimiter | null = null
+  private defaultJobOptions: JobOptions | undefined
 
   constructor(name: string, options?: QueueConfig) {
     this.name = name
     this.prefix = options?.prefix || config.prefix || 'queue'
     this.redisClient = getRedisClient(options)
     this.keyPrefix = `${this.prefix}:${this.name}`
     this.events = new JobEvents(name)
+    this.defaultJobOptions = options?.defaultJobOptions
 
     // Set logger level if specified
     if (options?.logLevel) {
@@ -51,6 +56,12 @@ export class Queue<T = any> {
       this.logger.debug(`Stalled job checker started for queue ${name}`)
     }
 
+    // Initialize rate limiter if provided
+    if (options?.limiter) {
+      this.limiter = new RateLimiter(this, options.limiter)
+      this.logger.debug(`Rate limiter configured for queue ${name}`)
+    }
+
     // Initialize scripts
     this.init()
   }
@@ -71,10 +82,29 @@ export class Queue<T = any> {
   }
 
   /**
-   * Add a job to the queue
+   * Add a job to the queue with rate limiting support
    */
   async add(data: T, options?: JobOptions): Promise<Job<T>> {
     try {
+      // Check rate limit if configured
+      if (this.limiter) {
+        // If we have keyPrefix in the limiter, check rate limit based on data
+        const limiterResult = await this.limiter.check(data)
+
+        if (limiterResult.limited) {
+          this.logger.warn(`Rate limit exceeded, retrying in ${limiterResult.resetIn}ms`)
+
+          // If rate limited, add to delayed queue with the reset time
+          const opts = {
+            ...this.defaultJobOptions,
+            ...options,
+            delay: limiterResult.resetIn
+          }
+
+          return this.add(data, opts)
+        }
+      }
+
       const opts = mergeOptions(options)
       const jobId = opts.jobId || generateId()
       const timestamp = Date.now()
@@ -492,4 +522,28 @@ export class Queue<T = any> {
   getJobKey(jobId: string): string {
     return `${this.keyPrefix}:job:${jobId}`
   }
+
+  /**
+   * Check if the queue is rate limited for a specific key
+   */
+  async isRateLimited(key?: string, data?: T): Promise<{ limited: boolean; resetIn: number }> {
+    if (!this.limiter) {
+      return { limited: false, resetIn: 0 }
+    }
+
+    let result: RateLimitResult
+
+    if (key) {
+      // Use explicit key if provided
+      result = await this.limiter.checkByKey(key)
+    } else {
+      // Use data with keyPrefix from limiter options
+      result = await this.limiter.check(data)
+    }
+
+    return {
+      limited: result.limited,
+      resetIn: result.resetIn,
+    }
+  }
 }

src/rate-limiter.ts

@@ -19,13 +19,53 @@ export class RateLimiter {
   /**
    * Check if the rate limit has been exceeded
    */
-  async check(): Promise<RateLimitResult> {
+  async check(data?: any): Promise<RateLimitResult> {
     const key = this.queue.getKey('limit')
     const now = Date.now()
 
+    // Determine the identifier based on keyPrefix if provided
+    let identifier = this.queue.name
+
+    if (this.options.keyPrefix && data) {
+      if (typeof this.options.keyPrefix === 'function') {
+        // Use the function to generate a key from the data
+        const keyValue = this.options.keyPrefix(data)
+        identifier = `${this.queue.name}:${keyValue}`
+      } else {
+        // Use the keyPrefix as a property path in the data object
+        const keyValue = data[this.options.keyPrefix]
+        if (keyValue) {
+          identifier = `${this.queue.name}:${keyValue}`
+        }
+      }
+    }
+
     const result = await this.queue.redisClient.send('rateLimit', [
       key,
-      this.queue.name,
+      identifier,
+      this.options.max.toString(),
+      this.options.duration.toString(),
+      now.toString(),
+    ])
+
+    return {
+      limited: result[0] === 1,
+      remaining: result[1],
+      resetIn: result[2],
+    }
+  }
+
+  /**
+   * Check if the rate limit has been exceeded for a specific key
+   */
+  async checkByKey(key: string): Promise<RateLimitResult> {
+    const identifier = `${this.queue.name}:${key}`
+    const limitKey = this.queue.getKey('limit')
+    const now = Date.now()
+
+    const result = await this.queue.redisClient.send('rateLimit', [
+      limitKey,
+      identifier,
       this.options.max.toString(),
       this.options.duration.toString(),
       now.toString(),

src/types.ts

@@ -53,6 +53,7 @@ export interface JobOptions {
 export interface RateLimiter {
   max: number
   duration: number
+  keyPrefix?: string | ((data: any) => string)
 }
 
 export interface Job<T = any> {