smartproxy/ts/core/utils/lifecycle-component.ts

/**
 * Base class for components that need proper resource lifecycle management
 * Provides automatic cleanup of timers and event listeners to prevent memory leaks
 */
export abstract class LifecycleComponent {
  private timers: Set<NodeJS.Timeout> = new Set();
  private intervals: Set<NodeJS.Timeout> = new Set();
  private listeners: Array<{ 
    target: any; 
    event: string; 
    handler: Function;
    once?: boolean;
  }> = [];
  private childComponents: Set<LifecycleComponent> = new Set();
  protected isShuttingDown = false;
  private cleanupPromise?: Promise<void>;

  /**
   * Create a managed setTimeout that will be automatically cleaned up
   */
  protected setTimeout(handler: Function, timeout: number): NodeJS.Timeout {
    if (this.isShuttingDown) {
      // Return a dummy timer if shutting down
      return setTimeout(() => {}, 0);
    }

    const wrappedHandler = () => {
      this.timers.delete(timer);
      if (!this.isShuttingDown) {
        handler();
      }
    };

    const timer = setTimeout(wrappedHandler, timeout);
    this.timers.add(timer);
    return timer;
  }

  /**
   * Create a managed setInterval that will be automatically cleaned up
   */
  protected setInterval(handler: Function, interval: number): NodeJS.Timeout {
    if (this.isShuttingDown) {
      // Return a dummy timer if shutting down
      return setInterval(() => {}, interval);
    }

    const wrappedHandler = () => {
      if (!this.isShuttingDown) {
        handler();
      }
    };

    const timer = setInterval(wrappedHandler, interval);
    this.intervals.add(timer);
    
    // Allow process to exit even with timer
    if (typeof timer.unref === 'function') {
      timer.unref();
    }
    
    return timer;
  }

  /**
   * Clear a managed timeout
   */
  protected clearTimeout(timer: NodeJS.Timeout): void {
    clearTimeout(timer);
    this.timers.delete(timer);
  }

  /**
   * Clear a managed interval
   */
  protected clearInterval(timer: NodeJS.Timeout): void {
    clearInterval(timer);
    this.intervals.delete(timer);
  }

  /**
   * Add a managed event listener that will be automatically removed on cleanup
   */
  protected addEventListener(
    target: any, 
    event: string, 
    handler: Function,
    options?: { once?: boolean }
  ): void {
    if (this.isShuttingDown) {
      return;
    }

    // For 'once' listeners, we need to wrap the handler to remove it from our tracking
    let actualHandler = handler;
    if (options?.once) {
      actualHandler = (...args: any[]) => {
        // Call the original handler
        handler(...args);
        
        // Remove from our internal tracking
        const index = this.listeners.findIndex(
          l => l.target === target && l.event === event && l.handler === handler
        );
        if (index !== -1) {
          this.listeners.splice(index, 1);
        }
      };
    }

    // Support both EventEmitter and DOM-style event targets
    if (typeof target.on === 'function') {
      if (options?.once) {
        target.once(event, actualHandler);
      } else {
        target.on(event, actualHandler);
      }
    } else if (typeof target.addEventListener === 'function') {
      target.addEventListener(event, actualHandler, options);
    } else {
      throw new Error('Target must support on() or addEventListener()');
    }

    // Store the original handler in our tracking (not the wrapped one)
    this.listeners.push({ 
      target, 
      event, 
      handler,
      once: options?.once 
    });
  }

  /**
   * Remove a specific event listener
   */
  protected removeEventListener(target: any, event: string, handler: Function): void {
    // Remove from target
    if (typeof target.removeListener === 'function') {
      target.removeListener(event, handler);
    } else if (typeof target.removeEventListener === 'function') {
      target.removeEventListener(event, handler);
    }

    // Remove from our tracking
    const index = this.listeners.findIndex(
      l => l.target === target && l.event === event && l.handler === handler
    );
    if (index !== -1) {
      this.listeners.splice(index, 1);
    }
  }

  /**
   * Register a child component that should be cleaned up when this component is cleaned up
   */
  protected registerChildComponent(component: LifecycleComponent): void {
    this.childComponents.add(component);
  }

  /**
   * Unregister a child component
   */
  protected unregisterChildComponent(component: LifecycleComponent): void {
    this.childComponents.delete(component);
  }

  /**
   * Override this method to implement component-specific cleanup logic
   */
  protected async onCleanup(): Promise<void> {
    // Override in subclasses
  }

  /**
   * Clean up all managed resources
   */
  public async cleanup(): Promise<void> {
    // Return existing cleanup promise if already cleaning up
    if (this.cleanupPromise) {
      return this.cleanupPromise;
    }

    this.cleanupPromise = this.performCleanup();
    return this.cleanupPromise;
  }

  private async performCleanup(): Promise<void> {
    this.isShuttingDown = true;
    
    // First, clean up child components
    const childCleanupPromises: Promise<void>[] = [];
    for (const child of this.childComponents) {
      childCleanupPromises.push(child.cleanup());
    }
    await Promise.all(childCleanupPromises);
    this.childComponents.clear();

    // Clear all timers
    for (const timer of this.timers) {
      clearTimeout(timer);
    }
    this.timers.clear();

    // Clear all intervals
    for (const timer of this.intervals) {
      clearInterval(timer);
    }
    this.intervals.clear();

    // Remove all event listeners
    for (const { target, event, handler } of this.listeners) {
      // All listeners need to be removed, including 'once' listeners that might not have fired
      if (typeof target.removeListener === 'function') {
        target.removeListener(event, handler);
      } else if (typeof target.removeEventListener === 'function') {
        target.removeEventListener(event, handler);
      }
    }
    this.listeners = [];

    // Call subclass cleanup
    await this.onCleanup();
  }

  /**
   * Check if the component is shutting down
   */
  protected isShuttingDownState(): boolean {
    return this.isShuttingDown;
  }
}