smom-dbis-138/services/relay/src/MessageQueue.js

/**
 * Message Queue for managing pending relay messages
 */

import { mkdir, readFile, rename, writeFile } from 'fs/promises';
import path from 'path';

export class MessageQueue {
  constructor(logger, options = {}) {
    this.logger = logger;
    this.persistencePath = options.persistencePath || '';
    this.queue = [];
    this.processed = new Set();
    this.failed = new Set();
    this.retryCounts = new Map();
    this.messageStore = new Map();
    this.inFlight = new Map();
    this.lastPersistedAt = null;
    this.persistenceEnabled = Boolean(this.persistencePath);
  }

  async init() {
    await this.loadSnapshot();
  }

  serialize() {
    return {
      version: 2,
      queue: this.queue,
      processedIds: Array.from(this.processed),
      failedIds: Array.from(this.failed),
      retryCounts: Object.fromEntries(this.retryCounts.entries()),
      savedAt: new Date().toISOString()
    };
  }

  static encodeForPersistence(value) {
    if (typeof value === 'bigint') {
      return { __type: 'bigint', value: value.toString() };
    }
    if (Array.isArray(value)) {
      return value.map((item) => MessageQueue.encodeForPersistence(item));
    }
    if (value && typeof value === 'object') {
      return Object.fromEntries(
        Object.entries(value).map(([key, entry]) => [key, MessageQueue.encodeForPersistence(entry)])
      );
    }
    return value;
  }

  static decodeFromPersistence(value) {
    if (Array.isArray(value)) {
      return value.map((item) => MessageQueue.decodeFromPersistence(item));
    }
    if (value && typeof value === 'object') {
      if (value.__type === 'bigint' && typeof value.value === 'string') {
        return BigInt(value.value);
      }
      return Object.fromEntries(
        Object.entries(value).map(([key, entry]) => [key, MessageQueue.decodeFromPersistence(entry)])
      );
    }
    return value;
  }

  async persistSnapshot() {
    if (!this.persistenceEnabled) return;

    const directory = path.dirname(this.persistencePath);
    const tmpPath = `${this.persistencePath}.tmp`;
    await mkdir(directory, { recursive: true });
    await writeFile(
      tmpPath,
      JSON.stringify(MessageQueue.encodeForPersistence(this.serialize()), null, 2)
    );
    await rename(tmpPath, this.persistencePath);
    this.lastPersistedAt = new Date().toISOString();
  }

  async loadSnapshot() {
    if (!this.persistenceEnabled) return;

    try {
      const raw = await readFile(this.persistencePath, 'utf8');
      const parsed = JSON.parse(raw);
      const restoredQueueRaw = Array.isArray(parsed.queue) ? parsed.queue : [];
      const restoredQueue = MessageQueue.decodeFromPersistence(restoredQueueRaw);
      const restoredRetryCounts = parsed.retryCounts && typeof parsed.retryCounts === 'object'
        ? Object.entries(parsed.retryCounts)
        : [];

      this.queue = restoredQueue;
      this.messageStore.clear();
      for (const message of restoredQueue) {
        if (message?.messageId) {
          this.messageStore.set(message.messageId, message);
        }
      }
      this.retryCounts = new Map(
        restoredRetryCounts.map(([messageId, count]) => [messageId, Number(count) || 0])
      );
      const restoredProcessed = Array.isArray(parsed.processedIds) ? parsed.processedIds : [];
      const restoredFailed = Array.isArray(parsed.failedIds) ? parsed.failedIds : [];
      for (const messageId of restoredProcessed) {
        if (messageId) this.processed.add(messageId);
      }
      for (const messageId of restoredFailed) {
        if (messageId) this.failed.add(messageId);
      }
      this.lastPersistedAt = parsed.savedAt || null;

      if (restoredQueue.length > 0 || restoredProcessed.length > 0) {
        this.logger.info(
          `Loaded relay snapshot from ${this.persistencePath}: ` +
            `queue=${restoredQueue.length} processed=${restoredProcessed.length} failed=${restoredFailed.length}`
        );
      }
    } catch (error) {
      if (error?.code === 'ENOENT') {
        return;
      }
      this.logger.warn(`Failed to load relay queue snapshot from ${this.persistencePath}`, error);
    }
  }
  
