diff options
Diffstat (limited to 'src/main/scala/xyz/driver/core/trace')
8 files changed, 208 insertions, 0 deletions
diff --git a/src/main/scala/xyz/driver/core/trace/GoogleServiceTracer.scala b/src/main/scala/xyz/driver/core/trace/GoogleServiceTracer.scala new file mode 100644 index 0000000..ead4c6f --- /dev/null +++ b/src/main/scala/xyz/driver/core/trace/GoogleServiceTracer.scala @@ -0,0 +1,14 @@ +package xyz.driver.core.trace + +import akka.http.scaladsl.model.headers.RawHeader +import com.google.cloud.trace.Tracer +import com.google.cloud.trace.core.{SpanContextFactory, TraceContext} + +final case class GoogleStackdriverTraceSpan(tracer: Tracer, context: TraceContext) extends CanMakeHeader { + def header: RawHeader = + RawHeader(TracingHeaderKey, SpanContextFactory.toHeader(context.getHandle.getCurrentSpanContext)) +} + +trait GoogleServiceTracer extends ServiceTracer { + type TracerSpanPayload = GoogleStackdriverTraceSpan +} diff --git a/src/main/scala/xyz/driver/core/trace/GoogleStackdriverTrace.scala b/src/main/scala/xyz/driver/core/trace/GoogleStackdriverTrace.scala new file mode 100644 index 0000000..1ff8d10 --- /dev/null +++ b/src/main/scala/xyz/driver/core/trace/GoogleStackdriverTrace.scala @@ -0,0 +1,47 @@ +package xyz.driver.core.trace + +import java.io.FileInputStream +import java.nio.file.{Files, Paths} +import java.util + +import akka.http.scaladsl.model.HttpRequest +import com.google.auth.oauth2.GoogleCredentials +import com.google.cloud.trace.grpc.v1.GrpcTraceConsumer +import com.google.cloud.trace.v1.consumer.TraceConsumer +import com.typesafe.scalalogging.Logger + +final class GoogleStackdriverTrace(projectId: String, + clientSecretsFile: String, + appName: String, + appEnvironment: String, + log: Logger) + extends GoogleServiceTracer { + + // initialize our various tracking storage systems + private val clientSecretsInputStreamOpt: Option[FileInputStream] = if (Files.exists(Paths.get(clientSecretsFile))) { + Some(new FileInputStream(clientSecretsFile)) + } else { + None + } + // if the google credentials are invalid, just log the traces + private val traceConsumer: TraceConsumer = clientSecretsInputStreamOpt.fold[TraceConsumer] { + log.warn(s"Google credentials not found in path: $clientSecretsFile") + new LoggingTraceConsumer(log) + } { clientSecretsInputStream => + GrpcTraceConsumer + .create( + "cloudtrace.googleapis.com", + GoogleCredentials + .fromStream(clientSecretsInputStream) + .createScoped(util.Arrays.asList("https://www.googleapis.com/auth/trace.append")) + ) + } + + private val googleServiceTracer = + new GoogleStackdriverTraceWithConsumer(projectId, appName, appEnvironment, traceConsumer) + + override def startSpan(httpRequest: HttpRequest): GoogleStackdriverTraceSpan = + googleServiceTracer.startSpan(httpRequest) + + override def endSpan(span: GoogleStackdriverTraceSpan): Unit = googleServiceTracer.endSpan(span) +} diff --git a/src/main/scala/xyz/driver/core/trace/GoogleStackdriverTraceWithConsumer.scala b/src/main/scala/xyz/driver/core/trace/GoogleStackdriverTraceWithConsumer.scala new file mode 100644 index 0000000..7fed3c7 --- /dev/null +++ b/src/main/scala/xyz/driver/core/trace/GoogleStackdriverTraceWithConsumer.scala @@ -0,0 +1,75 @@ +package xyz.driver.core.trace + +import akka.http.scaladsl.model.HttpRequest +import com.google.cloud.trace.core._ +import com.google.cloud.trace.sink.TraceSink +import com.google.cloud.trace.v1.TraceSinkV1 +import com.google.cloud.trace.v1.consumer.{SizedBufferingTraceConsumer, TraceConsumer} +import com.google.cloud.trace.v1.producer.TraceProducer +import com.google.cloud.trace.v1.util.RoughTraceSizer +import com.google.cloud.trace.{SpanContextHandler, SpanContextHandlerTracer, Tracer} + +import scala.compat.java8.OptionConverters._ + +final class GoogleStackdriverTraceWithConsumer(projectId: String, + appName: String, + appEnvironment: String, + traceConsumer: TraceConsumer) + extends GoogleServiceTracer { + + private val traceProducer: TraceProducer = new TraceProducer() + private val threadSafeBufferingTraceConsumer = + new SizedBufferingTraceConsumer(traceConsumer, new RoughTraceSizer(), 100) + + private val traceSink: TraceSink = new TraceSinkV1(projectId, traceProducer, threadSafeBufferingTraceConsumer) + + private val spanContextFactory: SpanContextFactory = new SpanContextFactory( + new ConstantTraceOptionsFactory(true, true)) + private val timestampFactory: TimestampFactory = new JavaTimestampFactory() + + override def startSpan(httpRequest: HttpRequest): TracerSpanPayload = { + val parentHeaderOption: Option[akka.http.javadsl.model.HttpHeader] = + httpRequest.getHeader(TracingHeaderKey).asScala + val (spanContext: SpanContext, spanKind: SpanKind) = parentHeaderOption.fold { + (spanContextFactory.initialContext(), SpanKind.RPC_CLIENT) + } { parentHeader => + (spanContextFactory.fromHeader(parentHeader.value()), SpanKind.RPC_SERVER) + } + + val contextHandler: SpanContextHandler = new SimpleSpanContextHandler(spanContext) + val httpMethod = httpRequest.method.value + val httpHost = httpRequest.uri.authority.host.address() + val httpRelative = httpRequest.uri.toRelative.toString() + val tracer: Tracer = new SpanContextHandlerTracer(traceSink, contextHandler, spanContextFactory, timestampFactory) + // Create a span using the given timestamps. + // https://cloud.google.com/trace/docs/reference/v1/rest/v1/projects.traces#TraceSpan + val spanOptions: StartSpanOptions = (new StartSpanOptions()).setSpanKind(spanKind) + + val spanLabelBuilder = Labels + .builder() + .add("/http/method", httpMethod) + .add("/http/url", httpRelative) + .add("/http/host", httpHost) + .add("/component", appName) + .add("/environment", appEnvironment) + + parentHeaderOption.foreach { parentHeader => + spanLabelBuilder.add("/span/parent", parentHeader.value()) + } + + // The cloudTrace analysis reporting UI makes it easy to query by name prefix. + // this spanName gives us the ability to grab things that are specific to a particular UDE/env, as well as all + // endpoints in that service, as well as a particular endpoint in a particular environment/service. + val spanName: String = s"($appEnvironment->$appName)$httpRelative" + + val context: TraceContext = tracer.startSpan(spanName, spanOptions) + tracer.annotateSpan(context, spanLabelBuilder.build()) + GoogleStackdriverTraceSpan(tracer, context) + } + + override def endSpan(span: TracerSpanPayload): Unit = { + span.tracer.endSpan(span.context) + threadSafeBufferingTraceConsumer.flush() // flush out the thread safe buffer + } + +} diff --git a/src/main/scala/xyz/driver/core/trace/LoggingTrace.scala b/src/main/scala/xyz/driver/core/trace/LoggingTrace.scala new file mode 100644 index 0000000..9db85b7 --- /dev/null +++ b/src/main/scala/xyz/driver/core/trace/LoggingTrace.scala @@ -0,0 +1,21 @@ +package xyz.driver.core.trace + +import akka.http.scaladsl.model.HttpRequest +import com.google.cloud.trace.v1.consumer.TraceConsumer +import com.typesafe.scalalogging.Logger + +final class LoggingTrace(appName: String, appEnvironment: String, log: Logger) extends GoogleServiceTracer { + + private val traceConsumer: TraceConsumer = new LoggingTraceConsumer(log) + private val googleServiceTracer = new GoogleStackdriverTraceWithConsumer( + "logging-tracer", + appName, + appEnvironment, + traceConsumer + ) + + override def startSpan(httpRequest: HttpRequest): GoogleStackdriverTraceSpan = + googleServiceTracer.startSpan(httpRequest) + + override def endSpan(span: GoogleStackdriverTraceSpan): Unit = googleServiceTracer.endSpan(span) +} diff --git a/src/main/scala/xyz/driver/core/trace/LoggingTraceConsumer.scala b/src/main/scala/xyz/driver/core/trace/LoggingTraceConsumer.scala new file mode 100644 index 0000000..24df94a --- /dev/null +++ b/src/main/scala/xyz/driver/core/trace/LoggingTraceConsumer.scala @@ -0,0 +1,11 @@ +package xyz.driver.core.trace + +import com.google.cloud.trace.v1.consumer.TraceConsumer +import com.google.devtools.cloudtrace.v1.Traces +import com.typesafe.scalalogging.Logger + +class LoggingTraceConsumer(log: Logger) extends TraceConsumer { + def receive(traces: Traces): Unit = { + log.trace(s"Received traces: $traces") + } +} diff --git a/src/main/scala/xyz/driver/core/trace/ServiceTracer.scala b/src/main/scala/xyz/driver/core/trace/ServiceTracer.scala new file mode 100644 index 0000000..25562cd --- /dev/null +++ b/src/main/scala/xyz/driver/core/trace/ServiceTracer.scala @@ -0,0 +1,17 @@ +package xyz.driver.core.trace + +import akka.http.scaladsl.model.HttpRequest +import akka.http.scaladsl.model.headers.RawHeader + +trait CanMakeHeader { + def header: RawHeader +} + +trait ServiceTracer { + + type TracerSpanPayload <: CanMakeHeader + + def startSpan(httpRequest: HttpRequest): TracerSpanPayload + + def endSpan(span: TracerSpanPayload): Unit +} diff --git a/src/main/scala/xyz/driver/core/trace/SimpleSpanContextHandler.scala b/src/main/scala/xyz/driver/core/trace/SimpleSpanContextHandler.scala new file mode 100644 index 0000000..4ea0e8c --- /dev/null +++ b/src/main/scala/xyz/driver/core/trace/SimpleSpanContextHandler.scala @@ -0,0 +1,17 @@ +package xyz.driver.core.trace + +import com.google.cloud.trace.SpanContextHandler +import com.google.cloud.trace.DetachedSpanContextHandle +import com.google.cloud.trace.core.{SpanContext, SpanContextHandle} + +@SuppressWarnings(Array("org.wartremover.warts.Var")) +class SimpleSpanContextHandler(rootSpan: SpanContext) extends SpanContextHandler { + private var currentSpanContext = rootSpan + + override def current(): SpanContext = currentSpanContext + + override def attach(context: SpanContext): SpanContextHandle = { + currentSpanContext = context + new DetachedSpanContextHandle(context) + } +} diff --git a/src/main/scala/xyz/driver/core/trace/package.scala b/src/main/scala/xyz/driver/core/trace/package.scala new file mode 100644 index 0000000..0dec833 --- /dev/null +++ b/src/main/scala/xyz/driver/core/trace/package.scala @@ -0,0 +1,6 @@ +package xyz.driver.core + +package object trace { + // this happens to be the same as the name google uses, but can be anything we want + val TracingHeaderKey: String = "X-Cloud-Trace-Context" +} |