Skip to content
Permalink
Browse files

Add span context to all queries when using VITESS over MySQL (#1200)

Signed-off-by: Andres Taylor <antaylor@squareup.com>
  • Loading branch information...
systay committed Sep 19, 2019
1 parent a734044 commit 88d297bf15afcf6c46cbda205c4ed4cfc378b6e6
@@ -3,6 +3,7 @@ dependencies {
compile project(':misk')
compile project(':misk-testing')
compile project(':misk-hibernate')
testCompile dep.openTracingDatadog
}

test {

This file was deleted.

@@ -0,0 +1,61 @@
package misk.jdbc

import datadog.opentracing.DDSpan
import datadog.opentracing.DDTracer
import datadog.trace.common.writer.Writer
import net.ttddyy.dsproxy.transform.TransformInfo
import org.junit.jupiter.api.Test
import org.mockito.Mockito
import io.opentracing.mock.MockTracer
import org.assertj.core.api.Assertions.assertThat
import javax.sql.DataSource

class SpanInjectorTest {
@Test
fun testNotDecorateIfItsVitess() {
val tracer = MockTracer()
val config = DataSourceConfig(DataSourceType.VITESS)
val injector = SpanInjector(tracer, config)
val ds = Mockito.mock(DataSource::class.java)

assertThat(injector.decorate(ds)).isSameAs(ds)
}

@Test
fun testDecoratesIfItsVitessMysql() {
val tracer = MockTracer()
val config = DataSourceConfig(DataSourceType.VITESS_MYSQL)
val injector = SpanInjector(tracer, config)
val ds = Mockito.mock(DataSource::class.java)

assertThat(injector.decorate(ds)).isNotSameAs(ds)
}

@Test
fun testDatadog() {
DDTracer(NoopWriter()).use { tracer ->
tracer.buildSpan("operation").startActive(true).use { scope ->
val config = DataSourceConfig(DataSourceType.VITESS_MYSQL)
val injector = SpanInjector(tracer, config)
val query = "SELECT * FROM table"
val transformInfo = TransformInfo(null, null, query, false, 0)
val result = injector.transformQuery(transformInfo)
assertThat(result).contains((scope.span() as DDSpan).traceId)
}
}
}
}

class NoopWriter : Writer {
override fun start() {
}

override fun write(trace: MutableList<DDSpan>?) {
}

override fun close() {
}

override fun incrementTraceCount() {
}
}
@@ -24,7 +24,6 @@ dependencies {
compile dep.openTracing
compile dep.openTracingUtil
compile dep.openTracingJdbc
compile dep.tracingJaeger
compile dep.vitess
compile dep.tink
compile project(':misk')
@@ -16,7 +16,7 @@ import misk.jdbc.DataSourceDecorator
import misk.jdbc.DataSourceService
import misk.jdbc.DataSourceType
import misk.jdbc.DatabasePool
import misk.jdbc.JaegerSpanInjector
import misk.jdbc.SpanInjector
import misk.jdbc.PingDatabaseService
import misk.jdbc.RealDatabasePool
import misk.metrics.Metrics
@@ -179,14 +179,14 @@ class HibernateModule(
.dependsOn<DataSourceService>(qualifier))

if (config.type == DataSourceType.VITESS_MYSQL) {
val jaegerSpanInjectorDecoratorKey = JaegerSpanInjector::class.toKey(qualifier)
val jaegerSpanInjectorDecoratorKey = SpanInjector::class.toKey(qualifier)
bind(jaegerSpanInjectorDecoratorKey)
.toProvider(object : Provider<JaegerSpanInjector> {
.toProvider(object : Provider<SpanInjector> {
@com.google.inject.Inject(optional = true)
var tracer: Tracer? = null

override fun get(): JaegerSpanInjector =
JaegerSpanInjector(tracer, config)
override fun get(): SpanInjector =
SpanInjector(tracer, config)
}).asSingleton()
}

This file was deleted.

@@ -0,0 +1,54 @@
package misk.jdbc

import io.opentracing.Tracer
import io.opentracing.propagation.Format
import net.ttddyy.dsproxy.support.ProxyDataSourceBuilder
import net.ttddyy.dsproxy.transform.QueryTransformer
import net.ttddyy.dsproxy.transform.TransformInfo
import javax.sql.DataSource
import io.opentracing.propagation.TextMap
import java.lang.RuntimeException

/***
* On the fly decorates queries with the span context, so the query call can be traced all the way through Vitess
*/
class SpanInjector(
val tracer: Tracer?,
val config: DataSourceConfig
) : QueryTransformer, DataSourceDecorator {
override fun decorate(dataSource: DataSource): DataSource {
if (config.type != DataSourceType.VITESS_MYSQL || tracer == null) return dataSource
return ProxyDataSourceBuilder(dataSource).queryTransformer(this).build()
}

override fun transformQuery(transformInfo: TransformInfo?): String {
val activeSpan = this.tracer?.activeSpan() ?: return transformInfo?.query!!

return "/*VT_SPAN_CONTEXT=${getContextString(activeSpan.context())}*/${transformInfo?.query!!}"
}

private fun getContextString(activeSpan: io.opentracing.SpanContext): String {
val carrier = StringBuilderCarrier()

this.tracer!!.inject(activeSpan, Format.Builtin.TEXT_MAP, carrier)

return carrier.toString()
}
}

class StringBuilderCarrier : TextMap {
val map = HashMap<String, String>()

override fun iterator(): MutableIterator<MutableMap.MutableEntry<String, String>> {
throw RuntimeException("not implemented")
}

override fun put(key: String?, value: String?) {
if (key != null && value != null)
map.put(key, value)
}

override fun toString(): String {
return map.map { e -> e.key + "=" + e.value }.joinToString(separator = ":")
}
}

0 comments on commit 88d297b

Please sign in to comment.
You can’t perform that action at this time.