  async add(messageData, options = {}) {
    const messageId = messageData.messageId;
    const forceReplay = options.forceReplay === true;

    // Skip if already processed
    if (this.processed.has(messageId)) {
      this.logger.debug(`Message ${messageId} already processed, skipping`);
      return;
    }

    if (this.failed.has(messageId)) {
      if (!forceReplay) {
        this.logger.debug(`Message ${messageId} already failed, skipping`);
        return;
      }
      this.failed.delete(messageId);
      this.logger.info(`Cleared failed status for replayed message ${messageId}`);
    }
    
    // Check if already in queue
    const existingIndex = this.queue.findIndex(m => m.messageId === messageId);
    if (existingIndex >= 0) {
      this.logger.debug(`Message ${messageId} already in queue`);
      return;
    }
    
    this.messageStore.set(messageId, messageData);
    // Add to queue
    this.queue.push(messageData);
    await this.persistSnapshot();
    this.logger.info(`Added message ${messageId} to queue. Queue size: ${this.queue.length}`);
  }
  
  async getNext() {
    if (this.queue.length === 0) {
      return null;
    }
    
    const messageData = this.queue.shift();
    if (messageData && messageData.messageId) {
      this.inFlight.set(messageData.messageId, messageData);
      await this.persistSnapshot();
    }
    return messageData;
  }
  
  async markProcessed(messageId) {
    this.processed.add(messageId);
    this.retryCounts.delete(messageId);
    this.inFlight.delete(messageId);
    this.messageStore.delete(messageId);
    this.queue = this.queue.filter((m) => m.messageId !== messageId);
    await this.persistSnapshot();
    this.logger.info(`Message ${messageId} marked as processed`);
  }
  
  async markFailed(messageId) {
    this.failed.add(messageId);
    this.retryCounts.delete(messageId);
    this.inFlight.delete(messageId);
    await this.persistSnapshot();
    this.logger.error(`Message ${messageId} marked as failed`);
  }
  
  async getRetryCount(messageId) {
    return this.retryCounts.get(messageId) || 0;
  }

  async resetRetryCount(messageId) {
    this.retryCounts.delete(messageId);
    await this.persistSnapshot();
  }
  
  async retry(messageId, options = {}) {
    const { increment = true } = options;
    const count = this.retryCounts.get(messageId) || 0;
    if (increment) {
      this.retryCounts.set(messageId, count + 1);
      await this.persistSnapshot();
    }

    const existingIndex = this.queue.findIndex(m => m.messageId === messageId);
    if (existingIndex >= 0) {
      const nextCount = increment ? count + 1 : count;
      this.logger.info(`Message ${messageId} retry count: ${nextCount}`);
      return;
    }

    const messageData = this.inFlight.get(messageId) || this.messageStore.get(messageId);
    if (!messageData) {
      this.logger.warn(`Cannot requeue ${messageId}; original message payload is unavailable`);
      return;
    }

    this.inFlight.delete(messageId);
    this.queue.push(messageData);
    await this.persistSnapshot();
    const nextCount = increment ? count + 1 : count;
    this.logger.info(`Message ${messageId} requeued. Retry count: ${nextCount}. Queue size: ${this.queue.length}`);
  }
  
  getStats() {
    return {
      queueSize: this.queue.length,
      processed: this.processed.size,
      failed: this.failed.size,
      inFlight: this.inFlight.size,
      persistenceEnabled: this.persistenceEnabled,
      lastPersistedAt: this.lastPersistedAt
    };
  }
}