promiscuous-poseidon_cluster 0.3.0

data/lib/poseidon_cluster.rb

@@ -0,0 +1 @@
+require 'poseidon/cluster'

data/poseidon_cluster.gemspec

@@ -0,0 +1,28 @@
+Gem::Specification.new do |s|
+  s.required_ruby_version = '>= 2.0.0'
+  s.required_rubygems_version = ">= 1.8.0"
+
+  s.name        = File.basename(__FILE__, '.gemspec')
+  s.summary     = "Poseidon cluster extensions"
+  s.description = "Cluster extensions for Poseidon, a producer and consumer implementation for Kafka >= 0.8"
+  s.version     = "0.3.0"
+
+  s.authors     = ["Black Square Media"]
+  s.email       = "info@blacksquaremedia.com"
+  s.homepage    = "https://github.com/promiscuous-io/poseidon_cluster"
+
+  s.require_path = 'lib'
+  s.files        = `git ls-files`.split("\n")
+  s.test_files   = `git ls-files -- {test,spec,features,scenario}/*`.split("\n")
+
+  s.add_dependency "poseidon", ">= 0.0.5.pre1"
+  s.add_dependency "zk"
+
+  s.add_development_dependency "rake"
+  s.add_development_dependency "bundler"
+  s.add_development_dependency "rspec"
+  s.add_development_dependency "rspec-its"
+  s.add_development_dependency "yard"
+  s.add_development_dependency "coveralls"
+
+end

data/scenario/.gitignore

@@ -0,0 +1 @@
+output.txt

data/scenario/consumer.rb

@@ -0,0 +1,17 @@
+#!/usr/bin/env ruby
+require 'bundler/setup'
+require 'poseidon_cluster'
+
+name     = ARGV[0].to_s
+output   = File.open(ARGV[1], "a")
+output.sync = true
+
+total    = 0
+consumer = Poseidon::ConsumerGroup.new "my-group", ["localhost:29092"], ["localhost:22181"], "my-topic", max_bytes: 256*1024
+consumer.fetch_loop do |n, messages|
+  break if name[0] > 'Q' && total > 0
+  messages.each do |m|
+    output.write "#{name},#{n},#{m.value}\n"
+  end
+  total += messages.size
+end

data/scenario/producer.rb

@@ -0,0 +1,23 @@
+#!/usr/bin/env ruby
+require 'bundler/setup'
+require 'poseidon'
+
+limit, offset = ARGV[0].to_i, ARGV[1].to_i
+producer = Poseidon::Producer.new ["localhost:29092"], "poseidon-producer"
+
+while limit > 0 do
+  batch  = limit > 10000 ? 10000 : limit
+  limit -= batch
+
+  messages = (0...batch).map do
+    num = offset.to_s.rjust(8, "0")
+    offset += 1
+    Poseidon::MessageToSend.new "my-topic", num, Time.now.to_s+num
+  end
+
+  10.times do
+    ok = producer.send_messages messages
+    break if ok
+    sleep(1)
+  end
+end

data/scenario/run.rb

@@ -0,0 +1,35 @@
+#!/usr/bin/env ruby
+
+require 'bundler/setup'
+require 'timeout'
+require File.expand_path("../scenario", __FILE__)
+
+# Start Zookeeper & Kafka
+Scenario.run do
+  5.times do
+    produce 1000
+  end
+  consume "A"
+  consume "B"
+  consume "C"
+  checkpoint!
+
+  15.times { produce 1000 }
+  consume "D"
+  10.times { produce 1000 }
+  consume "X"
+  10.times { produce 1000 }
+  checkpoint!
+
+  20.times { produce 1000 }
+  consume "E"
+  consume "F"
+  15.times { produce 1000 }
+  consume "Y"
+  50.times { produce 100 }
+  20.times { produce 1000 }
+
+  checkpoint!
+end
+
+

data/scenario/scenario.rb

@@ -0,0 +1,134 @@
+require 'fileutils'
+require 'pathname'
+
+module Scenario
+  extend self
+
+  ROOT    = Pathname.new(File.expand_path("../", __FILE__))
+  VERSION = "0.8.1.1"
+  SERVER  = ROOT.join "kafka_2.10-#{VERSION}"
+
+  TOPIC_NAME = "my-topic"
+  KAFKA_BIN  = SERVER.join("bin", "kafka-server-start.sh")
+  KAFKA_CFG  = SERVER.join("config", "server-poseidon.properties")
+  KAFKA_TMP  = "/tmp/kafka-logs-poseidon"
+  ZOOKP_BIN  = SERVER.join("bin", "zookeeper-server-start.sh")
+  ZOOKP_CFG  = SERVER.join("config", "zookeeper-poseidon.properties")
+  ZOOKP_TMP  = "/tmp/zookeeper-poseidon"
+  LOG4J_CFG  = SERVER.join("config", "log4j.properties")
+  OUTPUT     = Scenario::ROOT.join("output.txt")
+
+  @@pids     = {}
+  @@total    = 0
+
+  def run(&block)
+    setup
+    instance_eval(&block)
+  rescue => e
+    abort [e, *e.backtrace[0,20]].join("\n")
+  ensure
+    teardown
+  end
+
+  def setup
+    FileUtils.rm_rf OUTPUT.to_s
+    configure
+
+    # Ensure all required files are present
+    [KAFKA_BIN, ZOOKP_BIN, KAFKA_CFG, ZOOKP_CFG].each do |path|
+      abort "Unable to locate #{path}. File does not exist!" unless path.file?
+    end
+
+    Signal.trap("INT") { teardown }
+
+    spawn KAFKA_BIN, KAFKA_CFG
+    spawn ZOOKP_BIN, ZOOKP_CFG
+    sleep(2)
+  end
+
+  def teardown
+    @@pids.each do |_, pid|
+      Process.kill :TERM, pid
+    end
+    sleep(1)
+    FileUtils.rm_rf KAFKA_TMP.to_s
+    FileUtils.rm_rf ZOOKP_TMP.to_s
+
+    fail! unless numlines == @@total
+  end
+
+  def configure
+    download
+
+    KAFKA_CFG.open("w") do |f|
+      f.write SERVER.join("config", "server.properties").read.
+        sub("=9092", "=29092").
+        sub(":2181", ":22181").
+        sub("num.partitions=2", "num.partitions=12").
+        sub("log.flush.interval.ms=1000", "log.flush.interval.ms=10").
+        sub("/tmp/kafka-logs", KAFKA_TMP)
+    end
+    ZOOKP_CFG.open("w") do |f|
+      f.write SERVER.join("config", "zookeeper.properties").read.
+        sub("/tmp/zookeeper", ZOOKP_TMP).
+        sub("=2181", "=22181")
+    end
+    content = LOG4J_CFG.read
+    LOG4J_CFG.open("w") do |f|
+      f.write content.gsub("INFO", "FATAL")
+    end if content.include?("INFO")
+  end
+
+  def download
+    return if SERVER.directory?
+    sh "cd #{ROOT} && curl http://www.mirrorservice.org/sites/ftp.apache.org/kafka/#{VERSION}/kafka_2.10-#{VERSION}.tgz | tar xz"
+  end
+
+  def checkpoint!(timeout = 10)
+    puts "--> Verifying #{@@total}"
+    timeout.times do
+      if numlines > @@total
+        break
+      elsif numlines < @@total
+        sleep(1)
+      else
+        return
+      end
+    end
+    fail!
+  end
+
+  def consume(name)
+    puts "--> Launching consumer #{name}"
+    spawn ROOT.join("consumer.rb"), name, OUTPUT
+  end
+
+  def produce(count)
+    puts "--> Producing messages #{@@total}-#{@@total+count-1}"
+    sh ROOT.join("producer.rb"), count, @@total
+    @@total += count
+  end
+
+  def numlines
+    `wc -l #{OUTPUT} 2> /dev/null`.to_i
+  end
+
+  def abort(message)
+    Kernel.abort "ERROR: #{message}"
+  end
+
+  def fail!
+    Kernel.abort "FAILED: expected #{@@total} but was #{numlines}"
+  end
+
+  def sh(*bits)
+    cmd = bits.join(" ")
+    system(cmd) || abort(cmd)
+  end
+
+  def spawn(*args)
+    cmd = args.join(" ")
+    @@pids[cmd] = Process.spawn(cmd)
+  end
+
+end

data/spec/lib/poseidon/cluster_spec.rb

