kazoo-ruby 0.3.3 → 0.4.0

data/lib/kazoo/partition.rb

@@ -15,6 +15,10 @@ module Kazoo
       replicas.length
     end
 
+    def preferred_leader
+      @replicas.first
+    end
+
     def leader
       @mutex.synchronize do
         refresh_state if @leader.nil?
@@ -80,13 +84,14 @@ module Kazoo
     protected
 
     def refresh_state
-      state_json = cluster.zk.get(path: "/brokers/topics/#{topic.name}/partitions/#{id}/state")
-      raise Kazoo::Error, "Failed to get partition state. Error code: #{state_json.fetch(:rc)}" unless state_json.fetch(:rc) == Zookeeper::Constants::ZOK
+      state_result = cluster.zk.get(path: "/brokers/topics/#{topic.name}/partitions/#{id}/state")
+      raise Kazoo::Error, "Failed to get partition state. Error code: #{state_result.fetch(:rc)}" unless state_result.fetch(:rc) == Zookeeper::Constants::ZOK
 
-      set_state(JSON.parse(state_json.fetch(:data)))
+      set_state_from_json(state_result.fetch(:data))
     end
 
-    def set_state(json)
+    def set_state_from_json(json_payload)
+      json = JSON.parse(json_payload)
       raise Kazoo::VersionNotSupported unless json.fetch('version') == 1
 
       @leader = cluster.brokers.fetch(json.fetch('leader'))

data/lib/kazoo/replica_assigner.rb

@@ -0,0 +1,74 @@
+module Kazoo
+  # Helper class to assign replicas to brokers for new partitions.
+  # It tries to e=venly divide the number of leaders and replicas over the brokers,
+  # in order to get a comparable load on all the brokers in the cluster.
+  class ReplicaAssigner
+    attr_reader :cluster
+    attr_reader :broker_leaders, :broker_replicas
+
+    def initialize(cluster)
+      @cluster = cluster
+      retrieve_initial_counts
+    end
+
+    def brokers
+      @cluster.brokers
+    end
+
+    def retrieve_initial_counts
+      @broker_leaders, @broker_replicas = {}, {}
+
+      @cluster.brokers.each do |_, broker|
+        @broker_leaders[broker], @broker_replicas[broker] = 0, 0
+      end
+
+      cluster.partitions.each do |partition|
+        @broker_leaders[partition.preferred_leader] += 1
+        partition.replicas.each do |broker|
+          @broker_replicas[broker] += 1
+        end
+      end
+    end
+
+    def cluster_leader_count
+      broker_leaders.values.inject(0, &:+)
+    end
+
+    def cluster_replica_count
+      broker_replicas.values.inject(0, &:+)
+    end
+
+    def assign(replication_factor)
+      raise Kazoo::ValidationError, "replication_factor should be higher than 0 " if replication_factor <= 0
+      raise Kazoo::ValidationError, "replication_factor should not be higher than the number of brokers " if replication_factor > brokers.length
+
+      # Order all brokers by the current number of leaders (ascending).
+      # The first one will be the leader replica
+      leader = @broker_leaders
+        .to_a
+        .sort_by { |pair| [pair[1], pair[0].id] }
+        .map(&:first)
+        .first
+
+      # Update the current broker replica counts.
+      # The assigned leader replica counts as a leader, but as a replica as well.
+      @broker_leaders[leader] += 1
+      @broker_replicas[leader] += 1
+
+      # To assign the other replcias, we remove the broker that was selected as leader from
+      # the list of brokers, and sort the rest by the number of replicas they are currently hosting.
+      # Then, we take the number of remaining replcias to complete the replication factor.
+      other_replicas = @broker_replicas
+        .to_a
+        .reject { |(key, _)| key == leader }
+        .sort_by { |pair| [pair[1], pair[0].id] }
+        .map(&:first)
+        .take(replication_factor - 1)
+
+      # Update the current broker replica counts.
+      other_replicas.each { |broker| @broker_replicas[broker] += 1 }
+
+      [leader].concat(other_replicas)
+    end
+  end
+end

data/lib/kazoo/subscription.rb

