Replaced PushStage based Take with GraphStage #19834
Changes suggested by @viktorklang Added take to binary compatibility exclude Cleanup
This commit is contained in:
parent
65a9e0a0f9
commit
1bcc1310d2
6 changed files with 50 additions and 30 deletions
|
|
@ -124,19 +124,30 @@ private[akka] final case class Recover[T](pf: PartialFunction[Throwable, T]) ext
|
|||
/**
|
||||
* INTERNAL API
|
||||
*/
|
||||
private[akka] final case class Take[T](count: Long) extends PushPullStage[T, T] {
|
||||
private var left: Long = count
|
||||
private[akka] final case class Take[T](count: Long) extends SimpleLinearGraphStage[T] {
|
||||
override def initialAttributes: Attributes = DefaultAttributes.take
|
||||
|
||||
override def onPush(elem: T, ctx: Context[T]): SyncDirective = {
|
||||
left -= 1
|
||||
if (left > 0) ctx.push(elem)
|
||||
else if (left == 0) ctx.pushAndFinish(elem)
|
||||
else ctx.finish() //Handle negative take counts
|
||||
override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) with InHandler with OutHandler {
|
||||
private var left: Long = count
|
||||
|
||||
override def onPush(): Unit = {
|
||||
val leftBefore = left
|
||||
if (leftBefore >= 1) {
|
||||
left = leftBefore - 1
|
||||
push(out, grab(in))
|
||||
}
|
||||
if (leftBefore <= 1) completeStage()
|
||||
}
|
||||
|
||||
override def onPull(): Unit = {
|
||||
if (left > 0) pull(in)
|
||||
else completeStage()
|
||||
}
|
||||
|
||||
setHandlers(in, out, this)
|
||||
}
|
||||
|
||||
override def onPull(ctx: Context[T]): SyncDirective =
|
||||
if (left <= 0) ctx.finish()
|
||||
else ctx.pull()
|
||||
override def toString: String = "Take"
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue