Add Sampler type and implementation that aligns with Jaeger sampling (#…

…1107)

misk-jaeger/src/main/kotlin/misk/tracing/backends/jaeger/JaegerBackendModule.kt

@@ -7,6 +7,7 @@ import io.opentracing.Tracer
 import misk.ServiceModule
 import misk.config.AppName
 import misk.inject.KAbstractModule
+import misk.sampling.Sampler
 
 class JaegerBackendModule(val config: JaegerBackendConfig?) : KAbstractModule() {
   override fun configure() {
@@ -33,4 +34,13 @@ class JaegerBackendModule(val config: JaegerBackendConfig?) : KAbstractModule()
         .withScopeManager(MDCScopeManager())
         .build()
   }
+
+  @Provides
+  @Singleton
+  @Tracing
+  fun sampler(tracer: Tracer): Sampler {
+    val samplingRate = config?.sampler?.param ?: Configuration.SamplerConfiguration().param
+
+    return SpanSampler(tracer = tracer, samplingRate = samplingRate.toDouble())
+  }
 }

misk-jaeger/src/main/kotlin/misk/tracing/backends/jaeger/SpanSampler.kt

@@ -0,0 +1,35 @@
+package misk.tracing.backends.jaeger
+
+import com.uber.jaeger.samplers.ProbabilisticSampler
+import io.opentracing.Span
+import io.opentracing.Tracer
+import misk.sampling.Sampler
+import com.uber.jaeger.Span as JaegerSpan
+
+/**
+ * A sampler that deterministically samples based on the current tracing span ID.
+ * Sampling rate when no span is present is 1.0.
+ *
+ * @param samplingRate Sampling rate. 0 is never and 1 is always.
+ */
+class SpanSampler(
+  private val tracer: Tracer,
+  samplingRate: Double
+) : Sampler {
+
+  init {
+    check(samplingRate in 0.0..1.0) { "sampling rate must be in range [0, 1]" }
+  }
+
+  private val sampler = ProbabilisticSampler(samplingRate)
+
+  /** Returns true if the sampling threshold is met for the current Jaeger span or there is no current span */
+  override fun sample(): Boolean {
+    val spanId: Long? = when (val span: Span? = tracer.activeSpan()) {
+      is JaegerSpan -> span.context().spanId
+      else -> null
+    }
+
+    return spanId == null || sampler.sample("irrelevant", spanId).isSampled
+  }
+}

misk-jaeger/src/main/kotlin/misk/tracing/backends/jaeger/Tracing.kt

@@ -0,0 +1,7 @@
+package misk.tracing.backends.jaeger
+
+import javax.inject.Qualifier
+
+@Qualifier
+@Target(AnnotationTarget.FIELD, AnnotationTarget.FUNCTION, AnnotationTarget.VALUE_PARAMETER)
+annotation class Tracing

misk/src/main/kotlin/misk/config/MiskConfig.kt

@@ -17,7 +17,6 @@ import com.fasterxml.jackson.module.kotlin.KotlinModule
 import com.fasterxml.jackson.module.kotlin.MissingKotlinParameterException
 import com.google.common.base.Joiner
 import misk.environment.Environment
-import misk.logging.getLogger
 import misk.resources.ResourceLoader
 import okio.buffer
 import okio.source
@@ -27,8 +26,6 @@ import java.io.FilenameFilter
 import java.net.URL
 
 object MiskConfig {
-  internal val logger = getLogger<Config>()
-
   @JvmStatic
   inline fun <reified T : Config> load(
     appName: String,

misk/src/main/kotlin/misk/sampling/Sampler.kt

@@ -0,0 +1,34 @@
+package misk.sampling
+
+import java.util.concurrent.ThreadLocalRandom
+import javax.inject.Inject
+import javax.inject.Singleton
+
+interface Sampler {
+  /** If an action should be taken based on the implementation's policy, returns true */
+  fun sample(): Boolean
+
+  /** If [sample] returns true, runs the given lambda */
+  fun sampledCall(f: () -> Unit) {
+    if (sample()) {
+      f()
+    }
+  }
+}
+
+/** A [Sampler] randomly invokes an action based on a sample percentage */
+class PercentSampler(
+  val samplePercentage: () -> Int,
+  val random: () -> Int = { ThreadLocalRandom.current().nextInt(0, 100) }
+) : Sampler {
+  constructor(samplePercentage: Int, random: () -> Int) : this({ samplePercentage }, random)
+  constructor(samplePercentage: Int) : this({ samplePercentage })
+
+  override fun sample(): Boolean = random() < samplePercentage()
+}
+
+/** Sampler that always invokes an action */
+@Singleton
+class AlwaysSampler @Inject constructor() : Sampler {
+  override fun sample(): Boolean = true
+}