@@ -0,0 +1,215 @@
+module Kazoo
+
+  # A Kazoo::Subscription describes interest in a set of topics of a Kafka cluster.
+  #
+  # Use Kazoo::Subscription.build to instantiate a subscription. It will return one of
+  # the two known subclasses of Kazoo::Subscription: a Kazoo::StaticSubscription for
+  # a static list of topics, or a Kazoo::PatternSUbscription for a dynamic list based
+  # on a regular expression that serves as a white list or black list.
+  class Subscription
+    attr_reader :timestamp, :version
+
+    # Instantiates a whitelist subscription that matches every topic.
+    def self.everything
+      build(/.*/)
+    end
+
+    # Instantiates a Kazoo::Subscription based on the subscription argument.
+    #
+    # - If the subscription argument is the name of a topic, a Kazoo::Topic, or an array of those,
+    #   it will create a static subscription for the provided topic.
+    # - If the subscription argument is a regular expression, it will create a pattern subscription.
+    #   The `pattern` argument will determine whether it is a white_list (default), or black_list.
+    # - If the subscription argument is a Kazoo::Subscription, it will return the argument itself.
+    def self.build(subscription, pattern: :white_list, timestamp: Time.now)
+      case subscription
+      when Kazoo::Subscription
+        subscription
+      when String, Symbol, Kazoo::Topic, Array
+        topic_names = Array(subscription).map { |t| topic_name(t) }
+        Kazoo::StaticSubscription.new(topic_names, timestamp: timestamp)
+      when Regexp
+        Kazoo::PatternSubscription.new(subscription, pattern: pattern, timestamp: timestamp)
+      else
+        raise ArgumentError, "Don't know how to create a subscription from #{subscription.inspect}"
+      end
+    end
+
+    # Instantiates a Kazoo::Subscription based on a JSON payload as it is stored in Zookeeper.
+    #
+    # This method will raise Kazoo::InvalidSubscription if the JSON payload cannot be parsed.
+    # Only version 1 payloads are supported.
+    def self.from_json(json_payload)
+      json = JSON.parse(json_payload)
+      version, timestamp = json.fetch('version'), json.fetch('timestamp')
+      raise Kazoo::InvalidSubscription, "Only version 1 subscriptions are supported, found version #{version}!" unless version == 1
+
+      time = Time.at(BigDecimal.new(timestamp) / BigDecimal.new(1000))
+
+      pattern, subscription = json.fetch('pattern'), json.fetch('subscription')
+      raise Kazoo::InvalidSubscription, "Only subscriptions with a single stream are supported" unless subscription.values.all? { |streams| streams == 1 }
+
+      case pattern
+      when 'static'
+        topic_names = subscription.keys
+        Kazoo::StaticSubscription.new(topic_names, version: version, timestamp: time)
+
+      when 'white_list', 'black_list'
+        raise Kazoo::InvalidSubscription, "Only pattern subscriptions with a single expression are supported" unless subscription.keys.length == 1
+        regexp = Regexp.new(subscription.keys.first.tr(',', '|'))
+        Kazoo::PatternSubscription.new(regexp, pattern: pattern.to_sym, version: version, timestamp: time)
+
+      else
+        raise Kazoo::InvalidSubscription, "Unrecognized subscription pattern #{pattern.inspect}"
+      end
+
+    rescue JSON::ParserError, KeyError => e
+      raise Kazoo::InvalidSubscription.new(e.message)
+    end
+
+
+    # Returns a topic name based on various inputs.
+    # Helper method used by Kazoo::Subscription.build
+    def self.topic_name(topic)
+      case topic
+        when String, Symbol; topic.to_s
+        when Kazoo::Topic;   topic.name
+        else raise ArgumentError, "Cannot get topic name from #{topic.inspect}"
+      end
+    end
+
+    # Returns an array of all Kazoo::Topic instances in the given Kafka cluster
+    # that are matched by this subscription.
+    def topics(cluster)
+      cluster.topics.values.select { |topic| has_topic?(topic) }
+    end
+
+    # Returns an array of all Kazoo::Partition instances in the given Kafka cluster
+    # that are matched by this subscription.
+    def partitions(cluster)
+      topics(cluster).flat_map { |topic| topic.partitions }
+    end
+
+    # has_topic? should return true if a given Kazoo::Topic is part of this subscription.
+    def has_topic?(topic)
+      raise NotImplementedError
+    end
+
+    # Returns the JSON representation of this subscription that can be stored in Zookeeper.
+    def to_json(options = {})
+      JSON.dump(as_json(options))
+    end
+
+    def eql?(other)
+      other.kind_of?(Kazoo::Subscription) && other.pattern == pattern && other.subscription == subscription
+    end
+
+    alias_method :==, :eql?
+
+    def hash
+      [pattern, subscription].hash
+    end
+
+    def inspect
+      "#<#{self.class.name} pattern=#{pattern} subscription=#{subscription.inspect}>"
+    end
+
+    protected
+
+    # Subclasses should call super(**kwargs) in their initializer.
+    def initialize(timestamp: Time.now, version: 1)
+      @timestamp, @version = timestamp, version
+    end
+
+    # Subclasses should return a hash that can be converted to JSON to represent
+    # the subscription in Zookeeper.
+    def subscription
+      raise NotImplementedError
+    end
+
+    # Should return the name of the pattern, i.e. static, white_list, or black_list
+    def pattern
+      raise NotImplementedError
+    end
+
+    private
+
+    def as_json(options = {})
+      {
+        version:      version,
+        pattern:      pattern,
+        timestamp:    msec_timestamp,
+        subscription: subscription,
+      }
+    end
+
+    def msec_timestamp
+      (timestamp.to_i * 1000) + (timestamp.nsec / 1_000_000)
+    end
+  end
+
+  # StaticSubscription describes a subscription based on a static list of topic names.
+  class StaticSubscription < Kazoo::Subscription
+    attr_reader :topic_names
+
+    # Instantiates a static subscription instance. The topic_names argument must be
+    # an array of strings.
+    def initialize(topic_names, **kwargs)
+      super(**kwargs)
+      @topic_names = topic_names
+    end
+
+    def has_topic?(topic)
+      topic_names.include?(topic.name)
+    end
+
+    protected
+
+    def pattern
+      :static
+    end
+
+    def subscription
+      topic_names.inject({}) { |hash, topic_name| hash[topic_name] = 1; hash }
+    end
+  end
+
+  # PatternSubscription describes a subscription based on a regular expression that
+  # serves as either a whitelist or blacklist filter for all topics available in
+  # a Kafka cluster.
+  class PatternSubscription < Kazoo::Subscription
+    PATTERN_TYPES = [:black_list, :white_list].freeze
+
+    attr_reader :pattern, :regexp
+
+    def initialize(regexp, pattern: :white_list, **kwargs)
+      super(**kwargs)
+      raise ArgumentError, "#{pattern.inspect} is not a valid pattern type" unless PATTERN_TYPES.include?(pattern)
+      @regexp, @pattern = regexp, pattern
+    end
+
+    # Returns true if this subscription uses a whitelist pattern
+    def white_list?
+      pattern == :white_list
+    end
+
+    # Returns true if this subscription uses a blacklist pattern
+    def black_list?
+      pattern == :black_list
+    end
+
+    # Returns true if the whitelist or blacklist does not filter out the provided topic.
+    def has_topic?(topic)
+      case pattern
+        when :white_list; topic.name =~ regexp
+        when :black_list; topic.name !~ regexp
+      end
+    end
+
+    protected
+
+    def subscription
+      { regexp.inspect[1..-2] => 1 }
+    end
+  end
+end

