raboof/pekko
1
0
Fork 0
Code Issues Pull requests Projects Releases 6 Packages Wiki Activity Actions 11

=str Reduce allocation of throttle. (#31276)

Browse source
This commit is contained in:
kerr 2022-03-25 16:20:31 +08:00 • committed by GitHub
parent 3e1a43f913
commit 83c211b84f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 33 additions and 46 deletions
Download patch file Download diff file Expand all files Collapse all files

41
akka-stream/src/main/scala/akka/stream/impl/Throttle.scala
View file

@ -4,20 +4,21 @@
package akka.stream.impl package akka.stream.impl
import scala.concurrent.duration.{ FiniteDuration, _ }
import akka.annotation.InternalApi import akka.annotation.InternalApi
import akka.stream.ThrottleMode.Enforcing
import akka.stream._ import akka.stream._
import akka.stream.ThrottleMode.{ Enforcing, Shaping }
import akka.stream.impl.fusing.GraphStages.SimpleLinearGraphStage import akka.stream.impl.fusing.GraphStages.SimpleLinearGraphStage
import akka.stream.stage._ import akka.stream.stage._
import akka.util.NanoTimeTokenBucket import akka.util.NanoTimeTokenBucket
import scala.concurrent.duration.{ FiniteDuration, _ }
/** /**
* INTERNAL API * INTERNAL API
*/ */
@InternalApi private[akka] object Throttle { @InternalApi private[akka] object Throttle {
final val AutomaticMaximumBurst = -1 final val AutomaticMaximumBurst = -1
private case object TimerKey
} }
/** /**
@ -45,27 +46,17 @@ import akka.util.NanoTimeTokenBucket
else maximumBurst else maximumBurst
require(!(mode == ThrottleMode.Enforcing && effectiveMaximumBurst < 0), "maximumBurst must be > 0 in Enforcing mode") require(!(mode == ThrottleMode.Enforcing && effectiveMaximumBurst < 0), "maximumBurst must be > 0 in Enforcing mode")
private val timerName: String = "ThrottleTimer" override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
new TimerGraphStageLogic(shape) with InHandler with OutHandler {
override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new TimerGraphStageLogic(shape) {
private val tokenBucket = new NanoTimeTokenBucket(effectiveMaximumBurst, nanosBetweenTokens) private val tokenBucket = new NanoTimeTokenBucket(effectiveMaximumBurst, nanosBetweenTokens)
private var currentElement: T = _
var willStop = false
var currentElement: T = _
val enforcing = mode match {
case Enforcing => true
case Shaping => false
}
override def preStart(): Unit = tokenBucket.init() override def preStart(): Unit = tokenBucket.init()
// This scope is here just to not retain an extra reference to the handler below.
// We can't put this code into preRestart() because setHandler() must be called before that.
{
val handler = new InHandler with OutHandler {
override def onUpstreamFinish(): Unit = override def onUpstreamFinish(): Unit =
if (isAvailable(out) && isTimerActive(timerName)) willStop = true if (!(isAvailable(out) && isTimerActive(Throttle.TimerKey))) {
else completeStage() completeStage()
}
override def onPush(): Unit = { override def onPush(): Unit = {
val elem = grab(in) val elem = grab(in)
@ -74,27 +65,23 @@ import akka.util.NanoTimeTokenBucket
if (delayNanos == 0L) push(out, elem) if (delayNanos == 0L) push(out, elem)
else { else {
if (enforcing) failStage(new RateExceededException("Maximum throttle throughput exceeded.")) if (mode eq Enforcing) failStage(new RateExceededException("Maximum throttle throughput exceeded."))
else { else {
currentElement = elem currentElement = elem
scheduleOnce(timerName, delayNanos.nanos) scheduleOnce(Throttle.TimerKey, delayNanos.nanos)
} }
} }
} }
override def onPull(): Unit = pull(in) override def onPull(): Unit = pull(in)
}
setHandlers(in, out, handler)
// After this point, we no longer need the `handler` so it can just fall out of scope.
}
override protected def onTimer(key: Any): Unit = { override protected def onTimer(key: Any): Unit = {
push(out, currentElement) push(out, currentElement)
currentElement = null.asInstanceOf[T] currentElement = null.asInstanceOf[T]
if (willStop) completeStage() if (isClosed(in)) completeStage()
} }
setHandlers(in, out, this)
} }
override def toString = "Throttle" override def toString = "Throttle"