Merge pull request #284 from NthPortal/Tracer-mapK/PR

Add `mapK` to `Tracer`, `SpanBuilder`, etc.

core/common/src/main/scala/org/typelevel/otel4s/KindTransformer.scala

@@ -0,0 +1,107 @@
+/*
+ * Copyright 2022 Typelevel
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.typelevel.otel4s
+
+import cats.Functor
+import cats.Monad
+import cats.data.EitherT
+import cats.data.IorT
+import cats.data.Kleisli
+import cats.data.OptionT
+import cats.data.StateT
+import cats.effect.MonadCancelThrow
+import cats.effect.Resource
+import cats.~>
+
+/** A utility for transforming the higher-kinded type `F` to another
+  * higher-kinded type `G`.
+  */
+@annotation.implicitNotFound("No transformer defined from ${F} to ${G}")
+trait KindTransformer[F[_], G[_]] {
+
+  /** A higher-kinded function that lifts the kind `F` into a `G`.
+    *
+    * @note
+    *   This method is usually best implemented by a `liftK` method on `G`'s
+    *   companion object.
+    */
+  val liftK: F ~> G
+
+  /** Modify the context of `G[A]` using the natural transformation `f`.
+    *
+    * This method is "limited" in the sense that while most `mapK` methods can
+    * modify the context using arbitrary transformations, this method can only
+    * modify the context using natural transformations.
+    *
+    * @note
+    *   This method is usually best implemented by a `mapK` method on `G`.
+    */
+  def limitedMapK[A](ga: G[A])(f: F ~> F): G[A]
+
+  /** Lifts a natural transformation from `F` to `F` into a natural
+    * transformation from `G` to `G`.
+    */
+  final def liftFunctionK(f: F ~> F): G ~> G =
+    new (G ~> G) {
+      def apply[A](ga: G[A]): G[A] = limitedMapK(ga)(f)
+    }
+}
+
+object KindTransformer {
+  implicit def optionT[F[_]: Functor]: KindTransformer[F, OptionT[F, *]] =
+    new KindTransformer[F, OptionT[F, *]] {
+      val liftK: F ~> OptionT[F, *] = OptionT.liftK
+      def limitedMapK[A](ga: OptionT[F, A])(f: F ~> F): OptionT[F, A] =
+        ga.mapK(f)
+    }
+
+  implicit def eitherT[F[_]: Functor, L]: KindTransformer[F, EitherT[F, L, *]] =
+    new KindTransformer[F, EitherT[F, L, *]] {
+      val liftK: F ~> EitherT[F, L, *] = EitherT.liftK
+      def limitedMapK[R](ga: EitherT[F, L, R])(f: F ~> F): EitherT[F, L, R] =
+        ga.mapK(f)
+    }
+
+  implicit def iorT[F[_]: Functor, L]: KindTransformer[F, IorT[F, L, *]] =
+    new KindTransformer[F, IorT[F, L, *]] {
+      val liftK: F ~> IorT[F, L, *] = IorT.liftK
+      def limitedMapK[R](ga: IorT[F, L, R])(f: F ~> F): IorT[F, L, R] =
+        ga.mapK(f)
+    }
+
+  implicit def kleisli[F[_], A]: KindTransformer[F, Kleisli[F, A, *]] =
+    new KindTransformer[F, Kleisli[F, A, *]] {
+      val liftK: F ~> Kleisli[F, A, *] = Kleisli.liftK
+      def limitedMapK[B](ga: Kleisli[F, A, B])(f: F ~> F): Kleisli[F, A, B] =
+        ga.mapK(f)
+    }
+
+  implicit def stateT[F[_]: Monad, S]: KindTransformer[F, StateT[F, S, *]] =
+    new KindTransformer[F, StateT[F, S, *]] {
+      val liftK: F ~> StateT[F, S, *] = StateT.liftK
+      def limitedMapK[A](ga: StateT[F, S, A])(f: F ~> F): StateT[F, S, A] =
+        ga.mapK(f)
+    }
+
+  implicit def resource[F[_]: MonadCancelThrow]
+      : KindTransformer[F, Resource[F, *]] =
+    new KindTransformer[F, Resource[F, *]] {
+      val liftK: F ~> Resource[F, *] = Resource.liftK
+      def limitedMapK[A](ga: Resource[F, A])(f: F ~> F): Resource[F, A] =
+        ga.mapK(f)
+    }
+}

core/common/src/main/scala/org/typelevel/otel4s/meta/InstrumentMeta.scala

@@ -17,6 +17,7 @@
 package org.typelevel.otel4s.meta
 
 import cats.Applicative
