TraceRoot.scala

package mailroom.tracing

import cats.Id
import cats.effect.{IO, IOLocal, Resource}
import natchez._

import java.net.URI

/**
  * Similar to Trace[F] but  allows for the creation of a new root span, either completely fresh or continued
  * off of a kernel. Wraps a Trace[F] so we can delegate it's methods
  */
trait TraceRoot[F[_]] extends Trace[F] {
  def trace: Trace[F]

  def root[A](spanName:               String)(fa:     F[A]): F[A]
  def continue[A](spanName:           String, kernel: Kernel)(fa: F[A]): F[A]
  def continueOrElseRoot[A](spanName: String, kernel: Kernel)(fa: F[A]): F[A]

  override def put(fields: (String, TraceValue)*): F[Unit] = trace.put(fields: _*)
  override def kernel: F[Kernel] = trace.kernel
  override def span[A](name: String)(k: F[A]): F[A] = trace.span(name)(k)
  override def traceId:  F[Option[String]] = trace.traceId
  override def traceUri: F[Option[URI]]    = trace.traceUri
}

object TraceRoot {
  def apply[F[_]](implicit ev: TraceRoot[F]): ev.type = ev
}

class IOTraceRoot(ep: EntryPoint[IO], local: IOLocal[Span[IO]]) extends TraceRoot[IO] {

  override val trace: Trace[IO] =
    new Trace[IO] {

      def put(fields: (String, TraceValue)*): IO[Unit] =
        local.get.flatMap(_.put(fields: _*))

      def kernel: IO[Kernel] =
        local.get.flatMap(_.kernel)

      def span[A](name: String)(k: IO[A]): IO[A] =
        local.get.flatMap(parent => parent.span(name).flatMap(child => Resource.make(local.set(child))(_ => local.set(parent))).use(_ => k))

      def traceId: IO[Option[String]] =
        local.get.flatMap(_.traceId)

      def traceUri: IO[Option[URI]] =
        local.get.flatMap(_.traceUri)

    }

  def ioTraceOverride[A](span: Span[IO])(fa: IO[A]): IO[A] =
    Resource.make(local.set(span))(_ => local.reset).use(_ => fa)

  def root[A](spanName: String)(fa: IO[A]): IO[A] = ep.root(spanName).use(span => ioTraceOverride(span)(fa))

  def continue[A](spanName: String, kernel: Kernel)(fa: IO[A]): IO[A] =
    ep.continue(spanName, kernel).use(span => ioTraceOverride(span)(fa))

  def continueOrElseRoot[A](spanName: String, kernel: Kernel)(fa: IO[A]): IO[A] =
    ep.continueOrElseRoot(spanName, kernel).use(span => ioTraceOverride(span)(fa))
}

object IdTraceRoot extends TraceRoot[Id] {
  override val trace: Trace[Id] = Trace.Implicits.noop[Id]

  override def root[A](spanName:               String)(fa:     Id[A]): Id[A] = fa
  override def continue[A](spanName:           String, kernel: Kernel)(fa: Id[A]): Id[A] = fa
  override def continueOrElseRoot[A](spanName: String, kernel: Kernel)(fa: Id[A]): Id[A] = fa
}