<commit_before>require 'nerve/reporter/base'
require 'thread'
require 'zk'


class Nerve::Reporter
  class Zookeeper < Base
    @@zk_pool = {}
    @@zk_pool_lock = Mutex.new

    def initialize(service)
      %w{zk_hosts zk_path instance_id host port}.each do |required|
        raise ArgumentError, "missing required argument #{required} for new service watcher" unless service[required]
      end
      @path = service['zk_hosts'].shuffle.join(',')
      @data = parse_data({'host' => service['host'], 'port' => service['port'], 'name' => service['instance_id']})

      @key = service['zk_path'] + "/#{service['instance_id']}_"
      @full_key = nil
    end

    def start()
      log.info "nerve: waiting to connect to zookeeper at #{@path}"
      # Ensure that all Zookeeper reporters re-use a single zookeeper
      # connection to any given connection string. Note that you will
      # end up with a number of connections equal to the number of hosts in
      # the connection string because the randomization in initialize
      @@zk_pool_lock.synchronize {
        unless @@zk_pool.has_key?(@path)
          log.info "nerve: creating pooled connection at #{@path}"
          @@zk_pool[@path] = ZK.new(@path)
          log.info "nerve: successfully created zk connection to #{@path}"
        else
          log.info "nerve: re-using existing zookeeper connection at #{@path}"
        end
        @zk = @@zk_pool[@path]
        log.info "nerve: retrieved zk connection to #{@path}"
      }
    end

    def stop()
      log.info "nerve: closing zk connection at #{@path}"
      report_down
      @zk.close!
    end

    def report_up()
      zk_save
    end

    def report_down
      zk_delete
    end

    def update_data(new_data='')
      @data = parse_data(new_data)
      zk_save
    end

    def ping?
      return @zk.ping?
    end

    private

    def zk_delete
      if @full_key
        @zk.delete(@full_key, :ignore => :no_node)
        @full_key = nil<commit_msg>Use pooling instead of shuffle for connection distribution<commit_after>require 'nerve/reporter/base'
require 'thread'
require 'zk'


class Nerve::Reporter
  class Zookeeper < Base
    @@zk_pool = {}
    @@zk_pool_count = {}
    @@zk_pool_lock = Mutex.new

    def initialize(service)
      %w{zk_hosts zk_path instance_id host port}.each do |required|
        raise ArgumentError, "missing required argument #{required} for new service watcher" unless service[required]
      end
      # Since we pool we get one connection per zookeeper cluster
      @path = service['zk_hosts'].sort.join(',')
      @data = parse_data({'host' => service['host'], 'port' => service['port'], 'name' => service['instance_id']})

      @key = service['zk_path'] + "/#{service['instance_id']}_"
      @full_key = nil
    end

    def start()
      log.info "nerve: waiting to connect to zookeeper to #{@path}"
      # Ensure that all Zookeeper reporters re-use a single zookeeper
      # connection to any given set of zk hosts.
      @@zk_pool_lock.synchronize {
        unless @@zk_pool.has_key?(@path)
          log.info "nerve: creating pooled connection to #{@path}"
          @@zk_pool[@path] = ZK.new(@path)
          @@zk_pool_count[@path] = 1
          log.info "nerve: successfully created zk connection to #{@path}"
        else
          @@zk_pool_count[@path] += 1
          log.info "nerve: re-using existing zookeeper connection to #{@path}"
        end
        @zk = @@zk_pool[@path]
        log.info "nerve: retrieved zk connection to #{@path}"
      }
    end

    def stop()
      log.info "nerve: removing zk node at #{@full_key}" if @full_key
      report_down
      @@zk_pool_lock.synchronize {
        @@zk_pool_count[@path] -= 1
        # Last thread to use the connection closes it
        if @@zk_pool_count[@path] == 0
          log.info "nerve: closing zk connection to #{@path}"
          @zk.close!
        end
      }
    end

    def report_up()
      zk_save
    end

    def report_down
      zk_delete
    end

    def update_data(new_data='')
      @data = parse_data(new_data)
      zk_save
    end

    def ping?
      return @zk.ping?
    end

    private

    def zk_delete
      if @full_key
        @zk.delete(@full_key, :ignore => :no_node)
        @full_key = nil
{'fraction_non_alphanumeric': 0.11467353147671425, 'fraction_numerical': 0.000936110461034402, 'mean_word_length': 2.5322314049586776, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 11, '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': '3722152', 'n_tokens_mistral': 1379, 'n_tokens_neox': 1314, 'n_words': 389}