Skip to content

Commit

Permalink
preserve scopes (#568)
Browse files Browse the repository at this point in the history
  • Loading branch information
@adamgfraser
adamgfraser authored Aug 27, 2022
1 parent c02851c commit d66c28a
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 45 deletions.
101 changes: 56 additions & 45 deletions opentelemetry/src/main/scala/zio/telemetry/opentelemetry/Tracing.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,13 @@ trait Tracing {

private[opentelemetry] val currentContext: FiberRef[Context]

private[opentelemetry] def createRoot(spanName: String, spanKind: SpanKind): URIO[Scope, Context]
private[opentelemetry] def createRoot(spanName: String, spanKind: SpanKind): UIO[(UIO[Unit], Context)]

private[opentelemetry] def createChildOf(
parent: Context,
spanName: String,
spanKind: SpanKind
): URIO[Scope, Context]
): UIO[(UIO[Unit], Context)]

private[opentelemetry] def createChildOfUnsafe(parent: Context, spanName: String, spanKind: SpanKind): UIO[Context]

Expand All @@ -39,10 +39,14 @@ object Tracing {
private def currentContext: URIO[Tracing, FiberRef[Context]] =
ZIO.service[Tracing].map(_.currentContext)

private def createRoot(spanName: String, spanKind: SpanKind): URIO[Scope with Tracing, Context] =
private def createRoot(spanName: String, spanKind: SpanKind): URIO[Tracing, (UIO[Unit], Context)] =
ZIO.serviceWithZIO[Tracing](_.createRoot(spanName, spanKind))

private def createChildOf(parent: Context, spanName: String, spanKind: SpanKind): URIO[Scope with Tracing, Context] =
private def createChildOf(
parent: Context,
spanName: String,
spanKind: SpanKind
): URIO[Tracing, (UIO[Unit], Context)] =
ZIO.serviceWithZIO[Tracing](_.createChildOf(parent, spanName, spanKind))

private def createChildOfUnsafe(parent: Context, spanName: String, spanKind: SpanKind): URIO[Tracing, Context] =
Expand Down Expand Up @@ -91,12 +95,14 @@ object Tracing {
spanKind: SpanKind,
toErrorStatus: PartialFunction[E, StatusCode]
)(effect: ZIO[R, E, A]): ZIO[R with Tracing, E, A] =
ZIO.scoped[R with Tracing] {
for {
context <- extractContext(propagator, carrier, getter)
ctx <- createChildOf(context, spanName, spanKind)
r <- finalizeSpanUsingEffect(effect, ctx, toErrorStatus)
} yield r
extractContext(propagator, carrier, getter).flatMap { context =>
ZIO.acquireReleaseWith {
createChildOf(context, spanName, spanKind)
} { case (r, _) =>
r
} { case (_, ctx) =>
finalizeSpanUsingEffect(effect, ctx, toErrorStatus)
}
}

/**
Expand Down Expand Up @@ -127,8 +133,12 @@ object Tracing {
spanKind: SpanKind,
toErrorStatus: PartialFunction[E, StatusCode]
)(effect: ZIO[R, E, A]): ZIO[R with Tracing, E, A] =
ZIO.scoped[R with Tracing] {
createRoot(spanName, spanKind).flatMap(finalizeSpanUsingEffect(effect, _, toErrorStatus))
ZIO.acquireReleaseWith {
createRoot(spanName, spanKind)
} { case (r, _) =>
r
} { case (_, ctx) =>
finalizeSpanUsingEffect(effect, ctx, toErrorStatus)
}

/**
Expand All @@ -140,12 +150,14 @@ object Tracing {
spanKind: SpanKind,
toErrorStatus: PartialFunction[E, StatusCode]
)(effect: ZIO[R, E, A]): ZIO[R with Tracing, E, A] =
ZIO.scoped[R with Tracing] {
for {
old <- getCurrentContext
ctx <- createChildOf(old, spanName, spanKind)
r <- finalizeSpanUsingEffect(effect, ctx, toErrorStatus)
} yield r
getCurrentContext.flatMap { old =>
ZIO.acquireReleaseWith {
createChildOf(old, spanName, spanKind)
} { case (r, _) =>
r
} { case (_, ctx) =>
finalizeSpanUsingEffect(effect, ctx, toErrorStatus)
}
}

/**
Expand Down Expand Up @@ -237,9 +249,12 @@ object Tracing {
spanKind: SpanKind,
toErrorStatus: PartialFunction[E, StatusCode]
)(effect: ZIO[R, E, A]): ZIO[R with Tracing, E, A] =
ZIO.scoped[R with Tracing] {
ZIO.acquireReleaseWith {
createChildOf(Context.root().`with`(span), spanName, spanKind)
.flatMap(finalizeSpanUsingEffect(effect, _, toErrorStatus))
} { case (r, _) =>
r
} { case (_, ctx) =>
finalizeSpanUsingEffect(effect, ctx, toErrorStatus)
}

/**
Expand Down Expand Up @@ -346,35 +361,31 @@ object Tracing {

def currentNanos: UIO[Long] = Clock.currentTime(TimeUnit.NANOSECONDS)

def createRoot(spanName: String, spanKind: SpanKind): URIO[Scope, Context] =
def createRoot(spanName: String, spanKind: SpanKind): UIO[(UIO[Unit], Context)] =
for {
nanoSeconds <- currentNanos
span <- ZIO.acquireRelease(
ZIO.succeed(
tracer
.spanBuilder(spanName)
.setNoParent()
.setSpanKind(spanKind)
.setStartTimestamp(nanoSeconds, TimeUnit.NANOSECONDS)
.startSpan()
)
)(endSpan)
} yield span.storeInContext(Context.root())

def createChildOf(parent: Context, spanName: String, spanKind: SpanKind): URIO[Scope, Context] =
span <- ZIO.succeed(
tracer
.spanBuilder(spanName)
.setNoParent()
.setSpanKind(spanKind)
.setStartTimestamp(nanoSeconds, TimeUnit.NANOSECONDS)
.startSpan()
)
} yield (endSpan(span), span.storeInContext(Context.root()))

def createChildOf(parent: Context, spanName: String, spanKind: SpanKind): UIO[(UIO[Unit], Context)] =
for {
nanoSeconds <- currentNanos
span <- ZIO.acquireRelease(
ZIO.succeed(
tracer
.spanBuilder(spanName)
.setParent(parent)
.setSpanKind(spanKind)
.setStartTimestamp(nanoSeconds, TimeUnit.NANOSECONDS)
.startSpan()
)
)(endSpan)
} yield span.storeInContext(parent)
span <- ZIO.succeed(
tracer
.spanBuilder(spanName)
.setParent(parent)
.setSpanKind(spanKind)
.setStartTimestamp(nanoSeconds, TimeUnit.NANOSECONDS)
.startSpan()
)
} yield (endSpan(span), span.storeInContext(parent))

def createChildOfUnsafe(parent: Context, spanName: String, spanKind: SpanKind): UIO[Context] =
for {
Expand Down
9 changes: 9 additions & 0 deletions opentelemetry/src/test/scala/zio/telemetry/opentelemetry/TracingTest.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -304,6 +304,15 @@ object TracingTest extends ZIOSpecDefault {
baggage <- Tracing.getCurrentBaggage
entryValue = Option(baggage.getEntryValue("some"))
} yield assert(entryValue)(equalTo(Some("thing")))
},
test("resources") {
for {
ref <- Ref.make(false)
scope <- Scope.make
resource = ZIO.addFinalizer(ref.set(true))
_ <- scope.extend(resource.span("Resource"))
released <- ref.get
} yield assert(released)(isFalse)
}
).provideCustomLayer(tracingMockLayer)
)
Expand Down

0 comments on commit d66c28a

Please sign in to comment.