chore: Avoid lift SingleSource to InflightSource.

Author: He-Pin

akka-stream/src/main/scala/akka/stream/impl/fusing/StreamOfStreams.scala

@@ -159,19 +159,6 @@ private[akka] object FlattenConcat {
     def materialize(): Unit = ()
   }
 
-  private final class InflightSingleSource[T](elem: T) extends InflightSource[T] {
-    private var _hasNext = true
-    override def hasNext: Boolean = _hasNext
-    override def next(): T =
-      if (_hasNext) {
-        _hasNext = false
-        elem
-      } else throw new NoSuchElementException("next called after completion")
-    override def tryPull(): Unit = ()
-    override def cancel(cause: Throwable): Unit = ()
-    override def isClosed: Boolean = !hasNext
-  }
-
   private final class InflightIteratorSource[T](iterator: Iterator[T]) extends InflightSource[T] {
     override def hasNext: Boolean = iterator.hasNext
     override def next(): T = iterator.next()
@@ -232,16 +219,16 @@ private[akka] final class FlattenConcat[T, M](parallelism: Int)
 
   override def initialAttributes: Attributes = DefaultAttributes.flattenConcat
   override val shape: FlowShape[Graph[SourceShape[T], M], T] = FlowShape(in, out)
-  override def createLogic(enclosingAttributes: Attributes) =
-    new GraphStageLogic(shape) with InHandler with OutHandler {
+  override def createLogic(enclosingAttributes: Attributes) = {
+    final object FlattenConcatLogic extends GraphStageLogic(shape) with InHandler with OutHandler {
       import FlattenConcat._
-      private var queue: BufferImpl[InflightSource[T]] = _
+      // InflightSource[T] or SingleSource[T]
+      // AnyRef here to avoid lift the SingleSource[T] to InflightSource[T]
+      private var queue: BufferImpl[AnyRef] = _
       private val invokeCb: InflightSource[T] => Unit =
         getAsyncCallback[InflightSource[T]](futureSourceCompleted).invoke
 
-      override def preStart(): Unit = {
-        queue = BufferImpl(parallelism, enclosingAttributes)
-      }
+      override def preStart(): Unit = queue = BufferImpl(parallelism, enclosingAttributes)
 
       private def futureSourceCompleted(futureSource: InflightSource[T]): Unit = {
         if (queue.peek() eq futureSource) {
@@ -257,8 +244,7 @@ private[akka] final class FlattenConcat[T, M](parallelism: Int)
       }
 
       override def onPush(): Unit = {
-        val source = grab(in)
-        addSource(source)
+        addSource(grab(in))
         //must try pull after addSource to avoid queue overflow
         if (!queue.isFull) { // try to keep the maximum parallelism
           tryPull(in)
@@ -273,23 +259,32 @@ private[akka] final class FlattenConcat[T, M](parallelism: Int)
       }
 
       override def onPull(): Unit = {
-        if (queue.nonEmpty) {
-          val currentSource = queue.peek()
-          //purge if possible
-          if (currentSource.hasNext) {
-            push(out, currentSource.next())
-            if (currentSource.isClosed) {
-              handleCurrentSourceClosed(currentSource)
+        //purge if possible
+        queue.peek() match {
+          case src: SingleSource[T] @unchecked =>
+            push(out, src.elem)
+            removeSource()
+          case src: InflightSource[T] @unchecked => pushOut(src)
+          case null => //queue is empty
+            if (!hasBeenPulled(in)) {
+              tryPull(in)
+            } else if (isClosed(in)) {
+              completeStage()
             }
-          } else if (currentSource.isClosed) {
-            handleCurrentSourceClosed(currentSource)
-          } else {
-            currentSource.tryPull()
+          case _ => throw new IllegalStateException("Should not reach here.")
+        }
+      }
+
+      private def pushOut(src: InflightSource[T]): Unit = {
+        if (src.hasNext) {
+          push(out, src.next())
+          if (src.isClosed) {
+            handleCurrentSourceClosed(src)
           }
-        } else if (!hasBeenPulled(in)) {
-          tryPull(in)
-        } else if (isClosed(in)) {
-          completeStage()
+        } else if (src.isClosed) {
+          handleCurrentSourceClosed(src)
+        } else {
+          src.tryPull()
         }
       }
 
@@ -308,18 +303,18 @@ private[akka] final class FlattenConcat[T, M](parallelism: Int)
       private def cancelInflightSources(cause: Throwable): Unit = {
         if (queue.nonEmpty) {
           var source = queue.dequeue()
-          while (source ne null) {
-            source.cancel(cause)
+          while ((source ne null) && (source.isInstanceOf[InflightSource[T] @unchecked])) {
+            source.asInstanceOf[InflightSource[T]].cancel(cause)
             source = queue.dequeue()
           }
         }
       }
 
-      private def addSourceElem(elem: T): Unit = {
+      private def addSource(singleSource: SingleSource[T]): Unit = {
         if (isAvailable(out) && queue.isEmpty) {
-          push(out, elem)
+          push(out, singleSource.elem)
         } else {
-          queue.enqueue(new InflightSingleSource(elem))
+          queue.enqueue(singleSource)
         }
       }
 
@@ -397,7 +392,7 @@ private[akka] final class FlattenConcat[T, M](parallelism: Int)
         (TraversalBuilder.getValuePresentedSource(source): @nowarn("cat=lint-infer-any")) match {
           case OptionVal.Some(graph) =>
             graph match {
-              case single: SingleSource[T] @unchecked => addSourceElem(single.elem)
+              case single: SingleSource[T] @unchecked => addSource(single)
               case futureSource: FutureSource[T] @unchecked =>
                 val future = futureSource.future
                 future.value match {
@@ -415,37 +410,50 @@ private[akka] final class FlattenConcat[T, M](parallelism: Int)
 
       }
 
+      private def removeSource(): Unit = {
+        queue.dequeue()
+        pullIfNeeded()
+      }
+
       private def removeSource(source: InflightSource[T]): Unit = {
         if (queue.nonEmpty && (source eq queue.peek())) {
           //only dequeue if it's the current emitting source
-          val s = queue.dequeue()
-          if (s != source) {
-            throw new IllegalStateException("Should not reach here.")
-          }
-          if (isClosed(in)) {
-            if (queue.isEmpty) {
-              completeStage()
-            } else {
-              //pull the new emitting source
-              val nextSource = queue.peek()
-              nextSource.tryPull()
-            }
+          queue.dequeue()
+          pullIfNeeded()
+        } //not the head source, just ignore
+      }
+
+      private def pullIfNeeded(): Unit = {
+        if (isClosed(in)) {
+          if (queue.isEmpty) {
+            completeStage()
           } else {
-            if (queue.nonEmpty) {
-              //pull the new emitting source
-              val nextSource = queue.peek()
-              nextSource.tryPull()
-            }
-            if (!hasBeenPulled(in)) {
-              tryPull(in)
-            }
+            tryPullNextSourceInQueue()
           }
-        } //not the head source, just ignore
+        } else {
+          if (queue.nonEmpty) {
+            tryPullNextSourceInQueue()
+          }
+          if (!hasBeenPulled(in)) {
+            tryPull(in)
+          }
+        }
+      }
+
+      private def tryPullNextSourceInQueue(): Unit = {
+        //pull the new emitting source
+        val nextSource = queue.peek()
+        if (nextSource.isInstanceOf[InflightSource[T] @unchecked]) {
+          nextSource.asInstanceOf[InflightSource[T]].tryPull()
+        }
       }
 
       setHandlers(in, out, this)
     }
 
+    FlattenConcatLogic
+  }
+
   override def toString: String = s"FlattenConcat($parallelism)"
 }