update to new shutdown semantics

Fixes logstash-plugins#41.

lib/logstash/inputs/kafka.rb

@@ -133,31 +133,32 @@ def run(logstash_queue)
     @logger.info('Running kafka', :group_id => @group_id, :topic_id => @topic_id, :zk_connect => @zk_connect)
     begin
       @consumer_group.run(@consumer_threads,@kafka_client_queue)
-      begin
-        while true
-          event = @kafka_client_queue.pop
-          queue_event(event, logstash_queue)
-        end
-      rescue LogStash::ShutdownSignal
-        @logger.info('Kafka got shutdown signal')
-        @consumer_group.shutdown
+
+      while !stop?
+        event = @kafka_client_queue.pop
+        queue_event(event, logstash_queue)
       end
+
       until @kafka_client_queue.empty?
         queue_event(@kafka_client_queue.pop,logstash_queue)
       end
+
       @logger.info('Done running kafka input')
     rescue => e
       @logger.warn('kafka client threw exception, restarting',
                    :exception => e)
-      if @consumer_group.running?
-        @consumer_group.shutdown
-      end
+      stop
       sleep(Float(@consumer_restart_sleep_ms) * 1 / 1000)
       retry
     end
-    finished
   end # def run
 
+  public
+  def stop
+    super
+    @consumer_group.shutdown if @consumer_group.running?
+  end
+
   private
   def create_consumer_group(options)
     Kafka::Group.new(options)
@@ -182,5 +183,4 @@ def queue_event(message_and_metadata, output_queue)
                     :backtrace => e.backtrace)
     end # begin
   end # def queue_event
-
 end #class LogStash::Inputs::Kafka

spec/inputs/kafka_spec.rb

@@ -7,8 +7,7 @@ class LogStash::Inputs::TestKafka < LogStash::Inputs::Kafka
   private
   def queue_event(msg, output_queue)
     super(msg, output_queue)
-    # need to raise exception here to stop the infinite loop
-    raise LogStash::ShutdownSignal
+    stop
   end
 end
 
@@ -30,6 +29,25 @@ def run(a_num_threads, a_queue)
   end
 end
 
+class LogStash::Inputs::TestInfiniteKafka < LogStash::Inputs::Kafka
+  private
+  def queue_event(msg, output_queue)
+    super(msg, output_queue)
+  end
+end
+
+class TestInfiniteKafkaGroup < Kafka::Group
+  def run(a_num_threads, a_queue)
+    blah = TestMessageAndMetadata.new(@topic, 0, nil, 'Kafka message')
+    Thread.new do
+      while true
+        a_queue << blah
+        sleep 0.2
+      end
+    end
+  end
+end
+
 describe 'inputs/kafka' do
   let (:kafka_config) {{'topic_id' => 'test'}}
   let (:empty_config) {{}}
@@ -57,6 +75,30 @@ def run(a_num_threads, a_queue)
     expect {input.register}.to raise_error
   end
 
+  context "interrupted plugin" do
+    let(:plugin) { LogStash::Inputs::TestInfiniteKafka.new(kafka_config) }
+    let!(:queue) { SizedQueue.new(20) }
+    let!(:consumer_thread) { Thread.new(queue) { |queue| loop { queue.pop } } }
+    subject { Thread.new(plugin, queue) { |plugin, queue| plugin.run(queue) } }
+
+    before do
+      expect(plugin).to receive(:create_consumer_group) do |options|
+        TestInfiniteKafkaGroup.new(options)
+      end
+      plugin.register
+    end
+
+    after do
+      Thread.kill(consumer_thread)
+    end
+
+    it "should shutdown when stopped is called" do
+      expect(subject).to be_alive
+      plugin.stop
+      wait(3).for { subject }.to_not be_alive
+    end
+  end
+
   it 'should populate kafka config with default values' do
     kafka = LogStash::Inputs::TestKafka.new(kafka_config)
     insist {kafka.zk_connect} == 'localhost:2181'