smartproxy/ts/networkproxy/classes.np.connectionpool.ts

import * as plugins from '../plugins.js';
import { type INetworkProxyOptions, type IConnectionEntry, type ILogger, createLogger } from './classes.np.types.js';

/**
 * Manages a pool of backend connections for efficient reuse
 */
export class ConnectionPool {
  private connectionPool: Map<string, Array<IConnectionEntry>> = new Map();
  private roundRobinPositions: Map<string, number> = new Map();
  private logger: ILogger;
  
  constructor(private options: INetworkProxyOptions) {
    this.logger = createLogger(options.logLevel || 'info');
  }
  
  /**
   * Get a connection from the pool or create a new one
   */
  public getConnection(host: string, port: number): Promise<plugins.net.Socket> {
    return new Promise((resolve, reject) => {
      const poolKey = `${host}:${port}`;
      const connectionList = this.connectionPool.get(poolKey) || [];
      
      // Look for an idle connection
      const idleConnectionIndex = connectionList.findIndex(c => c.isIdle);
      
      if (idleConnectionIndex >= 0) {
        // Get existing connection from pool
        const connection = connectionList[idleConnectionIndex];
        connection.isIdle = false;
        connection.lastUsed = Date.now();
        this.logger.debug(`Reusing connection from pool for ${poolKey}`);
        
        // Update the pool
        this.connectionPool.set(poolKey, connectionList);
        
        resolve(connection.socket);
        return;
      }
      
      // No idle connection available, create a new one if pool isn't full
      const poolSize = this.options.connectionPoolSize || 50;
      if (connectionList.length < poolSize) {
        this.logger.debug(`Creating new connection to ${host}:${port}`);
        
        try {
          const socket = plugins.net.connect({
            host,
            port,
            keepAlive: true,
            keepAliveInitialDelay: 30000 // 30 seconds
          });
          
          socket.once('connect', () => {
            // Add to connection pool
            const connection = {
              socket,
              lastUsed: Date.now(),
              isIdle: false
            };
            
            connectionList.push(connection);
            this.connectionPool.set(poolKey, connectionList);
            
            // Setup cleanup when the connection is closed
            socket.once('close', () => {
              const idx = connectionList.findIndex(c => c.socket === socket);
              if (idx >= 0) {
                connectionList.splice(idx, 1);
                this.connectionPool.set(poolKey, connectionList);
                this.logger.debug(`Removed closed connection from pool for ${poolKey}`);
              }
            });
            
            resolve(socket);
          });
          
          socket.once('error', (err) => {
            this.logger.error(`Error creating connection to ${host}:${port}`, err);
            reject(err);
          });
        } catch (err) {
          this.logger.error(`Failed to create connection to ${host}:${port}`, err);
          reject(err);
        }
      } else {
        // Pool is full, wait for an idle connection or reject
        this.logger.warn(`Connection pool for ${poolKey} is full (${connectionList.length})`);
        reject(new Error(`Connection pool for ${poolKey} is full`));
      }
    });
  }
  
  /**
   * Return a connection to the pool for reuse
   */
  public returnConnection(socket: plugins.net.Socket, host: string, port: number): void {
    const poolKey = `${host}:${port}`;
    const connectionList = this.connectionPool.get(poolKey) || [];
    
    // Find this connection in the pool
    const connectionIndex = connectionList.findIndex(c => c.socket === socket);
    
    if (connectionIndex >= 0) {
      // Mark as idle and update last used time
      connectionList[connectionIndex].isIdle = true;
      connectionList[connectionIndex].lastUsed = Date.now();
      
      this.logger.debug(`Returned connection to pool for ${poolKey}`);
    } else {
      this.logger.warn(`Attempted to return unknown connection to pool for ${poolKey}`);
    }
  }
  
  /**
   * Cleanup the connection pool by removing idle connections
   * or reducing pool size if it exceeds the configured maximum
   */
  public cleanupConnectionPool(): void {
    const now = Date.now();
    const idleTimeout = this.options.keepAliveTimeout || 120000; // 2 minutes default
    
    for (const [host, connections] of this.connectionPool.entries()) {
      // Sort by last used time (oldest first)
      connections.sort((a, b) => a.lastUsed - b.lastUsed);
      
      // Remove idle connections older than the idle timeout
      let removed = 0;
      while (connections.length > 0) {
        const connection = connections[0];
        
        // Remove if idle and exceeds timeout, or if pool is too large
        if ((connection.isIdle && now - connection.lastUsed > idleTimeout) ||
            connections.length > (this.options.connectionPoolSize || 50)) {
          
          try {
            if (!connection.socket.destroyed) {
              connection.socket.end();
              connection.socket.destroy();
            }
          } catch (err) {
            this.logger.error(`Error destroying pooled connection to ${host}`, err);
          }
          
          connections.shift(); // Remove from pool
          removed++;
        } else {
          break; // Stop removing if we've reached active or recent connections
        }
      }
      
      if (removed > 0) {
        this.logger.debug(`Removed ${removed} idle connections from pool for ${host}, ${connections.length} remaining`);
      }
      
      // Update the pool with the remaining connections
      if (connections.length === 0) {
        this.connectionPool.delete(host);
      } else {
        this.connectionPool.set(host, connections);
      }
    }
  }
  
  /**
   * Close all connections in the pool
   */
  public closeAllConnections(): void {
    for (const [host, connections] of this.connectionPool.entries()) {
      this.logger.debug(`Closing ${connections.length} connections to ${host}`);
      
      for (const connection of connections) {
        try {
          if (!connection.socket.destroyed) {
            connection.socket.end();
            connection.socket.destroy();
          }
        } catch (error) {
          this.logger.error(`Error closing connection to ${host}:`, error);
        }
      }
    }
    
    this.connectionPool.clear();
    this.roundRobinPositions.clear();
  }
  
  /**
   * Get load balancing target using round-robin
   */
  public getNextTarget(targets: string[], port: number): { host: string, port: number } {
    const targetKey = targets.join(',');
    
    // Initialize position if not exists
    if (!this.roundRobinPositions.has(targetKey)) {
      this.roundRobinPositions.set(targetKey, 0);
    }
    
    // Get current position and increment for next time
    const currentPosition = this.roundRobinPositions.get(targetKey)!;
    const nextPosition = (currentPosition + 1) % targets.length;
    this.roundRobinPositions.set(targetKey, nextPosition);
    
    // Return the selected target
    return {
      host: targets[currentPosition],
      port
    };
  }
  
  /**
   * Gets the connection pool status
   */
  public getPoolStatus(): Record<string, { total: number, idle: number }> {
    return Object.fromEntries(
      Array.from(this.connectionPool.entries()).map(([host, connections]) => [
        host, 
        {
          total: connections.length, 
          idle: connections.filter(c => c.isIdle).length
        }
      ])
    );
  }
  
  /**
   * Setup a periodic cleanup task
   */
  public setupPeriodicCleanup(interval: number = 60000): NodeJS.Timeout {
    const timer = setInterval(() => {
      this.cleanupConnectionPool();
    }, interval);
    
    // Don't prevent process exit
    if (timer.unref) {
      timer.unref();
    }
    
    return timer;
  }
}