Fix mapConcat context propagation (#30289)

* Fix mapConcat context propagation

* Remove default ContextPropagation implementation to reduce possible overhead when no Telemetry is in use

* fix pullFirstSubElement onPull

* CR follow up
This commit is contained in:
Yury Gribkov 2021-06-08 06:09:47 -04:00 committed by GitHub
parent d69d04957b
commit 3da5c2bdd0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 11 deletions

View file

@ -26,11 +26,6 @@ import akka.annotation.InternalApi
} }
private[akka] final class ContextPropagationImpl extends ContextPropagation { private[akka] final class ContextPropagationImpl extends ContextPropagation {
private val buffer = Buffer[Unit](1, 1) def suspendContext(): Unit = ()
def suspendContext(): Unit = { def resumeContext(): Unit = ()
buffer.enqueue(())
}
def resumeContext(): Unit = {
buffer.dequeue()
}
} }

View file

@ -2212,15 +2212,20 @@ private[akka] final class StatefulMapConcat[In, Out](val f: () => In => Iterable
lazy val decider = inheritedAttributes.mandatoryAttribute[SupervisionStrategy].decider lazy val decider = inheritedAttributes.mandatoryAttribute[SupervisionStrategy].decider
var currentIterator: Iterator[Out] = _ var currentIterator: Iterator[Out] = _
var plainFun = f() var plainFun = f()
val contextPropagation = ContextPropagation()
def hasNext = if (currentIterator != null) currentIterator.hasNext else false def hasNext = if (currentIterator != null) currentIterator.hasNext else false
setHandlers(in, out, this) setHandlers(in, out, this)
def pushPull(): Unit = def pushPull(shouldResumeContext: Boolean): Unit =
if (hasNext) { if (hasNext) {
if (shouldResumeContext) contextPropagation.resumeContext()
push(out, currentIterator.next()) push(out, currentIterator.next())
if (!hasNext && isClosed(in)) completeStage() if (hasNext) {
// suspend context for the next element
contextPropagation.suspendContext()
} else if (isClosed(in)) completeStage()
} else if (!isClosed(in)) } else if (!isClosed(in))
pull(in) pull(in)
else completeStage() else completeStage()
@ -2230,13 +2235,13 @@ private[akka] final class StatefulMapConcat[In, Out](val f: () => In => Iterable
override def onPush(): Unit = override def onPush(): Unit =
try { try {
currentIterator = plainFun(grab(in)).iterator currentIterator = plainFun(grab(in)).iterator
pushPull() pushPull(shouldResumeContext = false)
} catch handleException } catch handleException
override def onUpstreamFinish(): Unit = onFinish() override def onUpstreamFinish(): Unit = onFinish()
override def onPull(): Unit = override def onPull(): Unit =
try pushPull() try pushPull(shouldResumeContext = true)
catch handleException catch handleException
private def handleException: Catcher[Unit] = { private def handleException: Catcher[Unit] = {