Improve riemann.kafka resilience to empty messages (#1016)

src/riemann/kafka.clj

@@ -5,8 +5,8 @@
             [riemann.test :as test])
   (:use [riemann.common        :only [event]]
         [riemann.core          :only [stream!]]
-        [riemann.service       :only [Service ServiceEquiv]]
-        [clojure.tools.logging :only [info error]]))
+        [riemann.service       :only [Service ServiceEquiv ServiceStatus]]
+        [clojure.tools.logging :only [debug info error]]))
 
 (defn kafka
   "Returns a function that is invoked with a topic name and an optional message key and returns a stream. That stream is a function which takes an event or a sequence of events and sends them to Kafka.
@@ -82,10 +82,15 @@
                 msgs-by-topic (get msgs :by-topic)]
             (doseq [records msgs-by-topic
                     record (last records)]
-              (let [event (event (get record :value))]
-                (stream! @core event)))))
+              (try
+                (let [event (event (get record :value))]
+                  (stream! @core event))
+                (catch java.lang.NullPointerException e
+                  (debug (str "Invalid message: " record)))))))
         (catch Exception e
-          (error e "Interrupted consumption"))
+          (do
+            (reset! running? false)
+            (error e "Interrupted consumption")))
         (finally
           (client/close! consumer))))))
 
@@ -104,6 +109,8 @@
       ServiceEquiv
       (equiv? [this other]
         (= opts (:opts other)))
+      ServiceStatus
+      (running? [this] @running?)
       Service
       (conflict? [this other]
         (= opts (:opts other)))

src/riemann/service.clj

@@ -44,6 +44,10 @@
   (equiv? [s1 s2]
     (nil? s2)))
 
+(defprotocol ServiceStatus
+  (running? [service]
+          "Queries the running state of the service."))
+
 (defrecord ScheduledTaskService [name equiv-key interval delay f core task]
   ServiceEquiv
   (equiv? [this other]

test/riemann/kafka_test.clj

@@ -21,9 +21,16 @@
                            :state "warning"
                            :time (unix-time)})
 
+(def consumer-timeout-ms 5000)  ; 5s seems necessary on an average developer workstation
+(def burst-event-count 5000000)
+(def burst-consumer-timeout-ms 10000)
+
 ; These tests assume that you have a single server kafka instance running on localhost:9092 (plain) 
-; and localhost:9093 (SSL) and the topics  "riemann" and "custom" configured:
-; ./bin/kafka-topics --create --topic riemann/custom --partitions 1 --replication-factor 1 --zookeeper localhost:2181
+; and localhost:9093 (SSL) and the topics  "riemann", "custom", and "burst" configured:
+; $KAFKA_HOME/bin/kafka-topics.sh --bootstrap-server localhost:9092 --create --partitions 1 --replication-factor 1 --topic riemann
+; $KAFKA_HOME/bin/kafka-topics.sh --bootstrap-server localhost:9092 --create --partitions 1 --replication-factor 1 --topic custom
+; $KAFKA_HOME/bin/kafka-topics.sh --bootstrap-server localhost:9092 --create --partitions 1 --replication-factor 1 --topic burst \
+;   --config segment.bytes=$((10*1024*1024)) --config retention.bytes=$((100*1024*1024)) --config compression.type=lz4
 ;
 ; Important Kafka server.properties settings:
 ; listeners=PLAINTEXT://localhost:9092,SSL://localhost:9093
@@ -49,7 +56,7 @@
         ; Producer writes to riemann topic
         (kafka-output first-event)
         ; Verify event arrives
-        (let [event (deref sink 1000 :timed-out)]
+        (let [event (deref sink consumer-timeout-ms :timed-out)]
           (is (= "firstservice"
                  (:service event)))
           (is (= "myhost"
@@ -86,7 +93,7 @@
     (try
       (testing "kafka ssl with edn serializer, custom topic and nil key"
         (kafka-output second-event)
-        (let [event (deref sink 1000 :timed-out)]
+        (let [event (deref sink consumer-timeout-ms :timed-out)]
           (is (= "secondservice" 
                  (:service event)))
           (is (= "myhost" 
@@ -100,3 +107,37 @@
         (Thread/sleep 1000)
         (core/stop! core)))))
 
+(defn counting-sink [consumer-counter]
+  (fn stream [event] (swap! consumer-counter inc)))
+
+(deftest ^:kafka ^:integration kafka-integration-burst-test
+  (let [consumer (kafka-consumer {:consumer.config {:bootstrap.servers "localhost:9092"
+                                                    :group.id "burstgroup"
+                                                    :auto.commit.interval.ms 1000}
+                                  :topics ["burst"]
+                                  :poll.timeout.ms 1000})
+
+        consumer-counter (atom 0)
+        time-elapsed-ms (atom 0)
+        time-interval-ms 1000
+        core (core/transition! (core/core)
+                               {:services [consumer]
+                                :streams [(counting-sink consumer-counter)]})
+        producer (kafka)
+        kafka-output (producer "burst")]
+    (try
+      (testing "kafka plain with json serializer and burst topic"
+        (dotimes [n burst-event-count] (kafka-output first-event))
+        (do
+          (while (and (< @time-elapsed-ms burst-consumer-timeout-ms) (< @consumer-counter burst-event-count))
+            (do
+              (Thread/sleep time-interval-ms)
+              (swap! time-elapsed-ms + time-interval-ms)
+              (println (str "Consumed events: " @consumer-counter "/" burst-event-count " in " @time-elapsed-ms "ms"
+                (if (.running? consumer) " (running)" " (stopped)")))))
+          (is (.running? consumer))))
+      (finally
+        ; Wait for kafka auto commit
+        (Thread/sleep 1000)
+        (core/stop! core)))))
+