<commit_before>_ = require "lodash"
genericPool = require "generic-pool"
debug = require("debug")("thrift-pool")

TIMEOUT_MESSAGE = "Thrift-pool: Connection timeout"
CLOSE_MESSAGE = "Thrift-pool: Connection closed"

class ClosedError extends Error
  closed: true

id = 0
# create_cb creates and initializes a connection
#  @param thrift, used to create connection
#  @param pool_options, host and port used to create connection
#  @param thrift_options, passed to thrift connection,
create_cb = (thrift, pool_options, thrift_options, cb) ->
  cb = _.once cb
  connection = thrift.createConnection pool_options.host, pool_options.port, thrift_options
  connection.__ended = false
  connection.__id = id += 1
  connection.on "connect", ->
    debug "in connect callback"
    connection.connection.setKeepAlive(true)
    cb null, connection
  connection.on "error", (err) ->
    debug "in error callback"
    connection.__ended = true
    debug {err, connection_id: connection.__id}
    cb err
  connection.on "close", ->
    debug "in close callback"
    connection.__ended = true
    cb new ClosedError CLOSE_MESSAGE
  # timeout listener only applies if timeout is passed into thrift_options
  if thrift_options.timeout?
    debug "adding timeout listener"
    connection.on "timeout", ->
      debug "in timeout callback"
      connection.__ended = true
      cb new Error TIMEOUT_MESSAGE

# create_pool initializes a generic-pool
#   @param thrift library to use to in create_cb
#   @param pool_options, host/port are used in create_cb
#          max, min, idleTimeouts are used by generic pool
#   @param thrift_options used in create_cb
create_pool = (thrift, pool_options = {}, thrift_options = {}) ->
  pool = genericPool.Pool
    name: "thrift"
    create: (cb) ->
      create_cb thrift, pool_options, thrift_options, cb
    destroy: (connection) ->
      debug "in destroy"
      connection.end()
    validate: (connection) ->
      debug "in validate"
      not connection.__ended
    max: pool_options.max_connections
    min: pool_options.min_connections
    idleTimeoutMillis: pool_options.idle_timeout

retry_fn = (num_retries, is_retryable, fn, cb) ->
  fn (err, result) ->
    if err
      debug {err}
      debug {num_retries, is_retryable: is_retryable err}
    if not err or not is_retryable(err) or num_retries <= 0
      debug 'calling callback'
      cb err, result
    else
      debug 'retrying'
      retry_fn num_retries - 1, is_retryable, fn, cb

module.exports = (thrift, service, pool_options = {}, thrift_options = {}) ->

  throw new Error "Thrift-pool: You must specify #{key}" for key in ["host", "port"] when not pool_options[key]

  pool_options = _.defaults pool_options,
    max_connections: 1 # Max number of connections to keep open at any given time
    min_connections: 0 # Min number of connections to keep open at any given time
    idle_timeout: 10000 # Time (ms) to wait until closing idle connections
    retries: 2

  pool = create_pool thrift, pool_options, thrift_options

  # add_listeners adds listeners for error, close, and timeout
  #   @param connection, connection to add listeners to
  #   @param cb_error, callback to attach to "error" listener
  #   @param cb_timeout, callback to attach to "timeout" listener
  #   @param cb_close, callback to attach to "close" listener
  add_listeners = (connection, cb_error, cb_timeout, cb_close) ->
    connection.on "error", cb_error
    connection.on "close", cb_close
    if thrift_options.timeout?
      connection.on "timeout", cb_timeout

  # remove_listeners removes error, timeout, and close listeners with given callbacks
  #   @param connection, connection to remove listeners from
  #   @param cb_error, error callback to remove from "error" listener
  #   @param cb_timeout, timeout callback to remove from "timeout" listener
  #   @param cb_close, close callback to remove from "close" listener
  remove_listeners = (connection, cb_error, cb_timeout, cb_close) ->
    connection.removeListener "error", cb_error
    connection.removeListener "close", cb_close
    if thrift_options.timeout?
      connection.removeListener "timeout", cb_timeout

  # wrap_thrift_attempt when called with a function and arguments/callback:
  #   - acquires a connection
  #   - adds additional connection event listeners
  #   - creates a client with the acquired connection
  #   - calls client with fn and passed args and callback
  #   - connection is released before results are returned
  #  @return, function that takes in arguments and a callback
  wrap_thrift_attempt = (fn) -> (args..., cb) ->
    pool.acquire (err, connection) ->
      debug "Connection acquired"
      debug {err, connection_id: connection?.__id} if err?
      return cb err if err?
      cb = _.once cb
      cb_error = (err) ->
        debug "in error callback, post-acquire listener"
        debug {err, connection_id: connection.__id}
        cb err
      cb_timeout = ->
        debug "in timeout callback, post-acquire listener"
        debug {err, connection_id: connection.__id}
        cb new Error TIMEOUT_MESSAGE
      cb_close = ->
        debug "in close callback, post-acquire listener"
        debug {err, connection_id: connection.__id}
        cb new ClosedError CLOSE_MESSAGE
      add_listeners connection, cb_error, cb_timeout, cb_close
      client = thrift.createClient service, connection
      debug "Client created"
      debug {client}
      client[fn] args..., (err, results...) ->
        debug "In client callback"
        remove_listeners connection, cb_error, cb_timeout, cb_close
        pool.release connection<commit_msg>Don't call debug {err, ...} when err doesn't exist<commit_after>_ = require "lodash"
genericPool = require "generic-pool"
debug = require("debug")("thrift-pool")

