Add tracing to client HTTP transport and improve tracing tags

3 files changed, 64 insertions, 20 deletions

diff --git a/src/main/scala/xyz/driver/core/reporting/GoogleReporter.scala b/src/main/scala/xyz/driver/core/reporting/GoogleReporter.scala
index 2f889f5..d4d20a4 100644
--- a/src/main/scala/xyz/driver/core/reporting/GoogleReporter.scala
+++ b/src/main/scala/xyz/driver/core/reporting/GoogleReporter.scala
@@ -10,6 +10,7 @@ import akka.stream.{Materializer, OverflowStrategy}
 import com.google.auth.oauth2.ServiceAccountCredentials
 import com.softwaremill.sttp._
 import com.typesafe.scalalogging.Logger
+import org.slf4j.MDC
 import spray.json.DerivedJsonProtocol._
 import spray.json._
 import xyz.driver.core.reporting.Reporter.CausalRelation
@@ -17,7 +18,7 @@ import xyz.driver.core.reporting.Reporter.CausalRelation
 import scala.async.Async._
 import scala.concurrent.duration._
 import scala.concurrent.{ExecutionContext, Future}
-import scala.util.Random
+import scala.util.{Failure, Random, Success, Try}
 import scala.util.control.NonFatal
 
 /** A reporter that collects traces and submits them to
@@ -117,7 +118,8 @@ class GoogleReporter(
     TruncatableString(displayName),
     Instant.now(),
     Instant.now(),
-    Attributes(attributes ++ Map("namespace" -> namespace))
+    Attributes(attributes ++ Map("service.namespace" -> namespace)),
+    Failure(new IllegalStateException("span status not set"))
   )
 
   def traceWithOptionalParent[A](
@@ -145,12 +147,20 @@ class GoogleReporter(
     }
     val span   = startSpan(child.traceId, child.spanId, parent.map(_._1.spanId), operationName, tags)
     val result = operation(child)
-    result.onComplete { _ =>
+    result.onComplete { result =>
       span.endTime = Instant.now()
+      span.status = result
       submit(span)
     }
     result
   }
+
+  override def log(severity: Reporter.Severity, message: String, reason: Option[Throwable])(
+      implicit ctx: SpanContext): Unit = {
+    MDC.put("trace", s"projects/${credentials.getProjectId}/traces/${ctx.traceId}")
+    super.log(severity, message, reason)
+  }
+
 }
 
 object GoogleReporter {
@@ -167,7 +177,8 @@ object GoogleReporter {
       displayName: TruncatableString,
       startTime: Instant,
       var endTime: Instant,
-      attributes: Attributes
+      var attributes: Attributes,
+      var status: Try[_]
   )
 
   private case class Spans(spans: Seq[Span])
@@ -185,9 +196,25 @@ object GoogleReporter {
     }
   }
 
+  private implicit val statusFormat = new RootJsonFormat[Try[_]] {
+    override def read(json: JsValue): Try[_] = sys.error("unimplemented")
+    override def write(obj: Try[_]) = {
+      // error codes defined at https://github.com/googleapis/googleapis/blob/master/google/rpc/code.proto
+      val (code, message) = obj match {
+        case Success(_) => (0, "success")
+        case Failure(_) => (2, "failure")
+      }
+      JsObject(
+        "code"    -> code.toJson,
+        "message" -> message.toJson,
+        "details" -> JsArray()
+      )
+    }
+  }
+
   private implicit val attributeFormat: RootJsonFormat[Attributes]                = jsonFormat1(Attributes)
   private implicit val truncatableStringFormat: RootJsonFormat[TruncatableString] = jsonFormat1(TruncatableString)
-  private implicit val spanFormat: RootJsonFormat[Span]                           = jsonFormat7(Span)
+  private implicit val spanFormat: RootJsonFormat[Span]                           = jsonFormat8(Span)
   private implicit val spansFormat: RootJsonFormat[Spans]                         = jsonFormat1(Spans)
 
 }
diff --git a/src/main/scala/xyz/driver/core/reporting/Reporter.scala b/src/main/scala/xyz/driver/core/reporting/Reporter.scala
index 9649ada..57e2310 100644
--- a/src/main/scala/xyz/driver/core/reporting/Reporter.scala
+++ b/src/main/scala/xyz/driver/core/reporting/Reporter.scala
@@ -2,7 +2,7 @@ package xyz.driver.core.reporting
 
 import com.typesafe.scalalogging.Logger
 import org.slf4j.helpers.NOPLogger
-import xyz.driver.core.reporting.Reporter.CausalRelation
+import xyz.driver.core.reporting.Reporter.{CausalRelation, Severity}
 
 import scala.concurrent.Future
 
@@ -128,20 +128,28 @@ trait Reporter {
       Some(ctx -> relation)
     )(op)
 
+  /** Log a message. */
+  def log(severity: Severity, message: String, reason: Option[Throwable])(implicit ctx: SpanContext): Unit
+
   /** Log a debug message. */
