@applitools/execution-grid-tunnel/src/tunnel-process-manager/tunnel-connection-pool.js

Allows user to run tests with exection-grid and navigate to private hosts and ips

const {TunnelConnection, ErrorCircuitBreaker} = require('../utils')
const {v4: uuid} = require('uuid')

const POOL_STATUS = {
  OFF: 'off', // start is not called yet
  INIT: 'init', // 
  INIT_ERROR: 'init_error',
  RUNNING: 'running', // The process connected to the server
  RECONNECT: 'reconnect', // The process tries to reconnect to server
  ERROR: 'error',
  STOPPING: 'stoping', // User asks to stop the tunnel
}

class TunnelConnectionPool{
  constructor(options){
    this._isPoolShuttingDown = false
    this._haveConnectionSucceeded = false
    this._options = options
    this._connectionMap = new Map() 
    this._createConnectionMap = new Map() 
    this._connectionStatusMap = new Map()
    this._availableConnectionMap = new Map() 
    this._logger = this._options.logger
    this._onSetStatusCallbacks = []
    this._poolStatus = POOL_STATUS.OFF

    this._circuitBreaker = new ErrorCircuitBreaker({logger: this._options.logger})
    this._circuitBreaker.onOpen(() => {
      if (this._poolStatus === POOL_STATUS.RUNNING){
        this.setStatus(POOL_STATUS.RECONNECT)
      }
    })
  }

  async _createNewConnection(connectionId){
    const {tunnelId, host, port, token, keepAliveMessage, connectedMessage, protocol, localProxyOptions, logger} = this._options
    const TunnelConnectionClass = this._options.TunnelConnection || TunnelConnection
    const connection = new TunnelConnectionClass({keepAliveMessage, connectedMessage, logger, connectionId})
    await connection.connect({tunnelId, host, port, protocol, token, localProxyOptions})
    return connection
  }

  async _handleConnectionLifecycle(){
    await this._circuitBreaker.waitForClosedState

    const id = uuid()
    try{
      this._createConnectionMap.set(id, id)
      
      const connection = await this._createNewConnection(id)

      this._createConnectionMap.delete(id)

      this._fireStatus('connected', id, connection)

      connection._remoteConnection?.on('error', e => {
        if (connection._remoteConnection?.closed && connection._localConnection.closed){
          this._fireStatus('end', id)
        }else {
          this._fireStatus('error', id, e)
        }
      })

      connection._localConnection?.on('close', () => {
          this._fireStatus('end', id)
      })

      const occupiedPromise = connection.waitForAttachingRequest()
        .then(()=> this._fireStatus('occupied', id))
      
      const endPromise = connection.waitForRemoteConnectionClosing()
        .then(() => this._fireStatus('end', id))
      
      this._circuitBreaker.resetErrorCounter()
      this.setStatus(POOL_STATUS.RUNNING)

      return
    }catch(e){
      this._createConnectionMap.delete(id)
      
      this._logger.warn({
        action: 'handle-connection-lifecycle',
        error: e.message,
        id
      })
      
      this._circuitBreaker.incErrorCounter()
      this._fireStatus('error', id, e)
      return e
    }
  }

  async _fireStatus(status, id, data = undefined){
    //console.log(status, id)
    this._logger.debug(status, id)

    if(status === 'connected'){
      this._connectionMap.set(id, data)
      this._connectionStatusMap.set(id, status)
      this._availableConnectionMap.set(id, id)

      if (!this._haveConnectionSucceeded){
        this._haveConnectionSucceeded = true
        const count = this._options.preAllocation - this._connectionMap.size
      
        for(let i=0; i< count; i++){
          this._handleConnectionLifecycle()
        }
      }
    }

    if (status === 'occupied'){
      if (!this._connectionStatusMap.has(id)) return
      this._connectionStatusMap.set(id, 'occupied')
      this._availableConnectionMap.delete(id)
      this._connectionStatusMap.delete(id)
    }

    if (status === 'error'){
      this._availableConnectionMap.delete(id)
      this._connectionStatusMap.set(id, "error")
    }
    
    if (status === 'end'){
      this._availableConnectionMap.delete(id)
      this._connectionStatusMap.delete(id)
      this._connectionMap.delete(id)
    }

    if(status === 'occupied' || status === 'end' || status === 'error'){
      // console.log(this._haveConnectionSucceeded, this._shouldCloseConnection)
      if (this._haveConnectionSucceeded && !this._isPoolShuttingDown){
        if ((this._availableConnectionMap.size + this._createConnectionMap.size) < this._options.preAllocation &&
            this._connectionMap.size < this._options.maxConnections){
              this._handleConnectionLifecycle()
          }
      }
    }
    this._logger.debug({
      action: 'pool-statistics',
      available: this._availableConnectionMap.size,
      inCreation: this._createConnectionMap.size,
      total: this._connectionMap.size
    })
  }

  async start(){
    this.setStatus(POOL_STATUS.INIT)
    const error = await this._handleConnectionLifecycle()
    
    if (error){
      this._logger.error({
        action: 'tunnel-pool-error',
        error: error.message
      })
      this.setStatus(POOL_STATUS.INIT_ERROR)

      throw error
    }
  }
  
  async end(timeout = 5000){
    this._isPoolShuttingDown = true

    for(const connection of this._connectionMap.values()){
      connection.end()
    }
    
    return new Promise(async (resolve, reject) => {
      const timeoutId = setTimeout(() => reject(new Error('Timeout Error')), timeout)
      while(this._connectionMap.size > 0){
        await new Promise(r => setTimeout(r,50))
      }
      clearTimeout(timeoutId)
      resolve()
    })
  }

  async onSetStatus(cb){
    this._onSetStatusCallbacks.push(cb)
  }
  
  async setStatus(status){
    if (this._poolStatus === status) return

    this._poolStatus = status

    for(const cb of this._onSetStatusCallbacks){
      cb(status)
    }
  }
  
  destroy(){
    this._isPoolShuttingDown = true
    this.setStatus(POOL_STATUS.STOPPING)
    for(const connection of this._connectionMap.values()){
      connection.destroy()
    }
  }
} 

module.exports = {TunnelConnectionPool, POOL_STATUS}