data/lib/kazoo/topic.rb

@@ -1,40 +1,35 @@
 module Kazoo
   class Topic
-    VALID_TOPIC_NAMES = %r{\A[a-zA-Z0-9\\._\\-]+\z}
+    ALL_PRELOAD_METHODS     = [:partitions, :config].freeze
+    DEFAULT_PRELOAD_METHODS = [:partitions].freeze
+    VALID_TOPIC_NAMES       = %r{\A[a-zA-Z0-9\\._\\-]+\z}
     BLACKLISTED_TOPIC_NAMES = %r{\A\.\.?\z}
 
     attr_reader :cluster, :name
-    attr_writer :partitions
 
-    def initialize(cluster, name)
+    def initialize(cluster, name, config: nil, partitions: nil)
       @cluster, @name = cluster, name
-    end
-
-    def self.from_json(cluster, name, json)
-      raise Kazoo::VersionNotSupported unless json.fetch('version') == 1
-
-      topic = new(cluster, name)
-      topic.partitions = json.fetch('partitions').map do |(id, replicas)|
-        topic.partition(id.to_i, replicas: replicas.map { |b| cluster.brokers[b] })
-      end.sort_by(&:id)
 
-      return topic
+      self.partitions = partitions
+      self.config = config
     end
 
     def partitions
-      @partitions ||= begin
-        result = cluster.zk.get(path: "/brokers/topics/#{name}")
-        raise Kazoo::Error, "Failed to get list of partitions for #{name}. Result code: #{result.fetch(:rc)}" if result.fetch(:rc) != Zookeeper::Constants::ZOK
+      @partitions ||= load_partitions_from_zookeeper
+    end
 
-        partition_json = JSON.parse(result.fetch(:data))
-        partition_json.fetch('partitions').map do |(id, replicas)|
-          partition(id.to_i, replicas: replicas.map { |b| cluster.brokers[b] })
-        end
-      end
+    def partitions=(ps)
+      @partitions = ps
+    end
+
+    def partition(index, **kwargs)
+      Kazoo::Partition.new(self, index, **kwargs)
     end
 
