RedditBot/src/utils/rate-limiter.ts

import PQueue from 'p-queue';
import type { Config } from '../config/index.js';
import type { Logger } from './logger.js';

const DEFAULT_HOUR_WINDOW_MS = 3600000;

/** Loads timestamps of successful sends made since the supplied epoch time. */
export type SentMessageHistoryProvider = (since: number) => Promise<readonly number[]>;

export interface RateLimiterDependencies {
  /**
   * Optional durable history used to restore the rolling hourly limit after a
   * process restart. The database adapter deliberately lives outside this
   * utility so the limiter remains independently testable.
   */
  sentMessageHistory?: SentMessageHistoryProvider;
  now?: () => number;
  sleep?: (ms: number, signal?: AbortSignal) => Promise<void>;
  random?: () => number;
}

export class RateLimiter {
  private queue: PQueue;
  private lastMessageTime = 0;
  private successfulMessageTimes: number[] = [];
  private initialization: Promise<void> | null = null;
  private readonly hourWindowMs: number;
  private readonly now: () => number;
  private readonly sleep: (ms: number, signal?: AbortSignal) => Promise<void>;
  private readonly random: () => number;

  constructor(
    private config: Config,
    private logger: Logger,
    hourWindowMs: number = DEFAULT_HOUR_WINDOW_MS,
    dependencies: RateLimiterDependencies = {},
  ) {
    this.queue = new PQueue({ concurrency: 1 });
    this.hourWindowMs = hourWindowMs;
    this.now = dependencies.now ?? Date.now;
    this.sleep = dependencies.sleep ?? abortableSleep;
    this.random = dependencies.random ?? Math.random;
    if (dependencies.sentMessageHistory) {
      this.initialization = this.loadHistory(dependencies.sentMessageHistory);
    }
  }

  /** Wait for durable rate-limit history to be loaded before accepting work. */
  async initialize(): Promise<void> {
    await this.initialization;
  }

  async throttle<T>(
    fn: () => Promise<T>,
    shouldCount: (result: T) => boolean = () => true,
    signal?: AbortSignal,
    priority = 0,
  ): Promise<T> {
    return this.queue.add(async () => {
      await this.initialize();
      throwIfAborted(signal);
      const now = this.now();
      this.pruneHistory(now);

      // Enforce the limit over a rolling hour so restarts cannot reset it.
      if (this.successfulMessageTimes.length >= this.config.rateLimit.maxMessagesPerHour) {
        const oldestSend = this.successfulMessageTimes[0];
        const waitTime = Math.max(0, oldestSend + this.hourWindowMs - now);
        this.logger.info(`Hourly limit reached, waiting ${Math.round(waitTime / 60000)} minutes`);
        await this.sleep(waitTime, signal);
        this.pruneHistory(this.now());
      }

      // Enforce minimum delay with jitter
      // We check against the previously completed message time
      const timeSinceLastMessage = now - this.lastMessageTime;
      const minDelay = this.config.rateLimit.minDelayBetweenMessagesMs;
      const jitter = this.random() * this.config.rateLimit.jitterMs;

      if (this.lastMessageTime > 0 && timeSinceLastMessage < minDelay) {
        const waitTime = minDelay - timeSinceLastMessage + jitter;
        this.logger.debug(`Rate limiting: waiting ${Math.round(waitTime / 1000)} seconds`);
        await this.sleep(waitTime, signal);
      }

      // Execute function
      const result = await fn();

      // Update stats only if shouldCount returns true
      if (shouldCount(result)) {
        this.lastMessageTime = this.now();
        this.successfulMessageTimes.push(this.lastMessageTime);
      }

      return result;
    }, { priority }) as Promise<T>;
  }

  private async loadHistory(provider: SentMessageHistoryProvider): Promise<void> {
    const now = this.now();
    const timestamps = await provider(now - this.hourWindowMs);
    this.successfulMessageTimes = timestamps
      .filter((timestamp) => Number.isFinite(timestamp) && timestamp > now - this.hourWindowMs && timestamp <= now)
      .sort((a, b) => a - b);
    this.lastMessageTime = this.successfulMessageTimes.at(-1) ?? 0;
  }

  private pruneHistory(now: number): void {
    const cutoff = now - this.hourWindowMs;
    const before = this.successfulMessageTimes.length;
    this.successfulMessageTimes = this.successfulMessageTimes.filter((timestamp) => timestamp > cutoff);
    if (before !== this.successfulMessageTimes.length) {
      this.logger.debug('Hourly rate limit counter reset');
    }
  }

  getStats(): { messagesThisHour: number; minutesSinceLastMessage: number } {
    this.pruneHistory(this.now());
    return {
      messagesThisHour: this.successfulMessageTimes.length,
      minutesSinceLastMessage: Math.round((this.now() - this.lastMessageTime) / 60000),
    };
  }

  /**
   * Get active timer information for the dashboard
   */
  getTimers(): {
    nextMessageAvailableAt: number | null;
    hourlyResetAt: number;
    isHourlyLimitReached: boolean;
    messagesRemaining: number;
  } {
    const now = this.now();
    this.pruneHistory(now);
    const minDelay = this.config.rateLimit.minDelayBetweenMessagesMs;
    const maxMessages = this.config.rateLimit.maxMessagesPerHour;

    // Calculate when the next message can be sent based on min delay
    let nextMessageAvailableAt: number | null = null;
    if (this.lastMessageTime > 0) {
      const nextAvailable = this.lastMessageTime + minDelay;
      if (nextAvailable > now) {
        nextMessageAvailableAt = nextAvailable;
      }
    }

    // Calculate when the hourly counter resets
    const hourlyResetAt = (this.successfulMessageTimes[0] ?? now) + this.hourWindowMs;

    // Check if hourly limit is reached
    const isHourlyLimitReached = this.successfulMessageTimes.length >= maxMessages;

    // If hourly limit is reached, next message is available when the hour resets
    if (isHourlyLimitReached) {
      nextMessageAvailableAt = hourlyResetAt;
    }

    return {
      nextMessageAvailableAt,
      hourlyResetAt,
      isHourlyLimitReached,
      messagesRemaining: Math.max(0, maxMessages - this.successfulMessageTimes.length),
    };
  }
}

function throwIfAborted(signal?: AbortSignal): void {
  if (signal?.aborted) throw new DOMException('Operation aborted', 'AbortError');
}

function abortableSleep(ms: number, signal?: AbortSignal): Promise<void> {
  if (signal?.aborted) return Promise.reject(new DOMException('Operation aborted', 'AbortError'));
  return new Promise((resolve, reject) => {
    let timer: NodeJS.Timeout;
    const onAbort = (): void => {
      clearTimeout(timer);
      cleanup();
      reject(new DOMException('Operation aborted', 'AbortError'));
    };
    const cleanup = (): void => signal?.removeEventListener('abort', onAbort);
    const cleanupAndResolve = (): void => {
      cleanup();
      resolve();
    };
    timer = setTimeout(cleanupAndResolve, ms);
    signal?.addEventListener('abort', onAbort, { once: true });
  });
}