Implement PubSub api

benchmarks/src/main/scala/zio/redis/benchmarks/BenchmarkRuntime.scala

@@ -35,6 +35,7 @@ object BenchmarkRuntime {
   private final val Layer =
     ZLayer.make[Redis](
       RedisExecutor.local,
+      RedisPubSub.local,
       ZLayer.succeed[BinaryCodec](ProtobufCodec),
       Redis.layer
     )

example/src/main/scala/example/Main.scala

@@ -21,7 +21,7 @@ import example.config.AppConfig
 import sttp.client3.httpclient.zio.HttpClientZioBackend
 import zhttp.service.Server
 import zio._
-import zio.redis.{Redis, RedisExecutor}
+import zio.redis.{Redis, RedisExecutor, RedisPubSub}
 import zio.schema.codec.{BinaryCodec, ProtobufCodec}
 
 object Main extends ZIOAppDefault {
@@ -33,6 +33,7 @@ object Main extends ZIOAppDefault {
         ContributorsCache.layer,
         HttpClientZioBackend.layer(),
         RedisExecutor.layer,
+        RedisPubSub.layer,
         Redis.layer,
         ZLayer.succeed[BinaryCodec](ProtobufCodec)
       )

redis/src/main/scala/zio/redis/ClusterExecutor.scala

@@ -150,9 +150,11 @@ object ClusterExecutor {
     } yield ExecutorScope(executor, closableScope)
 
   private def redis(address: RedisUri) = {
-    val executorLayer = ZLayer.succeed(RedisConfig(address.host, address.port)) >>> RedisExecutor.layer
-    val codecLayer    = ZLayer.succeed[BinaryCodec](StringUtf8Codec)
-    val redisLayer    = executorLayer ++ codecLayer >>> Redis.layer
+    val redisConfigLayer = ZLayer.succeed(RedisConfig(address.host, address.port))
+    val codecLayer       = ZLayer.succeed[BinaryCodec](StringUtf8Codec)
+    val executorLayer    = redisConfigLayer >>> RedisExecutor.layer
+    val pubSubLayer      = redisConfigLayer ++ codecLayer >>> RedisPubSub.layer
+    val redisLayer       = executorLayer ++ pubSubLayer ++ codecLayer >>> Redis.layer
     for {
       closableScope <- Scope.make
       layer         <- closableScope.extend[Any](redisLayer.memoize)

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

@@ -18,6 +18,7 @@ package zio.redis
 
 import zio._
 import zio.redis.options.Cluster.{Node, Partition, SlotRange}
+import zio.redis.options.PubSub.NumberOfSubscribers
 import zio.schema.Schema
 import zio.schema.codec.BinaryCodec
 
@@ -827,4 +828,48 @@ object Output {
         case other => throw ProtocolError(s"$other isn't an array")
       }
   }
+
+  case object PushProtocolOutput extends Output[PushProtocol] {
+    protected def tryDecode(respValue: RespValue)(implicit codec: BinaryCodec): 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)(implicit codec: BinaryCodec): Chunk[NumberOfSubscribers] =
+      respValue match {
+        case RespValue.Array(values) =>
+          Chunk.fromIterator(values.grouped(2).map { chunk =>
+            val channel           = MultiStringOutput.unsafeDecode(chunk(0))
+            val numOfSubscription = LongOutput.unsafeDecode(chunk(1))
+            NumberOfSubscribers(channel, numOfSubscription)
+          })
+        case other => throw ProtocolError(s"$other isn't an array")
+      }
+  }
 }

redis/src/main/scala/zio/redis/PubSubCommand.scala

@@ -0,0 +1,26 @@
+package zio.redis
+
+import zio.IO
+
+sealed trait PubSubCommand
+
+private[redis] object PubSubCommand {
+  case class Subscribe(
+    channel: String,
+    channels: List[String],
+    callback: PushProtocol => IO[RedisError, Unit]
+  ) extends PubSubCommand
+  case class PSubscribe(
+    pattern: String,
+    patterns: List[String],
+    callback: PushProtocol => IO[RedisError, Unit]
+  ) extends PubSubCommand
+  case class Unsubscribe(
+    channels: List[String],
+    callback: PushProtocol => IO[RedisError, Unit]
+  ) extends PubSubCommand
+  case class PUnsubscribe(
+    patterns: List[String],
+    callback: PushProtocol => IO[RedisError, Unit]
+  ) extends PubSubCommand
+}

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

@@ -32,15 +32,17 @@ trait Redis
     with api.Streams
     with api.Scripting
     with api.Cluster
+    with api.PubSub
 
 object Redis {
-  lazy val layer: URLayer[RedisExecutor with BinaryCodec, Redis] =
+  lazy val layer: URLayer[RedisExecutor with RedisPubSub with BinaryCodec, Redis] =
     ZLayer {
       for {
         executor <- ZIO.service[RedisExecutor]
+        pubSub   <- ZIO.service[RedisPubSub]
         codec    <- ZIO.service[BinaryCodec]
-      } yield Live(codec, executor)
+      } yield Live(codec, executor, pubSub)
     }
 
-  private final case class Live(codec: BinaryCodec, executor: RedisExecutor) extends Redis
+  private final case class Live(codec: BinaryCodec, executor: RedisExecutor, pubSub: RedisPubSub) extends Redis
 }

