Bound Artery stream shutdown

docs/src/main/paradox/remoting-artery.md

@@ -93,6 +93,14 @@ listening for connections and handling messages as not to interfere with other a
 The example above only illustrates the bare minimum of properties you have to add to enable remoting.
 All settings are described in @ref:[Remote Configuration](#remote-configuration-artery).
 
+### Shutdown
+
+Artery first flushes outstanding remote messages and then aborts its streams during `ActorSystem` termination.
+The `pekko.remote.artery.advanced.shutdown-streams-timeout` setting limits how long Artery waits for those streams to
+complete. If the timeout expires, Artery proceeds with transport shutdown so a stalled stream cannot indefinitely
+delay transport-specific shutdown handling. The separate
+`pekko.remote.artery.advanced.shutdown-flush-timeout` setting controls the preceding flush.
+
 ## Introduction
 
 We recommend @ref:[Pekko Cluster](cluster-usage.md) over using remoting directly. As remoting is the

remote/src/main/resources/reference.conf

@@ -1024,6 +1024,11 @@ pekko {
         # remote messages has been completed
         shutdown-flush-timeout = 1 second
 
+        # After flushing, remoting aborts its streams and waits this long for them
+        # to complete before proceeding with transport shutdown. This bounds the
+        # graceful drain independently of transport-specific liveness timeouts.
+        shutdown-streams-timeout = 10 seconds
+
         # Before sending notification of terminated actor (DeathWatchNotification) other messages
         # will be flushed to make sure that the Terminated message arrives after other messages.
         # It will wait this long for the flush acknowledgement before continuing.

remote/src/main/scala/org/apache/pekko/remote/artery/ArterySettings.scala

@@ -175,6 +175,10 @@ private[pekko] final class ArterySettings private (config: Config) {
       config
         .getMillisDuration("shutdown-flush-timeout")
         .requiring(timeout => timeout > Duration.Zero, "shutdown-flush-timeout must be more than zero")
+    val ShutdownStreamsTimeout: FiniteDuration =
+      config
+        .getMillisDuration("shutdown-streams-timeout")
+        .requiring(timeout => timeout > Duration.Zero, "shutdown-streams-timeout must be more than zero")
     val DeathWatchNotificationFlushTimeout: FiniteDuration = {
       toRootLowerCase(config.getString("death-watch-notification-flush-timeout")) match {
         case "off" => Duration.Zero

remote/src/main/scala/org/apache/pekko/remote/artery/ArteryTransport.scala

@@ -22,6 +22,7 @@ import java.util.concurrent.atomic.AtomicReference
 import scala.annotation.nowarn
 import scala.annotation.tailrec
 import scala.concurrent.Await
+import scala.concurrent.ExecutionContext
 import scala.concurrent.Future
 import scala.concurrent.Promise
 import scala.concurrent.duration._
@@ -36,9 +37,10 @@ import pekko.annotation.InternalStableApi
 import pekko.dispatch.Dispatchers
 import pekko.event.Logging
 import pekko.event.MarkerLoggingAdapter
+import pekko.pattern.after
+import pekko.protobufv3.internal.CodedInputStream
 import pekko.remote.AddressUidExtension
 import pekko.remote.ContainerFormats
-import pekko.protobufv3.internal.CodedInputStream
 import pekko.remote.RemoteActorRef
 import pekko.remote.RemoteActorRefProvider
 import pekko.remote.RemoteTransport
@@ -53,11 +55,11 @@ import pekko.remote.artery.compress.CompressionProtocol.CompressionMessage
 import pekko.remote.transport.ThrottlerTransportAdapter.Blackhole
 import pekko.remote.transport.ThrottlerTransportAdapter.SetThrottle
 import pekko.remote.transport.ThrottlerTransportAdapter.Unthrottled
+import pekko.serialization.SerializationExtension
 import pekko.stream._
 import pekko.stream.scaladsl.Flow
 import pekko.stream.scaladsl.Keep
 import pekko.stream.scaladsl.Sink
-import pekko.serialization.SerializationExtension
 import pekko.util.OptionVal
 import pekko.util.WildcardIndex
 
@@ -658,8 +660,30 @@ private[remote] abstract class ArteryTransport(_system: ExtendedActorSystem, _pr
 
     killSwitch.abort(ShutdownSignal)
     flightRecorder.transportKillSwitchPulled()
+    val streamsStopped = streamsCompleted.recover { case _ => Done }
+    val shutdownStreamsTimeout = settings.Advanced.ShutdownStreamsTimeout
+    val streamsStoppedOrTimedOut =
+      try {
+        val timeoutResult = scheduleShutdownStreamsTimeout(shutdownStreamsTimeout) {
+          if (streamsStopped.isCompleted) streamsStopped
+          else {
+            log.warning(
+              "Graceful shutdown of Artery streams timed out after [{}]. Proceeding with transport shutdown.",
+              shutdownStreamsTimeout.toCoarsest)
+            Future.successful(Done)
+          }
+        }
+        Future.firstCompletedOf(List(streamsStopped, timeoutResult))
+      } catch {
+        case _: IllegalStateException =>
+          log.warning(
+            "Could not schedule the graceful shutdown timeout for Artery streams because the system scheduler is " +
+            "shut down. Proceeding with transport shutdown.")
+          Future.successful(Done)
+      }
+
     for {
-      _ <- streamsCompleted.recover { case _ => Done }
+      _ <- streamsStoppedOrTimedOut
       _ <- shutdownTransport().recover { case _ => Done }
     } yield {
       // no need to explicitly shut down the contained access since it's lifecycle is bound to the Decoder
@@ -669,6 +693,10 @@ private[remote] abstract class ArteryTransport(_system: ExtendedActorSystem, _pr
     }
   }
 
+  protected def scheduleShutdownStreamsTimeout(timeout: FiniteDuration)(result: => Future[Done])(
+      implicit ec: ExecutionContext): Future[Done] =
+    after(timeout, system.scheduler)(result)
+
   protected def shutdownTransport(): Future[Done]
 
   @tailrec final protected def updateStreamMatValues(streamId: Int, values: InboundStreamMatValues[LifeCycle]): Unit = {

remote/src/test/scala/org/apache/pekko/remote/artery/ArteryTransportShutdownSpec.scala

@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.pekko.remote.artery
+
+import java.util.concurrent.atomic.AtomicInteger
+
+import scala.concurrent.duration.FiniteDuration
+import scala.concurrent.{ ExecutionContext, Future, Promise }
+
+import org.apache.pekko
+import pekko.Done
+import pekko.actor.ExtendedActorSystem
+import pekko.remote.{ RARP, RemoteActorRefProvider }
+import pekko.stream.scaladsl.Sink
+import pekko.testkit.PekkoSpec
+
+import com.typesafe.config.ConfigFactory
+
+class ArteryTransportShutdownSpec
+    extends PekkoSpec(
+      ConfigFactory
+        .parseString("pekko.remote.artery.advanced.shutdown-streams-timeout = 100 ms")
+        .withFallback(ArterySpecSupport.defaultConfig)) {
+
+  "ArteryTransport shutdown" must {
+    "proceed with transport shutdown when stream completion times out" in {
+      assertTransportShutsDown(new StalledStreamTransport(schedulerUnavailable = false))
+    }
+
+    "proceed with transport shutdown when the system scheduler is unavailable" in {
+      assertTransportShutsDown(new StalledStreamTransport(schedulerUnavailable = true))
+    }
+  }
+
+  private def assertTransportShutsDown(transport: StalledStreamTransport): Unit = {
+    transport.addStalledInboundStream()
+    val shutdown = transport.shutdown()
+
+    transport.transportShutdownStarted.future.futureValue should ===(Done)
+    shutdown.futureValue should ===(Done)
+    transport.shutdown().futureValue should ===(Done)
+    transport.transportShutdownCount.get should ===(1)
+  }
+
+  private class StalledStreamTransport(schedulerUnavailable: Boolean)
+      extends ArteryTransport(
+        system.asInstanceOf[ExtendedActorSystem],
+        RARP(system).provider.asInstanceOf[RemoteActorRefProvider]) {
+    override type LifeCycle = Unit
+
+    val transportShutdownStarted = Promise[Done]()
+    val transportShutdownCount = new AtomicInteger
+    private val streamCompleted = Promise[Done]()
+
+    override protected def startTransport(): Unit = ()
+
+    override protected def bindInboundStreams(): (Int, Int) = (0, 0)
+
+    override protected def runInboundStreams(port: Int, bindPort: Int): Unit = ()
+
+    override protected def scheduleShutdownStreamsTimeout(timeout: FiniteDuration)(result: => Future[Done])(
+        implicit ec: ExecutionContext): Future[Done] =
+      if (schedulerUnavailable) throw new IllegalStateException("scheduler unavailable")
+      else super.scheduleShutdownStreamsTimeout(timeout)(result)
+
+    override protected def shutdownTransport(): Future[Done] = {
+      transportShutdownCount.incrementAndGet()
+      transportShutdownStarted.trySuccess(Done)
+      Future.successful(Done)
+    }
+
+    def addStalledInboundStream(): Unit =
+      updateStreamMatValues(
+        ArteryTransport.ControlStreamId,
+        ArteryTransport.InboundStreamMatValues((), streamCompleted.future))
+
+    override protected def outboundTransportSink(
+        outboundContext: OutboundContext,
+        streamId: Int,
+        bufferPool: EnvelopeBufferPool): Sink[EnvelopeBuffer, Future[Done]] =
+      Sink.ignore
+  }
+}