WebSocket Improvements

zio-http-testkit/src/main/scala/zio/http/TestChannel.scala

@@ -1,5 +1,7 @@
 package zio.http
+
 import zio._
+import zio.stacktracer.TracingImplicits.disableAutoTrace
 
 import zio.http.ChannelEvent.{Unregistered, UserEvent, UserEventTriggered}
 
@@ -8,15 +10,15 @@ case class TestChannel(
   out: Queue[WebSocketChannelEvent],
   promise: Promise[Nothing, Unit],
 ) extends WebSocketChannel {
-  def awaitShutdown: UIO[Unit]                                 =
+  def awaitShutdown(implicit trace: Trace): UIO[Unit]                                 =
     promise.await
-  def receive: Task[WebSocketChannelEvent]                     =
+  def receive(implicit trace: Trace): Task[WebSocketChannelEvent]                     =
     in.take
-  def send(in: WebSocketChannelEvent): Task[Unit]              =
+  def send(in: WebSocketChannelEvent)(implicit trace: Trace): Task[Unit]              =
     out.offer(in).unit
-  def sendAll(in: Iterable[WebSocketChannelEvent]): Task[Unit] =
+  def sendAll(in: Iterable[WebSocketChannelEvent])(implicit trace: Trace): Task[Unit] =
     out.offerAll(in).unit
-  def shutdown: UIO[Unit]                                      =
+  def shutdown(implicit trace: Trace): UIO[Unit]                                      =
     in.offer(ChannelEvent.Unregistered) *>
       out.offer(ChannelEvent.Unregistered) *>
       promise.succeed(()).unit
@@ -27,7 +29,7 @@ object TestChannel {
     in: Queue[WebSocketChannelEvent],
     out: Queue[WebSocketChannelEvent],
     promise: Promise[Nothing, Unit],
-  ): ZIO[Any, Nothing, TestChannel] =
+  )(implicit trace: Trace): ZIO[Any, Nothing, TestChannel] =
     for {
       _ <- out.offer(UserEventTriggered(UserEvent.HandshakeComplete))
     } yield TestChannel(in, out, promise)

zio-http/src/main/scala/zio/http/Channel.scala