-  def debug(message: String)(implicit ctx: SpanContext): Unit
+  def debug(message: String)(implicit ctx: SpanContext): Unit = log(Severity.Debug, message, None)
+  def debug(message: String, reason: Throwable)(implicit ctx: SpanContext): Unit =
+    log(Severity.Debug, message, Some(reason))
 
   /** Log an informational message. */
-  def info(message: String)(implicit ctx: SpanContext): Unit
+  def info(message: String)(implicit ctx: SpanContext): Unit = log(Severity.Informational, message, None)
+  def info(message: String, reason: Throwable)(implicit ctx: SpanContext): Unit =
+    log(Severity.Informational, message, Some(reason))
 
   /** Log a warning message. */
-  def warn(message: String)(implicit ctx: SpanContext): Unit
-
-  /** Log a error message. */
-  def error(message: String)(implicit ctx: SpanContext): Unit
+  def warn(message: String)(implicit ctx: SpanContext): Unit = log(Severity.Warning, message, None)
+  def warn(message: String, reason: Throwable)(implicit ctx: SpanContext): Unit =
+    log(Severity.Warning, message, Some(reason))
 
-  /** Log a error message with an associated throwable that caused the error condition. */
-  def error(message: String, reason: Throwable)(implicit ctx: SpanContext): Unit
+  /** Log an error message. */
+  def error(message: String)(implicit ctx: SpanContext): Unit = log(Severity.Error, message, None)
+  def error(message: String, reason: Throwable)(implicit ctx: SpanContext): Unit =
+    log(Severity.Error, message, Some(reason))
 
 }
 
@@ -164,4 +172,12 @@ object Reporter {
     case object Follows extends CausalRelation
   }
 
+  sealed trait Severity
+  object Severity {
+    case object Debug         extends Severity
+    case object Informational extends Severity
+    case object Warning       extends Severity
+    case object Error         extends Severity
+  }
+
 }
diff --git a/src/main/scala/xyz/driver/core/reporting/ScalaLoggerLike.scala b/src/main/scala/xyz/driver/core/reporting/ScalaLoggerLike.scala
index c1131fb..eda81fb 100644
--- a/src/main/scala/xyz/driver/core/reporting/ScalaLoggerLike.scala
+++ b/src/main/scala/xyz/driver/core/reporting/ScalaLoggerLike.scala
@@ -5,12 +5,13 @@ trait ScalaLoggerLike extends Reporter {
 
   def logger: Logger
 
-  override def debug(message: String)(implicit ctx: SpanContext): Unit = logger.debug(message)
-  override def info(message: String)(implicit ctx: SpanContext): Unit  = logger.info(message)
-  override def warn(message: String)(implicit ctx: SpanContext): Unit  = logger.warn(message)
-  override def error(message: String)(implicit ctx: SpanContext): Unit = logger.error(message)
-  override def error(message: String, reason: Throwable)(implicit ctx: SpanContext): Unit =
-    logger.error(message, reason)
+  override def log(severity: Reporter.Severity, message: String, reason: Option[Throwable])(
+      implicit ctx: SpanContext): Unit = severity match {
+    case Reporter.Severity.Debug         => logger.debug(message, reason.orNull)
+    case Reporter.Severity.Informational => logger.info(message, reason.orNull)
+    case Reporter.Severity.Warning       => logger.warn(message, reason.orNull)
+    case Reporter.Severity.Error         => logger.error(message, reason.orNull)
+  }
 
 }