redis/src/main/scala/zio/redis/RedisEnvironment.scala

@@ -21,4 +21,5 @@ import zio.schema.codec.BinaryCodec
 private[redis] trait RedisEnvironment {
   protected def codec: BinaryCodec
   protected def executor: RedisExecutor
+  protected def pubSub: RedisPubSub
 }

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

@@ -46,5 +46,6 @@ 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 NoPubSubStream(key: String)       extends RedisError
+  final case class IOError(exception: IOException)   extends RedisError
 }

redis/src/main/scala/zio/redis/RedisExecutor.scala

@@ -24,8 +24,8 @@ trait RedisExecutor {
 
 object RedisExecutor {
   lazy val layer: ZLayer[RedisConfig, RedisError.IOError, RedisExecutor] =
-    RedisConnectionLive.layer >>> SingleNodeExecutor.layer
+    RedisConnectionLive.layer.fresh >>> SingleNodeExecutor.layer
 
   lazy val local: ZLayer[Any, RedisError.IOError, RedisExecutor] =
-    RedisConnectionLive.default >>> SingleNodeExecutor.layer
+    RedisConnectionLive.default.fresh >>> SingleNodeExecutor.layer
 }

redis/src/main/scala/zio/redis/RedisPubSub.scala

@@ -0,0 +1,26 @@
+package zio.redis
+
+import zio.schema.codec.BinaryCodec
+import zio.stream._
+import zio.{Chunk, IO, ZIO, ZLayer}
+
+trait RedisPubSub {
+  def execute(command: PubSubCommand): IO[RedisError, Chunk[Stream[RedisError, RespValue]]]
+}
+
+object RedisPubSub {
+  lazy val layer: ZLayer[RedisConfig with BinaryCodec, RedisError.IOError, RedisPubSub] =
+    RedisConnectionLive.layer.fresh >>> pubSublayer
+
+  lazy val local: ZLayer[BinaryCodec, RedisError.IOError, RedisPubSub] =
+    RedisConnectionLive.default.fresh >>> pubSublayer
+
+  private lazy val pubSublayer: ZLayer[RedisConnection with BinaryCodec, RedisError.IOError, RedisPubSub] =
+    ZLayer.scoped(
+      for {
+        conn   <- ZIO.service[RedisConnection]
+        codec  <- ZIO.service[BinaryCodec]
+        pubSub <- SingleNodeRedisPubSub.create(conn, codec)
+      } yield pubSub
+    )
+}

redis/src/main/scala/zio/redis/RedisPubSubCommand.scala

@@ -0,0 +1,22 @@
+package zio.redis
+
+import zio.redis.Output.ArbitraryOutput
+import zio.schema.Schema
+import zio.schema.codec.BinaryCodec
+import zio.stream._
+import zio.{Chunk, IO, ZIO}
+
+private[redis] final case class RedisPubSubCommand(command: PubSubCommand, codec: BinaryCodec, executor: RedisPubSub) {
+  def run[R: Schema](): IO[RedisError, Chunk[Stream[RedisError, R]]] =
+    executor
+      .execute(command)
+      .map(
+        _.map(
+          _.mapZIO(msg =>
+            ZIO
+              .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec))
+              .refineToOrDie[RedisError]
+          )
+        )
+      )
+}

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] = ???
@@ -46,4 +47,8 @@ object ResultBuilder {
   trait ResultOutputBuilder extends ResultBuilder {
     def returning[R: Output]: IO[RedisError, R]
   }
+
+  trait ResultStreamBuilder[+F[_]] {
+    def returning[R: Schema]: IO[RedisError, F[Stream[RedisError, R]]]
+  }
 }