Close the async ECS discovery clients on coordinated shutdown

discovery-aws-api-async/src/main/scala/org/apache/pekko/discovery/awsapi/ecs/AsyncEcsServiceDiscovery.scala

@@ -15,6 +15,7 @@ package org.apache.pekko.discovery.awsapi.ecs
 
 import java.net.InetAddress
 import java.util.concurrent.TimeoutException
+import java.util.concurrent.atomic.AtomicReference
 
 import scala.collection.immutable.Seq
 import scala.concurrent.duration._
@@ -24,7 +25,8 @@ import scala.jdk.FutureConverters._
 import scala.util.Try
 
 import org.apache.pekko
-import pekko.actor.ActorSystem
+import pekko.Done
+import pekko.actor.{ ActorSystem, CoordinatedShutdown }
 import pekko.annotation.ApiMayChange
 import pekko.discovery.ServiceDiscovery.{ Resolved, ResolvedTarget }
 import pekko.discovery.awsapi.ecs.AsyncEcsServiceDiscovery.{ resolveTasks, Tag }
@@ -51,14 +53,41 @@ class AsyncEcsServiceDiscovery(system: ActorSystem) extends ServiceDiscovery {
     }
     .toList
 
-  private lazy val ecsClient = {
+  // `httpClientBuilder` (rather than `httpClient`) hands ownership of the Netty client to the
+  // SDK, so that closing the ECS client also shuts down its event loop threads.
+  private[ecs] def createEcsClient(): EcsAsyncClient = {
     val conf = ClientOverrideConfiguration.builder().retryStrategy(DefaultRetryStrategy.doNotRetry()).build()
-    val httpClient = NettyNioAsyncHttpClient.create()
-    EcsAsyncClient.builder().overrideConfiguration(conf).httpClient(httpClient).build()
+    EcsAsyncClient
+      .builder()
+      .overrideConfiguration(conf)
+      .httpClientBuilder(NettyNioAsyncHttpClient.builder())
+      .build()
+  }
+
+  // holds the client once it has been successfully built, so that shutdown never forces
+  // (or re-attempts) the lazy initialisation
+  private val builtEcsClient = new AtomicReference[EcsAsyncClient]()
+
+  private lazy val ecsClient: EcsAsyncClient = {
+    val client = createEcsClient()
+    builtEcsClient.set(client)
+    client
   }
 
   private implicit val ec: ExecutionContext = system.dispatcher
 
+  CoordinatedShutdown(system).addTask(CoordinatedShutdown.PhaseServiceUnbind, "ecs-async-client-close") { () =>
+    builtEcsClient.getAndSet(null) match {
+      case null   => Future.successful(Done)
+      case client =>
+        // closing a Netty-backed client blocks until its event loops have shut down
+        Future {
+          client.close()
+          Done
+        }(system.dispatchers.lookup("pekko.actor.default-blocking-io-dispatcher"))
+    }
+  }
+
   override def lookup(lookup: Lookup, resolveTimeout: FiniteDuration): Future[Resolved] =
     Future.firstCompletedOf(
       Seq(

discovery-aws-api-async/src/main/scala/org/apache/pekko/discovery/awsapi/ecs/AsyncEcsTaskSetDiscovery.scala

@@ -15,6 +15,7 @@ package org.apache.pekko.discovery.awsapi.ecs
 
 import java.net.InetAddress
 import java.util.concurrent.TimeoutException
+import java.util.concurrent.atomic.AtomicReference
 
 import scala.collection.immutable.Seq
 import scala.concurrent.duration._
@@ -24,7 +25,8 @@ import scala.jdk.FutureConverters._
 import scala.util.Try
 
 import org.apache.pekko
-import pekko.actor.ActorSystem
+import pekko.Done
+import pekko.actor.{ ActorSystem, CoordinatedShutdown }
 import pekko.annotation.ApiMayChange
 import pekko.discovery.ServiceDiscovery.{ Resolved, ResolvedTarget }
 import pekko.discovery.awsapi.ecs.AsyncEcsTaskSetDiscovery.resolveTasks
@@ -55,16 +57,39 @@ class AsyncEcsTaskSetDiscovery(system: ActorSystem) extends ServiceDiscovery {
   private val config = system.settings.config.getConfig("pekko.discovery.aws-api-ecs-task-set-async")
   private val cluster = config.getString("cluster")
 
-  private lazy val ecsClient = {
+  private[ecs] def createEcsClient(): EcsAsyncClient = {
     val conf = ClientOverrideConfiguration.builder().retryStrategy(DefaultRetryStrategy.doNotRetry()).build()
     EcsAsyncClient.builder().overrideConfiguration(conf).build()
   }
 
+  // holds the client once it has been successfully built, so that shutdown never forces
+  // (or re-attempts) the lazy initialisation
+  private val builtEcsClient = new AtomicReference[EcsAsyncClient]()
+
+  private lazy val ecsClient: EcsAsyncClient = {
+    val client = createEcsClient()
+    builtEcsClient.set(client)
+    client
+  }
+
   private implicit val actorSystem: ActorSystem = system
   private implicit val ec: ExecutionContext = system.dispatcher
 
   private val httpClient: HttpExt = Http()
 
+  CoordinatedShutdown(system).addTask(CoordinatedShutdown.PhaseServiceUnbind, "ecs-task-set-async-client-close") {
+    () =>
+      builtEcsClient.getAndSet(null) match {
+        case null   => Future.successful(Done)
+        case client =>
+          // closing the client blocks until the underlying HTTP client has shut down
+          Future {
+            client.close()
+            Done
+          }(system.dispatchers.lookup("pekko.actor.default-blocking-io-dispatcher"))
+      }
+  }
+
   override def lookup(lookup: Lookup, resolveTimeout: FiniteDuration): Future[Resolved] =
     Future.firstCompletedOf(
       Seq(

discovery-aws-api-async/src/test/scala/org/apache/pekko/discovery/awsapi/ecs/AsyncEcsClientShutdownSpec.scala

@@ -0,0 +1,139 @@
+/*
+ * 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.discovery.awsapi.ecs
+
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.atomic.{ AtomicBoolean, AtomicInteger }
+
+import org.apache.pekko
+import pekko.actor.{ ActorSystem, CoordinatedShutdown }
+import pekko.discovery.Lookup
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.wordspec.AnyWordSpec
+import software.amazon.awssdk.services.ecs.EcsAsyncClient
+import software.amazon.awssdk.services.ecs.model.{
+  DescribeTasksRequest,
+  DescribeTasksResponse,
+  ListTasksRequest,
+  ListTasksResponse
+}
+
+import scala.concurrent.Await
+import scala.concurrent.duration._
+import scala.util.Try
+
+object AsyncEcsClientShutdownSpec {
+
+  class StubEcsAsyncClient extends EcsAsyncClient {
+    val closed = new AtomicBoolean(false)
+
+    override def serviceName(): String = EcsAsyncClient.SERVICE_NAME
+
+    override def close(): Unit = closed.set(true)
+
+    override def listTasks(request: ListTasksRequest): CompletableFuture[ListTasksResponse] =
+      CompletableFuture.completedFuture(ListTasksResponse.builder().build())
+
+    override def describeTasks(request: DescribeTasksRequest): CompletableFuture[DescribeTasksResponse] =
+      CompletableFuture.completedFuture(DescribeTasksResponse.builder().build())
+  }
+
+}
+
+class AsyncEcsClientShutdownSpec extends AnyWordSpec with Matchers {
+
+  import AsyncEcsClientShutdownSpec._
+
+  private def withSystem[T](name: String)(body: ActorSystem => T): T = {
+    val system = ActorSystem(name)
+    try body(system)
+    finally Await.ready(system.terminate(), 30.seconds)
+  }
+
+  private def shutdown(system: ActorSystem): Unit =
+    Await.result(CoordinatedShutdown(system).run(CoordinatedShutdown.UnknownReason), 30.seconds)
+
+  "AsyncEcsServiceDiscovery" should {
+
+    "close the ECS client on coordinated shutdown" in withSystem("AsyncEcsServiceDiscoverySpec") { system =>
+      val stub = new StubEcsAsyncClient
+      val discovery = new AsyncEcsServiceDiscovery(system) {
+        override private[ecs] def createEcsClient(): EcsAsyncClient = stub
+      }
+
+      Await.result(discovery.lookup(Lookup("my-service"), 10.seconds), 30.seconds).addresses should be(empty)
+      stub.closed.get() should ===(false)
+
+      shutdown(system)
+
+      stub.closed.get() should ===(true)
+    }
+
+    "not create an ECS client during shutdown when discovery was never used" in
+    withSystem("AsyncEcsServiceDiscoveryUnusedSpec") { system =>
+      val created = new AtomicInteger(0)
+      new AsyncEcsServiceDiscovery(system) {
+        override private[ecs] def createEcsClient(): EcsAsyncClient = {
+          created.incrementAndGet()
+          new StubEcsAsyncClient
+        }
+      }
+
+      shutdown(system)
+
+      created.get() should ===(0)
+    }
+
+  }
+
+  "AsyncEcsTaskSetDiscovery" should {
+
+    "close the ECS client on coordinated shutdown" in withSystem("AsyncEcsTaskSetDiscoverySpec") { system =>
+      val stub = new StubEcsAsyncClient
+      val discovery = new AsyncEcsTaskSetDiscovery(system) {
+        override private[ecs] def createEcsClient(): EcsAsyncClient = stub
+      }
+
+      // requires the ECS_CONTAINER_METADATA_URI environment variable, so it is expected to fail here,
+      // but not before the (lazily created) ECS client has been built
+      Try(Await.ready(discovery.lookup(Lookup("my-service"), 10.seconds), 30.seconds))
+      stub.closed.get() should ===(false)
+
+      shutdown(system)
+
+      stub.closed.get() should ===(true)
+    }
+
+    "not create an ECS client during shutdown when discovery was never used" in
+    withSystem("AsyncEcsTaskSetDiscoveryUnusedSpec") { system =>
+      val created = new AtomicInteger(0)
+      new AsyncEcsTaskSetDiscovery(system) {
+        override private[ecs] def createEcsClient(): EcsAsyncClient = {
+          created.incrementAndGet()
+          new StubEcsAsyncClient
+        }
+      }
+
+      shutdown(system)
+
+      created.get() should ===(0)
+    }
+
+  }
+
+}

docs/src/main/paradox/discovery/aws.md

@@ -199,6 +199,10 @@ pekko.discovery {
 This uses AWS SDK v2. The advantage here is that the SDK does
 non-blocking IO, which you probably want.
 
+Both of the discovery methods below create their `EcsAsyncClient` lazily, on the first lookup, and close it
+again during the `service-unbind` phase of Coordinated Shutdown. If discovery is never used, no client is
+created and nothing needs to be closed.
+
 @@dependency[sbt,Gradle,Maven] {
   symbol1=PekkoManagementVersion
   value1=$project.version$