<commit_before>require 'zk'

module DCell
  module Registry
    class ZkAdapter
      include Node
      include Global

      PREFIX  = "/dcell"
      DEFAULT_PORT = 2181

      # Create a new connection to Zookeeper
      #
      # servers: a list of Zookeeper servers to connect to. Each server in the
      #          list has a host/port configuration
      def initialize(options)
        # Stringify keys :/
        options = options.inject({}) { |h,(k,v)| h[k.to_s] = v; h }

        @env = options['env'] || 'production'
        @base_path = "#{PREFIX}/#{@env}"

        # Let them specify a single server instead of many
        server = options['server']
        if server
          servers = [server]
        else
          servers = options['servers']
          raise "no Zookeeper servers given" unless servers
        end

        # Add the default Zookeeper port unless specified
        servers.map! do |server|
          if server[/:\d+$/]
            server
          else
            "#{server}:#{DEFAULT_PORT}"
          end
        end

        @zk = ZK.new(*servers)
        @node_registry = Registry.new(@zk, @base_path, :nodes, true)
        @global_registry = Registry.new(@zk, @base_path, :globals, false)
      end

      class Registry
        def initialize(zk, base_path, name, ephemeral)
          @zk = zk
          @base_path = File.join(base_path, name.to_s)
          @ephemeral = ephemeral
          @zk.mkdir_p @base_path
        end

        def get(key)
          result, _ = @zk.get("#{@base_path}/#{key}", watch: true)
          MessagePack.unpack(result, options={:symbolize_keys => true}) if result
        rescue ZK::Exceptions::NoNode
        end

        def _set(key, value, unique)
          path = "#{@base_path}/#{key}"
          @zk.create path, value.to_msgpack, :ephemeral => @ephemeral
        rescue ZK::Exceptions::NodeExists
          raise KeyExists if unique
          @zk.set path, value.to_msgpack
        end

        def set(key, value)
          _set(key, value, false)
        end

        def all
          @zk.children @base_path
        end

        def remove(key)
          path = "#{@base_path}/#{key}"
          @zk.delete path
        end

        def clear_all
          all.each do |key|<commit_msg>registries:zk: properly handle removal of non-existant key<commit_after>require 'zk'

module DCell
  module Registry
    class ZkAdapter
      include Node
      include Global

      PREFIX  = "/dcell"
      DEFAULT_PORT = 2181

      # Create a new connection to Zookeeper
      #
      # servers: a list of Zookeeper servers to connect to. Each server in the
      #          list has a host/port configuration
      def initialize(options)
        # Stringify keys :/
        options = options.inject({}) { |h,(k,v)| h[k.to_s] = v; h }

        @env = options['env'] || 'production'
        @base_path = "#{PREFIX}/#{@env}"

        # Let them specify a single server instead of many
        server = options['server']
        if server
          servers = [server]
        else
          servers = options['servers']
          raise "no Zookeeper servers given" unless servers
        end

        # Add the default Zookeeper port unless specified
        servers.map! do |server|
          if server[/:\d+$/]
            server
          else
            "#{server}:#{DEFAULT_PORT}"
          end
        end

        @zk = ZK.new(*servers)
        @node_registry = Registry.new(@zk, @base_path, :nodes, true)
        @global_registry = Registry.new(@zk, @base_path, :globals, false)
      end

      class Registry
        def initialize(zk, base_path, name, ephemeral)
          @zk = zk
          @base_path = File.join(base_path, name.to_s)
          @ephemeral = ephemeral
          @zk.mkdir_p @base_path
        end

        def get(key)
          result, _ = @zk.get("#{@base_path}/#{key}", watch: true)
          MessagePack.unpack(result, options={:symbolize_keys => true}) if result
        rescue ZK::Exceptions::NoNode
        end

        def _set(key, value, unique)
          path = "#{@base_path}/#{key}"
          @zk.create path, value.to_msgpack, :ephemeral => @ephemeral
        rescue ZK::Exceptions::NodeExists
          raise KeyExists if unique
          @zk.set path, value.to_msgpack
        end

        def set(key, value)
          _set(key, value, false)
        end

        def all
          @zk.children @base_path
        end

        def remove(key)
          path = "#{@base_path}/#{key}"
          @zk.delete path
        rescue ZK::Exceptions::NoNode
        end

        def clear_all
          all.each do |key|
{'fraction_non_alphanumeric': 0.10597884739909345, 'fraction_numerical': 0.0017267429311461256, 'mean_word_length': 1.805084745762712, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3498987', 'n_tokens_mistral': 1408, 'n_tokens_neox': 1365, 'n_words': 381}