= core: fix NPE when finishing segments, fixes #88

2 files changed, 20 insertions, 5 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
index 6ea30511..a15a4326 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
@@ -26,6 +26,8 @@ import kamon.trace.TraceContextAware.DefaultTraceContextAware
 import kamon.trace.TraceContext.SegmentIdentity
 import kamon.metric.TraceMetrics.TraceMetricRecorder
 
+import scala.annotation.tailrec
+
 trait TraceContext {
   def name: String
   def token: String
@@ -136,10 +138,12 @@ class SimpleMetricCollectionContext(traceName: String, val token: String, metada
     }
   }
 
-  private def drainFinishedSegments(metricRecorder: TraceMetricRecorder): Unit = {
-    while (!finishedSegments.isEmpty) {
-      val segmentData = finishedSegments.poll()
-      metricRecorder.segmentRecorder(segmentData.identity).record(segmentData.duration)
+  @tailrec private def drainFinishedSegments(metricRecorder: TraceMetricRecorder): Unit = {
+    val segment = finishedSegments.poll()
+    if(segment != null) {
+      metricRecorder.segmentRecorder(segment.identity).record(segment.duration)
+      println("CLOSED")
+      drainFinishedSegments(metricRecorder)
     }
   }
 
diff --git a/kamon-playground/src/main/scala/test/SimpleRequestProcessor.scala b/kamon-playground/src/main/scala/test/SimpleRequestProcessor.scala
index 1db2db67..b6bcc677 100644
--- a/kamon-playground/src/main/scala/test/SimpleRequestProcessor.scala
+++ b/kamon-playground/src/main/scala/test/SimpleRequestProcessor.scala
@@ -21,6 +21,7 @@ import akka.routing.RoundRobinPool
 import akka.util.Timeout
 import kamon.Kamon
 import kamon.metric.Subscriptions.TickMetricSnapshot
+import kamon.metric.TraceMetrics.HttpClientRequest
 import kamon.metric._
 import kamon.spray.KamonTraceDirectives
 import kamon.trace.TraceRecorder
@@ -124,6 +125,16 @@ object SimpleRequestProcessor extends App with SimpleRoutingApp with RequestBuil
             throw new NullPointerException
             "okk"
           }
+        } ~
+        path("segment") {
+          complete {
+            val segment = TraceRecorder.startSegment(HttpClientRequest("hello-world"))
+            (replier ? "hello").mapTo[String].onComplete { t =>
+              segment.get.finish()
+            }
+
+            "segment"
+          }
         }
     }
   }
@@ -166,7 +177,7 @@ class Replier extends Actor with ActorLogging {
       if (TraceRecorder.currentContext.isEmpty)
         log.warning("PROCESSING A MESSAGE WITHOUT CONTEXT")
 
-      log.info("Processing at the Replier, and self is: {}", self)
+      //log.info("Processing at the Replier, and self is: {}", self)
       sender ! anything
   }
 }