+ core: add a couple helper methods for creating segments, related to #126.

1 files changed, 20 insertions, 1 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
index 25f3a0a3..85b7396f 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
@@ -18,7 +18,9 @@ package kamon.trace
 
 import java.io.ObjectStreamException
 import kamon.trace.TraceContextAware.DefaultTraceContextAware
-import kamon.util.{ Function, RelativeNanoTimestamp }
+import kamon.util.{ SameThreadExecutionContext, Supplier, Function, RelativeNanoTimestamp }
+
+import scala.concurrent.Future
 
 trait TraceContext {
   def name: String
@@ -45,6 +47,23 @@ trait TraceContext {
     if (nonEmpty)
       Some(f(this))
     else None
+
+  def withNewSegment[T](segmentName: String, category: String, library: String)(code: ⇒ T): T = {
+    val segment = startSegment(segmentName, category, library)
+    try code finally segment.finish()
+  }
+
+  // Java variant.
+  def withNewSegment[T](segmentName: String, category: String, library: String, code: Supplier[T]): T =
+    withNewSegment(segmentName, category, library)(code.get)
+
+  def withNewAsyncSegment[T](segmentName: String, category: String, library: String)(code: ⇒ Future[T]): Future[T] = {
+    val segment = startSegment(segmentName, category, library)
+    val result = code
+    code.onComplete(_ ⇒ segment.finish())(SameThreadExecutionContext)
+    result
+  }
+
 }
 
 trait Segment {