• Home
• io.janstenpickle
• trace4cats-rate-sampling_3
• 0.13.1+21-b3ad01df
• source code
• RateSpanSampler.scala

×

Project download

Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $

You can buy this project and download/modify it how often you want.

io.janstenpickle.trace4cats.rate.sampling.RateSpanSampler.scala Maven / Gradle / Ivy

package io.janstenpickle.trace4cats.rate.sampling

import cats.Applicative
import cats.effect.kernel.{Resource, Temporal}
import cats.syntax.functor._
import io.janstenpickle.trace4cats.kernel.SpanSampler
import io.janstenpickle.trace4cats.model.{SampleDecision, SpanContext, SpanKind, TraceId}
import io.janstenpickle.trace4cats.rate.{TokenBucket, TokenInterval}

import scala.concurrent.duration._

object RateSpanSampler {
  def apply[F[_]: Applicative: TokenBucket]: SpanSampler[F] = apply(rootSpansOnly = true)

  def apply[F[_]: Applicative: TokenBucket](rootSpansOnly: Boolean): SpanSampler[F] =
    new SpanSampler[F] {
      override def shouldSample(
        parentContext: Option[SpanContext],
        traceId: TraceId,
        spanName: String,
        spanKind: SpanKind
      ): F[SampleDecision] = {
        val sampleDecisionF: F[SampleDecision] = TokenBucket[F].request1.map {
          case false => SampleDecision.Drop
          case true => SampleDecision.Include
        }

        parentContext.map(_.traceFlags.sampled).fold(sampleDecisionF) {
          case SampleDecision.Include =>
            if (rootSpansOnly) Applicative[F].pure(SampleDecision.Include)
            else sampleDecisionF
          case SampleDecision.Drop => Applicative[F].pure(SampleDecision.Drop)
        }
      }

    }

  def create[F[_]: Temporal](bucketSize: Int, tokenRate: Double): Resource[F, SpanSampler[F]] =
    create(bucketSize, tokenRate, rootSpansOnly = true)

  def create[F[_]: Temporal](bucketSize: Int, tokenRate: Double, rootSpansOnly: Boolean): Resource[F, SpanSampler[F]] =
    TokenInterval(tokenRate)
      .fold(Resource.pure[F, SpanSampler[F]](SpanSampler.always[F]))(create[F](bucketSize, _, rootSpansOnly))

  def create[F[_]: Temporal](bucketSize: Int, tokenInterval: FiniteDuration): Resource[F, SpanSampler[F]] =
    create(bucketSize, tokenInterval, rootSpansOnly = true)

  def create[F[_]: Temporal](
    bucketSize: Int,
    tokenInterval: FiniteDuration,
    rootSpansOnly: Boolean
  ): Resource[F, SpanSampler[F]] =
    TokenBucket.create[F](bucketSize, tokenInterval).map(implicit tb => apply[F](rootSpansOnly))
}