Websocket: Update To Unsafe Signal for Close Rather than Pure Signal

blaze-core/src/main/scala/org/http4s/blazecore/websocket/Http4sWSStage.scala

@@ -2,25 +2,27 @@ package org.http4s
 package blazecore
 package websocket
 
+import java.util.concurrent.atomic.AtomicReference
+
 import cats.effect._
-import cats.effect.implicits._
 import cats.implicits._
 import fs2._
-import fs2.async.mutable.Signal
 import org.http4s.{websocket => ws4s}
 import org.http4s.blaze.pipeline.{Command, LeafBuilder, TailStage, TrunkBuilder}
 import org.http4s.blaze.pipeline.stages.SerializingStage
 import org.http4s.blaze.util.Execution.{directec, trampoline}
 import org.http4s.websocket.WebsocketBits._
+
 import scala.concurrent.ExecutionContext
 import scala.util.{Failure, Success}
+import cats.effect.IO
 
 class Http4sWSStage[F[_]](ws: ws4s.Websocket[F])(implicit F: Effect[F], val ec: ExecutionContext)
     extends TailStage[WebSocketFrame] {
 
   def name: String = "Http4s WebSocket Stage"
 
-  private val deadSignal: F[Signal[F, Boolean]] = async.signalOf[F, Boolean](false)
+  private val deadSignal: AtomicReference[Boolean] = new AtomicReference(false)
 
   //////////////////////// Source and Sink generators ////////////////////////
 
@@ -40,29 +42,22 @@ class Http4sWSStage[F[_]](ws: ws4s.Websocket[F])(implicit F: Effect[F], val ec:
         channelRead().onComplete {
           case Success(ws) =>
             ws match {
-              case Close(_) =>
-                for {
-                  t <- deadSignal.map(_.set(true))
-                } yield {
-                  t.runAsync(_ => IO.unit).unsafeRunSync()
-                  cb(Left(Command.EOF))
-                }
-
+              case c@Close(_) =>
+                deadSignal.set(true)
+                cb(Right(c)) // With Dead Signal Set, Return callback with the Close Frame
               case Ping(d) =>
                 channelWrite(Pong(d)).onComplete {
                   case Success(_) => go()
                   case Failure(Command.EOF) => cb(Left(Command.EOF))
                   case Failure(t) => cb(Left(t))
                 }(trampoline)
-
               case Pong(_) => go()
               case f => cb(Right(f))
             }
 
           case Failure(Command.EOF) => cb(Left(Command.EOF))
           case Failure(e) => cb(Left(e))
         }(trampoline)
-
       go()
     }
     Stream.repeatEval(t)
@@ -80,18 +75,18 @@ class Http4sWSStage[F[_]](ws: ws4s.Websocket[F])(implicit F: Effect[F], val ec:
     val onStreamFinalize: F[Unit] =
       for {
         dec <- F.delay(count.decrementAndGet())
-        _ <- deadSignal.map(signal => if (dec == 0) signal.set(true))
+        _ <- if (dec == 0) F.delay(deadSignal.set(true)) else ().pure[F]
       } yield ()
 
     // Effect to send a close to the other endpoint
     val sendClose: F[Unit] = F.delay(sendOutboundCommand(Command.Disconnect))
 
-    val wsStream = for {
-      dead <- deadSignal
-      in = inputstream.to(ws.receive).onFinalize(onStreamFinalize)
-      out = ws.send.onFinalize(onStreamFinalize).to(snk).drain
-      merged <- in.mergeHaltR(out).interruptWhen(dead).onFinalize(sendClose).compile.drain
-    } yield merged
+    val wsStream = inputstream.to(ws.receive).onFinalize(onStreamFinalize)
+      .mergeHaltR(ws.send.onFinalize(onStreamFinalize).to(snk).drain)
+      .interruptWhen(Stream.repeatEval(F.delay(deadSignal.get()))) // Check The Status of the deadSignal Atomic Ref
+      .onFinalize(sendClose)
+      .compile
+      .drain
 
     async.unsafeRunAsync {
       wsStream.attempt.flatMap {
@@ -108,7 +103,7 @@ class Http4sWSStage[F[_]](ws: ws4s.Websocket[F])(implicit F: Effect[F], val ec:
   }
 
   override protected def stageShutdown(): Unit = {
-    deadSignal.map(_.set(true)).runAsync(_ => IO.unit).unsafeRunSync()
+    deadSignal.set(true)
     super.stageShutdown()
   }
 }