TIMEOUT_MESSAGE = "Thrift-pool: Connection timeout"
CLOSE_MESSAGE = "Thrift-pool: Connection closed"

class ClosedError extends Error
  closed: true

id = 0
# create_cb creates and initializes a connection
#  @param thrift, used to create connection
#  @param pool_options, host and port used to create connection
#  @param thrift_options, passed to thrift connection,
create_cb = (thrift, pool_options, thrift_options, cb) ->
  cb = _.once cb
  connection = thrift.createConnection pool_options.host, pool_options.port, thrift_options
  connection.__ended = false
  connection.__id = id += 1
  connection.on "connect", ->
    debug "in connect callback"
    connection.connection.setKeepAlive(true)
    cb null, connection
  connection.on "error", (err) ->
    debug "in error callback"
    connection.__ended = true
    debug {err, connection_id: connection.__id}
    cb err
  connection.on "close", ->
    debug "in close callback"
    connection.__ended = true
    cb new ClosedError CLOSE_MESSAGE
  # timeout listener only applies if timeout is passed into thrift_options
  if thrift_options.timeout?
    debug "adding timeout listener"
    connection.on "timeout", ->
      debug "in timeout callback"
      connection.__ended = true
      cb new Error TIMEOUT_MESSAGE

# create_pool initializes a generic-pool
#   @param thrift library to use to in create_cb
#   @param pool_options, host/port are used in create_cb
#          max, min, idleTimeouts are used by generic pool
#   @param thrift_options used in create_cb
create_pool = (thrift, pool_options = {}, thrift_options = {}) ->
  pool = genericPool.Pool
    name: "thrift"
    create: (cb) ->
      create_cb thrift, pool_options, thrift_options, cb
    destroy: (connection) ->
      debug "in destroy"
      connection.end()
    validate: (connection) ->
      debug "in validate"
      not connection.__ended
    max: pool_options.max_connections
    min: pool_options.min_connections
    idleTimeoutMillis: pool_options.idle_timeout

retry_fn = (num_retries, is_retryable, fn, cb) ->
  fn (err, result) ->
    if err
      debug {err}
      debug {num_retries, is_retryable: is_retryable err}
    if not err or not is_retryable(err) or num_retries <= 0
      debug 'calling callback'
      cb err, result
    else
      debug 'retrying'
      retry_fn num_retries - 1, is_retryable, fn, cb

module.exports = (thrift, service, pool_options = {}, thrift_options = {}) ->

  throw new Error "Thrift-pool: You must specify #{key}" for key in ["host", "port"] when not pool_options[key]

  pool_options = _.defaults pool_options,
    max_connections: 1 # Max number of connections to keep open at any given time
    min_connections: 0 # Min number of connections to keep open at any given time
    idle_timeout: 10000 # Time (ms) to wait until closing idle connections
    retries: 2

  pool = create_pool thrift, pool_options, thrift_options

  # add_listeners adds listeners for error, close, and timeout
  #   @param connection, connection to add listeners to
  #   @param cb_error, callback to attach to "error" listener
  #   @param cb_timeout, callback to attach to "timeout" listener
  #   @param cb_close, callback to attach to "close" listener
  add_listeners = (connection, cb_error, cb_timeout, cb_close) ->
    connection.on "error", cb_error
    connection.on "close", cb_close
    if thrift_options.timeout?
      connection.on "timeout", cb_timeout

  # remove_listeners removes error, timeout, and close listeners with given callbacks
  #   @param connection, connection to remove listeners from
  #   @param cb_error, error callback to remove from "error" listener
  #   @param cb_timeout, timeout callback to remove from "timeout" listener
  #   @param cb_close, close callback to remove from "close" listener
  remove_listeners = (connection, cb_error, cb_timeout, cb_close) ->
    connection.removeListener "error", cb_error
    connection.removeListener "close", cb_close
    if thrift_options.timeout?
      connection.removeListener "timeout", cb_timeout

  # wrap_thrift_attempt when called with a function and arguments/callback:
  #   - acquires a connection
  #   - adds additional connection event listeners
  #   - creates a client with the acquired connection
  #   - calls client with fn and passed args and callback
  #   - connection is released before results are returned
  #  @return, function that takes in arguments and a callback
  wrap_thrift_attempt = (fn) -> (args..., cb) ->
    pool.acquire (err, connection) ->
      debug "Connection acquired"
      debug {err, connection_id: connection?.__id} if err?
      return cb err if err?
      cb = _.once cb
      cb_error = (err) ->
        debug "in error callback, post-acquire listener"
        debug {err, connection_id: connection.__id}
        cb err
      cb_timeout = ->
        debug "in timeout callback, post-acquire listener"
        debug {connection_id: connection.__id}
        cb new Error TIMEOUT_MESSAGE
      cb_close = ->
        debug "in close callback, post-acquire listener"
        debug {connection_id: connection.__id}
        cb new ClosedError CLOSE_MESSAGE
      add_listeners connection, cb_error, cb_timeout, cb_close
      client = thrift.createClient service, connection
      debug "Client created"
      debug {client}
      client[fn] args..., (err, results...) ->
        debug "In client callback"
        remove_listeners connection, cb_error, cb_timeout, cb_close
        pool.release connection
{'fraction_non_alphanumeric': 0.08245411827289653, 'fraction_numerical': 0.0021278482134941042, 'mean_word_length': 3.4131455399061035, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7611949', 'n_tokens_mistral': 3204, 'n_tokens_neox': 3191, 'n_words': 1140}