@@ -0,0 +1,19 @@
+require 'spec_helper'
+
+describe Poseidon::Cluster do
+
+  it 'should generate incremented numbers (atomically)' do
+    num = described_class.inc!
+    (described_class.inc! - num).should == 1
+
+    (0...5).map do
+      Thread.new { 100.times { described_class.inc! }}
+    end.each &:join
+    (described_class.inc! - num).should == 502
+  end
+
+  it 'should generate GUIDs' do
+    described_class.guid.should match(/\A[\w\-\.]+?\-\d{1,5}\-\d{10}\-\d{1,3}\z/)
+  end
+
+end

data/spec/lib/poseidon/consumer_group_spec.rb

@@ -0,0 +1,286 @@
+require 'spec_helper'
+
+describe Poseidon::ConsumerGroup do
+
+  def fetch_response(n)
+    set = Poseidon::MessageSet.new
+    n.times {|i| set << Poseidon::Message.new(value: "value", key: "key", offset: i) }
+    pfr = Poseidon::Protocol::PartitionFetchResponse.new(0, 0, 100, set)
+    tfr = Poseidon::Protocol::TopicFetchResponse.new("mytopic", [pfr])
+    Poseidon::Protocol::FetchResponse.new(nil, [tfr])
+  end
+
+  let :brokers do
+    [ Poseidon::Protocol::Broker.new(1, "localhost", 29092),   # id,host,port
+      Poseidon::Protocol::Broker.new(2, "localhost", 29091), ]
+  end
+
+  let :partitions do
+    [ Poseidon::Protocol::PartitionMetadata.new(0, 0, 1, [1,2], []), # err,id,leader,replicas,isr
+      Poseidon::Protocol::PartitionMetadata.new(0, 1, 2, [1,2], []), ]
+  end
+
+  let :topics do
+    [ Poseidon::TopicMetadata.new(Poseidon::Protocol::TopicMetadataStruct.new(0, "mytopic", partitions)) ]
+  end
+
+  let :metadata do
+    Poseidon::Protocol::MetadataResponse.new nil, brokers.dup, topics.dup
+  end
+
+  let :zk_client do
+    double "ZK", mkdir_p: nil, get: nil, set: nil, delete: nil, create: "/path", register: nil, children: ["my-group-UNIQUEID"], close: nil
+  end
+
+  let(:group) { described_class.new "my-group", ["localhost:29092", "localhost:29091"], ["localhost:22181"], "mytopic" }
+  subject     { group }
+
+  before do
+    allow(ZK).to receive_messages(new: zk_client)
+    allow(Poseidon::Cluster).to receive_messages(guid: "UNIQUEID")
+    allow_any_instance_of(Poseidon::ConsumerGroup).to receive(:sleep)
+    allow_any_instance_of(Poseidon::PartitionConsumer).to receive_messages(resolve_offset_if_necessary: 0)
+    allow_any_instance_of(Poseidon::BrokerPool).to receive_messages(fetch_metadata_from_broker: metadata)
+
+    allow_any_instance_of(Poseidon::Connection).to receive(:fetch).with(10000, 1, ->req { req[0].partition_fetches[0].partition == 0 }).and_return(fetch_response(10))
+    allow_any_instance_of(Poseidon::Connection).to receive(:fetch).with(10000, 1, ->req { req[0].partition_fetches[0].partition == 1 }).and_return(fetch_response(5))
+  end
+
+  it               { should be_registered }
+  its(:name)       { should == "my-group" }
+  its(:topic)      { should == "mytopic" }
+  its(:pool)       { should be_instance_of(Poseidon::BrokerPool) }
+  its(:id)         { should == "my-group-UNIQUEID" }
+  its(:zk)         { should be(zk_client) }
+
+  its(:claimed)        { should == [0, 1] }
+  its(:metadata)       { should be_instance_of(Poseidon::ClusterMetadata) }
+  its(:topic_metadata) { should be_instance_of(Poseidon::TopicMetadata) }
+  its(:registries)     { should == {
+    consumer: "/consumers/my-group/ids",
+    owner:    "/consumers/my-group/owners/mytopic",
+    offset:   "/consumers/my-group/offsets/mytopic",
+  }}
+
+  its("metadata.brokers.keys") { should =~ [1,2] }
+  its("topic_metadata.partition_count") { should == 2 }
+
+  it "should register with zookeeper and rebalance" do
+    zk_client.should_receive(:mkdir_p).with("/consumers/my-group/ids")
+    zk_client.should_receive(:mkdir_p).with("/consumers/my-group/owners/mytopic")
+    zk_client.should_receive(:mkdir_p).with("/consumers/my-group/offsets/mytopic")
+    zk_client.should_receive(:create).with("/consumers/my-group/ids/my-group-UNIQUEID", "{}", ephemeral: true)
+    zk_client.should_receive(:register).with("/consumers/my-group/ids")
+    described_class.any_instance.should_receive :rebalance!
+
+    subject
+  end
+
+  it "should sort partitions by leader address" do
+    subject.partitions.map(&:id).should == [1, 0]
+  end
+
+  it "should not fail if topic doesn't exist" do
+    no_topics = Poseidon::Protocol::MetadataResponse.new nil, brokers.dup, []
+    Poseidon::BrokerPool.any_instance.stub(:fetch_metadata_from_broker).and_return(no_topics)
+
+    subject.partitions.should == []
+    subject.claimed.should == []
+  end
+
+  it "should return the offset for each partition" do
+    zk_client.should_receive(:get).with("/consumers/my-group/offsets/mytopic/0", ignore: :no_node).and_return([nil])
+    subject.offset(0).should == 0
+
+    zk_client.should_receive(:get).with("/consumers/my-group/offsets/mytopic/1", ignore: :no_node).and_return(["21", nil])
+    subject.offset(1).should == 21
+
+    zk_client.should_receive(:get).with("/consumers/my-group/offsets/mytopic/2", ignore: :no_node).and_return(["0", nil])
+    subject.offset(2).should == 0
+  end
+
+  it "should return the leader for a partition" do
+    subject.leader(0).should == brokers[0]
+    subject.leader(1).should == brokers[1]
+    subject.leader(2).should be_nil
+  end
+
+  it "should checkout individual partition consumers (atomically)" do
+    subject.checkout {|c| c.partition.should == 1 }.should be_truthy
+    subject.checkout {|c| c.partition.should == 0 }.should be_truthy
+
+    n = 0
+    a = Thread.new do
+      100.times { subject.checkout {|_| n+=1 } }
+      Thread.pass
+      100.times { subject.checkout {|_| n+=1 } }
+    end
+    b = Thread.new do
+      100.times { subject.checkout {|_| n+=1 } }
+      Thread.pass
+      100.times { subject.checkout {|_| n+=1 } }
+    end
+    [a, b].each &:join
+    n.should == 400
+  end
+
+  describe "consumer" do
+    subject { described_class::Consumer.new group, 1 }
+    before  { group.stub(:offset).with(1).and_return(432) }
+
+    it { should be_a(Poseidon::PartitionConsumer) }
+    its(:offset) { should == 432 }
+
+    it 'should start with the earliest offset if none stored' do
+      group.unstub(:offset)
+      subject.offset.should == :earliest_offset
+    end
+
+    it 'should start with the latest offset if none stored and in trailing mode' do
+      group.unstub(:offset)
+      trailing_consumer = described_class::Consumer.new group, 1, {trail: true}
+      trailing_consumer.offset.should == :latest_offset
+    end
+
+  end
+
+  describe "rebalance" do
+
+    it "should watch out for new consumers joining/leaving" do
+      described_class.any_instance.should_receive(:rebalance!)
+      subject
+    end
+
+    it "should distribute available partitions between consumers" do
+      subject.claimed.should == [0, 1]
+      zk_client.stub children: ["my-group-UNIQUEID", "my-group-OTHERID"]
+      -> { subject.send :rebalance! }.should change { subject.claimed }.to([0])
+      zk_client.stub children: ["my-group-UNIQUEID", "my-group-OTHERID", "my-group-THIRDID"]
+      -> { subject.send :rebalance! }.should change { subject.claimed }.to([])
+    end
+
+    it "should allocate partitions correctly" do
+      subject.claimed.should == [0, 1]
+
+      zk_client.stub children: ["my-group-UNIQUEID", "my-group-ZID"]
+      zk_client.should_receive(:delete).with("/consumers/my-group/owners/mytopic/1", ignore: :no_node)
+      -> { subject.send :rebalance! }.should change { subject.claimed }.to([1])
+
+      zk_client.stub children: ["my-group-UNIQUEID", "my-group-ZID", "my-group-AID"]
+      -> { subject.send :rebalance! }.should change { subject.claimed }.to([0])
+    end
+
+  end
+
+  describe "fetch" do
+
+    it "should return messages from claimed partitions" do
+      subject.fetch do |n, msg|
+        n.should == 1
+        msg.size.should == 5
+      end.should be_truthy
+
+      subject.fetch do |n, msg|
+        n.should == 0
+        msg.size.should == 10
+      end.should be_truthy
+
+      subject.fetch do |n, msg|
+        n.should == 1
+        msg.size.should == 5
+      end.should be_truthy
+    end
+
+    it "should auto-commit fetched offset" do
+      zk_client.should_receive(:set).with("/consumers/my-group/offsets/mytopic/1", "5")
+      subject.fetch {|n, _| n.should == 1 }
+    end
+
+    it "should skip auto-commits if requested" do
+      zk_client.should_not_receive(:set)
+      subject.fetch(commit: false) {|n, _| n.should == 1 }
+    end
+
+    it "should skip auto-commits if block results in false" do
+      zk_client.should_not_receive(:set)
+      subject.fetch {|n, _| n.should == 1; false }
+    end
+
+    it "should return false when trying to fetch messages without a claim" do
+      no_topics = Poseidon::Protocol::MetadataResponse.new nil, brokers.dup, []
+      Poseidon::BrokerPool.any_instance.stub fetch_metadata_from_broker: no_topics
+
+      subject.claimed.should == []
+      subject.fetch {|*|  }.should be_falsey
+    end
+
+    it "should return true even when no messages were fetched" do
+      Poseidon::Connection.any_instance.stub fetch: fetch_response(0)
+      subject.fetch {|*|  }.should be_truthy
+    end
+
+  end
+
+  describe "fetch_loop" do
+
+    it "should fetch indefinitely" do
+      total, cycles = 0, 0
+      subject.fetch_loop do |_, m|
+        total += m.size
+        break if (cycles+=1) > 2
+      end
+      total.should == 20
+      cycles.should == 3
+    end
+
+    it "should delay fetch was unsuccessful" do
+      subject.stub fetch: false
+
+      cycles = 0
+      subject.should_receive(:sleep).with(1)
+      subject.fetch_loop do |n, m|
+        n.should == -1
+        m.should == []
+        break if (cycles+=1) > 1
+      end
+    end
+
+    it "should delay fetch didn't yield any results" do
+      subject.stub(:fetch).and_yield(3, []).and_return(true)
+
+      cycles = 0
+      subject.should_receive(:sleep).with(1)
+      subject.fetch_loop do |n, m|
+        n.should == 3
+        m.should == []
+        break if (cycles+=1) > 1
+      end
+    end
+
+  end
+
+  describe "pick" do
+
+    { [3, ["N1", "N2", "N3"], "N1"]       => (0..0),
+      [3, ["N1", "N2", "N3"], "N2"]       => (1..1),
+      [3, ["N1", "N2", "N3"], "N3"]       => (2..2),
+      [4, ["N2", "N4", "N3", "N1"], "N3"] => (2..2),
+      [3, ["N1", "N2", "N3"], "N4"]       => nil,
+      [5, ["N1", "N2", "N3"], "N1"]       => (0..1),
+      [5, ["N1", "N2", "N3"], "N2"]       => (2..3),
+      [5, ["N1", "N2", "N3"], "N3"]       => (4..4),
+      [5, ["N1", "N2", "N3"], "N4"]       => nil,
+      [2, ["N1", "N2"], "N9"]             => nil,
+      [1, ["N1", "N2", "N3"], "N1"]       => (0..0),
+      [1, ["N1", "N2", "N3"], "N2"]       => nil,
+      [1, ["N1", "N2", "N3"], "N3"]       => nil,
+      [5, ["N1", "N2"], "N1"]             => (0..2),
+      [5, ["N1", "N2"], "N2"]             => (3..4),
+    }.each do |args, expected|
+      it "should pick #{expected.inspect} from #{args.inspect}" do
+        described_class.pick(*args).should == expected
+      end
+    end
+
+  end
+end