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 2102ff03b1..895ab3aac3 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 9445ae796b..9219a012f3 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 b91d73e733..3309472875 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 ae9da629a1..2e5364dfad 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