Implement PubSub api

modules/redis/src/main/scala/zio/redis/Output.scala

@@ -19,6 +19,7 @@ package zio.redis
 import zio._
 import zio.redis.internal.RespValue
 import zio.redis.options.Cluster.{Node, Partition, SlotRange}
+import zio.redis.options.PubSub.{NumberOfSubscribers, PushProtocol}
 import zio.schema.Schema
 import zio.schema.codec.BinaryCodec
 
@@ -639,6 +640,50 @@ object Output {
       }
   }
 
+  case object PushProtocolOutput extends Output[PushProtocol] {
+    protected def tryDecode(respValue: RespValue): PushProtocol =
+      respValue match {
+        case RespValue.NullArray => throw ProtocolError(s"Array must not be empty")
+        case RespValue.Array(values) =>
+          val name = MultiStringOutput.unsafeDecode(values(0))
+          val key  = MultiStringOutput.unsafeDecode(values(1))
+          name match {
+            case "subscribe" =>
+              val num = LongOutput.unsafeDecode(values(2))
+              PushProtocol.Subscribe(key, num)
+            case "psubscribe" =>
+              val num = LongOutput.unsafeDecode(values(2))
+              PushProtocol.PSubscribe(key, num)
+            case "unsubscribe" =>
+              val num = LongOutput.unsafeDecode(values(2))
+              PushProtocol.Unsubscribe(key, num)
+            case "punsubscribe" =>
+              val num = LongOutput.unsafeDecode(values(2))
+              PushProtocol.PUnsubscribe(key, num)
+            case "message" =>
+              PushProtocol.Message(key, values(2))
+            case "pmessage" =>
+              val channel = MultiStringOutput.unsafeDecode(values(2))
+              PushProtocol.PMessage(key, channel, values(3))
+            case other => throw ProtocolError(s"$other isn't a pushed message")
+          }
+        case other => throw ProtocolError(s"$other isn't an array")
+      }
+  }
+
+  case object NumSubResponseOutput extends Output[Chunk[NumberOfSubscribers]] {
+    protected def tryDecode(respValue: RespValue): Chunk[NumberOfSubscribers] =
+      respValue match {
+        case RespValue.Array(values) =>
+          Chunk.fromIterator(values.grouped(2).map { chunk =>
+            val channel   = MultiStringOutput.unsafeDecode(chunk(0))
+            val numOfSubs = LongOutput.unsafeDecode(chunk(1))
+            NumberOfSubscribers(channel, numOfSubs)
+          })
+        case other => throw ProtocolError(s"$other isn't an array")
+      }
+  }
+
   private def decodeDouble(bytes: Chunk[Byte]): Double = {
     val text = new String(bytes.toArray, StandardCharsets.UTF_8)
     try text.toDouble

modules/redis/src/main/scala/zio/redis/Redis.scala

@@ -32,6 +32,7 @@ trait Redis
     with api.Streams
     with api.Scripting
     with api.Cluster
+    with api.Publishing
 
 object Redis {
   lazy val cluster: ZLayer[CodecSupplier & RedisClusterConfig, RedisError, Redis] =

modules/redis/src/main/scala/zio/redis/RedisError.scala

@@ -46,5 +46,8 @@ object RedisError {
   object Moved {
     def apply(slotAndAddress: (Slot, RedisUri)): Moved = Moved(slotAndAddress._1, slotAndAddress._2)
   }
-  final case class IOError(exception: IOException) extends RedisError
+  final case class IOError(exception: IOException)       extends RedisError
+  final case class CommandNameNotFound(message: String)  extends RedisError
+  sealed trait PubSubError                               extends RedisError
+  final case class InvalidPubSubCommand(command: String) extends PubSubError
 }

modules/redis/src/main/scala/zio/redis/RedisSubscription.scala

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2021 John A. De Goes and the ZIO contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package zio.redis
+
+import zio._
+
+trait RedisSubscription extends api.Subscription
+
+object RedisSubscription {
+  lazy val local: ZLayer[CodecSupplier, RedisError.IOError, RedisSubscription] =
+    SubscriptionExecutor.local >>> makeLayer
+
+  lazy val singleNode: ZLayer[CodecSupplier & RedisConfig, RedisError.IOError, RedisSubscription] =
+    SubscriptionExecutor.layer >>> makeLayer
+
+  private def makeLayer: URLayer[CodecSupplier & SubscriptionExecutor, RedisSubscription] =
+    ZLayer.fromFunction(Live.apply _)
+
+  private final case class Live(codecSupplier: CodecSupplier, executor: SubscriptionExecutor) extends RedisSubscription
+}

modules/redis/src/main/scala/zio/redis/ResultBuilder.scala

@@ -19,6 +19,7 @@ package zio.redis
 import zio.IO
 import zio.redis.ResultBuilder.NeedsReturnType
 import zio.schema.Schema
+import zio.stream.Stream
 
 sealed trait ResultBuilder {
   final def map(f: Nothing => Any)(implicit nrt: NeedsReturnType): IO[Nothing, Nothing] = ???
@@ -45,4 +46,8 @@ object ResultBuilder {
   trait ResultOutputBuilder extends ResultBuilder {
     def returning[R: Output]: IO[RedisError, R]
   }
+
+  trait ResultStreamBuilder1[+F[_]] extends ResultBuilder {
+    def returning[R: Schema]: Stream[RedisError, F[R]]
+  }
 }

modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala

@@ -0,0 +1,219 @@
+/*
+ * Copyright 2021 John A. De Goes and the ZIO contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package zio.redis
+
+import zio.redis.Input.{CommandNameInput, StringInput}
+import zio.redis.Output.PushProtocolOutput
+import zio.redis.SingleNodeSubscriptionExecutor.{Request, RequestQueueSize, True}
+import zio.redis.api.Subscription
+import zio.redis.internal._
+import zio.redis.options.PubSub.PushProtocol
+import zio.stream._
+import zio.{Chunk, ChunkBuilder, IO, Promise, Queue, Ref, Schedule, Scope, URIO, ZIO}
+
+private[redis] final class SingleNodeSubscriptionExecutor private (
+  channelSubsRef: Ref[Map[String, Chunk[Queue[Take[RedisError, PushProtocol]]]]],
+  patternSubsRef: Ref[Map[String, Chunk[Queue[Take[RedisError, PushProtocol]]]]],
+  reqQueue: Queue[Request],
+  connection: RedisConnection
+) extends SubscriptionExecutor {
+  def execute(command: RespCommand): Stream[RedisError, PushProtocol] =
+    ZStream
+      .fromZIO(
+        for {
+          commandName <-
+            ZIO
+              .fromOption(command.args.collectFirst { case RespCommandArgument.CommandName(name) => name })
+              .orElseFail(RedisError.CommandNameNotFound(command.args.toString()))
+          stream <- commandName match {
+                      case Subscription.Subscribe    => ZIO.succeed(subscribe(channelSubsRef, command))
+                      case Subscription.PSubscribe   => ZIO.succeed(subscribe(patternSubsRef, command))
+                      case Subscription.Unsubscribe  => ZIO.succeed(unsubscribe(command))
+                      case Subscription.PUnsubscribe => ZIO.succeed(unsubscribe(command))
+                      case other                     => ZIO.fail(RedisError.InvalidPubSubCommand(other))
+                    }
+        } yield stream
+      )
+      .flatten
+
+  private def subscribe(
+    subscriptionRef: Ref[Map[String, Chunk[Queue[Take[RedisError, PushProtocol]]]]],
+    command: RespCommand
+  ): Stream[RedisError, PushProtocol] =
+    ZStream
+      .fromZIO(
+        for {
+          queues <-
+            ZIO.foreach(command.args.collect { case key: RespCommandArgument.Key => key.value.asString })(key =>
+              Queue
+                .unbounded[Take[RedisError, PushProtocol]]
+                .tap(queue =>
+                  subscriptionRef.update(subscription =>
+                    subscription.updated(key, subscription.getOrElse(key, Chunk.empty) ++ Chunk.single(queue))
+                  )
+                )
+                .map(key -> _)
+            )
+          promise <- Promise.make[RedisError, Unit]
+          _ <- reqQueue.offer(
+                 Request(
+                   command.args.map(_.value),
+                   promise
+                 )
+               )
+          streams = queues.map { case (key, queue) =>
+                      ZStream
+                        .fromQueueWithShutdown(queue)
+                        .ensuring(
+                          subscriptionRef.update(subscription =>
+                            subscription.get(key) match {
+                              case Some(queues) => subscription.updated(key, queues.filterNot(_ == queue))
+                              case None         => subscription
+                            }
+                          )
+                        )
+                    }
+          _ <- promise.await.tapError(_ => ZIO.foreachDiscard(queues) { case (_, queue) => queue.shutdown })
+        } yield streams.fold(ZStream.empty)(_ merge _)
+      )
+      .flatten
+      .flattenTake
+
+  private def unsubscribe(command: RespCommand): Stream[RedisError, PushProtocol] =
+    ZStream
+      .fromZIO(
+        for {
+          promise <- Promise.make[RedisError, Unit]
+          _       <- reqQueue.offer(Request(command.args.map(_.value), promise))
+          _       <- promise.await
+        } yield ZStream.empty
+      )
+      .flatten
+
+  private def send =
+    reqQueue.takeBetween(1, RequestQueueSize).flatMap { reqs =>
+      val buffer = ChunkBuilder.make[Byte]()
+      val it     = reqs.iterator
+
+      while (it.hasNext) {
+        val req = it.next()
+        buffer ++= RespValue.Array(req.command).asBytes
+      }
+
+      val bytes = buffer.result()
+
+      connection
+        .write(bytes)
+        .mapError(RedisError.IOError(_))
+        .tapBoth(
+          e => ZIO.foreachDiscard(reqs.map(_.promise))(_.fail(e)),
+          _ => ZIO.foreachDiscard(reqs.map(_.promise))(_.succeed(()))
+        )
+    }
+
+  private def receive: IO[RedisError, Unit] = {
+    def offerMessage(
+      subscriptionRef: Ref[Map[String, Chunk[Queue[Take[RedisError, PushProtocol]]]]],
+      key: String,
+      msg: PushProtocol
+    ) = for {
+      subscription <- subscriptionRef.get
+      _ <- ZIO.foreachDiscard(subscription.get(key))(
+             ZIO.foreachDiscard(_)(queue => queue.offer(Take.single(msg)).unlessZIO(queue.isShutdown))
+           )
+    } yield ()
+
+    def releaseStream(subscriptionRef: Ref[Map[String, Chunk[Queue[Take[RedisError, PushProtocol]]]]], key: String) =
+      for {
+        subscription <- subscriptionRef.getAndUpdate(_ - key)
+        _ <- ZIO.foreachDiscard(subscription.get(key))(
+               ZIO.foreachDiscard(_)(queue => queue.offer(Take.end).unlessZIO(queue.isShutdown))
+             )
+      } yield ()
+
+    connection.read
+      .mapError(RedisError.IOError(_))
+      .via(RespValue.Decoder)
+      .collectSome
+      .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp)))
+      .refineToOrDie[RedisError]
+      .foreach {
+        case msg @ PushProtocol.Subscribe(channel, _) => offerMessage(channelSubsRef, channel, msg)
+        case msg @ PushProtocol.Unsubscribe(channel, _) =>
+          offerMessage(channelSubsRef, channel, msg) *> releaseStream(channelSubsRef, channel)
+        case msg @ PushProtocol.Message(channel, _)    => offerMessage(channelSubsRef, channel, msg)
+        case msg @ PushProtocol.PSubscribe(pattern, _) => offerMessage(patternSubsRef, pattern, msg)
+        case msg @ PushProtocol.PUnsubscribe(pattern, _) =>
+          offerMessage(patternSubsRef, pattern, msg) *> releaseStream(patternSubsRef, pattern)
+        case msg @ PushProtocol.PMessage(pattern, _, _) => offerMessage(patternSubsRef, pattern, msg)
+      }
+  }
+
+  private def resubscribe: IO[RedisError, Unit] = {
+    def makeCommand(name: String, keys: Chunk[String]) =
+      if (keys.isEmpty)
+        Chunk.empty
+      else
+        RespValue
+          .Array((CommandNameInput.encode(name) ++ Input.Varargs(StringInput).encode(keys)).args.map(_.value))
+          .asBytes
+
+    for {
+      channels <- channelSubsRef.get.map(_.keys)
+      patterns <- patternSubsRef.get.map(_.keys)
+      commands = makeCommand(Subscription.Subscribe, Chunk.fromIterable(channels)) ++
+                   makeCommand(Subscription.PSubscribe, Chunk.fromIterable(patterns))
+      _ <- connection
+             .write(commands)
+             .when(commands.nonEmpty)
+             .mapError(RedisError.IOError(_))
+             .retryWhile(True)
+    } yield ()
+  }
+
+  /**
+   * Opens a connection to the server and launches receive operations. All failures are retried by opening a new
+   * connection. Only exits by interruption or defect.
+   */
+  val run: IO[RedisError, AnyVal] =
+    ZIO.logTrace(s"$this PubSub sender and reader has been started") *>
+      (send.repeat(Schedule.forever) race receive)
+        .tapError(e => ZIO.logWarning(s"Reconnecting due to error: $e") *> resubscribe)
+        .retryWhile(True)
+        .tapError(e => ZIO.logError(s"Executor exiting: $e"))
+}
+
+private[redis] object SingleNodeSubscriptionExecutor {
+  private final case class Request(
+    command: Chunk[RespValue.BulkString],
+    promise: Promise[RedisError, Unit]
+  )
+
+  private final val True: Any => Boolean = _ => true
+
+  private final val RequestQueueSize = 16
+
+  def create(conn: RedisConnection): URIO[Scope, SubscriptionExecutor] =
+    for {
+      reqQueue   <- Queue.bounded[Request](RequestQueueSize)
+      channelRef <- Ref.make(Map.empty[String, Chunk[Queue[Take[RedisError, PushProtocol]]]])
+      patternRef <- Ref.make(Map.empty[String, Chunk[Queue[Take[RedisError, PushProtocol]]]])
+      pubSub      = new SingleNodeSubscriptionExecutor(channelRef, patternRef, reqQueue, conn)
+      _          <- pubSub.run.forkScoped
+      _          <- logScopeFinalizer(s"$pubSub Subscription Node is closed")
+    } yield pubSub
+}

modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2021 John A. De Goes and the ZIO contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package zio.redis
+
+import zio.redis.internal.{RedisConnection, RespCommand}
+import zio.redis.options.PubSub.PushProtocol
+import zio.stream._
+import zio.{Layer, ZIO, ZLayer}
+
+trait SubscriptionExecutor {
+  private[redis] def execute(command: RespCommand): Stream[RedisError, PushProtocol]
+}
+
+object SubscriptionExecutor {
+  lazy val layer: ZLayer[RedisConfig, RedisError.IOError, SubscriptionExecutor] =
+    RedisConnection.layer.fresh >>> pubSublayer
+
+  lazy val local: Layer[RedisError.IOError, SubscriptionExecutor] =
+    RedisConnection.local.fresh >>> pubSublayer
+
+  private lazy val pubSublayer: ZLayer[RedisConnection, RedisError.IOError, SubscriptionExecutor] =
+    ZLayer.scoped(
+      for {
+        conn   <- ZIO.service[RedisConnection]
+        pubSub <- SingleNodeSubscriptionExecutor.create(conn)
+      } yield pubSub
+    )
+}