Simple TypeScript+DynamoDB rate limiter

rate-limiter.ts

import { DateTime, DateTimeUnit, DurationLike } from "luxon";
import { deburr, kebabCase } from "lodash";
import { marshall, unmarshall } from "@aws-sdk/util-dynamodb";
import { 
  ConditionalCheckFailedException,
  DynamoDBClient,
  UpdateItemCommand,
} from "@aws-sdk/client-dynamodb";

const client = new DynamoDBClient({ region: 'sa-east-1' })

const TABLE_NAME = 'rate-limiting-test'
const TTL: DurationLike = { hours: 1 } // how long to keep rate limiting records

type Plan = {
  name: string,
  limitPerUnit: number, 
  unit: DateTimeUnit
}

type PlanCheck = {
  isBlocked: boolean,
  plan: Plan,
  remaining: number,
  nextReset: Date,
  latency: number,
}

type RateLimitCheck = {
  isBlocked: boolean,
  latency: number,
  planChecks: PlanCheck[],
}

const isConditionalCheckFailed = (err: any): err is ConditionalCheckFailedException =>
  err instanceof ConditionalCheckFailedException

const checkPlanUsage = async (
  resourceId: string, 
  now: DateTime, 
  plan: Plan,
): Promise<PlanCheck> => {
  const bucket = now.startOf(plan.unit)
  const nextReset = now.plus({ [plan.unit]: 1 }).startOf(plan.unit)
  const planId = kebabCase(deburr(plan.name.trim()))
  const key = `${resourceId}:${planId}`

  const command = new UpdateItemCommand({
    TableName: TABLE_NAME,
    Key: marshall({ key, bucket: bucket.toISO() }),
    UpdateExpression: 'SET #usage = if_not_exists(#usage, :zero) + :one, #limit = :limit, #ttl = :ttl',
    ConditionExpression: '(attribute_not_exists(#usage) or #usage < :limit) and :limit > :zero',
    ExpressionAttributeNames: {
      '#usage': 'usage',
      '#limit': 'limit',
      '#ttl': 'ttl'
    },
    ExpressionAttributeValues: marshall({
      ':zero': 0,
      ':one': 1,
      ':limit': plan.limitPerUnit,
      ':ttl': nextReset.plus(TTL).toUnixInteger(),
    }),
    ReturnValues: 'UPDATED_NEW'
  })

  let isBlocked = true
  let remaining = 0
  const start = new Date()
  
  try {
    const { Attributes: attributes = {} } = await client.send(command)
    const { usage } = unmarshall(attributes)

    isBlocked = false
    remaining = plan.limitPerUnit - usage
  } catch (err) {
    if (!isConditionalCheckFailed(err)) {
      throw err
    }
  } finally {
    const end = new Date()
    return {
      isBlocked,
      plan,
      remaining,
      nextReset: nextReset.toJSDate(),
      latency: end.valueOf() - start.valueOf()
    }
  }
}

const checkRateLimit = async (
  resourceId: string, 
  plans: Plan[],
): Promise<RateLimitCheck> => {
  const now = DateTime.now().toUTC()
  const planChecks = await Promise.all(
    plans.map(up => checkPlanUsage(resourceId, now, up))
  )

  return {
    isBlocked: planChecks.some(pc => pc.isBlocked),
    latency: now.diffNow().toMillis() * -1,
    planChecks,
  }
}

async function main() {
  const resourceId = 'user-123'

  const hourLimits: Plan = {
    name: 'Hourly usage',
    unit: 'hour',
    limitPerUnit: 3
  }

  const minuteLimits: Plan = {
    name: 'Minute usage',
    unit: 'minute',
    limitPerUnit: 3
  }

  const results = await Promise.all([
    checkRateLimit(resourceId, [hourLimits, minuteLimits]),
    checkRateLimit(resourceId, [hourLimits, minuteLimits]),
    checkRateLimit(resourceId, [hourLimits, minuteLimits]),
    checkRateLimit(resourceId, [hourLimits, minuteLimits]),
    checkRateLimit(resourceId, [hourLimits, minuteLimits]),
  ])

  console.log(
    JSON.stringify(
      results.map(r => ({ 
        isBlocked: r.isBlocked,
        latency: r.latency,
      }))
    , null, 2)
  )
}

main().catch(e => console.log(e))

CONS/PROS

1. CON: Two separate network calls
2. CON: If checking against multiple plans in case one of them fails the request is still counted against the user quota
3. PRO: Possible to retrieve the remaining requests info so you can attach it to the response
4. PRO: Can handle large number of requests