-    def partition(*args)
-      Kazoo::Partition.new(self, *args)
+    def append_partition(**kwargs)
+      new_partition = partition(partitions.length, **kwargs)
+      partitions << new_partition
+      new_partition
     end
 
     def replication_factor
@@ -42,7 +37,7 @@ module Kazoo
     end
 
     def under_replicated?
-      partitions.any?(:under_replicated?)
+      partitions.any?(&:under_replicated?)
     end
 
     def inspect
@@ -80,29 +75,34 @@ module Kazoo
       false
     end
 
-    def create
-      raise Kazoo::Error, "The topic #{name} already exists!" if exists?
-      validate
+    def add_partitions(partitions: nil, replication_factor: nil)
+      raise ArgumentError, "partitions must be a positive integer" if Integer(partitions) <= 0
+      raise ArgumentError, "replication_factor must be a positive integer" if Integer(replication_factor) <= 0
+
+      raise Kazoo::TopicNotFound, "The topic #{name} does not exists!" unless exists?
 
-      result = cluster.zk.create(
-        path: "/config/topics/#{name}",
-        data: JSON.generate(version: 1, config: {})
-      )
+      replica_assigner = Kazoo::ReplicaAssigner.new(cluster)
 
-      if result.fetch(:rc) != Zookeeper::Constants::ZOK
-        raise Kazoo::Error, "Failed to create topic config node for #{name}. Error code: #{result.fetch(:rc)}"
+      partitions.times do
+        replicas = replica_assigner.assign(replication_factor)
+        append_partition(replicas: replicas)
       end
 
-      result = cluster.zk.create(
-        path: "/brokers/topics/#{name}",
-        data: JSON.generate(version: 1, partitions: partitions_as_json)
-      )
+      validate
+      write_partitions_to_zookeeper
+      wait_for_partitions
+      cluster.reset_metadata
+    end
 
-      if result.fetch(:rc) != Zookeeper::Constants::ZOK
-        raise Kazoo::Error, "Failed to create topic #{name}. Error code: #{result.fetch(:rc)}"
-      end
+    def save
+      raise Kazoo::Error, "The topic #{name} already exists!" if exists?
+      validate
+
+      write_config_to_zookeeper
+      write_partitions_to_zookeeper
 
       wait_for_partitions
+      cluster.reset_metadata
     end
 
     def destroy
@@ -136,9 +136,63 @@ module Kazoo
       else
         raise Kazoo::Error, "Failed to delete topic #{name}. Error code: #{result.fetch(:rc)}"
       end
+
+      cluster.reset_metadata
     end
 
     def config
+      @config ||= load_config_from_zookeeper
+    end
+
+    def config=(hash)
+      return if hash.nil?
+      @config = hash.inject({}) { |h, (k, v)| h[k.to_s] = v.to_s; h }
+    end
+
+    def set_config(key, value)
+      config[key.to_s] = value.to_s
+      write_config_to_zookeeper
+    end
+
+    def delete_config(key)
+      config.delete(key.to_s)
+      write_config_to_zookeeper
+    end
+
+    def reset_default_config
+      @config = {}
+      write_config_to_zookeeper
+    end
+
+    def set_partitions_from_json(json_payload)
+      partition_json = JSON.parse(json_payload)
+      raise Kazoo::VersionNotSupported if partition_json.fetch('version') != 1
+
+      @partitions = partition_json.fetch('partitions').map do |(id, replicas)|
+        partition(id.to_i, replicas: replicas.map { |b| cluster.brokers[b] })
+      end
+
+      @partitions.sort_by!(&:id)
+
+      self
+    end
+
+    def set_config_from_json(json_payload)
+      config_json = JSON.parse(json_payload)
+      raise Kazoo::VersionNotSupported if config_json.fetch('version') != 1
+
+      @config = config_json.fetch('config')
+
+      self
+    end
+
+    def load_partitions_from_zookeeper
+      result = cluster.zk.get(path: "/brokers/topics/#{name}")
+      raise Kazoo::Error, "Failed to get list of partitions for #{name}. Result code: #{result.fetch(:rc)}" if result.fetch(:rc) != Zookeeper::Constants::ZOK
+      set_partitions_from_json(result.fetch(:data)).partitions
+    end
+
+    def load_config_from_zookeeper
       result = cluster.zk.get(path: "/config/topics/#{name}")
       case result.fetch(:rc)
       when Zookeeper::Constants::ZOK
@@ -149,33 +203,30 @@ module Kazoo
         raise Kazoo::Error, "Failed to retrieve topic config"
       end
 
