Keep sender-side batching for TLS ciphertext input

stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageIsolatedSpec.scala

@@ -25,21 +25,25 @@ import javax.net.ssl.SSLEngineResult.Status.{ BUFFER_OVERFLOW, BUFFER_UNDERFLOW,
 
 import scala.concurrent.Await
 import scala.concurrent.ExecutionContext
+import scala.concurrent.Future
 import scala.concurrent.Promise
 import scala.concurrent.duration._
 import scala.util.Success
 
 import com.typesafe.config.{ Config, ConfigFactory }
 
 import org.apache.pekko
+import pekko.Done
 import pekko.NotUsed
 import pekko.actor.Cancellable
 import pekko.event.LoggingAdapter
 import pekko.stream._
 import pekko.stream.TLSProtocol._
+import pekko.stream.impl.fusing.ActorGraphInterpreter.ActorOutputBoundary
 import pekko.stream.impl.io.TlsGraphStage
 import pekko.stream.scaladsl._
 import pekko.stream.scaladsl.GraphDSL.Implicits._
+import pekko.stream.stage.{ GraphStageLogic, GraphStageWithMaterializedValue, OutHandler }
 import pekko.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber }
 import pekko.stream.testkit.scaladsl.{ TestSink, TestSource }
 import pekko.stream.testkit.scaladsl.StreamTestKit.assertAllStagesStopped
@@ -246,6 +250,57 @@ class TlsGraphStageIsolatedSpec extends StreamSpec(TlsSpec.configOverrides) with
 
       outputs.foldLeft(ByteString.empty)(_ ++ _) shouldEqual expected
     }
+
+    "batch transport input across its async boundary but not plaintext input" in {
+      assertAllStagesStopped {
+        val stop = Promise[Done]()
+        val (plainBatched, cipherBatched) =
+          RunnableGraph
+            .fromGraph(GraphDSL.createGraph(
+              new BoundaryBatchingProbe[SslTlsOutbound](stop.future),
+              new BoundaryBatchingProbe[ByteString](stop.future))(Tuple2.apply) { implicit b => (plain, cipher) =>
+              val tls = b.add(stageFlow(initSslContext("TLSv1.2"), TLS12Ciphers, clientMode = true, IgnoreComplete))
+              plain    ~> tls.in1
+              tls.out1 ~> Sink.ignore
+              cipher   ~> tls.in2
+              tls.out2 ~> Sink.ignore
+              ClosedShape
+            })
+            .run()
+
+        // Plaintext must not overtake a later plaintext-output cancellation; ciphertext has no such constraint.
+        Await.result(plainBatched, 3.seconds) shouldBe false
+        Await.result(cipherBatched, 3.seconds) shouldBe true
+        stop.success(Done)
+      }
+    }
+  }
+}
+
+/** Reports whether the output boundary after this source sends sender-side batches. */
+private final class BoundaryBatchingProbe[T](stop: Future[Done])
+    extends GraphStageWithMaterializedValue[SourceShape[T], Future[Boolean]] {
+  private val out = Outlet[T]("BoundaryBatchingProbe.out")
+  override val shape: SourceShape[T] = SourceShape(out)
+
+  override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Future[Boolean]) = {
+    val batched = Promise[Boolean]()
+    val logic = new GraphStageLogic(shape) with OutHandler {
+      override def preStart(): Unit = {
+        val complete = getAsyncCallback[Done](_ => completeStage())
+        stop.foreach(complete.invoke)(ExecutionContext.parasitic)
+      }
+
+      // The first pull follows the downstream subscription, after which the boundary has chosen its mode.
+      override def onPull(): Unit =
+        portToConn(out.id + inCount).inOwner match {
+          case boundary: ActorOutputBoundary => batched.trySuccess(boundary.batchesToSubscriber)
+          case other                         => batched.tryFailure(new AssertionError(s"Unexpected downstream $other"))
+        }
+
+      setHandler(out, this)
+    }
+    (logic, batched.future)
   }
 }
 

stream/src/main/scala/org/apache/pekko/stream/impl/PhasedFusingActorMaterializer.scala

@@ -820,12 +820,30 @@ private[pekko] object GraphStageIsland {
     logics.add(boundary)
     boundary.stageId = logics.size() - 1
     boundary.attributes = connection.inOwner.attributes.and(DefaultAttributes.inputBoundary)
+    boundary.unbatched = isUnbatchedInput(connection)
 
     boundary.portToConn(boundary.out.id + boundary.inCount) = connection
     connection.outHandler = boundary.handlers(0).asInstanceOf[OutHandler]
     connection.outOwner = boundary
   }
 
