From 10f4c8cc2ecd058258e0672cfd09bd48fdf5715c Mon Sep 17 00:00:00 2001 From: PJ Fanning Date: Fri, 30 Oct 2026 09:55:28 +0000 Subject: [PATCH] Keep sender-side batching for TLS ciphertext input Motivation: #3727 opted the GraphStage TLS engine's input boundaries out of sender-side batching so plaintext is not overtaken by a later plaintext-output cancellation. The opt-out applied to every input boundary of the TLS island, so ciphertext input from the transport also crossed the async boundary as one actor message per element, although it has no such ordering constraint. Modification: - UnbatchedInput now names the inlets it applies to. The materializer marks a BatchingActorInputBoundary unbatched only when the inlet it feeds is listed. - TlsGraphStage lists only its plaintext input inlet. Result: Plaintext input keeps the #3727 ordering guarantee; ciphertext input is batched across the async boundary again. Tests: - stream-tests/testOnly TlsGraphStageIsolatedSpec -z "batch transport input": passes; fails with cipherIn added back to UnbatchedInput (false was not equal to true) - stream-tests/testOnly org.apache.pekko.stream.io.Tls*: 286 passed - stream/mimaReportBinaryIssues: no issues - scalafmt --mode diff-ref=upstream/main: clean References: Refs #3727 --- .../stream/io/TlsGraphStageIsolatedSpec.scala | 55 +++++++++++++++++++ .../impl/PhasedFusingActorMaterializer.scala | 18 ++++++ .../impl/fusing/ActorGraphInterpreter.scala | 14 +++-- .../pekko/stream/impl/io/TlsGraphStage.scala | 6 +- 4 files changed, 86 insertions(+), 7 deletions(-) diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageIsolatedSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageIsolatedSpec.scala index 2102ff03b1e..895ab3aac37 100644 --- a/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageIsolatedSpec.scala +++ b/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageIsolatedSpec.scala @@ -25,6 +25,7 @@ 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 @@ -32,14 +33,17 @@ 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) } } diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/PhasedFusingActorMaterializer.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/PhasedFusingActorMaterializer.scala index 9445ae796bb..9219a012f3c 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/impl/PhasedFusingActorMaterializer.scala +++ b/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 diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/ActorGraphInterpreter.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/ActorGraphInterpreter.scala index b91d73e7332..33094728757 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/ActorGraphInterpreter.scala +++ b/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) diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/io/TlsGraphStage.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/io/TlsGraphStage.scala index ae9da629a1d..2e5364dfade 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/impl/io/TlsGraphStage.scala +++ b/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