+import cats.~>
 
 trait InstrumentMeta[F[_]] {
 
@@ -28,6 +29,9 @@ trait InstrumentMeta[F[_]] {
     */
   def unit: F[Unit]
 
+  /** Modify the context `F` using the transformation `f`. */
+  def mapK[G[_]](f: F ~> G): InstrumentMeta[G] =
+    new InstrumentMeta.MappedK(this)(f)
 }
 
 object InstrumentMeta {
@@ -44,4 +48,9 @@ object InstrumentMeta {
       val unit: F[Unit] = Applicative[F].unit
     }
 
+  private class MappedK[F[_], G[_]](meta: InstrumentMeta[F])(f: F ~> G)
+      extends InstrumentMeta[G] {
+    def isEnabled: Boolean = meta.isEnabled
+    def unit: G[Unit] = f(meta.unit)
+  }
 }

core/trace/src/main/scala/org/typelevel/otel4s/trace/Span.scala

@@ -18,6 +18,7 @@ package org.typelevel.otel4s
 package trace
 
 import cats.Applicative
+import cats.~>
 import org.typelevel.otel4s.meta.InstrumentMeta
 
 import scala.concurrent.duration.FiniteDuration
@@ -108,6 +109,15 @@ trait Span[F[_]] extends SpanMacro[F] {
     */
   final def end(timestamp: FiniteDuration): F[Unit] =
     backend.end(timestamp)
+
+  /** Modify the context `F` using the transformation `f`. */
+  def mapK[G[_]](f: F ~> G): Span[G] = Span.fromBackend(backend.mapK(f))
+
+  /** Modify the context `F` using an implicit [[KindTransformer]] from `F` to
+    * `G`.
+    */
+  final def mapK[G[_]](implicit kt: KindTransformer[F, G]): Span[G] =
+    mapK(kt.liftK)
 }
 
 object Span {
@@ -137,6 +147,15 @@ object Span {
 
     private[otel4s] def end: F[Unit]
     private[otel4s] def end(timestamp: FiniteDuration): F[Unit]
+
+    /** Modify the context `F` using the transformation `f`. */
+    def mapK[G[_]](f: F ~> G): Backend[G] = new Backend.MappedK(this)(f)
+
+    /** Modify the context `F` using an implicit [[KindTransformer]] from `F` to
+      * `G`.
+      */
+    final def mapK[G[_]](implicit kt: KindTransformer[F, G]): Backend[G] =
+      mapK(kt.liftK)
   }
 
   object Backend {
@@ -169,6 +188,36 @@ object Span {
         private[otel4s] def end: F[Unit] = unit
         private[otel4s] def end(timestamp: FiniteDuration): F[Unit] = unit
       }
+
+    /** Implementation for [[Backend.mapK]]. */
+    private class MappedK[F[_], G[_]](backend: Backend[F])(f: F ~> G)
+        extends Backend[G] {
+      def meta: InstrumentMeta[G] =
+        backend.meta.mapK(f)
+      def context: SpanContext = backend.context
+      def addAttributes(attributes: Attribute[_]*): G[Unit] =
+        f(backend.addAttributes(attributes: _*))
+      def addEvent(name: String, attributes: Attribute[_]*): G[Unit] =
+        f(backend.addEvent(name, attributes: _*))
+      def addEvent(
+          name: String,
+          timestamp: FiniteDuration,
+          attributes: Attribute[_]*
+      ): G[Unit] =
+        f(backend.addEvent(name, timestamp, attributes: _*))
+      def recordException(
+          exception: Throwable,
+          attributes: Attribute[_]*
+      ): G[Unit] =
+        f(backend.recordException(exception, attributes: _*))
+      def setStatus(status: Status): G[Unit] =
+        f(backend.setStatus(status))
+      def setStatus(status: Status, description: String): G[Unit] =
+        f(backend.setStatus(status, description))
+      private[otel4s] def end: G[Unit] = f(backend.end)
+      private[otel4s] def end(timestamp: FiniteDuration): G[Unit] =
+        f(backend.end(timestamp))
+    }
   }
 
   private[otel4s] def fromBackend[F[_]](back: Backend[F]): Span[F] =

core/trace/src/main/scala/org/typelevel/otel4s/trace/SpanBuilder.scala

@@ -19,6 +19,7 @@ package trace
 
 import cats.Applicative
 import cats.arrow.FunctionK
+import cats.effect.MonadCancelThrow
 import cats.effect.Resource
 
 import scala.concurrent.duration.FiniteDuration
@@ -112,6 +113,15 @@ trait SpanBuilder[F[_]] {
   def withParent(parent: SpanContext): SpanBuilder[F]
 
   def build: SpanOps[F]
+
+  /** Modify the context `F` using an implicit [[KindTransformer]] from `F` to
+    * `G`.
+    */
+  def mapK[G[_]: MonadCancelThrow](implicit
+      F: MonadCancelThrow[F],
+      kt: KindTransformer[F, G]
+  ): SpanBuilder[G] =
+    new SpanBuilder.MappedK(this)
 }
 
 object SpanBuilder {
@@ -152,4 +162,31 @@ object SpanBuilder {
       }
     }
 
+  /** Implementation for [[SpanBuilder.mapK]]. */
+  private class MappedK[F[_]: MonadCancelThrow, G[_]: MonadCancelThrow](
+      builder: SpanBuilder[F]
+  )(implicit kt: KindTransformer[F, G])
+      extends SpanBuilder[G] {
+    def addAttribute[A](attribute: Attribute[A]): SpanBuilder[G] =
+      new MappedK(builder.addAttribute(attribute))
+    def addAttributes(attributes: Attribute[_]*): SpanBuilder[G] =
+      new MappedK(builder.addAttributes(attributes: _*))
+    def addLink(
+        spanContext: SpanContext,
+        attributes: Attribute[_]*
+    ): SpanBuilder[G] =
+      new MappedK(builder.addLink(spanContext, attributes: _*))
+    def withFinalizationStrategy(
+        strategy: SpanFinalizer.Strategy
+    ): SpanBuilder[G] =
+      new MappedK(builder.withFinalizationStrategy(strategy))
+    def withSpanKind(spanKind: SpanKind): SpanBuilder[G] =
+      new MappedK(builder.withSpanKind(spanKind))
+    def withStartTimestamp(timestamp: FiniteDuration): SpanBuilder[G] =
+      new MappedK(builder.withStartTimestamp(timestamp))
+    def root: SpanBuilder[G] = new MappedK(builder.root)
+    def withParent(parent: SpanContext): SpanBuilder[G] =
+      new MappedK(builder.withParent(parent))
+    def build: SpanOps[G] = builder.build.mapK[G]
+  }
 }

core/trace/src/main/scala/org/typelevel/otel4s/trace/SpanOps.scala

@@ -14,9 +14,12 @@
  * limitations under the License.
  */
 
-package org.typelevel.otel4s.trace
+package org.typelevel.otel4s
+package trace
 
+import cats.effect.MonadCancelThrow
 import cats.effect.Resource
+import cats.syntax.functor._
 import cats.~>
 
 trait SpanOps[F[_]] {
@@ -138,6 +141,15 @@ trait SpanOps[F[_]] {
     *   See [[use]] for more details regarding lifecycle strategy
     */
   final def surround[A](fa: F[A]): F[A] = use(_ => fa)
+
+  /** Modify the context `F` using an implicit [[KindTransformer]] from `F` to
+    * `G`.
+    */
+  def mapK[G[_]: MonadCancelThrow](implicit
+      F: MonadCancelThrow[F],
+      kt: KindTransformer[F, G]
+  ): SpanOps[G] =
+    new SpanOps.MappedK(this)
 }
 
 object SpanOps {
@@ -156,6 +168,12 @@ object SpanOps {
       * [[cats.arrow.FunctionK FunctionK]] will not be traced.
       */
     def trace: F ~> F
+
+    /** Modify the context `F` using an implicit [[KindTransformer]] from `F` to
+      * `G`.
+      */
+    def mapK[G[_]](implicit kt: KindTransformer[F, G]): Res[G] =
+      Res(span.mapK[G], kt.liftFunctionK(trace))
   }
 
   object Res {
@@ -168,4 +186,21 @@ object SpanOps {
     def apply[F[_]](span: Span[F], trace: F ~> F): Res[F] =
       Impl(span, trace)
   }
+
+  /** Implementation for [[SpanOps.mapK]]. */
+  private class MappedK[F[_]: MonadCancelThrow, G[_]: MonadCancelThrow](
+      ops: SpanOps[F]
+  )(implicit kt: KindTransformer[F, G])
+      extends SpanOps[G] {
+    def startUnmanaged: G[Span[G]] =
+      kt.liftK(ops.startUnmanaged).map(_.mapK[G])
+
+    def resource: Resource[G, Res[G]] =
+      ops.resource.mapK(kt.liftK).map(res => res.mapK[G])
+
+    def use[A](f: Span[G] => G[A]): G[A] =
+      resource.use { res => res.trace(f(res.span)) }
+
+    def use_ : G[Unit] = kt.liftK(ops.use_)
+  }
 }