@@ -28,44 +28,44 @@ trait Channel[-In, +Out] { self =>
   /**
    * Await shutdown of the channel.
    */
-  def awaitShutdown: UIO[Unit]
+  def awaitShutdown(implicit trace: Trace): UIO[Unit]
 
   /**
    * Read a message from the channel, suspending until the next message is
    * available.
    */
-  def receive: Task[Out]
+  def receive(implicit trace: Trace): Task[Out]
 
   /**
    * Send a message to the channel.
    */
-  def send(in: In): Task[Unit]
+  def send(in: In)(implicit trace: Trace): Task[Unit]
 
   /**
    * Send all messages to the channel.
    */
-  def sendAll(in: Iterable[In]): Task[Unit]
+  def sendAll(in: Iterable[In])(implicit trace: Trace): Task[Unit]
 
   /**
    * Shut down the channel.
    */
-  def shutdown: UIO[Unit]
+  def shutdown(implicit trace: Trace): UIO[Unit]
 
   /**
    * Constructs a new channel that automatically transforms messages sent to
    * this channel using the specified function.
    */
   final def contramap[In2](f: In2 => In): Channel[In2, Out] =
     new Channel[In2, Out] {
-      def awaitShutdown: UIO[Unit]               =
+      def awaitShutdown(implicit trace: Trace): UIO[Unit]               =
         self.awaitShutdown
-      def receive: Task[Out]                     =
+      def receive(implicit trace: Trace): Task[Out]                     =
         self.receive
-      def send(in: In2): Task[Unit]              =
+      def send(in: In2)(implicit trace: Trace): Task[Unit]              =
         self.send(f(in))
-      def sendAll(in: Iterable[In2]): Task[Unit] =
+      def sendAll(in: Iterable[In2])(implicit trace: Trace): Task[Unit] =
         self.sendAll(in.map(f))
-      def shutdown: UIO[Unit]                    =
+      def shutdown(implicit trace: Trace): UIO[Unit]                    =
         self.shutdown
     }
 
@@ -75,15 +75,15 @@ trait Channel[-In, +Out] { self =>
    */
   final def map[Out2](f: Out => Out2)(implicit trace: Trace): Channel[In, Out2] =
     new Channel[In, Out2] {
-      def awaitShutdown: UIO[Unit]              =
+      def awaitShutdown(implicit trace: Trace): UIO[Unit]              =
         self.awaitShutdown
-      def receive: Task[Out2]                   =
+      def receive(implicit trace: Trace): Task[Out2]                   =
         self.receive.map(f)
-      def send(in: In): Task[Unit]              =
+      def send(in: In)(implicit trace: Trace): Task[Unit]              =
         self.send(in)
-      def sendAll(in: Iterable[In]): Task[Unit] =
+      def sendAll(in: Iterable[In])(implicit trace: Trace): Task[Unit] =
         self.sendAll(in)
-      def shutdown: UIO[Unit]                   =
+      def shutdown(implicit trace: Trace): UIO[Unit]                   =
         self.shutdown
     }
 

zio-http/src/main/scala/zio/http/WebSocketChannel.scala

@@ -17,6 +17,7 @@
 package zio.http
 
 import zio._
+import zio.stacktracer.TracingImplicits.disableAutoTrace
 
 import zio.http.ChannelEvent.{ExceptionCaught, Read, Registered, Unregistered, UserEventTriggered}
 import zio.http.netty.NettyChannel
@@ -31,31 +32,30 @@ private[http] object WebSocketChannel {
     queue: Queue[WebSocketChannelEvent],
   ): WebSocketChannel =
     new WebSocketChannel {
-      def awaitShutdown: UIO[Unit] =
+      def awaitShutdown(implicit trace: Trace): UIO[Unit] =
         nettyChannel.awaitClose
 
-      def receive: Task[WebSocketChannelEvent] =
+      def receive(implicit trace: Trace): Task[WebSocketChannelEvent] =
         queue.take
 
-      def send(in: WebSocketChannelEvent): Task[Unit] = {
+      def send(in: WebSocketChannelEvent)(implicit trace: Trace): Task[Unit] = {
         in match {
           case Read(message) => nettyChannel.writeAndFlush(frameToNetty(message))
           case _             => ZIO.unit
         }
       }
 
-      def sendAll(in: Iterable[WebSocketChannelEvent]): Task[Unit] =
+      def sendAll(in: Iterable[WebSocketChannelEvent])(implicit trace: Trace): Task[Unit] =
         ZIO.suspendSucceed {
           val iterator = in.iterator.collect { case Read(message) => message }
 
-          println(s"sendAll")
           ZIO.whileLoop(iterator.hasNext) {
             val message = iterator.next()
             if (iterator.hasNext) nettyChannel.write(frameToNetty(message))
             else nettyChannel.writeAndFlush(frameToNetty(message))
           }(_ => ())
         }
-      def shutdown: UIO[Unit]                                      =
+      def shutdown(implicit trace: Trace): UIO[Unit]                                      =
         nettyChannel.close(false).orDie
     }
 

zio-http/src/main/scala/zio/http/WebSocketConfig.scala

@@ -17,6 +17,7 @@
 package zio.http
 
 import zio.Duration
+import zio.stacktracer.TracingImplicits.disableAutoTrace
 
 /**
  * Server side websocket configuration

zio-http/src/main/scala/zio/http/WebSocketFrame.scala

@@ -17,6 +17,7 @@
 package zio.http
 
 import zio.Chunk
+import zio.stacktracer.TracingImplicits.disableAutoTrace
 
 sealed trait WebSocketFrame extends Product with Serializable { self =>
   def isFinal: Boolean = true