Fix resource leak on cancellation

Author: Kamil-Lontkowski

armeria-backend/cats-ce2/src/main/scala/sttp/client3/armeria/cats/ArmeriaCatsBackend.scala

@@ -1,6 +1,6 @@
 package sttp.client3.armeria.cats
 
-import cats.effect.{Concurrent, Resource, Sync}
+import cats.effect.{Concurrent, Resource, Sync, ExitCase}
 import com.linecorp.armeria.client.WebClient
 import com.linecorp.armeria.common.HttpData
 import com.linecorp.armeria.common.stream.StreamMessage
@@ -17,6 +17,12 @@ private final class ArmeriaCatsBackend[F[_]: Concurrent](client: WebClient, clos
 
   override val streams: NoStreams = NoStreams
 
+  override protected def ensureOnAbnormal[T](effect: F[T])(finalizer: => F[Unit]): F[T] =
+    Concurrent[F].guaranteeCase(effect) { exit =>
+      if (exit == ExitCase.Completed) Concurrent[F].unit
+      else Concurrent[F].recoverWith(finalizer) { case t => Concurrent[F].delay(t.printStackTrace()) }
+    }
+
   override protected def bodyFromStreamMessage: BodyFromStreamMessage[F, Nothing] =
     new BodyFromStreamMessage[F, Nothing] {
 

armeria-backend/cats/src/main/scala/sttp/client3/armeria/cats/ArmeriaCatsBackend.scala

@@ -17,6 +17,12 @@ private final class ArmeriaCatsBackend[F[_]: Async](client: WebClient, closeFact
 
   override val streams: NoStreams = NoStreams
 
+  override protected def ensureOnAbnormal[T](effect: F[T])(finalizer: => F[Unit]): F[T] =
+    Async[F].guaranteeCase(effect) { outcome =>
+      if (outcome.isSuccess) Async[F].unit
+      else Async[F].onError(finalizer) { case t => Async[F].delay(t.printStackTrace()) }
+    }
+
   override protected def bodyFromStreamMessage: BodyFromStreamMessage[F, Nothing] =
     new BodyFromStreamMessage[F, Nothing] {
 

armeria-backend/fs2-ce2/src/main/scala/sttp/client3/armeria/fs2/ArmeriaFs2Backend.scala

@@ -1,6 +1,6 @@
 package sttp.client3.armeria.fs2
 
-import cats.effect.{ConcurrentEffect, Resource, Sync}
+import cats.effect.{ConcurrentEffect, Resource, Sync, ExitCase}
 import com.linecorp.armeria.client.WebClient
 import com.linecorp.armeria.common.HttpData
 import com.linecorp.armeria.common.stream.StreamMessage
@@ -19,6 +19,12 @@ private final class ArmeriaFs2Backend[F[_]: ConcurrentEffect](client: WebClient,
 
   override val streams: Fs2Streams[F] = Fs2Streams[F]
 
+  override protected def ensureOnAbnormal[T](effect: F[T])(finalizer: => F[Unit]): F[T] =
+    ConcurrentEffect[F].guaranteeCase(effect) { exitCase =>
+      if (exitCase == ExitCase.Completed) ConcurrentEffect[F].unit
+      else ConcurrentEffect[F].onError(finalizer) { case t => ConcurrentEffect[F].delay(t.printStackTrace()) }
+    }
+
   override protected def bodyFromStreamMessage: BodyFromStreamMessage[F, Fs2Streams[F]] =
     new BodyFromStreamMessage[F, Fs2Streams[F]] {
 

armeria-backend/fs2/src/main/scala/sttp/client3/armeria/fs2/ArmeriaFs2Backend.scala

@@ -20,6 +20,12 @@ private final class ArmeriaFs2Backend[F[_]: Async](client: WebClient, closeFacto
 
   override val streams: Fs2Streams[F] = Fs2Streams[F]
 
+  override protected def ensureOnAbnormal[T](effect: F[T])(finalizer: => F[Unit]): F[T] =
+    Async[F].guaranteeCase(effect) { outcome =>
+      if (outcome.isSuccess) Async[F].unit
+      else Async[F].onError(finalizer) { case t => Async[F].delay(t.printStackTrace()) }
+    }
+
   override protected def bodyFromStreamMessage: BodyFromStreamMessage[F, Fs2Streams[F]] =
     new BodyFromStreamMessage[F, Fs2Streams[F]] {
 

armeria-backend/monix/src/main/scala/sttp/client3/armeria/monix/ArmeriaMonixBackend.scala

@@ -1,5 +1,6 @@
 package sttp.client3.armeria.monix
 
+import cats.effect.ExitCase
 import com.linecorp.armeria.client.WebClient
 import com.linecorp.armeria.common.HttpData
 import com.linecorp.armeria.common.stream.StreamMessage
@@ -19,6 +20,11 @@ private final class ArmeriaMonixBackend(client: WebClient, closeFactory: Boolean
 
   override val streams: MonixStreams = MonixStreams
 
+  override protected def ensureOnAbnormal[T](effect: Task[T])(finalizer: => Task[Unit]): Task[T] =
+    effect.guaranteeCase { exit =>
+      if (exit == ExitCase.Completed) Task.unit else finalizer.onErrorHandleWith(t => Task.eval(t.printStackTrace()))
+    }
+
   override protected def bodyFromStreamMessage: BodyFromStreamMessage[Task, MonixStreams] =
     new BodyFromStreamMessage[Task, MonixStreams] {
 

armeria-backend/scalaz/src/main/scala/sttp/client3/armeria/scalaz/ArmeriaScalazBackend.scala

@@ -17,6 +17,11 @@ private final class ArmeriaScalazBackend(client: WebClient, closeFactory: Boolea
 
   override val streams: NoStreams = NoStreams
 
+  override protected def ensureOnAbnormal[T](effect: Task[T])(finalizer: => Task[Unit]): Task[T] =
+    effect.handleWith { case e =>
+      finalizer.handleWith { case e2 => Task(e.addSuppressed(e2)) }.flatMap(_ => Task.fail(e))
+    }
+
   override protected def bodyFromStreamMessage: BodyFromStreamMessage[Task, Nothing] =
     new BodyFromStreamMessage[Task, Nothing] {
 

armeria-backend/src/main/scala/sttp/client3/armeria/AbstractArmeriaBackend.scala

@@ -50,16 +50,32 @@ abstract class AbstractArmeriaBackend[F[_], S <: Streams[S]](
 
   type SE = S with Effect[F]
 
+  // #1987: see the comments in HttpClientAsyncBackend
+  protected def ensureOnAbnormal[T](effect: F[T])(finalizer: => F[Unit]): F[T]
+
   protected def bodyFromStreamMessage: BodyFromStreamMessage[F, S]
 
   protected def streamToPublisher(stream: streams.BinaryStream): Publisher[HttpData]
 
   override def responseMonad: MonadError[F] = monad
 
-  override def send[T, R >: SE](request: Request[T, R]): F[Response[T]] =
-    monad.suspend(adjustExceptions(request)(execute(request)))
+  override def send[T, R >: SE](request: Request[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, R >: SE](request: Request[T, R]): F[Response[T]] = {
+  private def execute[T, R >: SE](
+      request: Request[T, R],
+      armeriaCtx: AtomicReference[ClientRequestContext]
+  ): F[Response[T]] = {
     val captor = Clients.newContextCaptor()
     try {
       val armeriaRes = requestToArmeria(request).execute()
@@ -82,6 +98,7 @@ abstract class AbstractArmeriaBackend[F[_], S <: Streams[S]](
             noopCanceler
           }
         case Success(ctx) =>
+          armeriaCtx.set(ctx)
           fromArmeriaResponse(request, armeriaRes, ctx)
       }
     } catch {

armeria-backend/src/main/scala/sttp/client3/armeria/future/ArmeriaFutureBackend.scala

@@ -18,6 +18,11 @@ private final class ArmeriaFutureBackend(client: WebClient, closeFactory: Boolea
 
   override val streams: NoStreams = NoStreams
 
+  override protected def ensureOnAbnormal[T](effect: Future[T])(finalizer: => Future[Unit]): Future[T] =
+    effect.recoverWith { case e =>
+      finalizer.recoverWith { case e2 => e.addSuppressed(e2); Future.failed(e) }.flatMap(_ => Future.failed(e))
+    }
+
   override protected def bodyFromStreamMessage: BodyFromStreamMessage[Future, Nothing] =
     new BodyFromStreamMessage[Future, Nothing] {
 

armeria-backend/zio/src/main/scala/sttp/client3/armeria/zio/ArmeriaZioBackend.scala

@@ -22,6 +22,11 @@ private final class ArmeriaZioBackend(runtime: Runtime[Any], client: WebClient,
 
   override val streams: ZioStreams = ZioStreams
 
+  override protected def ensureOnAbnormal[T](effect: Task[T])(finalizer: => Task[Unit]): Task[T] = effect.onExit {
+    exit =>
+      if (exit.isSuccess) ZIO.unit else finalizer.catchAll(t => ZIO.logErrorCause("Error in finalizer", Cause.fail(t)))
+  }.resurrect
+
   override protected def bodyFromStreamMessage: BodyFromStreamMessage[Task, ZioStreams] =
     new BodyFromStreamMessage[Task, ZioStreams] {
 

armeria-backend/zio1/src/main/scala/sttp/client3/armeria/zio/ArmeriaZioBackend.scala

@@ -23,6 +23,11 @@ private final class ArmeriaZioBackend(runtime: Runtime[Any], client: WebClient,
 
   override val streams: ZioStreams = ZioStreams
 
+  override protected def ensureOnAbnormal[T](effect: Task[T])(finalizer: => Task[Unit]): Task[T] = effect.onExit {
+    exit =>
+      if (exit.succeeded) ZIO.unit else finalizer.catchAll(t => ZIO.effect(t.printStackTrace()).orDie)
+  }.resurrect
+
   override protected def bodyFromStreamMessage: BodyFromStreamMessage[Task, ZioStreams] =
     new BodyFromStreamMessage[Task, ZioStreams] {