Skip to content

Commit

Permalink
Add mapK to Tracer, SpanBuilder, etc.
Browse files Browse the repository at this point in the history
Add `mapK` to `Tracer`, `Tracer.Meta`, `SpanBuilder`, `SpanOps`,
`Span`, `Span.Backend`, and `InstrumentMeta`.

Add `KindTransformer` type used by `mapK` implementation, allowing
for a single `mapK` implementation that supports many
transformation target types (`OptionT`, `EitherT`, `IorT`, `StateT`,
`Kliesli`, `Resource`) with very little work needed to add
additional or custom ones.
  • Loading branch information
NthPortal committed Sep 18, 2023
1 parent 2c59bc1 commit b68cd63
Show file tree
Hide file tree
Showing 7 changed files with 569 additions and 4 deletions.
123 changes: 123 additions & 0 deletions core/common/src/main/scala/org/typelevel/otel4s/KindTransformer.scala
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
/*
* 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.Applicative
import cats.Functor
import cats.Monad
import cats.data.EitherT
import cats.data.IorT
import cats.data.Kleisli
import cats.data.Nested
import cats.data.OptionT
import cats.data.StateT
import cats.effect.MonadCancelThrow
import cats.effect.Resource
import cats.implicits.toFunctorOps
import cats.syntax.all._
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)
}

implicit def nested[F[_]: Functor, G[_]: Applicative]
: KindTransformer[F, Nested[F, G, *]] =
new KindTransformer[F, Nested[F, G, *]] {
val liftK: F ~> Nested[F, G, *] =
new (F ~> Nested[F, G, *]) {
def apply[A](fa: F[A]): Nested[F, G, A] =
fa.map(_.pure[G]).nested
}
def limitedMapK[A](ga: Nested[F, G, A])(f: F ~> F): Nested[F, G, A] =
ga.mapK(f)
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package org.typelevel.otel4s.meta

import cats.Applicative
import cats.~>

trait InstrumentMeta[F[_]] {

Expand All @@ -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 {
Expand All @@ -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)
}
}
49 changes: 49 additions & 0 deletions core/trace/src/main/scala/org/typelevel/otel4s/trace/Span.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -94,6 +95,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 {
Expand Down Expand Up @@ -121,6 +131,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 {
Expand Down Expand Up @@ -151,6 +170,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] =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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]
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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[_]] {
Expand Down Expand Up @@ -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 {
Expand All @@ -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 {
Expand All @@ -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_)
}
}
Loading

0 comments on commit b68cd63

Please sign in to comment.