From 8ee1e4a2fd7031c29cab4c755c81709e7571cfb2 Mon Sep 17 00:00:00 2001
From: adamw <adam@warski.org>
Date: Fri, 24 Jan 2025 14:08:37 +0100
Subject: [PATCH] Better armeria

---
 .../armeria/AbstractArmeriaBackend.scala      | 23 +++++++++++++++----
 .../httpclient/HttpClientAsyncBackend.scala   |  3 +++
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/armeria-backend/src/main/scala/sttp/client4/armeria/AbstractArmeriaBackend.scala b/armeria-backend/src/main/scala/sttp/client4/armeria/AbstractArmeriaBackend.scala
index dafd86756..40bcde96b 100644
--- a/armeria-backend/src/main/scala/sttp/client4/armeria/AbstractArmeriaBackend.scala
+++ b/armeria-backend/src/main/scala/sttp/client4/armeria/AbstractArmeriaBackend.scala
@@ -61,10 +61,23 @@ abstract class AbstractArmeriaBackend[F[_], S <: Streams[S]](
   // #1987: see the comments in HttpClientAsyncBackend
   protected def ensureOnAbnormal[T](effect: F[T])(finalizer: => F[Unit]): F[T]
 
-  override def send[T](request: GenericRequest[T, R]): F[Response[T]] =
-    monad.suspend(adjustExceptions(request)(execute(request)))
+  override def send[T](request: GenericRequest[T, R]): F[Response[T]] = {
+    // #1987: see the comments in HttpClientAsyncBackend
+    val armeriaCtx = new AtomicReference[ClientRequestContext]()
+    ensureOnAbnormal {
+      monad.suspend(adjustExceptions(request)(execute(request, armeriaCtx)))
+    } {
+      monad.eval {
+        val ctx = armeriaCtx.get()
+        if (ctx != null) ctx.cancel()
+      }
+    }
+  }
 
-  private def execute[T](request: GenericRequest[T, R]): F[Response[T]] = {
+  private def execute[T](
+      request: GenericRequest[T, R],
+      armeriaCtx: AtomicReference[ClientRequestContext]
+  ): F[Response[T]] = {
     val captor = Clients.newContextCaptor()
     try {
       val armeriaRes = requestToArmeria(request).execute()
@@ -87,8 +100,8 @@ abstract class AbstractArmeriaBackend[F[_], S <: Streams[S]](
             noopCanceler
           }
         case Success(ctx) =>
-          // #1987: see the comments in HttpClientAsyncBackend
-          ensureOnAbnormal(fromArmeriaResponse(request, armeriaRes, ctx))(monad.eval(ctx.cancel()))
+          armeriaCtx.set(ctx)
+          fromArmeriaResponse(request, armeriaRes, ctx)
       }
     } catch {
       case NonFatal(ex) => monad.error(ex)
diff --git a/core/src/main/scalajvm/sttp/client4/httpclient/HttpClientAsyncBackend.scala b/core/src/main/scalajvm/sttp/client4/httpclient/HttpClientAsyncBackend.scala
index 7dde62408..aad7bd7ef 100644
--- a/core/src/main/scalajvm/sttp/client4/httpclient/HttpClientAsyncBackend.scala
+++ b/core/src/main/scalajvm/sttp/client4/httpclient/HttpClientAsyncBackend.scala
@@ -58,6 +58,9 @@ abstract class HttpClientAsyncBackend[F[_], S <: Streams[S], BH, B](
     * is consumed. This is done only in case of failure, as in case of success the body is either fully consumed as
     * specified in the response description, or when an `...Unsafe` response description is used, it's up to the user to
     * consume it.
+    *
+    * Any exceptions that occur while running `finalizer` should be added as suppressed or logged, not impacting the
+    * outcome of `effect`. If possible, `finalizer` should not be run in a cancellable way.
     */
   protected def ensureOnAbnormal[T](effect: F[T])(finalizer: => F[Unit]): F[T]