+  private def isUnbatchedInput(connection: Connection): Boolean = {
+    val owner = connection.inOwner
+    owner.attributes.get[ActorGraphInterpreter.UnbatchedInput] match {
+      case Some(unbatched) =>
+        owner.originalStage match {
+          case OptionVal.Some(stage) =>
+            // Port ids are positions in the shape, see TraversalBuilder.initShape.
+            val inlets = stage.shape.inlets
+            var i = 0
+            while (i < owner.inCount && (owner.portToConn(i) ne connection)) i += 1
+            i < owner.inCount && unbatched.inletNames.contains(inlets(i).s)
+          case _ => false
+        }
+      case None => false
+    }
+  }
+
   override def onIslandReady(): Unit = {
 
     val outConnSize = if (outConnections ne null) outConnections.size() else 0

stream/src/main/scala/org/apache/pekko/stream/impl/fusing/ActorGraphInterpreter.scala

@@ -57,8 +57,8 @@ import org.reactivestreams.Subscription
     def batch(elements: Array[AnyRef]): Unit
   }
 
-  // Opt out of sender-side batching when ordering with signals on another port matters (TLS).
-  private[stream] case object UnbatchedInput extends Attributes.Attribute
+  // Opt the named inlets out of sender-side batching when ordering with signals on another port matters (TLS).
+  private[stream] final case class UnbatchedInput(inletNames: Set[String]) extends Attributes.Attribute
 
   sealed abstract class BoundaryEvent extends DeadLetterSuppression with NoSerializationVerificationNeeded {
     def shell: GraphInterpreterShell
@@ -174,6 +174,9 @@ import org.reactivestreams.Subscription
     private def requestBatchSize = math.max(1, inputBuffer.length / 2)
     private var batchRemaining = requestBatchSize
 
+    // Set by the materializer before preStart when the fed inlet is listed in UnbatchedInput.
+    private[stream] var unbatched = false
+
     val out: Outlet[Any] = Outlet[Any]("UpstreamBoundary:" + internalPortName)
     out.id = 0
 
@@ -438,6 +441,7 @@ import org.reactivestreams.Subscription
 
     private var subscriber: Subscriber[Any] = _
     private var boundarySubscriber: BoundarySubscriber = _
+    private[stream] def batchesToSubscriber: Boolean = boundarySubscriber ne null
     private var downstreamDemand: Long = 0L
     // This flag is only used if complete/fail is called externally since this op turns into a Finished one inside the
     // interpreter (i.e. inside this op this flag has no effects since if it is completed the op will not be invoked)
@@ -563,9 +567,9 @@ import org.reactivestreams.Subscription
         if (subscriber eq null) {
           subscriber = sub
           boundarySubscriber = sub match {
-            case boundary: BatchingActorInputBoundary if boundary.attributes.contains(UnbatchedInput) => null
-            case boundary: BoundarySubscriber                                                         => boundary
-            case _                                                                                    => null
+            case boundary: BatchingActorInputBoundary if boundary.unbatched => null
+            case boundary: BoundarySubscriber                               => boundary
+            case _                                                          => null
           }
           val subscription = new Subscription with SubscriptionWithCancelException {
             override def request(elements: Long): Unit = actor ! new RequestMore(ActorOutputBoundary.this, elements)

stream/src/main/scala/org/apache/pekko/stream/impl/io/TlsGraphStage.scala

@@ -52,7 +52,7 @@ import pekko.util.ByteString
     closing: TLSClosing)
     extends GraphStage[BidiShape[SslTlsOutbound, ByteString, ByteString, SslTlsInbound]] {
 
-  private val plainIn = Inlet[SslTlsOutbound]("TlsGraphStage.plainIn")
+  private val plainIn = Inlet[SslTlsOutbound](TlsGraphStage.PlainInName)
   private val plainOut = Outlet[SslTlsInbound]("TlsGraphStage.plainOut")
   private val cipherIn = Inlet[ByteString]("TlsGraphStage.cipherIn")
   private val cipherOut = Outlet[ByteString]("TlsGraphStage.cipherOut")
@@ -1201,6 +1201,7 @@ import pekko.util.ByteString
 }
 
 @InternalApi private[pekko] object TlsGraphStage {
+  private[pekko] final val PlainInName = "TlsGraphStage.plainIn"
   private[pekko] final val TransportFlushTimer = "TlsGraphStageTransportFlush"
   private[pekko] final val UserInputBatchTimer = "TlsGraphStageUserInputBatch"
   private[pekko] final val InitialOutputFlushTimer = "TlsGraphStageInitialOutputFlush"
@@ -1215,7 +1216,8 @@ import pekko.util.ByteString
   private[pekko] val StreamTlsAttributes: Attributes =
     Attributes.asyncBoundary and Attributes.inputBuffer(InputBufferInitialSize, InputBufferMaxSize) and
     // Send plaintext messages before a subsequent plaintext-output cancellation from the same handler.
-    Attributes(ActorGraphInterpreter.UnbatchedInput)
+    // Ciphertext input keeps sender-side batching; only plaintext ordering relative to cancellation matters.
+    Attributes(ActorGraphInterpreter.UnbatchedInput(Set(PlainInName)))
 
   private[pekko] final val PlainDataAllowedFlag = 1 << 0
   private[pekko] final val StageClosedFlag = 1 << 1