-      config = JSON.parse(result.fetch(:data))
-      raise Kazoo::VersionNotSupported if config.fetch('version') != 1
-
-      config.fetch('config')
+      set_config_from_json(result.fetch(:data)).config
     end
 
-    def set_config(key, value)
-      new_config = config
-      new_config[key.to_s] = value.to_s
-      write_config(new_config)
-    end
+    def write_partitions_to_zookeeper
+      path = "/brokers/topics/#{name}"
+      data = JSON.generate(version: 1, partitions: partitions_as_json)
 
-    def delete_config(key)
-      new_config = config
-      new_config.delete(key.to_s)
-      write_config(new_config)
-    end
+      result = cluster.zk.set(path: path, data: data)
+      case result.fetch(:rc)
+      when Zookeeper::Constants::ZOK
+        # continue
+      when Zookeeper::Constants::ZNONODE
+        result = cluster.zk.create(path: path, data: data)
+        raise Kazoo::Error, "Failed to write partitions to zookeeper. Result code: #{result.fetch(:rc)}" unless result.fetch(:rc) == Zookeeper::Constants::ZOK
+      else
+        raise Kazoo::Error, "Failed to write partitions to zookeeper. Result code: #{result.fetch(:rc)}"
+      end
 
-    def reset_default_config
-      write_config({})
+      cluster.reset_metadata
     end
 
-    def write_config(config_hash)
-      raise Kazoo::TopicNotFound, "Topic #{name.inspect} does not exist" unless exists?
-
-      config = config_hash.inject({}) { |h, (k,v)| h[k.to_s] = v.to_s; h }
-      config_json = JSON.generate(version: 1, config: config)
+    def write_config_to_zookeeper
+      config_hash = config.inject({}) { |h, (k,v)| h[k.to_s] = v.to_s; h }
+      config_json = JSON.generate(version: 1, config: config_hash)
 
       # Set topic config
       result = cluster.zk.set(path: "/config/topics/#{name}", data: config_json)
@@ -184,20 +235,30 @@ module Kazoo
         # continue
       when Zookeeper::Constants::ZNONODE
         result = cluster.zk.create(path: "/config/topics/#{name}", data: config_json)
-        raise Kazoo::Error, "Failed to set topic config" unless result.fetch(:rc) == Zookeeper::Constants::ZOK
+        raise Kazoo::Error, "Failed to write topic config to zookeeper. Result code: #{result.fetch(:rc)}" unless result.fetch(:rc) == Zookeeper::Constants::ZOK
       else
-        raise Kazoo::Error, "Failed to set topic config"
+        raise Kazoo::Error, "Failed to write topic config to zookeeper. Result code: #{result.fetch(:rc)}"
       end
 
       # Set config change notification
       result = cluster.zk.create(path: "/config/changes/config_change_", data: name.inspect, sequence: true)
       raise Kazoo::Error, "Failed to set topic config change notification" unless result.fetch(:rc) == Zookeeper::Constants::ZOK
+
+      cluster.reset_metadata
     end
 
-    def self.create(cluster, name, partitions: nil, replication_factor: nil)
-      topic = new(cluster, name)
-      topic.send(:sequentially_assign_partitions, partitions, replication_factor)
-      topic.create
+    def self.create(cluster, name, partitions: nil, replication_factor: nil, config: {})
+      topic = new(cluster, name, config: config, partitions: [])
+      raise Kazoo::Error, "Topic #{name} already exists" if topic.exists?
+
+      replica_assigner = Kazoo::ReplicaAssigner.new(cluster)
+
+      partitions.times do
+        replicas = replica_assigner.assign(replication_factor)
+        topic.append_partition(replicas: replicas)
+      end
+
+      topic.save
       topic
     end
 
@@ -214,21 +275,6 @@ module Kazoo
       threads.each(&:join)
     end
 
-    def sequentially_assign_partitions(partition_count, replication_factor, brokers: nil)
-      brokers = cluster.brokers.values if brokers.nil?
-      raise ArgumentError, "replication_factor should be smaller or equal to the number of brokers" if replication_factor > brokers.length
-
-      # Sequentially assign replicas to brokers. There might be a better way.
-      @partitions = 0.upto(partition_count - 1).map do |partition_index|
-        replicas = 0.upto(replication_factor - 1).map do |replica_index|
-          broker_index = (partition_index + replica_index) % brokers.length
-          brokers[broker_index]
-        end
-
-        self.partition(partition_index, replicas: replicas)
-      end
-    end
-
     def partitions_as_json
       partitions.inject({}) do |hash, partition|
         hash[partition.id] = partition.replicas.map(&:id)