Use Bracket in the Http1Writer.write

http4s · Nov 13, 2021 · ff40852 · ff40852
1 parent 7ade2a8
commit ff40852
Showing 1 changed file with 12 additions and 10 deletions.
diff --git a/blaze-core/src/main/scala/org/http4s/blazecore/util/Http1Writer.scala b/blaze-core/src/main/scala/org/http4s/blazecore/util/Http1Writer.scala
@@ -18,6 +18,8 @@ package org.http4s
 package blazecore
 package util
 
+import cats.effect.ExitCase
+import cats.effect.syntax.bracket._
 import cats.syntax.all._
 import org.http4s.util.StringWriter
 import org.log4s.getLogger
@@ -28,16 +30,16 @@ import scala.concurrent._
 
 private[http4s] trait Http1Writer[F[_]] extends EntityBodyWriter[F] {
   final def write(headerWriter: StringWriter, body: EntityBody[F]): F[Boolean] =
-    fromFutureNoShift(F.delay(writeHeaders(headerWriter))).attempt.flatMap {
-      case Right(()) =>
-        writeEntityBody(body)
-      case Left(t) =>
-        body.drain.compile.drain.handleError { t2 =>
-          // Don't lose this error when sending the other
-          // TODO implement with cats.effect.Bracket when we have it
-          Http1Writer.logger.error(t2)("Error draining body")
-        } *> F.raiseError(t)
-    }
+    fromFutureNoShift(F.delay(writeHeaders(headerWriter)))
+      .guaranteeCase {
+        case ExitCase.Completed =>
+          F.unit
+
+        case ExitCase.Error(_) | ExitCase.Canceled =>
+          body.drain.compile.drain.handleError { t2 =>
+            Http1Writer.logger.error(t2)("Error draining body")
+          }
+      } >> writeEntityBody(body)
 
   /* Writes the header.  It is up to the writer whether to flush immediately or to
    * buffer the header with a subsequent chunk. */