From a2805041734335972f593a33ec940e14ff3ce24c Mon Sep 17 00:00:00 2001 From: opg1 Date: Wed, 28 Dec 2022 21:19:14 +0900 Subject: [PATCH 01/51] Add redis pubsub layer Fix RedisExecutor structure Add PubSub api Fix broken compile Add PushProtocolOutput suite Fix message field type as generic Fix broken output Fix PushProtocol output spec Add key property in PushProtocol Add test implementation Add PubSub integration test Fix formatting Refactor RedisPubSub Apply RedisPubSub refactoring Apply RedisPubSub refactoring to t/c Remove unused file Fix logic bugs Fix broken t/c Fix unsubscribe process Fix pubSubSpec Add request message broker in SingleNodeRedisPubSub Simplify RedisPubSub's public api Revert unrelated changes --- .../redis/benchmarks/BenchmarkRuntime.scala | 1 + example/src/main/scala/example/Main.scala | 3 +- redis/src/main/scala/zio/redis/Output.scala | 46 ++++ redis/src/main/scala/zio/redis/Redis.scala | 13 +- .../main/scala/zio/redis/RedisExecutor.scala | 4 +- .../main/scala/zio/redis/RedisPubSub.scala | 25 +++ .../scala/zio/redis/RedisPubSubCommand.scala | 19 ++ .../main/scala/zio/redis/ResultBuilder.scala | 5 + .../zio/redis/SingleNodeRedisPubSub.scala | 169 ++++++++++++++ .../main/scala/zio/redis/TestExecutor.scala | 209 +++++++++++++++++- .../src/main/scala/zio/redis/api/PubSub.scala | 89 ++++++++ .../scala/zio/redis/options/Cluster.scala | 6 +- .../main/scala/zio/redis/options/PubSub.scala | 49 ++++ redis/src/main/scala/zio/redis/package.scala | 3 +- redis/src/test/scala/zio/redis/ApiSpec.scala | 18 +- .../scala/zio/redis/ClusterExecutorSpec.scala | 1 + redis/src/test/scala/zio/redis/KeysSpec.scala | 1 + .../src/test/scala/zio/redis/OutputSpec.scala | 88 ++++++++ .../src/test/scala/zio/redis/PubSubSpec.scala | 175 +++++++++++++++ 19 files changed, 903 insertions(+), 21 deletions(-) create mode 100644 redis/src/main/scala/zio/redis/RedisPubSub.scala create mode 100644 redis/src/main/scala/zio/redis/RedisPubSubCommand.scala create mode 100644 redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala create mode 100644 redis/src/main/scala/zio/redis/api/PubSub.scala create mode 100644 redis/src/main/scala/zio/redis/options/PubSub.scala create mode 100644 redis/src/test/scala/zio/redis/PubSubSpec.scala diff --git a/benchmarks/src/main/scala/zio/redis/benchmarks/BenchmarkRuntime.scala b/benchmarks/src/main/scala/zio/redis/benchmarks/BenchmarkRuntime.scala index 59d2b1ff7..c91c41fe9 100644 --- a/benchmarks/src/main/scala/zio/redis/benchmarks/BenchmarkRuntime.scala +++ b/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), RedisLive.layer ) diff --git a/example/src/main/scala/example/Main.scala b/example/src/main/scala/example/Main.scala index 55fb13398..5b07c4b95 100644 --- a/example/src/main/scala/example/Main.scala +++ b/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.{RedisExecutor, RedisLive} +import zio.redis.{RedisExecutor, RedisLive, 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, RedisLive.layer, ZLayer.succeed[BinaryCodec](ProtobufCodec) ) diff --git a/redis/src/main/scala/zio/redis/Output.scala b/redis/src/main/scala/zio/redis/Output.scala index 21a29b3a5..f72db1fa1 100644 --- a/redis/src/main/scala/zio/redis/Output.scala +++ b/redis/src/main/scala/zio/redis/Output.scala @@ -827,4 +827,50 @@ 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" => + val message = values(2) + PushProtocol.Message(key, message) + case "pmessage" => + val channel = MultiStringOutput.unsafeDecode(values(2)) + val message = values(3) + PushProtocol.PMessage(key, channel, message) + 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[NumSubResponse]] { + protected def tryDecode(respValue: RespValue)(implicit codec: BinaryCodec): Chunk[NumSubResponse] = + 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)) + NumSubResponse(channel, numOfSubscription) + }) + case other => throw ProtocolError(s"$other isn't an array") + } + } } diff --git a/redis/src/main/scala/zio/redis/Redis.scala b/redis/src/main/scala/zio/redis/Redis.scala index 282fde35d..4781ee32b 100644 --- a/redis/src/main/scala/zio/redis/Redis.scala +++ b/redis/src/main/scala/zio/redis/Redis.scala @@ -34,11 +34,18 @@ trait Redis with api.Cluster { def codec: BinaryCodec def executor: RedisExecutor + def pubSub: RedisPubSub } -final case class RedisLive(codec: BinaryCodec, executor: RedisExecutor) extends Redis +final case class RedisLive(codec: BinaryCodec, executor: RedisExecutor, pubSub: RedisPubSub) extends Redis object RedisLive { - lazy val layer: URLayer[RedisExecutor with BinaryCodec, Redis] = - ZLayer.fromFunction(RedisLive.apply _) + lazy val layer: URLayer[RedisPubSub with RedisExecutor with BinaryCodec, Redis] = + ZLayer.fromZIO( + for { + codec <- ZIO.service[BinaryCodec] + executor <- ZIO.service[RedisExecutor] + pubSub <- ZIO.service[RedisPubSub] + } yield RedisLive(codec, executor, pubSub) + ) } diff --git a/redis/src/main/scala/zio/redis/RedisExecutor.scala b/redis/src/main/scala/zio/redis/RedisExecutor.scala index 36e88794b..f4fcaba42 100644 --- a/redis/src/main/scala/zio/redis/RedisExecutor.scala +++ b/redis/src/main/scala/zio/redis/RedisExecutor.scala @@ -8,10 +8,10 @@ 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 lazy val test: ULayer[RedisExecutor] = TestExecutor.layer diff --git a/redis/src/main/scala/zio/redis/RedisPubSub.scala b/redis/src/main/scala/zio/redis/RedisPubSub.scala new file mode 100644 index 000000000..36f9951f9 --- /dev/null +++ b/redis/src/main/scala/zio/redis/RedisPubSub.scala @@ -0,0 +1,25 @@ +package zio.redis + +import zio.schema.codec.BinaryCodec +import zio.stream._ +import zio.{ULayer, ZIO, ZLayer} + +trait RedisPubSub { + def execute(command: RedisPubSubCommand): ZStream[BinaryCodec, RedisError, PushProtocol] +} + +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 + + lazy val test: ULayer[RedisPubSub] = + TestExecutor.layer + + private lazy val pubSublayer: ZLayer[RedisConnection with BinaryCodec, RedisError.IOError, RedisPubSub] = + ZLayer.scoped( + ZIO.service[RedisConnection].flatMap(SingleNodeRedisPubSub.create(_)) + ) +} diff --git a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala new file mode 100644 index 000000000..fcfef72e6 --- /dev/null +++ b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala @@ -0,0 +1,19 @@ +package zio.redis + +import zio.ZLayer +import zio.stream.ZStream + +sealed abstract class RedisPubSubCommand + +object RedisPubSubCommand { + case class Subscribe(channel: String, channels: List[String]) extends RedisPubSubCommand + case class PSubscribe(pattern: String, patterns: List[String]) extends RedisPubSubCommand + case class Unsubscribe(channels: List[String]) extends RedisPubSubCommand + case class PUnsubscribe(patterns: List[String]) extends RedisPubSubCommand + + def run(command: RedisPubSubCommand): ZStream[Redis, RedisError, PushProtocol] = + ZStream.serviceWithStream { redis => + val codecLayer = ZLayer.succeed(redis.codec) + redis.pubSub.execute(command).provideLayer(codecLayer) + } +} diff --git a/redis/src/main/scala/zio/redis/ResultBuilder.scala b/redis/src/main/scala/zio/redis/ResultBuilder.scala index 7f89f91e0..c3d1a533c 100644 --- a/redis/src/main/scala/zio/redis/ResultBuilder.scala +++ b/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.ZStream 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 ResultOutputStreamBuilder { + def returning[R: Schema]: ZStream[Redis, RedisError, R] + } } diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala new file mode 100644 index 000000000..cfdbd5915 --- /dev/null +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -0,0 +1,169 @@ +package zio.redis + +import zio.redis.Input.{NonEmptyList, StringInput, Varargs} +import zio.redis.Output.PushProtocolOutput +import zio.redis.SingleNodeRedisPubSub.{Request, RequestQueueSize, True} +import zio.redis.api.PubSub +import zio.schema.codec.BinaryCodec +import zio.stream._ +import zio.{Chunk, ChunkBuilder, Hub, Promise, Queue, Ref, Schedule, UIO, ZIO} + +import scala.reflect.ClassTag + +final class SingleNodeRedisPubSub( + pubSubHubsRef: Ref[Map[SubscriptionKey, Hub[PushProtocol]]], + reqQueue: Queue[Request], + connection: RedisConnection +) extends RedisPubSub { + + def execute(command: RedisPubSubCommand): ZStream[BinaryCodec, RedisError, PushProtocol] = + command match { + case RedisPubSubCommand.Subscribe(channel, channels) => subscribe(channel, channels) + case RedisPubSubCommand.PSubscribe(pattern, patterns) => pSubscribe(pattern, patterns) + case RedisPubSubCommand.Unsubscribe(channels) => unsubscribe(channels) + case RedisPubSubCommand.PUnsubscribe(patterns) => pUnsubscribe(patterns) + } + + private def subscribe( + channel: String, + channels: List[String] + ): ZStream[BinaryCodec, RedisError, PushProtocol] = + makeSubscriptionStream(PubSub.Subscribe, SubscriptionKey.Channel(channel), channels.map(SubscriptionKey.Channel(_))) + + private def pSubscribe( + pattern: String, + patterns: List[String] + ): ZStream[BinaryCodec, RedisError, PushProtocol] = + makeSubscriptionStream( + PubSub.PSubscribe, + SubscriptionKey.Pattern(pattern), + patterns.map(SubscriptionKey.Pattern(_)) + ) + + private def unsubscribe(channels: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = + makeUnsubscriptionStream(PubSub.Unsubscribe, channels.map(SubscriptionKey.Channel(_))) + + private def pUnsubscribe(patterns: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = + makeUnsubscriptionStream(PubSub.PUnsubscribe, patterns.map(SubscriptionKey.Pattern(_))) + + private def makeSubscriptionStream(command: String, key: SubscriptionKey, keys: List[SubscriptionKey]) = + ZStream.unwrap[BinaryCodec, RedisError, PushProtocol]( + ZIO.serviceWithZIO[BinaryCodec] { implicit codec => + for { + promise <- Promise.make[RedisError, Unit] + chunk = StringInput.encode(command) ++ NonEmptyList(StringInput).encode((key.value, keys.map(_.value))) + stream <- makeStream(key :: keys) + _ <- reqQueue.offer(Request(chunk, promise)) + _ <- promise.await + } yield stream + } + ) + + private def makeUnsubscriptionStream[T <: SubscriptionKey: ClassTag](command: String, keys: List[T]) = + ZStream.unwrap[BinaryCodec, RedisError, PushProtocol]( + ZIO.serviceWithZIO[BinaryCodec] { implicit codec => + for { + targets <- if (keys.isEmpty) pubSubHubsRef.get.map(_.keys.collect { case t: T => t }.toList) + else ZIO.succeedNow(keys) + chunk = StringInput.encode(command) ++ Varargs(StringInput).encode(keys.map(_.value)) + promise <- Promise.make[RedisError, Unit] + stream <- makeStream(targets) + _ <- reqQueue.offer(Request(chunk, promise)) + _ <- promise.await + } yield stream + } + ) + + private def makeStream(keys: List[SubscriptionKey]): UIO[Stream[RedisError, PushProtocol]] = + for { + streams <- ZIO.foreach(keys)(getHub(_).map(ZStream.fromHub(_))) + stream = streams.fold(ZStream.empty)(_ merge _) + } yield stream + + private def getHub(key: SubscriptionKey) = { + def makeNewHub = + Hub + .unbounded[PushProtocol] + .tap(hub => pubSubHubsRef.update(_ + (key -> hub))) + + for { + hubs <- pubSubHubsRef.get + hub <- ZIO.fromOption(hubs.get(key)).orElse(makeNewHub) + } yield hub + } + + 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).serialize + } + + val bytes = buffer.result() + + connection + .write(bytes) + .mapError(RedisError.IOError(_)) + .tapBoth( + e => ZIO.foreachDiscard(reqs)(_.promise.fail(e)), + _ => ZIO.foreachDiscard(reqs)(_.promise.succeed(())) + ) + } + + private def receive: ZIO[BinaryCodec, RedisError, Unit] = + ZIO.serviceWithZIO[BinaryCodec] { implicit codec => + connection.read + .mapError(RedisError.IOError(_)) + .via(RespValue.decoder) + .collectSome + .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))) + .refineToOrDie[RedisError] + .foreach(push => getHub(push.key).flatMap(_.offer(push))) + } + + private def resubscribe: ZIO[BinaryCodec, RedisError, Unit] = + ZIO.serviceWithZIO[BinaryCodec] { implicit codec => + def makeCommand(name: String, keys: Set[String]) = + RespValue.Array(StringInput.encode(name) ++ Varargs(StringInput).encode(keys)).serialize + + for { + keySet <- pubSubHubsRef.get.map(_.keySet) + (channels, patterns) = keySet.partition(_.isChannelKey) + _ <- (connection.write(makeCommand(PubSub.Subscribe, channels.map(_.value))).when(channels.nonEmpty) *> + connection.write(makeCommand(PubSub.PSubscribe, patterns.map(_.value))).when(patterns.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: ZIO[BinaryCodec, RedisError, AnyVal] = + ZIO.logTrace(s"$this Executable sender and reader has been started") *> + (send.repeat[BinaryCodec, Long](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")) +} + +object SingleNodeRedisPubSub { + 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) = + for { + hubRef <- Ref.make(Map.empty[SubscriptionKey, Hub[PushProtocol]]) + reqQueue <- Queue.bounded[Request](RequestQueueSize) + pubSub = new SingleNodeRedisPubSub(hubRef, reqQueue, conn) + _ <- pubSub.run.forkScoped + _ <- logScopeFinalizer(s"$pubSub Node PubSub is closed") + } yield pubSub +} diff --git a/redis/src/main/scala/zio/redis/TestExecutor.scala b/redis/src/main/scala/zio/redis/TestExecutor.scala index 1dee64973..38ed91fdc 100644 --- a/redis/src/main/scala/zio/redis/TestExecutor.scala +++ b/redis/src/main/scala/zio/redis/TestExecutor.scala @@ -17,10 +17,15 @@ package zio.redis import zio._ +import zio.redis.Input.StringInput +import zio.redis.Output.PushProtocolOutput import zio.redis.RedisError.ProtocolError import zio.redis.RespValue.{BulkString, bulkString} import zio.redis.TestExecutor._ +import zio.redis.api.PubSub +import zio.schema.codec.BinaryCodec import zio.stm._ +import zio.stream.ZStream import java.nio.file.{FileSystems, Paths} import java.time.Instant @@ -38,8 +43,75 @@ final class TestExecutor private ( randomPick: Int => USTM[Int], hyperLogLogs: TMap[String, Set[String]], hashes: TMap[String, Map[String, String]], - sortedSets: TMap[String, Map[String, Double]] -) extends RedisExecutor { + sortedSets: TMap[String, Map[String, Double]], + pubSubs: TMap[SubscriptionKey, THub[RespValue]] +) extends RedisExecutor + with RedisPubSub { + + def execute(command: RedisPubSubCommand): ZStream[BinaryCodec, RedisError, PushProtocol] = + command match { + case RedisPubSubCommand.Subscribe(channel, channels) => subscribe(channel, channels) + case RedisPubSubCommand.PSubscribe(pattern, patterns) => pSubscribe(pattern, patterns) + case RedisPubSubCommand.Unsubscribe(channels) => unsubscribe(channels) + case RedisPubSubCommand.PUnsubscribe(patterns) => pUnsubscribe(patterns) + } + + private def subscribe(channel: String, channels: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = + pubSubStream(PubSub.Subscribe, (channel :: channels).map(SubscriptionKey.Channel(_)), false) + + private def unsubscribe(channels: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = + ZStream + .unwrap( + (for { + keys <- if (channels.isEmpty) pubSubs.keys.map(_.collect { case t: SubscriptionKey.Channel => t }) + else ZSTM.succeed(channels.map(SubscriptionKey.Channel(_))) + stream = pubSubStream(PubSub.Unsubscribe, keys, true) + } yield stream).commit + ) + + private def pSubscribe(pattern: String, patterns: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = + pubSubStream(PubSub.PSubscribe, (pattern :: patterns).map(SubscriptionKey.Pattern(_)), false) + + private def pUnsubscribe(patterns: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = + ZStream + .unwrap( + (for { + keys <- if (patterns.isEmpty) pubSubs.keys.map(_.collect { case t: SubscriptionKey.Pattern => t }) + else ZSTM.succeed(patterns.map(SubscriptionKey.Pattern(_))) + stream = pubSubStream(PubSub.PUnsubscribe, keys, true) + } yield stream).commit + ) + private def pubSubStream(cmd: String, keys: List[SubscriptionKey], isUnSubs: Boolean) = + ZStream + .unwrap[BinaryCodec, RedisError, PushProtocol]( + ZIO + .clockWith(_.instant) + .flatMap(now => + ZSTM + .serviceWithSTM[BinaryCodec] { implicit codec => + for { + streams <- + ZSTM.foreach(keys) { key => + for { + resp <- runCommand(cmd, StringInput.encode(key.value), now) + hub <- getPubSubHub(key) + queue <- hub.subscribe + _ <- resp match { + case RespValue.ArrayValues(value) => + hub.offer(value) + case other => + ZSTM.fail(RedisError.ProtocolError(s"Invalid pubsub command response $other")) + } + _ <- pubSubs.delete(key).when(isUnSubs) + } yield ZStream + .fromTQueue(queue) + .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp)).refineToOrDie[RedisError]) + } + } yield streams.fold(ZStream.empty)(_ merge _) + } + .commit + ) + ) override def execute(command: Chunk[RespValue.BulkString]): IO[RedisError, RespValue] = for { @@ -71,7 +143,7 @@ final class TestExecutor private ( val timeout = command.tail.last.asString.toInt runBlockingCommand(name.asString, command.tail, timeout, RespValue.NullBulkString, now) - case "CLIENT" | "STRALGO" => + case "CLIENT" | "STRALGO" | "PUBSUB" => val command1 = command.tail val name1 = command1.head runCommand(name.asString + " " + name1.asString, command1.tail, now).commit @@ -104,6 +176,123 @@ final class TestExecutor private ( private[this] def runCommand(name: String, input: Chunk[RespValue.BulkString], now: Instant): USTM[RespValue] = { name match { + case api.PubSub.Publish => + for { + (channels, patterns) <- pubSubs.keys.map(_.partition(_.isChannelKey)) + keyString = input(0).asString + msg = input(1) + targetChannels = channels.filter(_.value == keyString) + messages = + targetChannels.map { ch => + ( + ch, + RespValue.array(RespValue.bulkString("message"), RespValue.bulkString(ch.value), msg) + ) + } + targetPatterns = patterns.filter { pattern => + val matcher = FileSystems.getDefault.getPathMatcher("glob:" + pattern.value) + matcher.matches(Paths.get(keyString)) + } + pMessages = + targetPatterns.map { pattern => + ( + pattern, + RespValue + .array(RespValue.bulkString("pmessage"), RespValue.bulkString(pattern.value), input(0), msg) + ) + } + _ <- ZSTM.foreachDiscard(messages ++ pMessages) { case (key, message) => + getPubSubHub(key).flatMap(_.offer(message)) + } + } yield RespValue.Integer(targetChannels.size + targetPatterns.size.toLong) + + case api.PubSub.PubSubChannels => + for { + channels <- pubSubs.keys.map(_.collect { case t: SubscriptionKey.Channel => t.value }) + pattern = input(0).asString + matcher = FileSystems.getDefault.getPathMatcher("glob:" + pattern) + matchedChannels = channels + .filter(ch => matcher.matches(Paths.get(ch))) + .map(channel => RespValue.bulkString(channel)) + } yield RespValue.Array(Chunk.fromIterable(matchedChannels)) + + case api.PubSub.PubSubNumSub => + for { + channels <- pubSubs.keys.map(_.collect { case t: SubscriptionKey.Channel => t.value }) + keys = input.map(_.asString) + targets = keys.map(key => + RespValue.bulkString(key) -> RespValue.Integer( + if (channels contains key) 1L + else 0L + ) + ) + } yield RespValue.Array(targets.foldLeft(Chunk.empty[RespValue]) { case (acc, (bulk, num)) => + acc.appended(bulk).appended(num) + }) + + case api.PubSub.PubSubNumPat => + pubSubs.keys + .map(_.collect { case t: SubscriptionKey.Pattern => t.value }) + .map(patterns => RespValue.Integer(patterns.size.toLong)) + + case api.PubSub.Unsubscribe => + for { + channels <- pubSubs.keys.map(_.collect { case t: SubscriptionKey.Channel => t }) + keys = input.map(_.asString) + subsCount = channels.size.toLong - 1 + messages = keys.zipWithIndex.map { case (key, idx) => + RespValue.array( + RespValue.bulkString("unsubscribe"), + RespValue.bulkString(key), + RespValue.Integer((subsCount - idx) max 0L) + ) + } + } yield RespValue.Array(messages) + + case api.PubSub.PUnsubscribe => + for { + patterns <- pubSubs.keys.map(_.collect { case t: SubscriptionKey.Pattern => t }) + keys = input.map(_.asString) + subsCount = patterns.size.toLong - 1 + messages = keys.zipWithIndex.map { case (key, idx) => + RespValue.array( + RespValue.bulkString("punsubscribe"), + RespValue.bulkString(key), + RespValue.Integer((subsCount - idx) max 0L) + ) + } + } yield RespValue.Array(messages) + + case api.PubSub.Subscribe => + for { + subsCount <- pubSubs.keys.map(_.size + 1L) + keys = input.map(_.asString) + messages = + keys.zipWithIndex.map { case (key, idx) => + RespValue.array( + RespValue.bulkString("subscribe"), + RespValue.bulkString(key), + RespValue.Integer(subsCount + idx) + ) + } + _ <- ZSTM.foreachDiscard(keys.map(SubscriptionKey.Channel(_)))(getPubSubHub(_)) + } yield RespValue.Array(messages) + + case api.PubSub.PSubscribe => + for { + subsCount <- pubSubs.keys.map(_.size + 1L) + keys = input.map(_.asString) + subsMessages = + keys.zipWithIndex.map { case (key, idx) => + RespValue.array( + RespValue.bulkString("psubscribe"), + RespValue.bulkString(key), + RespValue.Integer(subsCount + idx) + ) + } + _ <- ZSTM.foreachDiscard(keys.map(SubscriptionKey.Pattern(_)))(getPubSubHub(_)) + } yield RespValue.Array(subsMessages) + case api.Connection.Auth => onConnection(name, input)(RespValue.bulkString("OK")) @@ -3582,6 +3771,13 @@ final class TestExecutor private ( } } + private def getPubSubHub(key: SubscriptionKey) = + for { + hubOpt <- pubSubs.get(key) + hub <- ZSTM.fromOption(hubOpt).orElse(THub.unbounded[RespValue]) + _ <- pubSubs.putIfAbsent(key, hub) + } yield hub + private[this] def orWrongType(predicate: USTM[Boolean])( program: => USTM[RespValue] ): USTM[RespValue] = @@ -3922,7 +4118,7 @@ object TestExecutor { lazy val redisType: RedisType = KeyType.toRedisType(`type`) } - lazy val layer: ULayer[RedisExecutor] = + lazy val layer: ULayer[RedisExecutor with RedisPubSub] = ZLayer { for { seed <- ZIO.randomWith(_.nextInt) @@ -3940,6 +4136,7 @@ object TestExecutor { clientTInfo = ClientTrackingInfo(ClientTrackingFlags(clientSideCaching = false), ClientTrackingRedirect.NotEnabled) clientTrackingInfo <- TRef.make(clientTInfo).commit + pubSubs <- TMap.empty[SubscriptionKey, THub[RespValue]].commit } yield new TestExecutor( clientInfo, clientTrackingInfo, @@ -3950,8 +4147,8 @@ object TestExecutor { randomPick, hyperLogLogs, hashes, - sortedSets + sortedSets, + pubSubs ) } - } diff --git a/redis/src/main/scala/zio/redis/api/PubSub.scala b/redis/src/main/scala/zio/redis/api/PubSub.scala new file mode 100644 index 000000000..9e47fcefa --- /dev/null +++ b/redis/src/main/scala/zio/redis/api/PubSub.scala @@ -0,0 +1,89 @@ +package zio.redis.api + +import zio.redis.Input._ +import zio.redis.Output._ +import zio.redis.ResultBuilder.ResultOutputStreamBuilder +import zio.redis._ +import zio.redis.api.PubSub._ +import zio.schema.Schema +import zio.stream._ +import zio.{Chunk, ZIO} + +trait PubSub { + final def subscribeStreamBuilder(channel: String, channels: String*): ResultOutputStreamBuilder = + new ResultOutputStreamBuilder { + override def returning[R: Schema]: ZStream[Redis, RedisError, R] = + ZStream.serviceWithStream[Redis] { redis => + RedisPubSubCommand + .run(RedisPubSubCommand.Subscribe(channel, channels.toList)) + .collect { case t: PushProtocol.Message => t.message } + .mapZIO(resp => + ZIO + .attempt(ArbitraryOutput[R]().unsafeDecode(resp)(redis.codec)) + .refineToOrDie[RedisError] + ) + } + } + + final def subscribe(channel: String, channels: String*): ZStream[Redis, RedisError, PushProtocol] = + RedisPubSubCommand.run(RedisPubSubCommand.Subscribe(channel, channels.toList)) + + final def unsubscribe(channels: String*): ZStream[Redis, RedisError, PushProtocol.Unsubscribe] = + RedisPubSubCommand.run(RedisPubSubCommand.Unsubscribe(channels.toList)).collect { + case t: PushProtocol.Unsubscribe => t + } + + final def pSubscribe(pattern: String, patterns: String*): ZStream[Redis, RedisError, PushProtocol] = + RedisPubSubCommand.run(RedisPubSubCommand.PSubscribe(pattern, patterns.toList)) + + final def pSubscribeStreamBuilder(pattern: String, patterns: String*): ResultOutputStreamBuilder = + new ResultOutputStreamBuilder { + override def returning[R: Schema]: ZStream[Redis, RedisError, R] = + ZStream.serviceWithStream[Redis] { redis => + RedisPubSubCommand + .run(RedisPubSubCommand.PSubscribe(pattern, patterns.toList)) + .collect { case t: PushProtocol.PMessage => t.message } + .mapZIO(resp => + ZIO + .attempt(ArbitraryOutput[R]().unsafeDecode(resp)(redis.codec)) + .refineToOrDie[RedisError] + ) + } + } + + final def pUnsubscribe(patterns: String*): ZStream[Redis, RedisError, PushProtocol.PUnsubscribe] = + RedisPubSubCommand.run(RedisPubSubCommand.PUnsubscribe(patterns.toList)).collect { + case t: PushProtocol.PUnsubscribe => t + } + + final def publish[A: Schema](channel: String, message: A): ZIO[Redis, RedisError, Long] = { + val command = RedisCommand(Publish, Tuple2(StringInput, ArbitraryInput[A]()), LongOutput) + command.run((channel, message)) + } + + final def pubSubChannels(pattern: String): ZIO[Redis, RedisError, Chunk[String]] = { + val command = RedisCommand(PubSubChannels, StringInput, ChunkOutput(MultiStringOutput)) + command.run(pattern) + } + + final def pubSubNumPat: ZIO[Redis, RedisError, Long] = { + val command = RedisCommand(PubSubNumPat, NoInput, LongOutput) + command.run(()) + } + + final def pubSubNumSub(channel: String, channels: String*): ZIO[Redis, RedisError, Chunk[NumSubResponse]] = { + val command = RedisCommand(PubSubNumSub, NonEmptyList(StringInput), NumSubResponseOutput) + command.run((channel, channels.toList)) + } +} + +private[redis] object PubSub { + final val Subscribe = "SUBSCRIBE" + final val Unsubscribe = "UNSUBSCRIBE" + final val PSubscribe = "PSUBSCRIBE" + final val PUnsubscribe = "PUNSUBSCRIBE" + final val Publish = "PUBLISH" + final val PubSubChannels = "PUBSUB CHANNELS" + final val PubSubNumPat = "PUBSUB NUMPAT" + final val PubSubNumSub = "PUBSUB NUMSUB" +} diff --git a/redis/src/main/scala/zio/redis/options/Cluster.scala b/redis/src/main/scala/zio/redis/options/Cluster.scala index 5c11752c4..35bf83aa8 100644 --- a/redis/src/main/scala/zio/redis/options/Cluster.scala +++ b/redis/src/main/scala/zio/redis/options/Cluster.scala @@ -15,14 +15,14 @@ */ package zio.redis.options -import zio.redis.{RedisExecutor, RedisUri} +import zio.redis.{RedisExecutor, RedisPubSub, RedisUri} import zio.{Chunk, Scope} object Cluster { private[redis] final val SlotsAmount = 16384 - final case class ExecutorScope(executor: RedisExecutor, scope: Scope.Closeable) + final case class ExecutorScope(executor: RedisExecutor, pubSub: RedisPubSub, scope: Scope.Closeable) final case class ClusterConnection( partitions: Chunk[Partition], @@ -31,6 +31,8 @@ object Cluster { ) { def executor(slot: Slot): Option[RedisExecutor] = executors.get(slots(slot)).map(_.executor) + def pubSub(slot: Slot): Option[RedisPubSub] = executors.get(slots(slot)).map(_.pubSub) + def addExecutor(uri: RedisUri, es: ExecutorScope): ClusterConnection = copy(executors = executors + (uri -> es)) } diff --git a/redis/src/main/scala/zio/redis/options/PubSub.scala b/redis/src/main/scala/zio/redis/options/PubSub.scala new file mode 100644 index 000000000..f8915607c --- /dev/null +++ b/redis/src/main/scala/zio/redis/options/PubSub.scala @@ -0,0 +1,49 @@ +package zio.redis.options + +import zio.redis.RespValue + +trait PubSub { + + sealed trait SubscriptionKey { self => + def value: String + + def isChannelKey = self match { + case _: SubscriptionKey.Channel => true + case _: SubscriptionKey.Pattern => false + } + + def isPatternKey = !isChannelKey + } + + object SubscriptionKey { + case class Channel(value: String) extends SubscriptionKey + case class Pattern(value: String) extends SubscriptionKey + } + + case class NumSubResponse(channel: String, subscriberCount: Long) + + sealed trait PushProtocol { + def key: SubscriptionKey + } + + object PushProtocol { + case class Subscribe(channel: String, numOfSubscription: Long) extends PushProtocol { + def key: SubscriptionKey = SubscriptionKey.Channel(channel) + } + case class PSubscribe(pattern: String, numOfSubscription: Long) extends PushProtocol { + def key: SubscriptionKey = SubscriptionKey.Pattern(pattern) + } + case class Unsubscribe(channel: String, numOfSubscription: Long) extends PushProtocol { + def key: SubscriptionKey = SubscriptionKey.Channel(channel) + } + case class PUnsubscribe(pattern: String, numOfSubscription: Long) extends PushProtocol { + def key: SubscriptionKey = SubscriptionKey.Pattern(pattern) + } + case class Message(channel: String, message: RespValue) extends PushProtocol { + def key: SubscriptionKey = SubscriptionKey.Channel(channel) + } + case class PMessage(pattern: String, channel: String, message: RespValue) extends PushProtocol { + def key: SubscriptionKey = SubscriptionKey.Pattern(pattern) + } + } +} diff --git a/redis/src/main/scala/zio/redis/package.scala b/redis/src/main/scala/zio/redis/package.scala index 948ec1623..6dd5ef9f6 100644 --- a/redis/src/main/scala/zio/redis/package.scala +++ b/redis/src/main/scala/zio/redis/package.scala @@ -25,7 +25,8 @@ package object redis with options.Strings with options.Lists with options.Streams - with options.Scripting { + with options.Scripting + with options.PubSub { type Id[+A] = A diff --git a/redis/src/test/scala/zio/redis/ApiSpec.scala b/redis/src/test/scala/zio/redis/ApiSpec.scala index 58fe1d456..4a445edf7 100644 --- a/redis/src/test/scala/zio/redis/ApiSpec.scala +++ b/redis/src/test/scala/zio/redis/ApiSpec.scala @@ -16,7 +16,8 @@ object ApiSpec with HashSpec with StreamsSpec with ScriptingSpec - with ClusterSpec { + with ClusterSpec + with PubSubSpec { def spec: Spec[TestEnvironment, Any] = suite("Redis commands")( @@ -38,10 +39,12 @@ object ApiSpec hyperLogLogSuite, hashSuite, streamsSuite, - scriptingSpec + scriptingSpec, + pubSubSuite ) - val Layer: Layer[Any, Redis] = ZLayer.make[Redis](RedisExecutor.local.orDie, ZLayer.succeed(codec), RedisLive.layer) + val Layer: Layer[Any, Redis] = + ZLayer.make[Redis](RedisExecutor.local.orDie, RedisPubSub.local.orDie, ZLayer.succeed(codec), RedisLive.layer) } private object Test { @@ -55,10 +58,12 @@ object ApiSpec hashSuite, sortedSetsSuite, geoSuite, - stringsSuite + stringsSuite, + pubSubSuite ).filterAnnotations(TestAnnotation.tagged)(t => !t.contains(BaseSpec.TestExecutorUnsupported)).get - val Layer: Layer[Any, Redis] = ZLayer.make[Redis](RedisExecutor.test, ZLayer.succeed(codec), RedisLive.layer) + val Layer: Layer[Any, Redis] = + ZLayer.make[Redis](RedisExecutor.test, RedisPubSub.test, ZLayer.succeed(codec), RedisLive.layer) } private object Cluster { @@ -75,7 +80,8 @@ object ApiSpec geoSuite, streamsSuite @@ clusterExecutorUnsupported, scriptingSpec @@ clusterExecutorUnsupported, - clusterSpec + clusterSpec, + pubSubSuite ).filterAnnotations(TestAnnotation.tagged)(t => !t.contains(BaseSpec.ClusterExecutorUnsupported)).get val Layer: Layer[Any, Redis] = diff --git a/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala b/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala index ab5b9a79f..502584458 100644 --- a/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala +++ b/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala @@ -69,6 +69,7 @@ object ClusterExecutorSpec extends BaseSpec { ZLayer.make[Redis]( ZLayer.succeed(RedisConfig(uri.host, uri.port)), RedisExecutor.layer, + RedisPubSub.layer, ZLayer.succeed(codec), RedisLive.layer ) diff --git a/redis/src/test/scala/zio/redis/KeysSpec.scala b/redis/src/test/scala/zio/redis/KeysSpec.scala index 5e3105082..c5ce10c2a 100644 --- a/redis/src/test/scala/zio/redis/KeysSpec.scala +++ b/redis/src/test/scala/zio/redis/KeysSpec.scala @@ -464,6 +464,7 @@ object KeysSpec { ZLayer.succeed(RedisConfig("localhost", 6380)), RedisConnectionLive.layer, SingleNodeExecutor.layer, + RedisPubSub.layer, ZLayer.succeed[BinaryCodec](ProtobufCodec), RedisLive.layer ) diff --git a/redis/src/test/scala/zio/redis/OutputSpec.scala b/redis/src/test/scala/zio/redis/OutputSpec.scala index e305130cb..eceae9120 100644 --- a/redis/src/test/scala/zio/redis/OutputSpec.scala +++ b/redis/src/test/scala/zio/redis/OutputSpec.scala @@ -1035,6 +1035,94 @@ object OutputSpec extends BaseSpec { res <- ZIO.attempt(ClientTrackingRedirectOutput.unsafeDecode(resp)).either } yield assert(res)(isLeft(isSubtype[ProtocolError](anything))) } + ), + suite("PushProtocol")( + test("subscribe") { + val channel = "foo" + val numOfSubscription = 1L + val input = + RespValue.array( + RespValue.bulkString("subscribe"), + RespValue.bulkString(channel), + RespValue.Integer(numOfSubscription) + ) + val expected = PushProtocol.Subscribe(channel, numOfSubscription) + assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + equalTo(expected) + ) + }, + test("psubscribe") { + val pattern = "f*" + val numOfSubscription = 1L + val input = + RespValue.array( + RespValue.bulkString("psubscribe"), + RespValue.bulkString(pattern), + RespValue.Integer(numOfSubscription) + ) + val expected = PushProtocol.PSubscribe(pattern, numOfSubscription) + assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + equalTo(expected) + ) + }, + test("unsubscribe") { + val channel = "foo" + val numOfSubscription = 1L + val input = + RespValue.array( + RespValue.bulkString("unsubscribe"), + RespValue.bulkString(channel), + RespValue.Integer(numOfSubscription) + ) + val expected = PushProtocol.Unsubscribe(channel, numOfSubscription) + assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + equalTo(expected) + ) + }, + test("punsubscribe") { + val pattern = "f*" + val numOfSubscription = 1L + val input = + RespValue.array( + RespValue.bulkString("punsubscribe"), + RespValue.bulkString(pattern), + RespValue.Integer(numOfSubscription) + ) + val expected = PushProtocol.PUnsubscribe(pattern, numOfSubscription) + assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + equalTo(expected) + ) + }, + test("message") { + val channel = "foo" + val message = RespValue.bulkString("bar") + val input = + RespValue.array( + RespValue.bulkString("message"), + RespValue.bulkString(channel), + message + ) + val expected = PushProtocol.Message(channel, message) + assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + equalTo(expected) + ) + }, + test("pmessage") { + val pattern = "f*" + val channel = "foo" + val message = RespValue.bulkString("bar") + val input = + RespValue.array( + RespValue.bulkString("pmessage"), + RespValue.bulkString(pattern), + RespValue.bulkString(channel), + message + ) + val expected = PushProtocol.PMessage(pattern, channel, message) + assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + equalTo(expected) + ) + } ) ) diff --git a/redis/src/test/scala/zio/redis/PubSubSpec.scala b/redis/src/test/scala/zio/redis/PubSubSpec.scala new file mode 100644 index 000000000..f1dcb38d0 --- /dev/null +++ b/redis/src/test/scala/zio/redis/PubSubSpec.scala @@ -0,0 +1,175 @@ +package zio.redis + +import zio.test.Assertion._ +import zio.test._ +import zio.{Chunk, ZIO} + +import scala.util.Random + +trait PubSubSpec extends BaseSpec { + def pubSubSuite: Spec[Redis, RedisError] = + suite("pubSubs")( + suite("subscribe")( + test("subscribe response") { + for { + channel <- generateRandomString() + res <- subscribe(channel).runHead + } yield assertTrue(res.get.key == SubscriptionKey.Channel(channel)) + }, + test("message response") { + for { + channel <- generateRandomString() + message = "bar" + stream <- subscribeStreamBuilder(channel) + .returning[String] + .runHead + .fork + _ <- pubSubChannels(channel) + .repeatUntil(_ contains channel) + _ <- publish(channel, message) + res <- stream.join + } yield assertTrue(res.get == message) + }, + test("multiple subscribe") { + val numOfPublish = 20 + for { + prefix <- generateRandomString(5) + channel1 <- generateRandomString(prefix) + channel2 <- generateRandomString(prefix) + pattern = prefix + '*' + message <- generateRandomString(5) + stream1 <- subscribe(channel1) + .runFoldWhile(Chunk.empty[PushProtocol])( + _.forall(_.isInstanceOf[PushProtocol.Unsubscribe] == false) + )(_ appended _) + .fork + stream2 <- subscribe(channel2) + .runFoldWhile(Chunk.empty[PushProtocol])( + _.forall(_.isInstanceOf[PushProtocol.Unsubscribe] == false) + )(_ appended _) + .fork + _ <- pubSubChannels(pattern) + .repeatUntil(channels => channels.size >= 2) + ch1SubsCount <- publish(channel1, message).replicateZIO(numOfPublish).map(_.head) + ch2SubsCount <- publish(channel2, message).replicateZIO(numOfPublish).map(_.head) + _ <- unsubscribe().runDrain.fork + res1 <- stream1.join + res2 <- stream2.join + } yield assertTrue(ch1SubsCount == 1L) && + assertTrue(ch2SubsCount == 1L) && + assertTrue(res1.size == numOfPublish + 2) && + assertTrue(res2.size == numOfPublish + 2) + }, + test("psubscribe response") { + for { + pattern <- generateRandomString() + res <- pSubscribe(pattern).runHead + } yield assertTrue(res.get.key.value == pattern) + }, + test("pmessage response") { + for { + prefix <- generateRandomString(5) + pattern = prefix + '*' + channel <- generateRandomString(prefix) + message <- generateRandomString(prefix) + stream <- pSubscribeStreamBuilder(pattern) + .returning[String] + .runHead + .fork + _ <- pubSubNumPat.repeatUntil(_ > 0) + _ <- publish(channel, message) + res <- stream.join + } yield assertTrue(res.get == message) + } + ), + suite("publish")(test("publish long type message") { + val message = 1L + assertZIO( + for { + channel <- generateRandomString() + stream <- subscribeStreamBuilder(channel) + .returning[Long] + .runFoldWhile(0L)(_ < 10L) { case (sum, message) => + sum + message + } + .fork + _ <- pubSubChannels(channel).repeatUntil(_ contains channel) + _ <- ZIO.replicateZIO(10)(publish(channel, message)) + res <- stream.join + } yield res + )(equalTo(10L)) + }), + suite("unsubscribe")( + test("don't receive message type after unsubscribe") { + val numOfPublished = 5 + for { + prefix <- generateRandomString(5) + pattern = prefix + '*' + channel <- generateRandomString(prefix) + message <- generateRandomString() + stream <- subscribe(channel) + .runFoldWhile(Chunk.empty[PushProtocol])(_.size < 2)(_ appended _) + .fork + _ <- pubSubChannels(pattern) + .repeatUntil(_ contains channel) + _ <- unsubscribe(channel).runHead + receiverCount <- publish(channel, message).replicateZIO(numOfPublished).map(_.head) + res <- stream.join + } yield assertTrue( + res.size == 2 + ) && assertTrue(receiverCount == 0L) + }, + test("unsubscribe response") { + for { + channel <- generateRandomString() + res <- unsubscribe(channel).runHead + } yield assertTrue(res.get.key.value == channel) + }, + test("punsubscribe response") { + for { + pattern <- generateRandomString() + res <- pUnsubscribe(pattern).runHead + } yield assertTrue(res.get.key.value == pattern) + }, + test("unsubscribe with empty param") { + for { + prefix <- generateRandomString(5) + pattern = prefix + '*' + channel1 <- generateRandomString(prefix) + channel2 <- generateRandomString(prefix) + stream1 <- + subscribe(channel1) + .runFoldWhile(Chunk.empty[PushProtocol])(_.forall(_.isInstanceOf[PushProtocol.Unsubscribe] == false))( + _ appended _ + ) + .fork + stream2 <- + subscribe(channel2) + .runFoldWhile(Chunk.empty[PushProtocol])(_.forall(_.isInstanceOf[PushProtocol.Unsubscribe] == false))( + _ appended _ + ) + .fork + _ <- pubSubChannels(pattern) + .repeatUntil(_.size >= 2) + _ <- unsubscribe().runDrain.fork + unsubscribeMessages <- stream1.join zip stream2.join + (result1, result2) = unsubscribeMessages + numSubResponses <- pubSubNumSub(channel1, channel2) + } yield assertTrue( + result1.size == 2 && result2.size == 2 + ) && assertTrue( + numSubResponses == Chunk( + NumSubResponse(channel1, 0L), + NumSubResponse(channel2, 0L) + ) + ) + } + ) + ) + + private def generateRandomString(prefix: String = "") = + ZIO.succeed(Random.alphanumeric.take(15).mkString).map(prefix + _.substring((prefix.length - 1) max 0)) + + private def generateRandomString(len: Int) = + ZIO.succeed(Random.alphanumeric.take(len).mkString) +} From cd4a3f3d75206cf9b47b0fac49e8356ef6845d99 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 7 Feb 2023 16:33:42 +0900 Subject: [PATCH 02/51] Fix PubSub api --- example/src/main/scala/example/Main.scala | 1 - .../scala/zio/redis/ClusterExecutor.scala | 8 +- .../main/scala/zio/redis/PubSubCommand.scala | 21 ++ redis/src/main/scala/zio/redis/Redis.scala | 3 +- .../scala/zio/redis/RedisEnvironment.scala | 1 + .../src/main/scala/zio/redis/RedisError.scala | 1 + .../main/scala/zio/redis/RedisPubSub.scala | 2 +- .../scala/zio/redis/RedisPubSubCommand.scala | 19 -- .../main/scala/zio/redis/ResultBuilder.scala | 6 +- .../zio/redis/SingleNodeRedisPubSub.scala | 108 +++++---- .../main/scala/zio/redis/TestExecutor.scala | 132 ++++++----- .../src/main/scala/zio/redis/api/PubSub.scala | 213 ++++++++++++++---- .../scala/zio/redis/options/Cluster.scala | 6 +- .../main/scala/zio/redis/options/PubSub.scala | 4 +- 14 files changed, 340 insertions(+), 185 deletions(-) create mode 100644 redis/src/main/scala/zio/redis/PubSubCommand.scala delete mode 100644 redis/src/main/scala/zio/redis/RedisPubSubCommand.scala diff --git a/example/src/main/scala/example/Main.scala b/example/src/main/scala/example/Main.scala index 5b07c4b95..261536984 100644 --- a/example/src/main/scala/example/Main.scala +++ b/example/src/main/scala/example/Main.scala @@ -33,7 +33,6 @@ object Main extends ZIOAppDefault { ContributorsCache.layer, HttpClientZioBackend.layer(), RedisExecutor.layer, - RedisPubSub.layer, RedisLive.layer, ZLayer.succeed[BinaryCodec](ProtobufCodec) ) diff --git a/redis/src/main/scala/zio/redis/ClusterExecutor.scala b/redis/src/main/scala/zio/redis/ClusterExecutor.scala index 2b530c357..4003a06c4 100644 --- a/redis/src/main/scala/zio/redis/ClusterExecutor.scala +++ b/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 >>> RedisLive.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 >>> RedisLive.layer for { closableScope <- Scope.make layer <- closableScope.extend[Any](redisLayer.memoize) diff --git a/redis/src/main/scala/zio/redis/PubSubCommand.scala b/redis/src/main/scala/zio/redis/PubSubCommand.scala new file mode 100644 index 000000000..54b4dfdc3 --- /dev/null +++ b/redis/src/main/scala/zio/redis/PubSubCommand.scala @@ -0,0 +1,21 @@ +package zio.redis + +import zio.schema.codec.BinaryCodec +import zio.stream._ +import zio.{IO, ZLayer} + +final case class RedisPubSubCommand(command: PubSubCommand, codec: BinaryCodec, executor: RedisPubSub) { + def run: IO[RedisError, List[Stream[RedisError, PushProtocol]]] = { + val codecLayer = ZLayer.succeed(codec) + executor.execute(command).provideLayer(codecLayer) + } +} + +sealed trait PubSubCommand + +object PubSubCommand { + case class Subscribe(channel: String, channels: List[String]) extends PubSubCommand + case class PSubscribe(pattern: String, patterns: List[String]) extends PubSubCommand + case class Unsubscribe(channels: List[String]) extends PubSubCommand + case class PUnsubscribe(patterns: List[String]) extends PubSubCommand +} diff --git a/redis/src/main/scala/zio/redis/Redis.scala b/redis/src/main/scala/zio/redis/Redis.scala index 4781ee32b..be20da62c 100644 --- a/redis/src/main/scala/zio/redis/Redis.scala +++ b/redis/src/main/scala/zio/redis/Redis.scala @@ -31,7 +31,8 @@ trait Redis with api.SortedSets with api.Streams with api.Scripting - with api.Cluster { + with api.Cluster + with api.PubSub { def codec: BinaryCodec def executor: RedisExecutor def pubSub: RedisPubSub diff --git a/redis/src/main/scala/zio/redis/RedisEnvironment.scala b/redis/src/main/scala/zio/redis/RedisEnvironment.scala index ece66e6c8..e3aa0c6e9 100644 --- a/redis/src/main/scala/zio/redis/RedisEnvironment.scala +++ b/redis/src/main/scala/zio/redis/RedisEnvironment.scala @@ -5,4 +5,5 @@ import zio.schema.codec.BinaryCodec private[redis] trait RedisEnvironment { def codec: BinaryCodec def executor: RedisExecutor + def pubSub: RedisPubSub } diff --git a/redis/src/main/scala/zio/redis/RedisError.scala b/redis/src/main/scala/zio/redis/RedisError.scala index a09eebca0..a0d042b73 100644 --- a/redis/src/main/scala/zio/redis/RedisError.scala +++ b/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 NoPubSubStream(key: String) extends RedisError final case class IOError(exception: IOException) extends RedisError } diff --git a/redis/src/main/scala/zio/redis/RedisPubSub.scala b/redis/src/main/scala/zio/redis/RedisPubSub.scala index 36f9951f9..734756bf9 100644 --- a/redis/src/main/scala/zio/redis/RedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/RedisPubSub.scala @@ -5,7 +5,7 @@ import zio.stream._ import zio.{ULayer, ZIO, ZLayer} trait RedisPubSub { - def execute(command: RedisPubSubCommand): ZStream[BinaryCodec, RedisError, PushProtocol] + def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] } object RedisPubSub { diff --git a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala deleted file mode 100644 index fcfef72e6..000000000 --- a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala +++ /dev/null @@ -1,19 +0,0 @@ -package zio.redis - -import zio.ZLayer -import zio.stream.ZStream - -sealed abstract class RedisPubSubCommand - -object RedisPubSubCommand { - case class Subscribe(channel: String, channels: List[String]) extends RedisPubSubCommand - case class PSubscribe(pattern: String, patterns: List[String]) extends RedisPubSubCommand - case class Unsubscribe(channels: List[String]) extends RedisPubSubCommand - case class PUnsubscribe(patterns: List[String]) extends RedisPubSubCommand - - def run(command: RedisPubSubCommand): ZStream[Redis, RedisError, PushProtocol] = - ZStream.serviceWithStream { redis => - val codecLayer = ZLayer.succeed(redis.codec) - redis.pubSub.execute(command).provideLayer(codecLayer) - } -} diff --git a/redis/src/main/scala/zio/redis/ResultBuilder.scala b/redis/src/main/scala/zio/redis/ResultBuilder.scala index c3d1a533c..db22b22ab 100644 --- a/redis/src/main/scala/zio/redis/ResultBuilder.scala +++ b/redis/src/main/scala/zio/redis/ResultBuilder.scala @@ -19,7 +19,7 @@ package zio.redis import zio.IO import zio.redis.ResultBuilder.NeedsReturnType import zio.schema.Schema -import zio.stream.ZStream +import zio.stream.Stream sealed trait ResultBuilder { final def map(f: Nothing => Any)(implicit nrt: NeedsReturnType): IO[Nothing, Nothing] = ??? @@ -48,7 +48,7 @@ object ResultBuilder { def returning[R: Output]: IO[RedisError, R] } - trait ResultOutputStreamBuilder { - def returning[R: Schema]: ZStream[Redis, RedisError, R] + trait ResultStreamBuilder[+F[_]] { + def returning[R: Schema]: IO[RedisError, F[Stream[RedisError, R]]] } } diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala index cfdbd5915..8d40fcce5 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -8,77 +8,99 @@ import zio.schema.codec.BinaryCodec import zio.stream._ import zio.{Chunk, ChunkBuilder, Hub, Promise, Queue, Ref, Schedule, UIO, ZIO} -import scala.reflect.ClassTag - final class SingleNodeRedisPubSub( pubSubHubsRef: Ref[Map[SubscriptionKey, Hub[PushProtocol]]], reqQueue: Queue[Request], connection: RedisConnection ) extends RedisPubSub { - def execute(command: RedisPubSubCommand): ZStream[BinaryCodec, RedisError, PushProtocol] = + def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = command match { - case RedisPubSubCommand.Subscribe(channel, channels) => subscribe(channel, channels) - case RedisPubSubCommand.PSubscribe(pattern, patterns) => pSubscribe(pattern, patterns) - case RedisPubSubCommand.Unsubscribe(channels) => unsubscribe(channels) - case RedisPubSubCommand.PUnsubscribe(patterns) => pUnsubscribe(patterns) + case PubSubCommand.Subscribe(channel, channels) => subscribe(channel, channels) + case PubSubCommand.PSubscribe(pattern, patterns) => pSubscribe(pattern, patterns) + case PubSubCommand.Unsubscribe(channels) => unsubscribe(channels) + case PubSubCommand.PUnsubscribe(patterns) => pUnsubscribe(patterns) } private def subscribe( channel: String, channels: List[String] - ): ZStream[BinaryCodec, RedisError, PushProtocol] = + ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = makeSubscriptionStream(PubSub.Subscribe, SubscriptionKey.Channel(channel), channels.map(SubscriptionKey.Channel(_))) private def pSubscribe( pattern: String, patterns: List[String] - ): ZStream[BinaryCodec, RedisError, PushProtocol] = + ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = makeSubscriptionStream( PubSub.PSubscribe, SubscriptionKey.Pattern(pattern), patterns.map(SubscriptionKey.Pattern(_)) ) - private def unsubscribe(channels: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = - makeUnsubscriptionStream(PubSub.Unsubscribe, channels.map(SubscriptionKey.Channel(_))) + private def unsubscribe( + channels: List[String] + ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = + makeUnsubscriptionStream( + PubSub.Unsubscribe, + if (channels.nonEmpty) + ZIO.succeedNow(channels.map(SubscriptionKey.Channel(_))) + else + pubSubHubsRef.get.map(_.keys.filter(_.isChannelKey).toList) + ) - private def pUnsubscribe(patterns: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = - makeUnsubscriptionStream(PubSub.PUnsubscribe, patterns.map(SubscriptionKey.Pattern(_))) + private def pUnsubscribe( + patterns: List[String] + ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = + makeUnsubscriptionStream( + PubSub.PUnsubscribe, + if (patterns.nonEmpty) + ZIO.succeedNow(patterns.map(SubscriptionKey.Pattern(_))) + else + pubSubHubsRef.get.map(_.keys.filter(_.isPatternKey).toList) + ) private def makeSubscriptionStream(command: String, key: SubscriptionKey, keys: List[SubscriptionKey]) = - ZStream.unwrap[BinaryCodec, RedisError, PushProtocol]( - ZIO.serviceWithZIO[BinaryCodec] { implicit codec => - for { - promise <- Promise.make[RedisError, Unit] - chunk = StringInput.encode(command) ++ NonEmptyList(StringInput).encode((key.value, keys.map(_.value))) - stream <- makeStream(key :: keys) - _ <- reqQueue.offer(Request(chunk, promise)) - _ <- promise.await - } yield stream - } - ) + ZIO.serviceWithZIO[BinaryCodec] { implicit codec => + for { + promise <- Promise.make[RedisError, Unit] + chunk = StringInput.encode(command) ++ NonEmptyList(StringInput).encode((key.value, keys.map(_.value))) + streams <- makeStreams(key :: keys) + _ <- reqQueue.offer(Request(chunk, promise)) + _ <- promise.await + } yield streams + } - private def makeUnsubscriptionStream[T <: SubscriptionKey: ClassTag](command: String, keys: List[T]) = - ZStream.unwrap[BinaryCodec, RedisError, PushProtocol]( - ZIO.serviceWithZIO[BinaryCodec] { implicit codec => - for { - targets <- if (keys.isEmpty) pubSubHubsRef.get.map(_.keys.collect { case t: T => t }.toList) - else ZIO.succeedNow(keys) - chunk = StringInput.encode(command) ++ Varargs(StringInput).encode(keys.map(_.value)) - promise <- Promise.make[RedisError, Unit] - stream <- makeStream(targets) - _ <- reqQueue.offer(Request(chunk, promise)) - _ <- promise.await - } yield stream - } - ) + private def makeUnsubscriptionStream(command: String, keys: UIO[List[SubscriptionKey]]) = { + def releaseHub(key: SubscriptionKey) = + for { + pubSubs <- pubSubHubsRef.get + hubOpt = pubSubs.get(key) + _ <- ZIO.fromOption(hubOpt).flatMap(_.shutdown).ignore + _ <- pubSubHubsRef.update(_ - key) + } yield () - private def makeStream(keys: List[SubscriptionKey]): UIO[Stream[RedisError, PushProtocol]] = - for { - streams <- ZIO.foreach(keys)(getHub(_).map(ZStream.fromHub(_))) - stream = streams.fold(ZStream.empty)(_ merge _) - } yield stream + ZIO.serviceWithZIO[BinaryCodec] { implicit codec => + for { + targets <- keys + chunk = StringInput.encode(command) ++ Varargs(StringInput).encode(targets.map(_.value)) + promise <- Promise.make[RedisError, Unit] + streams <- makeStreams(targets) + targetSet = targets.map(_.value).toSet + _ <- reqQueue.offer(Request(chunk, promise)) + _ <- promise.await + } yield streams.map(_.tap { + case PushProtocol.Unsubscribe(channel, _) if targetSet contains channel => + releaseHub(SubscriptionKey.Channel(channel)) + case PushProtocol.PUnsubscribe(pattern, _) if targetSet contains pattern => + releaseHub(SubscriptionKey.Pattern(pattern)) + case _ => ZIO.unit + }) + } + } + + private def makeStreams(keys: List[SubscriptionKey]): UIO[List[Stream[RedisError, PushProtocol]]] = + ZIO.foreach(keys)(getHub(_).map(ZStream.fromHub(_))) private def getHub(key: SubscriptionKey) = { def makeNewHub = diff --git a/redis/src/main/scala/zio/redis/TestExecutor.scala b/redis/src/main/scala/zio/redis/TestExecutor.scala index 38ed91fdc..93af54973 100644 --- a/redis/src/main/scala/zio/redis/TestExecutor.scala +++ b/redis/src/main/scala/zio/redis/TestExecutor.scala @@ -25,7 +25,7 @@ import zio.redis.TestExecutor._ import zio.redis.api.PubSub import zio.schema.codec.BinaryCodec import zio.stm._ -import zio.stream.ZStream +import zio.stream._ import java.nio.file.{FileSystems, Paths} import java.time.Instant @@ -48,70 +48,84 @@ final class TestExecutor private ( ) extends RedisExecutor with RedisPubSub { - def execute(command: RedisPubSubCommand): ZStream[BinaryCodec, RedisError, PushProtocol] = + def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = command match { - case RedisPubSubCommand.Subscribe(channel, channels) => subscribe(channel, channels) - case RedisPubSubCommand.PSubscribe(pattern, patterns) => pSubscribe(pattern, patterns) - case RedisPubSubCommand.Unsubscribe(channels) => unsubscribe(channels) - case RedisPubSubCommand.PUnsubscribe(patterns) => pUnsubscribe(patterns) + case PubSubCommand.Subscribe(channel, channels) => subscribe(channel, channels) + case PubSubCommand.PSubscribe(pattern, patterns) => pSubscribe(pattern, patterns) + case PubSubCommand.Unsubscribe(channels) => unsubscribe(channels) + case PubSubCommand.PUnsubscribe(patterns) => pUnsubscribe(patterns) } - private def subscribe(channel: String, channels: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = - pubSubStream(PubSub.Subscribe, (channel :: channels).map(SubscriptionKey.Channel(_)), false) - - private def unsubscribe(channels: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = - ZStream - .unwrap( - (for { - keys <- if (channels.isEmpty) pubSubs.keys.map(_.collect { case t: SubscriptionKey.Channel => t }) - else ZSTM.succeed(channels.map(SubscriptionKey.Channel(_))) - stream = pubSubStream(PubSub.Unsubscribe, keys, true) - } yield stream).commit + private def subscribe( + channel: String, + channels: List[String] + ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = + ZIO.clockWith( + _.instant.flatMap( + pubSubStream(PubSub.Subscribe, (channel :: channels).map(SubscriptionKey.Channel(_)), _, false).commit ) + ) - private def pSubscribe(pattern: String, patterns: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = - pubSubStream(PubSub.PSubscribe, (pattern :: patterns).map(SubscriptionKey.Pattern(_)), false) - - private def pUnsubscribe(patterns: List[String]): ZStream[BinaryCodec, RedisError, PushProtocol] = - ZStream - .unwrap( - (for { - keys <- if (patterns.isEmpty) pubSubs.keys.map(_.collect { case t: SubscriptionKey.Pattern => t }) - else ZSTM.succeed(patterns.map(SubscriptionKey.Pattern(_))) - stream = pubSubStream(PubSub.PUnsubscribe, keys, true) - } yield stream).commit - ) - private def pubSubStream(cmd: String, keys: List[SubscriptionKey], isUnSubs: Boolean) = - ZStream - .unwrap[BinaryCodec, RedisError, PushProtocol]( - ZIO - .clockWith(_.instant) - .flatMap(now => - ZSTM - .serviceWithSTM[BinaryCodec] { implicit codec => - for { - streams <- - ZSTM.foreach(keys) { key => - for { - resp <- runCommand(cmd, StringInput.encode(key.value), now) - hub <- getPubSubHub(key) - queue <- hub.subscribe - _ <- resp match { - case RespValue.ArrayValues(value) => - hub.offer(value) - case other => - ZSTM.fail(RedisError.ProtocolError(s"Invalid pubsub command response $other")) - } - _ <- pubSubs.delete(key).when(isUnSubs) - } yield ZStream - .fromTQueue(queue) - .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp)).refineToOrDie[RedisError]) - } - } yield streams.fold(ZStream.empty)(_ merge _) - } - .commit - ) + private def unsubscribe( + channels: List[String] + ) = + ZIO.clockWith { clock => + for { + now <- clock.instant + streams <- (for { + keys <- if (channels.isEmpty) pubSubs.keys.map(_.collect { case t: SubscriptionKey.Channel => t }) + else ZSTM.succeed(channels.map(SubscriptionKey.Channel(_))) + stream <- pubSubStream(PubSub.Unsubscribe, keys, now, true) + } yield stream).commit + } yield streams + } + + private def pSubscribe( + pattern: String, + patterns: List[String] + ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = + ZIO.clockWith( + _.instant.flatMap( + pubSubStream(PubSub.PSubscribe, (pattern :: patterns).map(SubscriptionKey.Pattern(_)), _, false).commit ) + ) + + private def pUnsubscribe( + patterns: List[String] + ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = + ZIO.clockWith { clock => + for { + now <- clock.instant + streams <- (for { + keys <- if (patterns.isEmpty) pubSubs.keys.map(_.collect { case t: SubscriptionKey.Pattern => t }) + else ZSTM.succeed(patterns.map(SubscriptionKey.Pattern(_))) + stream <- pubSubStream(PubSub.PUnsubscribe, keys, now, true) + } yield stream).commit + } yield streams + } + private def pubSubStream(cmd: String, keys: List[SubscriptionKey], now: Instant, isUnSubs: Boolean) = + ZSTM + .serviceWithSTM[BinaryCodec] { implicit codec => + for { + streams <- + ZSTM.foreach(keys) { key => + for { + resp <- runCommand(cmd, StringInput.encode(key.value), now) + hub <- getPubSubHub(key) + queue <- hub.subscribe + _ <- resp match { + case RespValue.ArrayValues(value) => + hub.offer(value) + case other => + ZSTM.fail(RedisError.ProtocolError(s"Invalid pubsub command response $other")) + } + _ <- (pubSubs.delete(key)).when(isUnSubs) + } yield ZStream + .fromTQueue(queue) + .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp)).refineToOrDie[RedisError]) + } + } yield streams + } override def execute(command: Chunk[RespValue.BulkString]): IO[RedisError, RespValue] = for { diff --git a/redis/src/main/scala/zio/redis/api/PubSub.scala b/redis/src/main/scala/zio/redis/api/PubSub.scala index 9e47fcefa..92ab0434d 100644 --- a/redis/src/main/scala/zio/redis/api/PubSub.scala +++ b/redis/src/main/scala/zio/redis/api/PubSub.scala @@ -2,82 +2,195 @@ package zio.redis.api import zio.redis.Input._ import zio.redis.Output._ -import zio.redis.ResultBuilder.ResultOutputStreamBuilder +import zio.redis.ResultBuilder.ResultStreamBuilder import zio.redis._ -import zio.redis.api.PubSub._ import zio.schema.Schema import zio.stream._ -import zio.{Chunk, ZIO} - -trait PubSub { - final def subscribeStreamBuilder(channel: String, channels: String*): ResultOutputStreamBuilder = - new ResultOutputStreamBuilder { - override def returning[R: Schema]: ZStream[Redis, RedisError, R] = - ZStream.serviceWithStream[Redis] { redis => - RedisPubSubCommand - .run(RedisPubSubCommand.Subscribe(channel, channels.toList)) - .collect { case t: PushProtocol.Message => t.message } - .mapZIO(resp => - ZIO - .attempt(ArbitraryOutput[R]().unsafeDecode(resp)(redis.codec)) - .refineToOrDie[RedisError] - ) - } - } +import zio.{Chunk, IO, Promise, ZIO} + +trait PubSub extends RedisEnvironment { + import PubSub._ + + final def subscribe(channel: String): ResultStreamBuilder[Id] = + subscribeWithCallback(channel)(emptyCallback) + + final def subscribe(channel: String, channels: List[String]): ResultStreamBuilder[List] = + subscribeWithCallback(channel, channels)(emptyCallback) - final def subscribe(channel: String, channels: String*): ZStream[Redis, RedisError, PushProtocol] = - RedisPubSubCommand.run(RedisPubSubCommand.Subscribe(channel, channels.toList)) + final def subscribeWithCallback(channel: String)(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = + new ResultStreamBuilder[Id] { + def returning[R: Schema]: IO[RedisError, Id[Stream[RedisError, R]]] = + getSubscribeStreams(channel, List.empty)(onSubscribe).flatMap(extractOne(channel, _)) + } - final def unsubscribe(channels: String*): ZStream[Redis, RedisError, PushProtocol.Unsubscribe] = - RedisPubSubCommand.run(RedisPubSubCommand.Unsubscribe(channels.toList)).collect { - case t: PushProtocol.Unsubscribe => t + final def subscribeWithCallback( + channel: String, + channels: List[String] + )(onSubscribe: PubSubCallback): ResultStreamBuilder[List] = + new ResultStreamBuilder[List] { + def returning[R: Schema]: IO[RedisError, List[Stream[RedisError, R]]] = + getSubscribeStreams(channel, channels)(onSubscribe) } - final def pSubscribe(pattern: String, patterns: String*): ZStream[Redis, RedisError, PushProtocol] = - RedisPubSubCommand.run(RedisPubSubCommand.PSubscribe(pattern, patterns.toList)) - - final def pSubscribeStreamBuilder(pattern: String, patterns: String*): ResultOutputStreamBuilder = - new ResultOutputStreamBuilder { - override def returning[R: Schema]: ZStream[Redis, RedisError, R] = - ZStream.serviceWithStream[Redis] { redis => - RedisPubSubCommand - .run(RedisPubSubCommand.PSubscribe(pattern, patterns.toList)) - .collect { case t: PushProtocol.PMessage => t.message } - .mapZIO(resp => - ZIO - .attempt(ArbitraryOutput[R]().unsafeDecode(resp)(redis.codec)) - .refineToOrDie[RedisError] + final def unsubscribe(channel: String): IO[RedisError, Promise[RedisError, Unit]] = + unsubscribeWithCallback(List(channel))(emptyCallback).flatMap(extractOne(channel, _)) + + final def unsubscribe(channels: List[String]): IO[RedisError, List[Promise[RedisError, Unit]]] = + unsubscribeWithCallback(channels)(emptyCallback) + + final def unsubscribeWithCallback(channel: String)( + onUnsubscribe: PubSubCallback + ): IO[RedisError, Promise[RedisError, Unit]] = + unsubscribeWithCallback(List(channel))(onUnsubscribe).flatMap(extractOne(channel, _)) + + final def unsubscribeWithCallback(channels: List[String])( + onUnsubscribe: PubSubCallback + ): IO[RedisError, List[Promise[RedisError, Unit]]] = + RedisPubSubCommand(PubSubCommand.Unsubscribe(channels), codec, pubSub).run + .flatMap(streams => + ZIO.foreach(streams)(stream => + for { + promise <- Promise.make[RedisError, Unit] + _ <- stream + .interruptWhen(promise) + .mapZIO { + case PushProtocol.Unsubscribe(key, numOfSubscription) => + onUnsubscribe(key, numOfSubscription) <* promise.succeed(()) + case _ => ZIO.unit + } + .runDrain + .fork + } yield promise + ) + ) + + private def getSubscribeStreams[R: Schema](channel: String, channels: List[String])(onSubscribe: PubSubCallback) = + RedisPubSubCommand(PubSubCommand.Subscribe(channel, channels), codec, pubSub).run + .flatMap(streams => + ZIO.foreach(streams)(stream => + Promise + .make[RedisError, Unit] + .map(promise => + stream + .interruptWhen(promise) + .mapZIO { + case PushProtocol.Subscribe(key, numOfSubscription) => + onSubscribe(key, numOfSubscription).as(None) + case _: PushProtocol.Unsubscribe => + promise.succeed(()).as(None) + case PushProtocol.Message(_, msg) => + ZIO + .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) + .refineToOrDie[RedisError] + .asSome + case _ => ZIO.none + } + .collectSome ) - } + ) + ) + + final def pSubscribe(pattern: String): ResultStreamBuilder[Id] = pSubscribeWithCallback(pattern)(emptyCallback) + + final def pSubscribeWithCallback(pattern: String)(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = + new ResultStreamBuilder[Id] { + def returning[R: Schema]: IO[RedisError, Id[Stream[RedisError, R]]] = + getPSubscribeStreams(pattern, List.empty)(onSubscribe).flatMap(extractOne(pattern, _)) } - final def pUnsubscribe(patterns: String*): ZStream[Redis, RedisError, PushProtocol.PUnsubscribe] = - RedisPubSubCommand.run(RedisPubSubCommand.PUnsubscribe(patterns.toList)).collect { - case t: PushProtocol.PUnsubscribe => t + final def pSubscribe(pattern: String, patterns: List[String]): ResultStreamBuilder[List] = + pSubscribeWithCallback(pattern, patterns)(emptyCallback) + + final def pSubscribeWithCallback(pattern: String, patterns: List[String])( + onSubscribe: PubSubCallback + ): ResultStreamBuilder[List] = + new ResultStreamBuilder[List] { + def returning[R: Schema]: IO[RedisError, List[Stream[RedisError, R]]] = + getPSubscribeStreams(pattern, patterns)(onSubscribe) } - final def publish[A: Schema](channel: String, message: A): ZIO[Redis, RedisError, Long] = { - val command = RedisCommand(Publish, Tuple2(StringInput, ArbitraryInput[A]()), LongOutput) + private def getPSubscribeStreams[R: Schema](pattern: String, patterns: List[String])(onSubscribe: PubSubCallback) = + RedisPubSubCommand(PubSubCommand.PSubscribe(pattern, patterns), codec, pubSub).run + .flatMap(streams => + ZIO.foreach(streams)(stream => + Promise + .make[RedisError, Unit] + .map(promise => + stream + .interruptWhen(promise) + .mapZIO { + case PushProtocol.PSubscribe(key, numOfSubscription) => + onSubscribe(key, numOfSubscription).as(None) + case _: PushProtocol.PUnsubscribe => + promise.succeed(()).as(None) + case PushProtocol.PMessage(_, _, msg) => + ZIO + .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) + .refineToOrDie[RedisError] + .asSome + case _ => ZIO.none + } + .collectSome + ) + ) + ) + + final def pUnsubscribe(pattern: String): IO[RedisError, Promise[RedisError, Unit]] = + pUnsubscribeWithCallback(List(pattern))(emptyCallback).flatMap(extractOne(pattern, _)) + + final def pUnsubscribeWithCallback(pattern: String)( + onUnsubscribe: PubSubCallback + ): IO[RedisError, Promise[RedisError, Unit]] = + pUnsubscribeWithCallback(List(pattern))(onUnsubscribe).flatMap(extractOne(pattern, _)) + + final def pUnsubscribeWithCallback( + patterns: List[String] + )(onUnsubscribe: PubSubCallback): IO[RedisError, List[Promise[RedisError, Unit]]] = + RedisPubSubCommand(PubSubCommand.PUnsubscribe(patterns), codec, pubSub).run + .flatMap(streams => + ZIO.foreach(streams)(stream => + for { + promise <- Promise.make[RedisError, Unit] + _ <- stream + .interruptWhen(promise) + .mapZIO { + case PushProtocol.PUnsubscribe(key, numOfSubscription) => + onUnsubscribe(key, numOfSubscription) <* promise.succeed(()) + case _ => ZIO.unit + } + .runDrain + .fork + } yield promise + ) + ) + + private def extractOne[A](key: String, elements: List[A]) = + ZIO.fromOption(elements.headOption).orElseFail(RedisError.NoPubSubStream(key)) + + final def publish[A: Schema](channel: String, message: A): IO[RedisError, Long] = { + val command = RedisCommand(Publish, Tuple2(StringInput, ArbitraryInput[A]()), LongOutput, codec, executor) command.run((channel, message)) } - final def pubSubChannels(pattern: String): ZIO[Redis, RedisError, Chunk[String]] = { - val command = RedisCommand(PubSubChannels, StringInput, ChunkOutput(MultiStringOutput)) + final def pubSubChannels(pattern: String): IO[RedisError, Chunk[String]] = { + val command = RedisCommand(PubSubChannels, StringInput, ChunkOutput(MultiStringOutput), codec, executor) command.run(pattern) } - final def pubSubNumPat: ZIO[Redis, RedisError, Long] = { - val command = RedisCommand(PubSubNumPat, NoInput, LongOutput) + final def pubSubNumPat: IO[RedisError, Long] = { + val command = RedisCommand(PubSubNumPat, NoInput, LongOutput, codec, executor) command.run(()) } - final def pubSubNumSub(channel: String, channels: String*): ZIO[Redis, RedisError, Chunk[NumSubResponse]] = { - val command = RedisCommand(PubSubNumSub, NonEmptyList(StringInput), NumSubResponseOutput) + final def pubSubNumSub(channel: String, channels: String*): IO[RedisError, Chunk[NumSubResponse]] = { + val command = RedisCommand(PubSubNumSub, NonEmptyList(StringInput), NumSubResponseOutput, codec, executor) command.run((channel, channels.toList)) } } private[redis] object PubSub { + private lazy val emptyCallback = (_: String, _: Long) => ZIO.unit + final val Subscribe = "SUBSCRIBE" final val Unsubscribe = "UNSUBSCRIBE" final val PSubscribe = "PSUBSCRIBE" diff --git a/redis/src/main/scala/zio/redis/options/Cluster.scala b/redis/src/main/scala/zio/redis/options/Cluster.scala index 35bf83aa8..5c11752c4 100644 --- a/redis/src/main/scala/zio/redis/options/Cluster.scala +++ b/redis/src/main/scala/zio/redis/options/Cluster.scala @@ -15,14 +15,14 @@ */ package zio.redis.options -import zio.redis.{RedisExecutor, RedisPubSub, RedisUri} +import zio.redis.{RedisExecutor, RedisUri} import zio.{Chunk, Scope} object Cluster { private[redis] final val SlotsAmount = 16384 - final case class ExecutorScope(executor: RedisExecutor, pubSub: RedisPubSub, scope: Scope.Closeable) + final case class ExecutorScope(executor: RedisExecutor, scope: Scope.Closeable) final case class ClusterConnection( partitions: Chunk[Partition], @@ -31,8 +31,6 @@ object Cluster { ) { def executor(slot: Slot): Option[RedisExecutor] = executors.get(slots(slot)).map(_.executor) - def pubSub(slot: Slot): Option[RedisPubSub] = executors.get(slots(slot)).map(_.pubSub) - def addExecutor(uri: RedisUri, es: ExecutorScope): ClusterConnection = copy(executors = executors + (uri -> es)) } diff --git a/redis/src/main/scala/zio/redis/options/PubSub.scala b/redis/src/main/scala/zio/redis/options/PubSub.scala index f8915607c..78a98f657 100644 --- a/redis/src/main/scala/zio/redis/options/PubSub.scala +++ b/redis/src/main/scala/zio/redis/options/PubSub.scala @@ -1,8 +1,10 @@ package zio.redis.options -import zio.redis.RespValue +import zio.IO +import zio.redis.{RedisError, RespValue} trait PubSub { + type PubSubCallback = (String, Long) => IO[RedisError, Unit] sealed trait SubscriptionKey { self => def value: String From a01213db562c25e4cdef1c8bd826b7660caebded Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 7 Feb 2023 16:34:18 +0900 Subject: [PATCH 03/51] Apply changes to t/c --- redis/src/test/scala/zio/redis/ApiSpec.scala | 1 + .../scala/zio/redis/ClusterExecutorSpec.scala | 1 + .../src/test/scala/zio/redis/PubSubSpec.scala | 168 ++++++++++-------- 3 files changed, 98 insertions(+), 72 deletions(-) diff --git a/redis/src/test/scala/zio/redis/ApiSpec.scala b/redis/src/test/scala/zio/redis/ApiSpec.scala index 4a445edf7..39d046515 100644 --- a/redis/src/test/scala/zio/redis/ApiSpec.scala +++ b/redis/src/test/scala/zio/redis/ApiSpec.scala @@ -88,6 +88,7 @@ object ApiSpec ZLayer.make[Redis]( ZLayer.succeed(RedisClusterConfig(Chunk(RedisUri("localhost", 5000)))), ClusterExecutor.layer, + RedisPubSub.local, ZLayer.succeed(codec), RedisLive.layer ) diff --git a/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala b/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala index 502584458..21c4cdf25 100644 --- a/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala +++ b/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala @@ -80,6 +80,7 @@ object ClusterExecutorSpec extends BaseSpec { ZLayer.make[Redis]( ZLayer.succeed(RedisClusterConfig(Chunk(address1, address2))), ClusterExecutor.layer.orDie, + RedisPubSub.local.orDie, ZLayer.succeed(codec), RedisLive.layer ) diff --git a/redis/src/test/scala/zio/redis/PubSubSpec.scala b/redis/src/test/scala/zio/redis/PubSubSpec.scala index f1dcb38d0..0a33d2611 100644 --- a/redis/src/test/scala/zio/redis/PubSubSpec.scala +++ b/redis/src/test/scala/zio/redis/PubSubSpec.scala @@ -2,7 +2,7 @@ package zio.redis import zio.test.Assertion._ import zio.test._ -import zio.{Chunk, ZIO} +import zio.{Chunk, Promise, ZIO} import scala.util.Random @@ -12,72 +12,85 @@ trait PubSubSpec extends BaseSpec { suite("subscribe")( test("subscribe response") { for { - channel <- generateRandomString() - res <- subscribe(channel).runHead - } yield assertTrue(res.get.key == SubscriptionKey.Channel(channel)) + redis <- ZIO.service[Redis] + channel <- generateRandomString() + promise <- Promise.make[RedisError, String] + resBuilder = redis.subscribeWithCallback(channel)((key: String, _: Long) => promise.succeed(key).unit) + stream <- resBuilder.returning[String] + _ <- stream.interruptWhen(promise).runDrain.fork + res <- promise.await + } yield assertTrue(res == channel) }, test("message response") { for { + redis <- ZIO.service[Redis] channel <- generateRandomString() message = "bar" - stream <- subscribeStreamBuilder(channel) - .returning[String] - .runHead - .fork - _ <- pubSubChannels(channel) + promise <- Promise.make[RedisError, String] + stream <- redis.subscribe(channel).returning[String] + fiber <- stream.interruptWhen(promise).runHead.fork + _ <- redis + .pubSubChannels(channel) .repeatUntil(_ contains channel) - _ <- publish(channel, message) - res <- stream.join + _ <- redis.publish(channel, message) + res <- fiber.join } yield assertTrue(res.get == message) }, test("multiple subscribe") { val numOfPublish = 20 for { + redis <- ZIO.service[Redis] prefix <- generateRandomString(5) channel1 <- generateRandomString(prefix) channel2 <- generateRandomString(prefix) pattern = prefix + '*' message <- generateRandomString(5) - stream1 <- subscribe(channel1) - .runFoldWhile(Chunk.empty[PushProtocol])( - _.forall(_.isInstanceOf[PushProtocol.Unsubscribe] == false) - )(_ appended _) + stream1 <- redis + .subscribe(channel1) + .returning[String] .fork - stream2 <- subscribe(channel2) - .runFoldWhile(Chunk.empty[PushProtocol])( - _.forall(_.isInstanceOf[PushProtocol.Unsubscribe] == false) - )(_ appended _) + stream2 <- redis + .subscribe(channel2) + .returning[String] .fork - _ <- pubSubChannels(pattern) + _ <- redis + .pubSubChannels(pattern) .repeatUntil(channels => channels.size >= 2) - ch1SubsCount <- publish(channel1, message).replicateZIO(numOfPublish).map(_.head) - ch2SubsCount <- publish(channel2, message).replicateZIO(numOfPublish).map(_.head) - _ <- unsubscribe().runDrain.fork - res1 <- stream1.join - res2 <- stream2.join - } yield assertTrue(ch1SubsCount == 1L) && - assertTrue(ch2SubsCount == 1L) && - assertTrue(res1.size == numOfPublish + 2) && - assertTrue(res2.size == numOfPublish + 2) + ch1SubsCount <- redis.publish(channel1, message).replicateZIO(numOfPublish).map(_.head) + ch2SubsCount <- redis.publish(channel2, message).replicateZIO(numOfPublish).map(_.head) + promises <- redis.unsubscribe(List.empty) + _ <- ZIO.foreachDiscard(promises)(_.await) + _ <- stream1.join + _ <- stream2.join + } yield assertTrue(ch1SubsCount == 1L) && assertTrue(ch2SubsCount == 1L) }, test("psubscribe response") { for { + redis <- ZIO.service[Redis] pattern <- generateRandomString() - res <- pSubscribe(pattern).runHead - } yield assertTrue(res.get.key.value == pattern) + promise <- Promise.make[RedisError, String] + _ <- redis + .pSubscribeWithCallback(pattern)((key: String, _: Long) => promise.succeed(key).unit) + .returning[String] + .flatMap(_.interruptWhen(promise).runHead) + .fork + res <- promise.await + } yield assertTrue(res == pattern) }, test("pmessage response") { for { + redis <- ZIO.service[Redis] prefix <- generateRandomString(5) pattern = prefix + '*' channel <- generateRandomString(prefix) message <- generateRandomString(prefix) - stream <- pSubscribeStreamBuilder(pattern) + stream <- redis + .pSubscribe(pattern) .returning[String] - .runHead + .flatMap(_.runHead) .fork - _ <- pubSubNumPat.repeatUntil(_ > 0) - _ <- publish(channel, message) + _ <- redis.pubSubNumPat.repeatUntil(_ > 0) + _ <- redis.publish(channel, message) res <- stream.join } yield assertTrue(res.get == message) } @@ -86,15 +99,16 @@ trait PubSubSpec extends BaseSpec { val message = 1L assertZIO( for { + redis <- ZIO.service[Redis] channel <- generateRandomString() - stream <- subscribeStreamBuilder(channel) + stream <- redis + .subscribe(channel) .returning[Long] - .runFoldWhile(0L)(_ < 10L) { case (sum, message) => + .flatMap(_.runFoldWhile(0L)(_ < 10L) { case (sum, message) => sum + message - } - .fork - _ <- pubSubChannels(channel).repeatUntil(_ contains channel) - _ <- ZIO.replicateZIO(10)(publish(channel, message)) + }.fork) + _ <- redis.pubSubChannels(channel).repeatUntil(_ contains channel) + _ <- ZIO.replicateZIO(10)(redis.publish(channel, message)) res <- stream.join } yield res )(equalTo(10L)) @@ -103,61 +117,71 @@ trait PubSubSpec extends BaseSpec { test("don't receive message type after unsubscribe") { val numOfPublished = 5 for { + redis <- ZIO.service[Redis] prefix <- generateRandomString(5) pattern = prefix + '*' channel <- generateRandomString(prefix) message <- generateRandomString() - stream <- subscribe(channel) - .runFoldWhile(Chunk.empty[PushProtocol])(_.size < 2)(_ appended _) - .fork - _ <- pubSubChannels(pattern) + _ <- redis + .subscribe(channel) + .returning[String] + .flatMap(_.runCollect) + .fork + _ <- redis + .pubSubChannels(pattern) .repeatUntil(_ contains channel) - _ <- unsubscribe(channel).runHead - receiverCount <- publish(channel, message).replicateZIO(numOfPublished).map(_.head) - res <- stream.join - } yield assertTrue( - res.size == 2 - ) && assertTrue(receiverCount == 0L) + promise <- redis.unsubscribe(channel) + _ <- promise.await + receiverCount <- redis.publish(channel, message).replicateZIO(numOfPublished).map(_.head) + } yield assertTrue(receiverCount == 0L) }, test("unsubscribe response") { for { + redis <- ZIO.service[Redis] channel <- generateRandomString() - res <- unsubscribe(channel).runHead - } yield assertTrue(res.get.key.value == channel) + promise <- Promise.make[RedisError, String] + _ <- redis + .unsubscribeWithCallback(channel)((key: String, _: Long) => promise.succeed(key).unit) + .flatMap(_.await) + res <- promise.await + } yield assertTrue(res == channel) }, test("punsubscribe response") { for { + redis <- ZIO.service[Redis] pattern <- generateRandomString() - res <- pUnsubscribe(pattern).runHead - } yield assertTrue(res.get.key.value == pattern) + promise <- Promise.make[RedisError, String] + _ <- redis + .pUnsubscribeWithCallback(pattern)((key: String, _: Long) => promise.succeed(key).unit) + .flatMap(_.await) + res <- promise.await + } yield assertTrue(res == pattern) }, test("unsubscribe with empty param") { for { + redis <- ZIO.service[Redis] prefix <- generateRandomString(5) pattern = prefix + '*' channel1 <- generateRandomString(prefix) channel2 <- generateRandomString(prefix) - stream1 <- - subscribe(channel1) - .runFoldWhile(Chunk.empty[PushProtocol])(_.forall(_.isInstanceOf[PushProtocol.Unsubscribe] == false))( - _ appended _ - ) + _ <- + redis + .subscribe(channel1) + .returning[String] + .flatMap(_.runCollect) .fork - stream2 <- - subscribe(channel2) - .runFoldWhile(Chunk.empty[PushProtocol])(_.forall(_.isInstanceOf[PushProtocol.Unsubscribe] == false))( - _ appended _ - ) + _ <- + redis + .subscribe(channel2) + .returning[String] + .flatMap(_.runCollect) .fork - _ <- pubSubChannels(pattern) + _ <- redis + .pubSubChannels(pattern) .repeatUntil(_.size >= 2) - _ <- unsubscribe().runDrain.fork - unsubscribeMessages <- stream1.join zip stream2.join - (result1, result2) = unsubscribeMessages - numSubResponses <- pubSubNumSub(channel1, channel2) + _ <- redis.unsubscribe(List.empty).flatMap(ZIO.foreach(_)(_.await)) + numSubResponses <- redis.pubSubNumSub(channel1, channel2) } yield assertTrue( - result1.size == 2 && result2.size == 2 - ) && assertTrue( numSubResponses == Chunk( NumSubResponse(channel1, 0L), NumSubResponse(channel2, 0L) From 954151993aa74002b4ba88a607cb2b8d512ad853 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 7 Feb 2023 16:39:12 +0900 Subject: [PATCH 04/51] Set private accessor --- redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala index 8d40fcce5..0c31abae2 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -174,7 +174,7 @@ final class SingleNodeRedisPubSub( } object SingleNodeRedisPubSub { - final case class Request(command: Chunk[RespValue.BulkString], promise: Promise[RedisError, Unit]) + private final case class Request(command: Chunk[RespValue.BulkString], promise: Promise[RedisError, Unit]) private final val True: Any => Boolean = _ => true From 9723182ac006934398ee15ac0d89aa9b64a455f9 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 7 Feb 2023 18:00:40 +0900 Subject: [PATCH 05/51] Fix cleanup code --- redis/src/main/scala/zio/redis/RedisPubSub.scala | 5 +---- redis/src/main/scala/zio/redis/TestExecutor.scala | 0 2 files changed, 1 insertion(+), 4 deletions(-) delete mode 100644 redis/src/main/scala/zio/redis/TestExecutor.scala diff --git a/redis/src/main/scala/zio/redis/RedisPubSub.scala b/redis/src/main/scala/zio/redis/RedisPubSub.scala index 734756bf9..b117c0264 100644 --- a/redis/src/main/scala/zio/redis/RedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/RedisPubSub.scala @@ -2,7 +2,7 @@ package zio.redis import zio.schema.codec.BinaryCodec import zio.stream._ -import zio.{ULayer, ZIO, ZLayer} +import zio.{ZIO, ZLayer} trait RedisPubSub { def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] @@ -15,9 +15,6 @@ object RedisPubSub { lazy val local: ZLayer[BinaryCodec, RedisError.IOError, RedisPubSub] = RedisConnectionLive.default.fresh >>> pubSublayer - lazy val test: ULayer[RedisPubSub] = - TestExecutor.layer - private lazy val pubSublayer: ZLayer[RedisConnection with BinaryCodec, RedisError.IOError, RedisPubSub] = ZLayer.scoped( ZIO.service[RedisConnection].flatMap(SingleNodeRedisPubSub.create(_)) diff --git a/redis/src/main/scala/zio/redis/TestExecutor.scala b/redis/src/main/scala/zio/redis/TestExecutor.scala deleted file mode 100644 index e69de29bb..000000000 From 65b524b5cc07bf54f952ce9153cc784cdd43deee Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 7 Feb 2023 22:07:08 +0900 Subject: [PATCH 06/51] Modify callback and unsubscribe logic in SingleNodeRedisPubSub --- .../main/scala/zio/redis/PubSubCommand.scala | 16 +- .../src/main/scala/zio/redis/RedisError.scala | 5 +- .../zio/redis/SingleNodeRedisPubSub.scala | 168 +++++++++++++----- .../src/main/scala/zio/redis/api/PubSub.scala | 129 +++++--------- 4 files changed, 187 insertions(+), 131 deletions(-) diff --git a/redis/src/main/scala/zio/redis/PubSubCommand.scala b/redis/src/main/scala/zio/redis/PubSubCommand.scala index 54b4dfdc3..fa86c10fb 100644 --- a/redis/src/main/scala/zio/redis/PubSubCommand.scala +++ b/redis/src/main/scala/zio/redis/PubSubCommand.scala @@ -14,8 +14,16 @@ final case class RedisPubSubCommand(command: PubSubCommand, codec: BinaryCodec, sealed trait PubSubCommand object PubSubCommand { - case class Subscribe(channel: String, channels: List[String]) extends PubSubCommand - case class PSubscribe(pattern: String, patterns: List[String]) extends PubSubCommand - case class Unsubscribe(channels: List[String]) extends PubSubCommand - case class PUnsubscribe(patterns: List[String]) extends PubSubCommand + case class Subscribe( + channel: String, + channels: List[String], + onSubscribe: PubSubCallback + ) extends PubSubCommand + case class PSubscribe( + pattern: String, + patterns: List[String], + onSubscribe: PubSubCallback + ) extends PubSubCommand + case class Unsubscribe(channels: List[String]) extends PubSubCommand + case class PUnsubscribe(patterns: List[String]) extends PubSubCommand } diff --git a/redis/src/main/scala/zio/redis/RedisError.scala b/redis/src/main/scala/zio/redis/RedisError.scala index a0d042b73..20c99f459 100644 --- a/redis/src/main/scala/zio/redis/RedisError.scala +++ b/redis/src/main/scala/zio/redis/RedisError.scala @@ -46,6 +46,7 @@ object RedisError { object Moved { def apply(slotAndAddress: (Slot, RedisUri)): Moved = Moved(slotAndAddress._1, slotAndAddress._2) } - final case class NoPubSubStream(key: String) extends RedisError - final case class IOError(exception: IOException) extends RedisError + final case class NoPubSubStream(key: String) extends RedisError + final case class NoUnsubscribeRequest(key: String) extends RedisError + final case class IOError(exception: IOException) extends RedisError } diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala index 0c31abae2..c8bf82417 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -6,36 +6,48 @@ import zio.redis.SingleNodeRedisPubSub.{Request, RequestQueueSize, True} import zio.redis.api.PubSub import zio.schema.codec.BinaryCodec import zio.stream._ -import zio.{Chunk, ChunkBuilder, Hub, Promise, Queue, Ref, Schedule, UIO, ZIO} +import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Schedule, UIO, ZIO} final class SingleNodeRedisPubSub( pubSubHubsRef: Ref[Map[SubscriptionKey, Hub[PushProtocol]]], + callbacksRef: Ref[Map[SubscriptionKey, Chunk[PubSubCallback]]], + unsubscribedRef: Ref[Map[SubscriptionKey, Promise[RedisError, PushProtocol]]], reqQueue: Queue[Request], connection: RedisConnection ) extends RedisPubSub { def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = command match { - case PubSubCommand.Subscribe(channel, channels) => subscribe(channel, channels) - case PubSubCommand.PSubscribe(pattern, patterns) => pSubscribe(pattern, patterns) - case PubSubCommand.Unsubscribe(channels) => unsubscribe(channels) - case PubSubCommand.PUnsubscribe(patterns) => pUnsubscribe(patterns) + case PubSubCommand.Subscribe(channel, channels, onSubscribe) => + subscribe(channel, channels, onSubscribe) + case PubSubCommand.PSubscribe(pattern, patterns, onSubscribe) => + pSubscribe(pattern, patterns, onSubscribe) + case PubSubCommand.Unsubscribe(channels) => unsubscribe(channels) + case PubSubCommand.PUnsubscribe(patterns) => pUnsubscribe(patterns) } private def subscribe( channel: String, - channels: List[String] + channels: List[String], + onSubscribe: PubSubCallback ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = - makeSubscriptionStream(PubSub.Subscribe, SubscriptionKey.Channel(channel), channels.map(SubscriptionKey.Channel(_))) + makeSubscriptionStream( + PubSub.Subscribe, + SubscriptionKey.Channel(channel), + channels.map(SubscriptionKey.Channel(_)), + onSubscribe + ) private def pSubscribe( pattern: String, - patterns: List[String] + patterns: List[String], + onSubscribe: PubSubCallback ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = makeSubscriptionStream( PubSub.PSubscribe, SubscriptionKey.Pattern(pattern), - patterns.map(SubscriptionKey.Pattern(_)) + patterns.map(SubscriptionKey.Pattern(_)), + onSubscribe ) private def unsubscribe( @@ -60,44 +72,42 @@ final class SingleNodeRedisPubSub( pubSubHubsRef.get.map(_.keys.filter(_.isPatternKey).toList) ) - private def makeSubscriptionStream(command: String, key: SubscriptionKey, keys: List[SubscriptionKey]) = + private def makeSubscriptionStream( + command: String, + key: SubscriptionKey, + keys: List[SubscriptionKey], + onSubscribe: PubSubCallback + ) = ZIO.serviceWithZIO[BinaryCodec] { implicit codec => for { promise <- Promise.make[RedisError, Unit] chunk = StringInput.encode(command) ++ NonEmptyList(StringInput).encode((key.value, keys.map(_.value))) streams <- makeStreams(key :: keys) - _ <- reqQueue.offer(Request(chunk, promise)) + _ <- reqQueue.offer(Request(chunk, promise, Some((key, onSubscribe)))) _ <- promise.await } yield streams } - private def makeUnsubscriptionStream(command: String, keys: UIO[List[SubscriptionKey]]) = { - def releaseHub(key: SubscriptionKey) = - for { - pubSubs <- pubSubHubsRef.get - hubOpt = pubSubs.get(key) - _ <- ZIO.fromOption(hubOpt).flatMap(_.shutdown).ignore - _ <- pubSubHubsRef.update(_ - key) - } yield () - + private def makeUnsubscriptionStream( + command: String, + keys: UIO[List[SubscriptionKey]] + ) = ZIO.serviceWithZIO[BinaryCodec] { implicit codec => for { - targets <- keys - chunk = StringInput.encode(command) ++ Varargs(StringInput).encode(targets.map(_.value)) - promise <- Promise.make[RedisError, Unit] - streams <- makeStreams(targets) - targetSet = targets.map(_.value).toSet - _ <- reqQueue.offer(Request(chunk, promise)) - _ <- promise.await - } yield streams.map(_.tap { - case PushProtocol.Unsubscribe(channel, _) if targetSet contains channel => - releaseHub(SubscriptionKey.Channel(channel)) - case PushProtocol.PUnsubscribe(pattern, _) if targetSet contains pattern => - releaseHub(SubscriptionKey.Pattern(pattern)) - case _ => ZIO.unit - }) + targets <- keys + chunk = StringInput.encode(command) ++ Varargs(StringInput).encode(targets.map(_.value)) + promise <- Promise.make[RedisError, Unit] + _ <- reqQueue.offer(Request(chunk, promise, None)) + _ <- promise.await + streams <- + ZIO.foreach(targets)(key => + for { + promise <- Promise.make[RedisError, PushProtocol] + _ <- unsubscribedRef.update(_ + (key -> promise)) + } yield ZStream.fromZIO(promise.await) + ) + } yield streams } - } private def makeStreams(keys: List[SubscriptionKey]): UIO[List[Stream[RedisError, PushProtocol]]] = ZIO.foreach(keys)(getHub(_).map(ZStream.fromHub(_))) @@ -114,7 +124,19 @@ final class SingleNodeRedisPubSub( } yield hub } - private def send = + private def send = { + def registerCallbacks(request: Request) = + ZIO + .fromOption(request.callbacks) + .flatMap { case (key, additionalCallbacks) => + for { + callbackMap <- callbacksRef.get + callbacks = callbackMap.getOrElse(key, Chunk.empty) + _ <- callbacksRef.update(_.updated(key, callbacks appended additionalCallbacks)) + } yield () + } + .orElse(ZIO.unit) + reqQueue.takeBetween(1, RequestQueueSize).flatMap { reqs => val buffer = ChunkBuilder.make[Byte]() val it = reqs.iterator @@ -131,11 +153,60 @@ final class SingleNodeRedisPubSub( .mapError(RedisError.IOError(_)) .tapBoth( e => ZIO.foreachDiscard(reqs)(_.promise.fail(e)), - _ => ZIO.foreachDiscard(reqs)(_.promise.succeed(())) + _ => ZIO.foreachDiscard(reqs)(req => registerCallbacks(req) *> req.promise.succeed(())) ) } + } + + private def receive: ZIO[BinaryCodec, RedisError, Unit] = { + def applySubscriptionCallback(protocol: PushProtocol): IO[RedisError, Unit] = { + def runAndCleanup(key: SubscriptionKey, numOfSubscription: Long) = + for { + callbackMap <- callbacksRef.get + callbacks = callbackMap.getOrElse(key, Chunk.empty) + _ <- ZIO.foreachDiscard(callbacks)(_(key.value, numOfSubscription)) + _ <- callbacksRef.update(_.updated(key, Chunk.empty)) + } yield () + + protocol match { + case PushProtocol.Subscribe(channel, numOfSubscription) => + runAndCleanup(SubscriptionKey.Channel(channel), numOfSubscription) + case PushProtocol.PSubscribe(pattern, numOfSubscription) => + runAndCleanup(SubscriptionKey.Channel(pattern), numOfSubscription) + case _ => ZIO.unit + } + } + + def releaseHub(key: SubscriptionKey) = + for { + pubSubs <- pubSubHubsRef.get + hubOpt = pubSubs.get(key) + _ <- ZIO.fromOption(hubOpt).flatMap(_.shutdown).orElse(ZIO.unit) + _ <- pubSubHubsRef.update(_ - key) + } yield () + + def handlePushProtocolMessage(msg: PushProtocol) = msg match { + case msg @ PushProtocol.Unsubscribe(channel, _) => + for { + _ <- releaseHub(SubscriptionKey.Channel(channel)) + map <- unsubscribedRef.get + promise <- ZIO + .fromOption(map.get(SubscriptionKey.Channel(channel))) + .orElseFail(RedisError.NoUnsubscribeRequest(channel)) + _ <- promise.succeed(msg) + } yield () + case msg @ PushProtocol.PUnsubscribe(pattern, _) => + for { + _ <- releaseHub(SubscriptionKey.Pattern(pattern)) + map <- unsubscribedRef.get + promise <- ZIO + .fromOption(map.get(SubscriptionKey.Pattern(pattern))) + .orElseFail(RedisError.NoUnsubscribeRequest(pattern)) + _ <- promise.succeed(msg) + } yield () + case other => getHub(other.key).flatMap(_.offer(other)) + } - private def receive: ZIO[BinaryCodec, RedisError, Unit] = ZIO.serviceWithZIO[BinaryCodec] { implicit codec => connection.read .mapError(RedisError.IOError(_)) @@ -143,8 +214,9 @@ final class SingleNodeRedisPubSub( .collectSome .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))) .refineToOrDie[RedisError] - .foreach(push => getHub(push.key).flatMap(_.offer(push))) + .foreach(push => applySubscriptionCallback(push) *> handlePushProtocolMessage(push)) } + } private def resubscribe: ZIO[BinaryCodec, RedisError, Unit] = ZIO.serviceWithZIO[BinaryCodec] { implicit codec => @@ -174,7 +246,11 @@ final class SingleNodeRedisPubSub( } object SingleNodeRedisPubSub { - private final case class Request(command: Chunk[RespValue.BulkString], promise: Promise[RedisError, Unit]) + private final case class Request( + command: Chunk[RespValue.BulkString], + promise: Promise[RedisError, Unit], + callbacks: Option[(SubscriptionKey, PubSubCallback)] + ) private final val True: Any => Boolean = _ => true @@ -182,10 +258,12 @@ object SingleNodeRedisPubSub { def create(conn: RedisConnection) = for { - hubRef <- Ref.make(Map.empty[SubscriptionKey, Hub[PushProtocol]]) - reqQueue <- Queue.bounded[Request](RequestQueueSize) - pubSub = new SingleNodeRedisPubSub(hubRef, reqQueue, conn) - _ <- pubSub.run.forkScoped - _ <- logScopeFinalizer(s"$pubSub Node PubSub is closed") + hubRef <- Ref.make(Map.empty[SubscriptionKey, Hub[PushProtocol]]) + callbackRef <- Ref.make(Map.empty[SubscriptionKey, Chunk[PubSubCallback]]) + unsubscribedRef <- Ref.make(Map.empty[SubscriptionKey, Promise[RedisError, PushProtocol]]) + reqQueue <- Queue.bounded[Request](RequestQueueSize) + pubSub = new SingleNodeRedisPubSub(hubRef, callbackRef, unsubscribedRef, reqQueue, conn) + _ <- pubSub.run.forkScoped + _ <- logScopeFinalizer(s"$pubSub Node PubSub is closed") } yield pubSub } diff --git a/redis/src/main/scala/zio/redis/api/PubSub.scala b/redis/src/main/scala/zio/redis/api/PubSub.scala index 92ab0434d..af162630b 100644 --- a/redis/src/main/scala/zio/redis/api/PubSub.scala +++ b/redis/src/main/scala/zio/redis/api/PubSub.scala @@ -17,7 +17,9 @@ trait PubSub extends RedisEnvironment { final def subscribe(channel: String, channels: List[String]): ResultStreamBuilder[List] = subscribeWithCallback(channel, channels)(emptyCallback) - final def subscribeWithCallback(channel: String)(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = + final def subscribeWithCallback( + channel: String + )(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = new ResultStreamBuilder[Id] { def returning[R: Schema]: IO[RedisError, Id[Stream[RedisError, R]]] = getSubscribeStreams(channel, List.empty)(onSubscribe).flatMap(extractOne(channel, _)) @@ -32,67 +34,43 @@ trait PubSub extends RedisEnvironment { getSubscribeStreams(channel, channels)(onSubscribe) } - final def unsubscribe(channel: String): IO[RedisError, Promise[RedisError, Unit]] = - unsubscribeWithCallback(List(channel))(emptyCallback).flatMap(extractOne(channel, _)) - - final def unsubscribe(channels: List[String]): IO[RedisError, List[Promise[RedisError, Unit]]] = - unsubscribeWithCallback(channels)(emptyCallback) - - final def unsubscribeWithCallback(channel: String)( - onUnsubscribe: PubSubCallback - ): IO[RedisError, Promise[RedisError, Unit]] = - unsubscribeWithCallback(List(channel))(onUnsubscribe).flatMap(extractOne(channel, _)) + final def unsubscribe(channel: String): IO[RedisError, Promise[RedisError, Long]] = + unsubscribe(List(channel)).flatMap(extractOne(channel, _)) - final def unsubscribeWithCallback(channels: List[String])( - onUnsubscribe: PubSubCallback - ): IO[RedisError, List[Promise[RedisError, Unit]]] = + final def unsubscribe(channels: List[String]): IO[RedisError, List[Promise[RedisError, Long]]] = RedisPubSubCommand(PubSubCommand.Unsubscribe(channels), codec, pubSub).run - .flatMap(streams => - ZIO.foreach(streams)(stream => + .flatMap( + ZIO.foreach(_)(stream => for { - promise <- Promise.make[RedisError, Unit] - _ <- stream - .interruptWhen(promise) - .mapZIO { - case PushProtocol.Unsubscribe(key, numOfSubscription) => - onUnsubscribe(key, numOfSubscription) <* promise.succeed(()) - case _ => ZIO.unit - } - .runDrain - .fork + promise <- Promise.make[RedisError, Long] + _ <- stream.mapZIO { + case PushProtocol.Unsubscribe(_, numOfSubscription) => promise.succeed(numOfSubscription) + case _ => promise.fail(RedisError.WrongType(s"Cannot handle message except Unsubscribe")) + }.runDrain.fork } yield promise ) ) - private def getSubscribeStreams[R: Schema](channel: String, channels: List[String])(onSubscribe: PubSubCallback) = - RedisPubSubCommand(PubSubCommand.Subscribe(channel, channels), codec, pubSub).run - .flatMap(streams => - ZIO.foreach(streams)(stream => - Promise - .make[RedisError, Unit] - .map(promise => - stream - .interruptWhen(promise) - .mapZIO { - case PushProtocol.Subscribe(key, numOfSubscription) => - onSubscribe(key, numOfSubscription).as(None) - case _: PushProtocol.Unsubscribe => - promise.succeed(()).as(None) - case PushProtocol.Message(_, msg) => - ZIO - .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) - .refineToOrDie[RedisError] - .asSome - case _ => ZIO.none - } - .collectSome - ) - ) - ) - - final def pSubscribe(pattern: String): ResultStreamBuilder[Id] = pSubscribeWithCallback(pattern)(emptyCallback) - - final def pSubscribeWithCallback(pattern: String)(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = + private def getSubscribeStreams[R: Schema]( + channel: String, + channels: List[String] + )(onSubscribe: PubSubCallback) = + RedisPubSubCommand(PubSubCommand.Subscribe(channel, channels, onSubscribe), codec, pubSub).run + .map(_.map(_.mapZIO { + case PushProtocol.Message(_, msg) => + ZIO + .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) + .refineToOrDie[RedisError] + .asSome + case _ => ZIO.none + }.collectSome)) + + final def pSubscribe(pattern: String): ResultStreamBuilder[Id] = + pSubscribeWithCallback(pattern)(emptyCallback) + + final def pSubscribeWithCallback( + pattern: String + )(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = new ResultStreamBuilder[Id] { def returning[R: Schema]: IO[RedisError, Id[Stream[RedisError, R]]] = getPSubscribeStreams(pattern, List.empty)(onSubscribe).flatMap(extractOne(pattern, _)) @@ -109,8 +87,11 @@ trait PubSub extends RedisEnvironment { getPSubscribeStreams(pattern, patterns)(onSubscribe) } - private def getPSubscribeStreams[R: Schema](pattern: String, patterns: List[String])(onSubscribe: PubSubCallback) = - RedisPubSubCommand(PubSubCommand.PSubscribe(pattern, patterns), codec, pubSub).run + private def getPSubscribeStreams[R: Schema]( + pattern: String, + patterns: List[String] + )(onSubscribe: PubSubCallback) = + RedisPubSubCommand(PubSubCommand.PSubscribe(pattern, patterns, onSubscribe), codec, pubSub).run .flatMap(streams => ZIO.foreach(streams)(stream => Promise @@ -135,31 +116,19 @@ trait PubSub extends RedisEnvironment { ) ) - final def pUnsubscribe(pattern: String): IO[RedisError, Promise[RedisError, Unit]] = - pUnsubscribeWithCallback(List(pattern))(emptyCallback).flatMap(extractOne(pattern, _)) + final def pUnsubscribe(pattern: String): IO[RedisError, Promise[RedisError, Long]] = + pUnsubscribe(List(pattern)).flatMap(extractOne(pattern, _)) - final def pUnsubscribeWithCallback(pattern: String)( - onUnsubscribe: PubSubCallback - ): IO[RedisError, Promise[RedisError, Unit]] = - pUnsubscribeWithCallback(List(pattern))(onUnsubscribe).flatMap(extractOne(pattern, _)) - - final def pUnsubscribeWithCallback( - patterns: List[String] - )(onUnsubscribe: PubSubCallback): IO[RedisError, List[Promise[RedisError, Unit]]] = + final def pUnsubscribe(patterns: List[String]): IO[RedisError, List[Promise[RedisError, Long]]] = RedisPubSubCommand(PubSubCommand.PUnsubscribe(patterns), codec, pubSub).run - .flatMap(streams => - ZIO.foreach(streams)(stream => + .flatMap( + ZIO.foreach(_)(stream => for { - promise <- Promise.make[RedisError, Unit] - _ <- stream - .interruptWhen(promise) - .mapZIO { - case PushProtocol.PUnsubscribe(key, numOfSubscription) => - onUnsubscribe(key, numOfSubscription) <* promise.succeed(()) - case _ => ZIO.unit - } - .runDrain - .fork + promise <- Promise.make[RedisError, Long] + _ <- stream.mapZIO { + case PushProtocol.PUnsubscribe(_, numOfSubscription) => promise.succeed(numOfSubscription) + case _ => promise.fail(RedisError.WrongType(s"Cannot handle message except PUnsubscribe")) + }.runDrain.fork } yield promise ) ) @@ -189,7 +158,7 @@ trait PubSub extends RedisEnvironment { } private[redis] object PubSub { - private lazy val emptyCallback = (_: String, _: Long) => ZIO.unit + lazy val emptyCallback = (_: String, _: Long) => ZIO.unit final val Subscribe = "SUBSCRIBE" final val Unsubscribe = "UNSUBSCRIBE" From 892ba208ba374955852440069c22a2e32b13de38 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 7 Feb 2023 22:07:30 +0900 Subject: [PATCH 07/51] Apply changes to t/c --- .../src/test/scala/zio/redis/PubSubSpec.scala | 40 +++++-------------- 1 file changed, 10 insertions(+), 30 deletions(-) diff --git a/redis/src/test/scala/zio/redis/PubSubSpec.scala b/redis/src/test/scala/zio/redis/PubSubSpec.scala index 0a33d2611..ae0d123b0 100644 --- a/redis/src/test/scala/zio/redis/PubSubSpec.scala +++ b/redis/src/test/scala/zio/redis/PubSubSpec.scala @@ -2,7 +2,7 @@ package zio.redis import zio.test.Assertion._ import zio.test._ -import zio.{Chunk, Promise, ZIO} +import zio.{Chunk, Promise, Ref, ZIO} import scala.util.Random @@ -12,13 +12,15 @@ trait PubSubSpec extends BaseSpec { suite("subscribe")( test("subscribe response") { for { - redis <- ZIO.service[Redis] - channel <- generateRandomString() - promise <- Promise.make[RedisError, String] - resBuilder = redis.subscribeWithCallback(channel)((key: String, _: Long) => promise.succeed(key).unit) - stream <- resBuilder.returning[String] - _ <- stream.interruptWhen(promise).runDrain.fork - res <- promise.await + redis <- ZIO.service[Redis] + channel <- generateRandomString() + promise <- Promise.make[RedisError, String] + ref <- Ref.make(promise) + resBuilder = + redis.subscribeWithCallback(channel)((key: String, _: Long) => ref.get.flatMap(_.succeed(key)).unit) + stream <- resBuilder.returning[String] + _ <- stream.interruptWhen(promise).runDrain.fork + res <- ref.get.flatMap(_.await) } yield assertTrue(res == channel) }, test("message response") { @@ -135,28 +137,6 @@ trait PubSubSpec extends BaseSpec { receiverCount <- redis.publish(channel, message).replicateZIO(numOfPublished).map(_.head) } yield assertTrue(receiverCount == 0L) }, - test("unsubscribe response") { - for { - redis <- ZIO.service[Redis] - channel <- generateRandomString() - promise <- Promise.make[RedisError, String] - _ <- redis - .unsubscribeWithCallback(channel)((key: String, _: Long) => promise.succeed(key).unit) - .flatMap(_.await) - res <- promise.await - } yield assertTrue(res == channel) - }, - test("punsubscribe response") { - for { - redis <- ZIO.service[Redis] - pattern <- generateRandomString() - promise <- Promise.make[RedisError, String] - _ <- redis - .pUnsubscribeWithCallback(pattern)((key: String, _: Long) => promise.succeed(key).unit) - .flatMap(_.await) - res <- promise.await - } yield assertTrue(res == pattern) - }, test("unsubscribe with empty param") { for { redis <- ZIO.service[Redis] From 43edb47850f656e7402e34bbf7791531b7444438 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 7 Feb 2023 22:14:51 +0900 Subject: [PATCH 08/51] Modify PubSub api --- .../src/main/scala/zio/redis/api/PubSub.scala | 53 +++++++------------ 1 file changed, 20 insertions(+), 33 deletions(-) diff --git a/redis/src/main/scala/zio/redis/api/PubSub.scala b/redis/src/main/scala/zio/redis/api/PubSub.scala index af162630b..575879c34 100644 --- a/redis/src/main/scala/zio/redis/api/PubSub.scala +++ b/redis/src/main/scala/zio/redis/api/PubSub.scala @@ -34,18 +34,19 @@ trait PubSub extends RedisEnvironment { getSubscribeStreams(channel, channels)(onSubscribe) } - final def unsubscribe(channel: String): IO[RedisError, Promise[RedisError, Long]] = + final def unsubscribe(channel: String): IO[RedisError, Promise[RedisError, (String, Long)]] = unsubscribe(List(channel)).flatMap(extractOne(channel, _)) - final def unsubscribe(channels: List[String]): IO[RedisError, List[Promise[RedisError, Long]]] = + final def unsubscribe(channels: List[String]): IO[RedisError, List[Promise[RedisError, (String, Long)]]] = RedisPubSubCommand(PubSubCommand.Unsubscribe(channels), codec, pubSub).run .flatMap( ZIO.foreach(_)(stream => for { - promise <- Promise.make[RedisError, Long] + promise <- Promise.make[RedisError, (String, Long)] _ <- stream.mapZIO { - case PushProtocol.Unsubscribe(_, numOfSubscription) => promise.succeed(numOfSubscription) - case _ => promise.fail(RedisError.WrongType(s"Cannot handle message except Unsubscribe")) + case PushProtocol.Unsubscribe(channel, numOfSubscription) => + promise.succeed((channel, numOfSubscription)) + case _ => promise.fail(RedisError.WrongType(s"Cannot handle message except Unsubscribe")) }.runDrain.fork } yield promise ) @@ -92,42 +93,28 @@ trait PubSub extends RedisEnvironment { patterns: List[String] )(onSubscribe: PubSubCallback) = RedisPubSubCommand(PubSubCommand.PSubscribe(pattern, patterns, onSubscribe), codec, pubSub).run - .flatMap(streams => - ZIO.foreach(streams)(stream => - Promise - .make[RedisError, Unit] - .map(promise => - stream - .interruptWhen(promise) - .mapZIO { - case PushProtocol.PSubscribe(key, numOfSubscription) => - onSubscribe(key, numOfSubscription).as(None) - case _: PushProtocol.PUnsubscribe => - promise.succeed(()).as(None) - case PushProtocol.PMessage(_, _, msg) => - ZIO - .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) - .refineToOrDie[RedisError] - .asSome - case _ => ZIO.none - } - .collectSome - ) - ) - ) + .map(_.map(_.mapZIO { + case PushProtocol.PMessage(_, _, msg) => + ZIO + .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) + .refineToOrDie[RedisError] + .asSome + case _ => ZIO.none + }.collectSome)) - final def pUnsubscribe(pattern: String): IO[RedisError, Promise[RedisError, Long]] = + final def pUnsubscribe(pattern: String): IO[RedisError, Promise[RedisError, (String, Long)]] = pUnsubscribe(List(pattern)).flatMap(extractOne(pattern, _)) - final def pUnsubscribe(patterns: List[String]): IO[RedisError, List[Promise[RedisError, Long]]] = + final def pUnsubscribe(patterns: List[String]): IO[RedisError, List[Promise[RedisError, (String, Long)]]] = RedisPubSubCommand(PubSubCommand.PUnsubscribe(patterns), codec, pubSub).run .flatMap( ZIO.foreach(_)(stream => for { - promise <- Promise.make[RedisError, Long] + promise <- Promise.make[RedisError, (String, Long)] _ <- stream.mapZIO { - case PushProtocol.PUnsubscribe(_, numOfSubscription) => promise.succeed(numOfSubscription) - case _ => promise.fail(RedisError.WrongType(s"Cannot handle message except PUnsubscribe")) + case PushProtocol.PUnsubscribe(pattern, numOfSubscription) => + promise.succeed((pattern, numOfSubscription)) + case _ => promise.fail(RedisError.WrongType(s"Cannot handle message except PUnsubscribe")) }.runDrain.fork } yield promise ) From a6f74fafcd83ec398f3130e0ce5e137928e69f13 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 7 Feb 2023 22:14:55 +0900 Subject: [PATCH 09/51] Apply changes to t/c --- .../src/test/scala/zio/redis/PubSubSpec.scala | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/redis/src/test/scala/zio/redis/PubSubSpec.scala b/redis/src/test/scala/zio/redis/PubSubSpec.scala index ae0d123b0..ff1114f49 100644 --- a/redis/src/test/scala/zio/redis/PubSubSpec.scala +++ b/redis/src/test/scala/zio/redis/PubSubSpec.scala @@ -137,6 +137,24 @@ trait PubSubSpec extends BaseSpec { receiverCount <- redis.publish(channel, message).replicateZIO(numOfPublished).map(_.head) } yield assertTrue(receiverCount == 0L) }, + test("unsubscribe response") { + for { + redis <- ZIO.service[Redis] + channel <- generateRandomString() + res <- redis + .unsubscribe(channel) + .flatMap(_.await) + } yield assertTrue(res._1 == channel) + }, + test("punsubscribe response") { + for { + redis <- ZIO.service[Redis] + pattern <- generateRandomString() + res <- redis + .pUnsubscribe(pattern) + .flatMap(_.await) + } yield assertTrue(res._1 == pattern) + }, test("unsubscribe with empty param") { for { redis <- ZIO.service[Redis] From c7decb7915ad1993468260c5ca2fa818b52ec429 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 7 Feb 2023 22:20:16 +0900 Subject: [PATCH 10/51] Fix typo --- redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala | 2 +- redis/src/test/scala/zio/redis/PubSubSpec.scala | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala index c8bf82417..9385148f3 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -172,7 +172,7 @@ final class SingleNodeRedisPubSub( case PushProtocol.Subscribe(channel, numOfSubscription) => runAndCleanup(SubscriptionKey.Channel(channel), numOfSubscription) case PushProtocol.PSubscribe(pattern, numOfSubscription) => - runAndCleanup(SubscriptionKey.Channel(pattern), numOfSubscription) + runAndCleanup(SubscriptionKey.Pattern(pattern), numOfSubscription) case _ => ZIO.unit } } diff --git a/redis/src/test/scala/zio/redis/PubSubSpec.scala b/redis/src/test/scala/zio/redis/PubSubSpec.scala index ff1114f49..60c692574 100644 --- a/redis/src/test/scala/zio/redis/PubSubSpec.scala +++ b/redis/src/test/scala/zio/redis/PubSubSpec.scala @@ -2,7 +2,7 @@ package zio.redis import zio.test.Assertion._ import zio.test._ -import zio.{Chunk, Promise, Ref, ZIO} +import zio.{Chunk, Promise, ZIO} import scala.util.Random @@ -15,12 +15,11 @@ trait PubSubSpec extends BaseSpec { redis <- ZIO.service[Redis] channel <- generateRandomString() promise <- Promise.make[RedisError, String] - ref <- Ref.make(promise) resBuilder = - redis.subscribeWithCallback(channel)((key: String, _: Long) => ref.get.flatMap(_.succeed(key)).unit) + redis.subscribeWithCallback(channel)((key: String, _: Long) => promise.succeed(key).unit) stream <- resBuilder.returning[String] _ <- stream.interruptWhen(promise).runDrain.fork - res <- ref.get.flatMap(_.await) + res <- promise.await } yield assertTrue(res == channel) }, test("message response") { From b6438f05bb0e28f51cef3f0d7aafae1aabfead11 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 7 Feb 2023 22:21:23 +0900 Subject: [PATCH 11/51] Modify accessor --- redis/src/main/scala/zio/redis/api/PubSub.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/redis/src/main/scala/zio/redis/api/PubSub.scala b/redis/src/main/scala/zio/redis/api/PubSub.scala index 575879c34..c1a2c509c 100644 --- a/redis/src/main/scala/zio/redis/api/PubSub.scala +++ b/redis/src/main/scala/zio/redis/api/PubSub.scala @@ -145,7 +145,7 @@ trait PubSub extends RedisEnvironment { } private[redis] object PubSub { - lazy val emptyCallback = (_: String, _: Long) => ZIO.unit + private lazy val emptyCallback = (_: String, _: Long) => ZIO.unit final val Subscribe = "SUBSCRIBE" final val Unsubscribe = "UNSUBSCRIBE" From 5fb75e44b76721b841550419613967c87fd8ed78 Mon Sep 17 00:00:00 2001 From: opg1 Date: Thu, 9 Feb 2023 22:27:02 +0900 Subject: [PATCH 12/51] Update redis/src/main/scala/zio/redis/Output.scala MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Aleksandar Novaković --- redis/src/main/scala/zio/redis/Output.scala | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/redis/src/main/scala/zio/redis/Output.scala b/redis/src/main/scala/zio/redis/Output.scala index f72db1fa1..e11461fb5 100644 --- a/redis/src/main/scala/zio/redis/Output.scala +++ b/redis/src/main/scala/zio/redis/Output.scala @@ -849,8 +849,7 @@ object Output { val num = LongOutput.unsafeDecode(values(2)) PushProtocol.PUnsubscribe(key, num) case "message" => - val message = values(2) - PushProtocol.Message(key, message) + PushProtocol.Message(key, values(2)) case "pmessage" => val channel = MultiStringOutput.unsafeDecode(values(2)) val message = values(3) From ce97f95a0c2634de637b9135e324254c7cc695da Mon Sep 17 00:00:00 2001 From: opg1 Date: Thu, 9 Feb 2023 22:27:13 +0900 Subject: [PATCH 13/51] Update redis/src/main/scala/zio/redis/Output.scala MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Aleksandar Novaković --- redis/src/main/scala/zio/redis/Output.scala | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/redis/src/main/scala/zio/redis/Output.scala b/redis/src/main/scala/zio/redis/Output.scala index e11461fb5..b5ee57329 100644 --- a/redis/src/main/scala/zio/redis/Output.scala +++ b/redis/src/main/scala/zio/redis/Output.scala @@ -852,8 +852,7 @@ object Output { PushProtocol.Message(key, values(2)) case "pmessage" => val channel = MultiStringOutput.unsafeDecode(values(2)) - val message = values(3) - PushProtocol.PMessage(key, channel, message) + 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") From a25805bba5aa90d9c705ffc2c818a5d38155f348 Mon Sep 17 00:00:00 2001 From: 0pg Date: Fri, 10 Feb 2023 02:34:23 +0900 Subject: [PATCH 14/51] Modify responsibility to transform streams --- redis/src/main/scala/zio/redis/Output.scala | 7 +- .../main/scala/zio/redis/PubSubCommand.scala | 17 +- .../scala/zio/redis/RedisPubSubCommand.scala | 35 ++++ .../zio/redis/SingleNodeRedisPubSub.scala | 78 +++------ .../src/main/scala/zio/redis/api/PubSub.scala | 157 ++++++++---------- .../main/scala/zio/redis/options/PubSub.scala | 21 +-- 6 files changed, 141 insertions(+), 174 deletions(-) create mode 100644 redis/src/main/scala/zio/redis/RedisPubSubCommand.scala diff --git a/redis/src/main/scala/zio/redis/Output.scala b/redis/src/main/scala/zio/redis/Output.scala index b5ee57329..72224361e 100644 --- a/redis/src/main/scala/zio/redis/Output.scala +++ b/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 @@ -859,14 +860,14 @@ object Output { } } - case object NumSubResponseOutput extends Output[Chunk[NumSubResponse]] { - protected def tryDecode(respValue: RespValue)(implicit codec: BinaryCodec): Chunk[NumSubResponse] = + 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)) - NumSubResponse(channel, numOfSubscription) + NumberOfSubscribers(channel, numOfSubscription) }) case other => throw ProtocolError(s"$other isn't an array") } diff --git a/redis/src/main/scala/zio/redis/PubSubCommand.scala b/redis/src/main/scala/zio/redis/PubSubCommand.scala index fa86c10fb..f040f399e 100644 --- a/redis/src/main/scala/zio/redis/PubSubCommand.scala +++ b/redis/src/main/scala/zio/redis/PubSubCommand.scala @@ -1,28 +1,15 @@ package zio.redis -import zio.schema.codec.BinaryCodec -import zio.stream._ -import zio.{IO, ZLayer} - -final case class RedisPubSubCommand(command: PubSubCommand, codec: BinaryCodec, executor: RedisPubSub) { - def run: IO[RedisError, List[Stream[RedisError, PushProtocol]]] = { - val codecLayer = ZLayer.succeed(codec) - executor.execute(command).provideLayer(codecLayer) - } -} - sealed trait PubSubCommand object PubSubCommand { case class Subscribe( channel: String, - channels: List[String], - onSubscribe: PubSubCallback + channels: List[String] ) extends PubSubCommand case class PSubscribe( pattern: String, - patterns: List[String], - onSubscribe: PubSubCallback + patterns: List[String] ) extends PubSubCommand case class Unsubscribe(channels: List[String]) extends PubSubCommand case class PUnsubscribe(patterns: List[String]) extends PubSubCommand diff --git a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala new file mode 100644 index 000000000..03d471e81 --- /dev/null +++ b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala @@ -0,0 +1,35 @@ +package zio.redis + +import zio.redis.Output.ArbitraryOutput +import zio.schema.Schema +import zio.schema.codec.BinaryCodec +import zio.stream._ +import zio.{IO, ZIO, ZLayer} + +final case class RedisPubSubCommand(command: PubSubCommand, codec: BinaryCodec, executor: RedisPubSub) { + def run[R: Schema]( + callback: PushProtocol => IO[RedisError, Unit] + ): IO[RedisError, List[Stream[RedisError, R]]] = { + val codecLayer = ZLayer.succeed(codec) + executor + .execute(command) + .provideLayer(codecLayer) + .map( + _.map( + _.tap(callback(_)).mapZIO { + case PushProtocol.Message(_, msg) => + ZIO + .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) + .refineToOrDie[RedisError] + .asSome + case PushProtocol.PMessage(_, _, msg) => + ZIO + .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) + .refineToOrDie[RedisError] + .asSome + case _ => ZIO.none + }.collectSome + ) + ) + } +} diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala index 9385148f3..e64fd41a7 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -6,11 +6,10 @@ import zio.redis.SingleNodeRedisPubSub.{Request, RequestQueueSize, True} import zio.redis.api.PubSub import zio.schema.codec.BinaryCodec import zio.stream._ -import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Schedule, UIO, ZIO} +import zio.{Chunk, ChunkBuilder, Hub, Promise, Queue, Ref, Schedule, UIO, ZIO} final class SingleNodeRedisPubSub( pubSubHubsRef: Ref[Map[SubscriptionKey, Hub[PushProtocol]]], - callbacksRef: Ref[Map[SubscriptionKey, Chunk[PubSubCallback]]], unsubscribedRef: Ref[Map[SubscriptionKey, Promise[RedisError, PushProtocol]]], reqQueue: Queue[Request], connection: RedisConnection @@ -18,36 +17,32 @@ final class SingleNodeRedisPubSub( def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = command match { - case PubSubCommand.Subscribe(channel, channels, onSubscribe) => - subscribe(channel, channels, onSubscribe) - case PubSubCommand.PSubscribe(pattern, patterns, onSubscribe) => - pSubscribe(pattern, patterns, onSubscribe) + case PubSubCommand.Subscribe(channel, channels) => + subscribe(channel, channels) + case PubSubCommand.PSubscribe(pattern, patterns) => + pSubscribe(pattern, patterns) case PubSubCommand.Unsubscribe(channels) => unsubscribe(channels) case PubSubCommand.PUnsubscribe(patterns) => pUnsubscribe(patterns) } private def subscribe( channel: String, - channels: List[String], - onSubscribe: PubSubCallback + channels: List[String] ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = makeSubscriptionStream( PubSub.Subscribe, SubscriptionKey.Channel(channel), - channels.map(SubscriptionKey.Channel(_)), - onSubscribe + channels.map(SubscriptionKey.Channel(_)) ) private def pSubscribe( pattern: String, - patterns: List[String], - onSubscribe: PubSubCallback + patterns: List[String] ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = makeSubscriptionStream( PubSub.PSubscribe, SubscriptionKey.Pattern(pattern), - patterns.map(SubscriptionKey.Pattern(_)), - onSubscribe + patterns.map(SubscriptionKey.Pattern(_)) ) private def unsubscribe( @@ -75,15 +70,14 @@ final class SingleNodeRedisPubSub( private def makeSubscriptionStream( command: String, key: SubscriptionKey, - keys: List[SubscriptionKey], - onSubscribe: PubSubCallback + keys: List[SubscriptionKey] ) = ZIO.serviceWithZIO[BinaryCodec] { implicit codec => for { promise <- Promise.make[RedisError, Unit] chunk = StringInput.encode(command) ++ NonEmptyList(StringInput).encode((key.value, keys.map(_.value))) streams <- makeStreams(key :: keys) - _ <- reqQueue.offer(Request(chunk, promise, Some((key, onSubscribe)))) + _ <- reqQueue.offer(Request(chunk, promise)) _ <- promise.await } yield streams } @@ -97,7 +91,7 @@ final class SingleNodeRedisPubSub( targets <- keys chunk = StringInput.encode(command) ++ Varargs(StringInput).encode(targets.map(_.value)) promise <- Promise.make[RedisError, Unit] - _ <- reqQueue.offer(Request(chunk, promise, None)) + _ <- reqQueue.offer(Request(chunk, promise)) _ <- promise.await streams <- ZIO.foreach(targets)(key => @@ -124,19 +118,7 @@ final class SingleNodeRedisPubSub( } yield hub } - private def send = { - def registerCallbacks(request: Request) = - ZIO - .fromOption(request.callbacks) - .flatMap { case (key, additionalCallbacks) => - for { - callbackMap <- callbacksRef.get - callbacks = callbackMap.getOrElse(key, Chunk.empty) - _ <- callbacksRef.update(_.updated(key, callbacks appended additionalCallbacks)) - } yield () - } - .orElse(ZIO.unit) - + private def send = reqQueue.takeBetween(1, RequestQueueSize).flatMap { reqs => val buffer = ChunkBuilder.make[Byte]() val it = reqs.iterator @@ -153,30 +135,11 @@ final class SingleNodeRedisPubSub( .mapError(RedisError.IOError(_)) .tapBoth( e => ZIO.foreachDiscard(reqs)(_.promise.fail(e)), - _ => ZIO.foreachDiscard(reqs)(req => registerCallbacks(req) *> req.promise.succeed(())) + _ => ZIO.foreachDiscard(reqs)(_.promise.succeed(())) ) } - } private def receive: ZIO[BinaryCodec, RedisError, Unit] = { - def applySubscriptionCallback(protocol: PushProtocol): IO[RedisError, Unit] = { - def runAndCleanup(key: SubscriptionKey, numOfSubscription: Long) = - for { - callbackMap <- callbacksRef.get - callbacks = callbackMap.getOrElse(key, Chunk.empty) - _ <- ZIO.foreachDiscard(callbacks)(_(key.value, numOfSubscription)) - _ <- callbacksRef.update(_.updated(key, Chunk.empty)) - } yield () - - protocol match { - case PushProtocol.Subscribe(channel, numOfSubscription) => - runAndCleanup(SubscriptionKey.Channel(channel), numOfSubscription) - case PushProtocol.PSubscribe(pattern, numOfSubscription) => - runAndCleanup(SubscriptionKey.Pattern(pattern), numOfSubscription) - case _ => ZIO.unit - } - } - def releaseHub(key: SubscriptionKey) = for { pubSubs <- pubSubHubsRef.get @@ -204,7 +167,10 @@ final class SingleNodeRedisPubSub( .orElseFail(RedisError.NoUnsubscribeRequest(pattern)) _ <- promise.succeed(msg) } yield () - case other => getHub(other.key).flatMap(_.offer(other)) + case msg @ PushProtocol.Subscribe(channel, _) => getHub(SubscriptionKey.Channel(channel)).flatMap(_.offer(msg)) + case msg @ PushProtocol.PSubscribe(pattern, _) => getHub(SubscriptionKey.Pattern(pattern)).flatMap(_.offer(msg)) + case msg @ PushProtocol.Message(channel, _) => getHub(SubscriptionKey.Channel(channel)).flatMap(_.offer(msg)) + case msg @ PushProtocol.PMessage(pattern, _, _) => getHub(SubscriptionKey.Pattern(pattern)).flatMap(_.offer(msg)) } ZIO.serviceWithZIO[BinaryCodec] { implicit codec => @@ -214,7 +180,7 @@ final class SingleNodeRedisPubSub( .collectSome .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))) .refineToOrDie[RedisError] - .foreach(push => applySubscriptionCallback(push) *> handlePushProtocolMessage(push)) + .foreach(push => handlePushProtocolMessage(push)) } } @@ -248,8 +214,7 @@ final class SingleNodeRedisPubSub( object SingleNodeRedisPubSub { private final case class Request( command: Chunk[RespValue.BulkString], - promise: Promise[RedisError, Unit], - callbacks: Option[(SubscriptionKey, PubSubCallback)] + promise: Promise[RedisError, Unit] ) private final val True: Any => Boolean = _ => true @@ -259,10 +224,9 @@ object SingleNodeRedisPubSub { def create(conn: RedisConnection) = for { hubRef <- Ref.make(Map.empty[SubscriptionKey, Hub[PushProtocol]]) - callbackRef <- Ref.make(Map.empty[SubscriptionKey, Chunk[PubSubCallback]]) unsubscribedRef <- Ref.make(Map.empty[SubscriptionKey, Promise[RedisError, PushProtocol]]) reqQueue <- Queue.bounded[Request](RequestQueueSize) - pubSub = new SingleNodeRedisPubSub(hubRef, callbackRef, unsubscribedRef, reqQueue, conn) + pubSub = new SingleNodeRedisPubSub(hubRef, unsubscribedRef, reqQueue, conn) _ <- pubSub.run.forkScoped _ <- logScopeFinalizer(s"$pubSub Node PubSub is closed") } yield pubSub diff --git a/redis/src/main/scala/zio/redis/api/PubSub.scala b/redis/src/main/scala/zio/redis/api/PubSub.scala index c1a2c509c..9b8e493c2 100644 --- a/redis/src/main/scala/zio/redis/api/PubSub.scala +++ b/redis/src/main/scala/zio/redis/api/PubSub.scala @@ -4,9 +4,10 @@ import zio.redis.Input._ import zio.redis.Output._ import zio.redis.ResultBuilder.ResultStreamBuilder import zio.redis._ +import zio.redis.options.PubSub.NumberOfSubscribers import zio.schema.Schema import zio.stream._ -import zio.{Chunk, IO, Promise, ZIO} +import zio.{Chunk, IO, Promise, Ref, ZIO} trait PubSub extends RedisEnvironment { import PubSub._ @@ -14,57 +15,19 @@ trait PubSub extends RedisEnvironment { final def subscribe(channel: String): ResultStreamBuilder[Id] = subscribeWithCallback(channel)(emptyCallback) - final def subscribe(channel: String, channels: List[String]): ResultStreamBuilder[List] = - subscribeWithCallback(channel, channels)(emptyCallback) - - final def subscribeWithCallback( - channel: String - )(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = + final def subscribeWithCallback(channel: String)(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = new ResultStreamBuilder[Id] { def returning[R: Schema]: IO[RedisError, Id[Stream[RedisError, R]]] = - getSubscribeStreams(channel, List.empty)(onSubscribe).flatMap(extractOne(channel, _)) + runSubscription(PubSubCommand.Subscribe(channel, List.empty), onSubscribe).flatMap(extractOne(channel, _)) } - final def subscribeWithCallback( - channel: String, - channels: List[String] - )(onSubscribe: PubSubCallback): ResultStreamBuilder[List] = - new ResultStreamBuilder[List] { - def returning[R: Schema]: IO[RedisError, List[Stream[RedisError, R]]] = - getSubscribeStreams(channel, channels)(onSubscribe) - } - - final def unsubscribe(channel: String): IO[RedisError, Promise[RedisError, (String, Long)]] = - unsubscribe(List(channel)).flatMap(extractOne(channel, _)) + final def subscribe(channel: String, channels: String*): ResultStreamBuilder[List] = + createStreamListBuilder(channel, channels.toList, emptyCallback, isPatterned = false) - final def unsubscribe(channels: List[String]): IO[RedisError, List[Promise[RedisError, (String, Long)]]] = - RedisPubSubCommand(PubSubCommand.Unsubscribe(channels), codec, pubSub).run - .flatMap( - ZIO.foreach(_)(stream => - for { - promise <- Promise.make[RedisError, (String, Long)] - _ <- stream.mapZIO { - case PushProtocol.Unsubscribe(channel, numOfSubscription) => - promise.succeed((channel, numOfSubscription)) - case _ => promise.fail(RedisError.WrongType(s"Cannot handle message except Unsubscribe")) - }.runDrain.fork - } yield promise - ) - ) - - private def getSubscribeStreams[R: Schema]( - channel: String, - channels: List[String] - )(onSubscribe: PubSubCallback) = - RedisPubSubCommand(PubSubCommand.Subscribe(channel, channels, onSubscribe), codec, pubSub).run - .map(_.map(_.mapZIO { - case PushProtocol.Message(_, msg) => - ZIO - .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) - .refineToOrDie[RedisError] - .asSome - case _ => ZIO.none - }.collectSome)) + final def subscribeWithCallback(channel: String, channels: String*)( + onSubscribe: PubSubCallback + ): ResultStreamBuilder[List] = + createStreamListBuilder(channel, channels.toList, onSubscribe, isPatterned = false) final def pSubscribe(pattern: String): ResultStreamBuilder[Id] = pSubscribeWithCallback(pattern)(emptyCallback) @@ -74,54 +37,22 @@ trait PubSub extends RedisEnvironment { )(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = new ResultStreamBuilder[Id] { def returning[R: Schema]: IO[RedisError, Id[Stream[RedisError, R]]] = - getPSubscribeStreams(pattern, List.empty)(onSubscribe).flatMap(extractOne(pattern, _)) + runSubscription(PubSubCommand.PSubscribe(pattern, List.empty), onSubscribe).flatMap(extractOne(pattern, _)) } - final def pSubscribe(pattern: String, patterns: List[String]): ResultStreamBuilder[List] = - pSubscribeWithCallback(pattern, patterns)(emptyCallback) + final def pSubscribe(pattern: String, patterns: String*): ResultStreamBuilder[List] = + createStreamListBuilder(pattern, patterns.toList, emptyCallback, isPatterned = true) - final def pSubscribeWithCallback(pattern: String, patterns: List[String])( + final def pSubscribeWithCallback(pattern: String, patterns: String*)( onSubscribe: PubSubCallback ): ResultStreamBuilder[List] = - new ResultStreamBuilder[List] { - def returning[R: Schema]: IO[RedisError, List[Stream[RedisError, R]]] = - getPSubscribeStreams(pattern, patterns)(onSubscribe) - } + createStreamListBuilder(pattern, patterns.toList, onSubscribe, isPatterned = true) - private def getPSubscribeStreams[R: Schema]( - pattern: String, - patterns: List[String] - )(onSubscribe: PubSubCallback) = - RedisPubSubCommand(PubSubCommand.PSubscribe(pattern, patterns, onSubscribe), codec, pubSub).run - .map(_.map(_.mapZIO { - case PushProtocol.PMessage(_, _, msg) => - ZIO - .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) - .refineToOrDie[RedisError] - .asSome - case _ => ZIO.none - }.collectSome)) - - final def pUnsubscribe(pattern: String): IO[RedisError, Promise[RedisError, (String, Long)]] = - pUnsubscribe(List(pattern)).flatMap(extractOne(pattern, _)) - - final def pUnsubscribe(patterns: List[String]): IO[RedisError, List[Promise[RedisError, (String, Long)]]] = - RedisPubSubCommand(PubSubCommand.PUnsubscribe(patterns), codec, pubSub).run - .flatMap( - ZIO.foreach(_)(stream => - for { - promise <- Promise.make[RedisError, (String, Long)] - _ <- stream.mapZIO { - case PushProtocol.PUnsubscribe(pattern, numOfSubscription) => - promise.succeed((pattern, numOfSubscription)) - case _ => promise.fail(RedisError.WrongType(s"Cannot handle message except PUnsubscribe")) - }.runDrain.fork - } yield promise - ) - ) + final def unsubscribe(channels: String*): IO[RedisError, Promise[RedisError, Chunk[(String, Long)]]] = + runUnsubscription(PubSubCommand.Unsubscribe(channels.toList)) - private def extractOne[A](key: String, elements: List[A]) = - ZIO.fromOption(elements.headOption).orElseFail(RedisError.NoPubSubStream(key)) + final def pUnsubscribe(patterns: String*): IO[RedisError, Promise[RedisError, Chunk[(String, Long)]]] = + runUnsubscription(PubSubCommand.PUnsubscribe(patterns.toList)) final def publish[A: Schema](channel: String, message: A): IO[RedisError, Long] = { val command = RedisCommand(Publish, Tuple2(StringInput, ArbitraryInput[A]()), LongOutput, codec, executor) @@ -138,10 +69,58 @@ trait PubSub extends RedisEnvironment { command.run(()) } - final def pubSubNumSub(channel: String, channels: String*): IO[RedisError, Chunk[NumSubResponse]] = { + final def pubSubNumSub(channel: String, channels: String*): IO[RedisError, Chunk[NumberOfSubscribers]] = { val command = RedisCommand(PubSubNumSub, NonEmptyList(StringInput), NumSubResponseOutput, codec, executor) command.run((channel, channels.toList)) } + + private def createStreamListBuilder( + key: String, + keys: List[String], + callback: PubSubCallback, + isPatterned: Boolean + ): ResultStreamBuilder[List] = + new ResultStreamBuilder[List] { + def returning[R: Schema]: IO[RedisError, List[Stream[RedisError, R]]] = + if (isPatterned) runSubscription(PubSubCommand.PSubscribe(key, keys), callback) + else runSubscription(PubSubCommand.Subscribe(key, keys), callback) + } + + private def runUnsubscription( + command: PubSubCommand + ): IO[RedisError, Promise[RedisError, Chunk[(String, Long)]]] = + for { + promise <- Promise.make[RedisError, Chunk[(String, Long)]] + ref <- Ref.make(Chunk.empty[(String, Long)]) + streams <- RedisPubSubCommand(command, codec, pubSub).run[Unit] { + case PushProtocol.Unsubscribe(channel, numOfSubscription) => + ref.update(_ appended (channel, numOfSubscription)) + case PushProtocol.PUnsubscribe(pattern, numOfSubscription) => + ref.update(_ appended (pattern, numOfSubscription)) + case _ => ZIO.unit + } + _ <- streams + .fold(ZStream.empty)(_ merge _) + .runDrain + .onDone( + e => promise.fail(e), + _ => ref.get.flatMap(promise.succeed(_)) + ) + .fork + } yield promise + + private def runSubscription[R: Schema]( + command: PubSubCommand, + onSubscribe: PubSubCallback + ): IO[RedisError, List[Stream[RedisError, R]]] = + RedisPubSubCommand(command, codec, pubSub).run[R] { + case PushProtocol.Subscribe(key, numOfSubscription) => onSubscribe(key, numOfSubscription) + case PushProtocol.PSubscribe(key, numOfSubscription) => onSubscribe(key, numOfSubscription) + case _ => ZIO.unit + } + + private def extractOne[A](key: String, elements: List[A]) = + ZIO.fromOption(elements.headOption).orElseFail(RedisError.NoPubSubStream(key)) } private[redis] object PubSub { diff --git a/redis/src/main/scala/zio/redis/options/PubSub.scala b/redis/src/main/scala/zio/redis/options/PubSub.scala index 78a98f657..ef0170abc 100644 --- a/redis/src/main/scala/zio/redis/options/PubSub.scala +++ b/redis/src/main/scala/zio/redis/options/PubSub.scala @@ -2,7 +2,6 @@ package zio.redis.options import zio.IO import zio.redis.{RedisError, RespValue} - trait PubSub { type PubSubCallback = (String, Long) => IO[RedisError, Unit] @@ -22,30 +21,32 @@ trait PubSub { case class Pattern(value: String) extends SubscriptionKey } - case class NumSubResponse(channel: String, subscriberCount: Long) - sealed trait PushProtocol { - def key: SubscriptionKey + def key: String } object PushProtocol { case class Subscribe(channel: String, numOfSubscription: Long) extends PushProtocol { - def key: SubscriptionKey = SubscriptionKey.Channel(channel) + def key: String = channel } case class PSubscribe(pattern: String, numOfSubscription: Long) extends PushProtocol { - def key: SubscriptionKey = SubscriptionKey.Pattern(pattern) + def key: String = pattern } case class Unsubscribe(channel: String, numOfSubscription: Long) extends PushProtocol { - def key: SubscriptionKey = SubscriptionKey.Channel(channel) + def key: String = channel } case class PUnsubscribe(pattern: String, numOfSubscription: Long) extends PushProtocol { - def key: SubscriptionKey = SubscriptionKey.Pattern(pattern) + def key: String = pattern } case class Message(channel: String, message: RespValue) extends PushProtocol { - def key: SubscriptionKey = SubscriptionKey.Channel(channel) + def key: String = channel } case class PMessage(pattern: String, channel: String, message: RespValue) extends PushProtocol { - def key: SubscriptionKey = SubscriptionKey.Pattern(pattern) + def key: String = pattern } } } + +object PubSub { + final case class NumberOfSubscribers(channel: String, subscriberCount: Long) +} From fd0b303bb5ad2043fdb4e274f45f2f286a83bbd5 Mon Sep 17 00:00:00 2001 From: 0pg Date: Fri, 10 Feb 2023 02:39:28 +0900 Subject: [PATCH 15/51] Apply changes to t/c --- .../src/test/scala/zio/redis/PubSubSpec.scala | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/redis/src/test/scala/zio/redis/PubSubSpec.scala b/redis/src/test/scala/zio/redis/PubSubSpec.scala index 60c692574..76bfd4166 100644 --- a/redis/src/test/scala/zio/redis/PubSubSpec.scala +++ b/redis/src/test/scala/zio/redis/PubSubSpec.scala @@ -1,5 +1,6 @@ package zio.redis +import zio.redis.options.PubSub.NumberOfSubscribers import zio.test.Assertion._ import zio.test._ import zio.{Chunk, Promise, ZIO} @@ -12,14 +13,13 @@ trait PubSubSpec extends BaseSpec { suite("subscribe")( test("subscribe response") { for { - redis <- ZIO.service[Redis] - channel <- generateRandomString() - promise <- Promise.make[RedisError, String] - resBuilder = - redis.subscribeWithCallback(channel)((key: String, _: Long) => promise.succeed(key).unit) - stream <- resBuilder.returning[String] - _ <- stream.interruptWhen(promise).runDrain.fork - res <- promise.await + redis <- ZIO.service[Redis] + channel <- generateRandomString() + promise <- Promise.make[RedisError, String] + resBuilder = redis.subscribeWithCallback(channel)((key: String, _: Long) => promise.succeed(key).unit) + stream <- resBuilder.returning[String] + _ <- stream.interruptWhen(promise).runDrain.fork + res <- promise.await } yield assertTrue(res == channel) }, test("message response") { @@ -59,8 +59,8 @@ trait PubSubSpec extends BaseSpec { .repeatUntil(channels => channels.size >= 2) ch1SubsCount <- redis.publish(channel1, message).replicateZIO(numOfPublish).map(_.head) ch2SubsCount <- redis.publish(channel2, message).replicateZIO(numOfPublish).map(_.head) - promises <- redis.unsubscribe(List.empty) - _ <- ZIO.foreachDiscard(promises)(_.await) + promises <- redis.unsubscribe() + _ <- promises.await _ <- stream1.join _ <- stream2.join } yield assertTrue(ch1SubsCount == 1L) && assertTrue(ch2SubsCount == 1L) @@ -143,7 +143,7 @@ trait PubSubSpec extends BaseSpec { res <- redis .unsubscribe(channel) .flatMap(_.await) - } yield assertTrue(res._1 == channel) + } yield assertTrue(res.head._1 == channel) }, test("punsubscribe response") { for { @@ -152,7 +152,7 @@ trait PubSubSpec extends BaseSpec { res <- redis .pUnsubscribe(pattern) .flatMap(_.await) - } yield assertTrue(res._1 == pattern) + } yield assertTrue(res.head._1 == pattern) }, test("unsubscribe with empty param") { for { @@ -176,12 +176,12 @@ trait PubSubSpec extends BaseSpec { _ <- redis .pubSubChannels(pattern) .repeatUntil(_.size >= 2) - _ <- redis.unsubscribe(List.empty).flatMap(ZIO.foreach(_)(_.await)) + _ <- redis.unsubscribe().flatMap(_.await) numSubResponses <- redis.pubSubNumSub(channel1, channel2) } yield assertTrue( numSubResponses == Chunk( - NumSubResponse(channel1, 0L), - NumSubResponse(channel2, 0L) + NumberOfSubscribers(channel1, 0L), + NumberOfSubscribers(channel2, 0L) ) ) } From 7cb149142f1efd4652d9132448d7de97285302ee Mon Sep 17 00:00:00 2001 From: 0pg Date: Fri, 10 Feb 2023 13:27:10 +0900 Subject: [PATCH 16/51] Move SubscriptionKey into only used place --- .../zio/redis/SingleNodeRedisPubSub.scala | 42 ++++++++++-------- .../main/scala/zio/redis/options/PubSub.scala | 44 +++---------------- 2 files changed, 31 insertions(+), 55 deletions(-) diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala index e64fd41a7..17499f31c 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -2,7 +2,7 @@ package zio.redis import zio.redis.Input.{NonEmptyList, StringInput, Varargs} import zio.redis.Output.PushProtocolOutput -import zio.redis.SingleNodeRedisPubSub.{Request, RequestQueueSize, True} +import zio.redis.SingleNodeRedisPubSub.{Request, RequestQueueSize, SubscriptionKey, True} import zio.redis.api.PubSub import zio.schema.codec.BinaryCodec import zio.stream._ @@ -31,8 +31,8 @@ final class SingleNodeRedisPubSub( ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = makeSubscriptionStream( PubSub.Subscribe, - SubscriptionKey.Channel(channel), - channels.map(SubscriptionKey.Channel(_)) + channelKey(channel), + channels.map(channelKey(_)) ) private def pSubscribe( @@ -41,8 +41,8 @@ final class SingleNodeRedisPubSub( ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = makeSubscriptionStream( PubSub.PSubscribe, - SubscriptionKey.Pattern(pattern), - patterns.map(SubscriptionKey.Pattern(_)) + patternKey(pattern), + patterns.map(patternKey(_)) ) private def unsubscribe( @@ -51,9 +51,9 @@ final class SingleNodeRedisPubSub( makeUnsubscriptionStream( PubSub.Unsubscribe, if (channels.nonEmpty) - ZIO.succeedNow(channels.map(SubscriptionKey.Channel(_))) + ZIO.succeedNow(channels.map(channelKey(_))) else - pubSubHubsRef.get.map(_.keys.filter(_.isChannelKey).toList) + pubSubHubsRef.get.map(_.keys.filter(_.isChannel).toList) ) private def pUnsubscribe( @@ -62,9 +62,9 @@ final class SingleNodeRedisPubSub( makeUnsubscriptionStream( PubSub.PUnsubscribe, if (patterns.nonEmpty) - ZIO.succeedNow(patterns.map(SubscriptionKey.Pattern(_))) + ZIO.succeedNow(patterns.map(patternKey(_))) else - pubSubHubsRef.get.map(_.keys.filter(_.isPatternKey).toList) + pubSubHubsRef.get.map(_.keys.filter(_.isPattern).toList) ) private def makeSubscriptionStream( @@ -151,26 +151,26 @@ final class SingleNodeRedisPubSub( def handlePushProtocolMessage(msg: PushProtocol) = msg match { case msg @ PushProtocol.Unsubscribe(channel, _) => for { - _ <- releaseHub(SubscriptionKey.Channel(channel)) + _ <- releaseHub(channelKey(channel)) map <- unsubscribedRef.get promise <- ZIO - .fromOption(map.get(SubscriptionKey.Channel(channel))) + .fromOption(map.get(channelKey(channel))) .orElseFail(RedisError.NoUnsubscribeRequest(channel)) _ <- promise.succeed(msg) } yield () case msg @ PushProtocol.PUnsubscribe(pattern, _) => for { - _ <- releaseHub(SubscriptionKey.Pattern(pattern)) + _ <- releaseHub(patternKey(pattern)) map <- unsubscribedRef.get promise <- ZIO - .fromOption(map.get(SubscriptionKey.Pattern(pattern))) + .fromOption(map.get(patternKey(pattern))) .orElseFail(RedisError.NoUnsubscribeRequest(pattern)) _ <- promise.succeed(msg) } yield () - case msg @ PushProtocol.Subscribe(channel, _) => getHub(SubscriptionKey.Channel(channel)).flatMap(_.offer(msg)) - case msg @ PushProtocol.PSubscribe(pattern, _) => getHub(SubscriptionKey.Pattern(pattern)).flatMap(_.offer(msg)) - case msg @ PushProtocol.Message(channel, _) => getHub(SubscriptionKey.Channel(channel)).flatMap(_.offer(msg)) - case msg @ PushProtocol.PMessage(pattern, _, _) => getHub(SubscriptionKey.Pattern(pattern)).flatMap(_.offer(msg)) + case msg @ PushProtocol.Subscribe(channel, _) => getHub(channelKey(channel)).flatMap(_.offer(msg)) + case msg @ PushProtocol.PSubscribe(pattern, _) => getHub(patternKey(pattern)).flatMap(_.offer(msg)) + case msg @ PushProtocol.Message(channel, _) => getHub(channelKey(channel)).flatMap(_.offer(msg)) + case msg @ PushProtocol.PMessage(pattern, _, _) => getHub(patternKey(pattern)).flatMap(_.offer(msg)) } ZIO.serviceWithZIO[BinaryCodec] { implicit codec => @@ -191,7 +191,7 @@ final class SingleNodeRedisPubSub( for { keySet <- pubSubHubsRef.get.map(_.keySet) - (channels, patterns) = keySet.partition(_.isChannelKey) + (channels, patterns) = keySet.partition(_.isChannel) _ <- (connection.write(makeCommand(PubSub.Subscribe, channels.map(_.value))).when(channels.nonEmpty) *> connection.write(makeCommand(PubSub.PSubscribe, patterns.map(_.value))).when(patterns.nonEmpty)) .mapError(RedisError.IOError(_)) @@ -199,6 +199,9 @@ final class SingleNodeRedisPubSub( } yield () } + private def patternKey(key: String) = SubscriptionKey(key, true) + private def channelKey(key: String) = SubscriptionKey(key, false) + /** * 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. @@ -212,6 +215,9 @@ final class SingleNodeRedisPubSub( } object SingleNodeRedisPubSub { + private final case class SubscriptionKey(value: String, isPattern: Boolean) { + def isChannel: Boolean = isPattern == false + } private final case class Request( command: Chunk[RespValue.BulkString], promise: Promise[RedisError, Unit] diff --git a/redis/src/main/scala/zio/redis/options/PubSub.scala b/redis/src/main/scala/zio/redis/options/PubSub.scala index ef0170abc..d06ebc88b 100644 --- a/redis/src/main/scala/zio/redis/options/PubSub.scala +++ b/redis/src/main/scala/zio/redis/options/PubSub.scala @@ -5,45 +5,15 @@ import zio.redis.{RedisError, RespValue} trait PubSub { type PubSubCallback = (String, Long) => IO[RedisError, Unit] - sealed trait SubscriptionKey { self => - def value: String - - def isChannelKey = self match { - case _: SubscriptionKey.Channel => true - case _: SubscriptionKey.Pattern => false - } - - def isPatternKey = !isChannelKey - } - - object SubscriptionKey { - case class Channel(value: String) extends SubscriptionKey - case class Pattern(value: String) extends SubscriptionKey - } - - sealed trait PushProtocol { - def key: String - } + sealed trait PushProtocol object PushProtocol { - case class Subscribe(channel: String, numOfSubscription: Long) extends PushProtocol { - def key: String = channel - } - case class PSubscribe(pattern: String, numOfSubscription: Long) extends PushProtocol { - def key: String = pattern - } - case class Unsubscribe(channel: String, numOfSubscription: Long) extends PushProtocol { - def key: String = channel - } - case class PUnsubscribe(pattern: String, numOfSubscription: Long) extends PushProtocol { - def key: String = pattern - } - case class Message(channel: String, message: RespValue) extends PushProtocol { - def key: String = channel - } - case class PMessage(pattern: String, channel: String, message: RespValue) extends PushProtocol { - def key: String = pattern - } + case class Subscribe(channel: String, numOfSubscription: Long) extends PushProtocol + case class PSubscribe(pattern: String, numOfSubscription: Long) extends PushProtocol + case class Unsubscribe(channel: String, numOfSubscription: Long) extends PushProtocol + case class PUnsubscribe(pattern: String, numOfSubscription: Long) extends PushProtocol + case class Message(channel: String, message: RespValue) extends PushProtocol + case class PMessage(pattern: String, channel: String, message: RespValue) extends PushProtocol } } From 87eee71cfabcba4e689c9408e8f9556ccee73670 Mon Sep 17 00:00:00 2001 From: 0pg Date: Sun, 12 Feb 2023 02:24:56 +0900 Subject: [PATCH 17/51] Fix onSubscribe callback race condition bug --- .../main/scala/zio/redis/RedisPubSub.scala | 4 +- .../scala/zio/redis/RedisPubSubCommand.scala | 4 +- .../zio/redis/SingleNodeRedisPubSub.scala | 124 +++++++++--------- .../src/main/scala/zio/redis/api/PubSub.scala | 18 +-- 4 files changed, 73 insertions(+), 77 deletions(-) diff --git a/redis/src/main/scala/zio/redis/RedisPubSub.scala b/redis/src/main/scala/zio/redis/RedisPubSub.scala index b117c0264..377b8c7a3 100644 --- a/redis/src/main/scala/zio/redis/RedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/RedisPubSub.scala @@ -2,10 +2,10 @@ package zio.redis import zio.schema.codec.BinaryCodec import zio.stream._ -import zio.{ZIO, ZLayer} +import zio.{Chunk, ZIO, ZLayer} trait RedisPubSub { - def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] + def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] } object RedisPubSub { diff --git a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala index 03d471e81..f5894a044 100644 --- a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala +++ b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala @@ -4,12 +4,12 @@ import zio.redis.Output.ArbitraryOutput import zio.schema.Schema import zio.schema.codec.BinaryCodec import zio.stream._ -import zio.{IO, ZIO, ZLayer} +import zio.{Chunk, IO, ZIO, ZLayer} final case class RedisPubSubCommand(command: PubSubCommand, codec: BinaryCodec, executor: RedisPubSub) { def run[R: Schema]( callback: PushProtocol => IO[RedisError, Unit] - ): IO[RedisError, List[Stream[RedisError, R]]] = { + ): IO[RedisError, Chunk[Stream[RedisError, R]]] = { val codecLayer = ZLayer.succeed(codec) executor .execute(command) diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala index 17499f31c..e40c66f80 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -9,13 +9,13 @@ import zio.stream._ import zio.{Chunk, ChunkBuilder, Hub, Promise, Queue, Ref, Schedule, UIO, ZIO} final class SingleNodeRedisPubSub( - pubSubHubsRef: Ref[Map[SubscriptionKey, Hub[PushProtocol]]], - unsubscribedRef: Ref[Map[SubscriptionKey, Promise[RedisError, PushProtocol]]], + pubSubHubsRef: Ref[Map[SubscriptionKey, Hub[Take[RedisError, PushProtocol]]]], reqQueue: Queue[Request], + resQueue: Queue[Promise[RedisError, PushProtocol]], connection: RedisConnection ) extends RedisPubSub { - def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = + def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] = command match { case PubSubCommand.Subscribe(channel, channels) => subscribe(channel, channels) @@ -28,7 +28,7 @@ final class SingleNodeRedisPubSub( private def subscribe( channel: String, channels: List[String] - ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = + ): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] = makeSubscriptionStream( PubSub.Subscribe, channelKey(channel), @@ -38,7 +38,7 @@ final class SingleNodeRedisPubSub( private def pSubscribe( pattern: String, patterns: List[String] - ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = + ): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] = makeSubscriptionStream( PubSub.PSubscribe, patternKey(pattern), @@ -47,7 +47,7 @@ final class SingleNodeRedisPubSub( private def unsubscribe( channels: List[String] - ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = + ): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] = makeUnsubscriptionStream( PubSub.Unsubscribe, if (channels.nonEmpty) @@ -58,7 +58,7 @@ final class SingleNodeRedisPubSub( private def pUnsubscribe( patterns: List[String] - ): ZIO[BinaryCodec, RedisError, List[Stream[RedisError, PushProtocol]]] = + ): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] = makeUnsubscriptionStream( PubSub.PUnsubscribe, if (patterns.nonEmpty) @@ -74,12 +74,16 @@ final class SingleNodeRedisPubSub( ) = ZIO.serviceWithZIO[BinaryCodec] { implicit codec => for { - promise <- Promise.make[RedisError, Unit] - chunk = StringInput.encode(command) ++ NonEmptyList(StringInput).encode((key.value, keys.map(_.value))) - streams <- makeStreams(key :: keys) - _ <- reqQueue.offer(Request(chunk, promise)) - _ <- promise.await - } yield streams + promises <- Promise.make[RedisError, PushProtocol].replicateZIO(keys.size + 1).map(Chunk.fromIterable(_)) + chunk = StringInput.encode(command) ++ NonEmptyList(StringInput).encode((key.value, keys.map(_.value))) + _ <- reqQueue.offer(Request(chunk, promises)) + streams <- ZIO.foreach((key :: keys) zip promises) { case (key, promise) => + for { + hub <- getHub(key) + stream = ZStream.fromHub(hub).flattenTake + } yield ZStream.fromZIO(promise.await) concat stream + } + } yield Chunk.fromIterable(streams) } private def makeUnsubscriptionStream( @@ -88,28 +92,18 @@ final class SingleNodeRedisPubSub( ) = ZIO.serviceWithZIO[BinaryCodec] { implicit codec => for { - targets <- keys - chunk = StringInput.encode(command) ++ Varargs(StringInput).encode(targets.map(_.value)) - promise <- Promise.make[RedisError, Unit] - _ <- reqQueue.offer(Request(chunk, promise)) - _ <- promise.await - streams <- - ZIO.foreach(targets)(key => - for { - promise <- Promise.make[RedisError, PushProtocol] - _ <- unsubscribedRef.update(_ + (key -> promise)) - } yield ZStream.fromZIO(promise.await) - ) + targets <- keys + chunk = StringInput.encode(command) ++ Varargs(StringInput).encode(targets.map(_.value)) + promises <- Promise.make[RedisError, PushProtocol].replicateZIO(targets.size).map(Chunk.fromIterable(_)) + _ <- reqQueue.offer(Request(chunk, promises)) + streams = promises.map(promise => ZStream.fromZIO(promise.await)) } yield streams } - private def makeStreams(keys: List[SubscriptionKey]): UIO[List[Stream[RedisError, PushProtocol]]] = - ZIO.foreach(keys)(getHub(_).map(ZStream.fromHub(_))) - private def getHub(key: SubscriptionKey) = { def makeNewHub = Hub - .unbounded[PushProtocol] + .unbounded[Take[RedisError, PushProtocol]] .tap(hub => pubSubHubsRef.update(_ + (key -> hub))) for { @@ -134,43 +128,45 @@ final class SingleNodeRedisPubSub( .write(bytes) .mapError(RedisError.IOError(_)) .tapBoth( - e => ZIO.foreachDiscard(reqs)(_.promise.fail(e)), - _ => ZIO.foreachDiscard(reqs)(_.promise.succeed(())) + e => ZIO.foreachDiscard(reqs.flatMap(_.promises))(_.fail(e)), + _ => ZIO.foreachDiscard(reqs.map(_.promises))(resQueue.offerAll(_)) ) } private def receive: ZIO[BinaryCodec, RedisError, Unit] = { - def releaseHub(key: SubscriptionKey) = - for { - pubSubs <- pubSubHubsRef.get - hubOpt = pubSubs.get(key) - _ <- ZIO.fromOption(hubOpt).flatMap(_.shutdown).orElse(ZIO.unit) - _ <- pubSubHubsRef.update(_ - key) - } yield () + def handlePushProtocolMessage(msg: PushProtocol): UIO[Unit] = { + def releasePendingPromise(msg: PushProtocol): UIO[Unit] = + resQueue.take.flatMap(_.succeed(msg)).unit - def handlePushProtocolMessage(msg: PushProtocol) = msg match { - case msg @ PushProtocol.Unsubscribe(channel, _) => + def handleUnsubscription(key: SubscriptionKey, msg: PushProtocol): UIO[Unit] = for { - _ <- releaseHub(channelKey(channel)) - map <- unsubscribedRef.get - promise <- ZIO - .fromOption(map.get(channelKey(channel))) - .orElseFail(RedisError.NoUnsubscribeRequest(channel)) - _ <- promise.succeed(msg) + _ <- releasePendingPromise(msg) + pubSubs <- pubSubHubsRef.get + hubOpt = pubSubs.get(key) + _ <- ZIO.fromOption(hubOpt).flatMap(_.shutdown).orElse(ZIO.unit) + _ <- pubSubHubsRef.update(_ - key) } yield () - case msg @ PushProtocol.PUnsubscribe(pattern, _) => + + def handleSubscription(key: SubscriptionKey, msg: PushProtocol): UIO[Unit] = for { - _ <- releaseHub(patternKey(pattern)) - map <- unsubscribedRef.get - promise <- ZIO - .fromOption(map.get(patternKey(pattern))) - .orElseFail(RedisError.NoUnsubscribeRequest(pattern)) - _ <- promise.succeed(msg) + _ <- resQueue.take.flatMap(_.succeed(msg)) + _ <- getHub(key).flatMap(_.offer(Take.single(msg))) } yield () - case msg @ PushProtocol.Subscribe(channel, _) => getHub(channelKey(channel)).flatMap(_.offer(msg)) - case msg @ PushProtocol.PSubscribe(pattern, _) => getHub(patternKey(pattern)).flatMap(_.offer(msg)) - case msg @ PushProtocol.Message(channel, _) => getHub(channelKey(channel)).flatMap(_.offer(msg)) - case msg @ PushProtocol.PMessage(pattern, _, _) => getHub(patternKey(pattern)).flatMap(_.offer(msg)) + + msg match { + case msg @ PushProtocol.Unsubscribe(channel, _) => + handleUnsubscription(channelKey(channel), msg) + case msg @ PushProtocol.PUnsubscribe(pattern, _) => + handleUnsubscription(patternKey(pattern), msg) + case msg @ PushProtocol.Subscribe(channel, _) => + handleSubscription(channelKey(channel), msg) + case msg @ PushProtocol.PSubscribe(pattern, _) => + handleSubscription(patternKey(pattern), msg) + case msg @ PushProtocol.Message(channel, _) => + getHub(channelKey(channel)).flatMap(_.offer(Take.single(msg))).unit + case msg @ PushProtocol.PMessage(pattern, _, _) => + getHub(patternKey(pattern)).flatMap(_.offer(Take.single(msg))).unit + } } ZIO.serviceWithZIO[BinaryCodec] { implicit codec => @@ -220,7 +216,7 @@ object SingleNodeRedisPubSub { } private final case class Request( command: Chunk[RespValue.BulkString], - promise: Promise[RedisError, Unit] + promises: Chunk[Promise[RedisError, PushProtocol]] ) private final val True: Any => Boolean = _ => true @@ -229,11 +225,11 @@ object SingleNodeRedisPubSub { def create(conn: RedisConnection) = for { - hubRef <- Ref.make(Map.empty[SubscriptionKey, Hub[PushProtocol]]) - unsubscribedRef <- Ref.make(Map.empty[SubscriptionKey, Promise[RedisError, PushProtocol]]) - reqQueue <- Queue.bounded[Request](RequestQueueSize) - pubSub = new SingleNodeRedisPubSub(hubRef, unsubscribedRef, reqQueue, conn) - _ <- pubSub.run.forkScoped - _ <- logScopeFinalizer(s"$pubSub Node PubSub is closed") + hubRef <- Ref.make(Map.empty[SubscriptionKey, Hub[Take[RedisError, PushProtocol]]]) + reqQueue <- Queue.bounded[Request](RequestQueueSize) + resQueue <- Queue.unbounded[Promise[RedisError, PushProtocol]] + pubSub = new SingleNodeRedisPubSub(hubRef, reqQueue, resQueue, conn) + _ <- pubSub.run.forkScoped + _ <- logScopeFinalizer(s"$pubSub Node PubSub is closed") } yield pubSub } diff --git a/redis/src/main/scala/zio/redis/api/PubSub.scala b/redis/src/main/scala/zio/redis/api/PubSub.scala index 9b8e493c2..df6fb410c 100644 --- a/redis/src/main/scala/zio/redis/api/PubSub.scala +++ b/redis/src/main/scala/zio/redis/api/PubSub.scala @@ -21,12 +21,12 @@ trait PubSub extends RedisEnvironment { runSubscription(PubSubCommand.Subscribe(channel, List.empty), onSubscribe).flatMap(extractOne(channel, _)) } - final def subscribe(channel: String, channels: String*): ResultStreamBuilder[List] = + final def subscribe(channel: String, channels: String*): ResultStreamBuilder[Chunk] = createStreamListBuilder(channel, channels.toList, emptyCallback, isPatterned = false) final def subscribeWithCallback(channel: String, channels: String*)( onSubscribe: PubSubCallback - ): ResultStreamBuilder[List] = + ): ResultStreamBuilder[Chunk] = createStreamListBuilder(channel, channels.toList, onSubscribe, isPatterned = false) final def pSubscribe(pattern: String): ResultStreamBuilder[Id] = @@ -40,12 +40,12 @@ trait PubSub extends RedisEnvironment { runSubscription(PubSubCommand.PSubscribe(pattern, List.empty), onSubscribe).flatMap(extractOne(pattern, _)) } - final def pSubscribe(pattern: String, patterns: String*): ResultStreamBuilder[List] = + final def pSubscribe(pattern: String, patterns: String*): ResultStreamBuilder[Chunk] = createStreamListBuilder(pattern, patterns.toList, emptyCallback, isPatterned = true) final def pSubscribeWithCallback(pattern: String, patterns: String*)( onSubscribe: PubSubCallback - ): ResultStreamBuilder[List] = + ): ResultStreamBuilder[Chunk] = createStreamListBuilder(pattern, patterns.toList, onSubscribe, isPatterned = true) final def unsubscribe(channels: String*): IO[RedisError, Promise[RedisError, Chunk[(String, Long)]]] = @@ -79,9 +79,9 @@ trait PubSub extends RedisEnvironment { keys: List[String], callback: PubSubCallback, isPatterned: Boolean - ): ResultStreamBuilder[List] = - new ResultStreamBuilder[List] { - def returning[R: Schema]: IO[RedisError, List[Stream[RedisError, R]]] = + ): ResultStreamBuilder[Chunk] = + new ResultStreamBuilder[Chunk] { + def returning[R: Schema]: IO[RedisError, Chunk[Stream[RedisError, R]]] = if (isPatterned) runSubscription(PubSubCommand.PSubscribe(key, keys), callback) else runSubscription(PubSubCommand.Subscribe(key, keys), callback) } @@ -112,14 +112,14 @@ trait PubSub extends RedisEnvironment { private def runSubscription[R: Schema]( command: PubSubCommand, onSubscribe: PubSubCallback - ): IO[RedisError, List[Stream[RedisError, R]]] = + ): IO[RedisError, Chunk[Stream[RedisError, R]]] = RedisPubSubCommand(command, codec, pubSub).run[R] { case PushProtocol.Subscribe(key, numOfSubscription) => onSubscribe(key, numOfSubscription) case PushProtocol.PSubscribe(key, numOfSubscription) => onSubscribe(key, numOfSubscription) case _ => ZIO.unit } - private def extractOne[A](key: String, elements: List[A]) = + private def extractOne[A](key: String, elements: Chunk[A]) = ZIO.fromOption(elements.headOption).orElseFail(RedisError.NoPubSubStream(key)) } From f532e4b454aef460d5f25a30682de2c393b5e55e Mon Sep 17 00:00:00 2001 From: 0pg Date: Sun, 12 Feb 2023 02:25:47 +0900 Subject: [PATCH 18/51] Add pacakge private accessor to PubSubCommand --- redis/src/main/scala/zio/redis/PubSubCommand.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/redis/src/main/scala/zio/redis/PubSubCommand.scala b/redis/src/main/scala/zio/redis/PubSubCommand.scala index f040f399e..b1632c16c 100644 --- a/redis/src/main/scala/zio/redis/PubSubCommand.scala +++ b/redis/src/main/scala/zio/redis/PubSubCommand.scala @@ -1,8 +1,8 @@ package zio.redis -sealed trait PubSubCommand +private[redis] sealed trait PubSubCommand -object PubSubCommand { +private[redis] object PubSubCommand { case class Subscribe( channel: String, channels: List[String] From e9e2a3cabfd317aac1d00e4371ee5c422409187b Mon Sep 17 00:00:00 2001 From: 0pg Date: Sun, 12 Feb 2023 02:32:08 +0900 Subject: [PATCH 19/51] Modify unsubscrption release logic --- redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala index e40c66f80..5eca2d699 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -143,7 +143,7 @@ final class SingleNodeRedisPubSub( _ <- releasePendingPromise(msg) pubSubs <- pubSubHubsRef.get hubOpt = pubSubs.get(key) - _ <- ZIO.fromOption(hubOpt).flatMap(_.shutdown).orElse(ZIO.unit) + _ <- ZIO.fromOption(hubOpt).flatMap(_.offer(Take.end)).orElse(ZIO.unit) _ <- pubSubHubsRef.update(_ - key) } yield () @@ -176,7 +176,7 @@ final class SingleNodeRedisPubSub( .collectSome .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))) .refineToOrDie[RedisError] - .foreach(push => handlePushProtocolMessage(push)) + .foreach(handlePushProtocolMessage(_)) } } From 49c141b383dfd3d81b28eed6048ae656483c9284 Mon Sep 17 00:00:00 2001 From: 0pg Date: Wed, 15 Feb 2023 02:00:48 +0900 Subject: [PATCH 20/51] Modify RedisPubSubCommand interface --- .../main/scala/zio/redis/PubSubCommand.scala | 20 ++++-- .../src/main/scala/zio/redis/RedisError.scala | 1 - .../main/scala/zio/redis/RedisPubSub.scala | 2 +- .../scala/zio/redis/RedisPubSubCommand.scala | 24 ++----- .../zio/redis/SingleNodeRedisPubSub.scala | 31 +++++++-- .../src/main/scala/zio/redis/api/PubSub.scala | 67 ++++++++++++------- 6 files changed, 90 insertions(+), 55 deletions(-) diff --git a/redis/src/main/scala/zio/redis/PubSubCommand.scala b/redis/src/main/scala/zio/redis/PubSubCommand.scala index b1632c16c..c118e6565 100644 --- a/redis/src/main/scala/zio/redis/PubSubCommand.scala +++ b/redis/src/main/scala/zio/redis/PubSubCommand.scala @@ -1,16 +1,26 @@ package zio.redis -private[redis] sealed trait PubSubCommand +import zio.IO + +sealed trait PubSubCommand private[redis] object PubSubCommand { case class Subscribe( channel: String, - channels: List[String] + channels: List[String], + callback: PushProtocol => IO[RedisError, Unit] ) extends PubSubCommand case class PSubscribe( pattern: String, - patterns: List[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 - case class Unsubscribe(channels: List[String]) extends PubSubCommand - case class PUnsubscribe(patterns: List[String]) extends PubSubCommand } diff --git a/redis/src/main/scala/zio/redis/RedisError.scala b/redis/src/main/scala/zio/redis/RedisError.scala index 20c99f459..ccca7efc7 100644 --- a/redis/src/main/scala/zio/redis/RedisError.scala +++ b/redis/src/main/scala/zio/redis/RedisError.scala @@ -47,6 +47,5 @@ object RedisError { def apply(slotAndAddress: (Slot, RedisUri)): Moved = Moved(slotAndAddress._1, slotAndAddress._2) } final case class NoPubSubStream(key: String) extends RedisError - final case class NoUnsubscribeRequest(key: String) extends RedisError final case class IOError(exception: IOException) extends RedisError } diff --git a/redis/src/main/scala/zio/redis/RedisPubSub.scala b/redis/src/main/scala/zio/redis/RedisPubSub.scala index 377b8c7a3..244070e51 100644 --- a/redis/src/main/scala/zio/redis/RedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/RedisPubSub.scala @@ -5,7 +5,7 @@ import zio.stream._ import zio.{Chunk, ZIO, ZLayer} trait RedisPubSub { - def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] + def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, RespValue]]] } object RedisPubSub { diff --git a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala index f5894a044..c297e8cb5 100644 --- a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala +++ b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala @@ -6,29 +6,19 @@ import zio.schema.codec.BinaryCodec import zio.stream._ import zio.{Chunk, IO, ZIO, ZLayer} -final case class RedisPubSubCommand(command: PubSubCommand, codec: BinaryCodec, executor: RedisPubSub) { - def run[R: Schema]( - callback: PushProtocol => IO[RedisError, Unit] - ): IO[RedisError, Chunk[Stream[RedisError, R]]] = { +private[redis] final case class RedisPubSubCommand(command: PubSubCommand, codec: BinaryCodec, executor: RedisPubSub) { + def run[R: Schema](): IO[RedisError, Chunk[Stream[RedisError, R]]] = { val codecLayer = ZLayer.succeed(codec) executor .execute(command) .provideLayer(codecLayer) .map( _.map( - _.tap(callback(_)).mapZIO { - case PushProtocol.Message(_, msg) => - ZIO - .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) - .refineToOrDie[RedisError] - .asSome - case PushProtocol.PMessage(_, _, msg) => - ZIO - .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) - .refineToOrDie[RedisError] - .asSome - case _ => ZIO.none - }.collectSome + _.mapZIO(msg => + ZIO + .attempt(ArbitraryOutput[R]().unsafeDecode(msg)(codec)) + .refineToOrDie[RedisError] + ) ) ) } diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala index 5eca2d699..d10ea1f95 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -6,7 +6,7 @@ import zio.redis.SingleNodeRedisPubSub.{Request, RequestQueueSize, SubscriptionK import zio.redis.api.PubSub import zio.schema.codec.BinaryCodec import zio.stream._ -import zio.{Chunk, ChunkBuilder, Hub, Promise, Queue, Ref, Schedule, UIO, ZIO} +import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Schedule, UIO, ZIO} final class SingleNodeRedisPubSub( pubSubHubsRef: Ref[Map[SubscriptionKey, Hub[Take[RedisError, PushProtocol]]]], @@ -15,15 +15,34 @@ final class SingleNodeRedisPubSub( connection: RedisConnection ) extends RedisPubSub { - def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] = + def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, RespValue]]] = { + def applyCallback( + streams: Chunk[Stream[RedisError, PushProtocol]], + callback: PushProtocol => IO[RedisError, Unit] + ) = + streams.map( + _.tap(callback(_)).map { + case PushProtocol.Message(_, msg) => Some(msg) + case PushProtocol.PMessage(_, _, msg) => Some(msg) + case _ => None + }.collectSome + ) + command match { - case PubSubCommand.Subscribe(channel, channels) => + case PubSubCommand.Subscribe(channel, channels, callback) => subscribe(channel, channels) - case PubSubCommand.PSubscribe(pattern, patterns) => + .map(applyCallback(_, callback)) + case PubSubCommand.PSubscribe(pattern, patterns, callback) => pSubscribe(pattern, patterns) - case PubSubCommand.Unsubscribe(channels) => unsubscribe(channels) - case PubSubCommand.PUnsubscribe(patterns) => pUnsubscribe(patterns) + .map(applyCallback(_, callback)) + case PubSubCommand.Unsubscribe(channels, callback) => + unsubscribe(channels) + .map(applyCallback(_, callback)) + case PubSubCommand.PUnsubscribe(patterns, callback) => + pUnsubscribe(patterns) + .map(applyCallback(_, callback)) } + } private def subscribe( channel: String, diff --git a/redis/src/main/scala/zio/redis/api/PubSub.scala b/redis/src/main/scala/zio/redis/api/PubSub.scala index df6fb410c..dd7c37b3e 100644 --- a/redis/src/main/scala/zio/redis/api/PubSub.scala +++ b/redis/src/main/scala/zio/redis/api/PubSub.scala @@ -7,7 +7,7 @@ import zio.redis._ import zio.redis.options.PubSub.NumberOfSubscribers import zio.schema.Schema import zio.stream._ -import zio.{Chunk, IO, Promise, Ref, ZIO} +import zio.{Chunk, IO, Promise, Ref, UIO, ZIO} trait PubSub extends RedisEnvironment { import PubSub._ @@ -18,7 +18,8 @@ trait PubSub extends RedisEnvironment { final def subscribeWithCallback(channel: String)(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = new ResultStreamBuilder[Id] { def returning[R: Schema]: IO[RedisError, Id[Stream[RedisError, R]]] = - runSubscription(PubSubCommand.Subscribe(channel, List.empty), onSubscribe).flatMap(extractOne(channel, _)) + runSubscription(PubSubCommand.Subscribe(channel, List.empty, createSubscriptionCallback(onSubscribe))) + .flatMap(extractOne(channel, _)) } final def subscribe(channel: String, channels: String*): ResultStreamBuilder[Chunk] = @@ -37,7 +38,8 @@ trait PubSub extends RedisEnvironment { )(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = new ResultStreamBuilder[Id] { def returning[R: Schema]: IO[RedisError, Id[Stream[RedisError, R]]] = - runSubscription(PubSubCommand.PSubscribe(pattern, List.empty), onSubscribe).flatMap(extractOne(pattern, _)) + runSubscription(PubSubCommand.PSubscribe(pattern, List.empty, createSubscriptionCallback(onSubscribe))) + .flatMap(extractOne(pattern, _)) } final def pSubscribe(pattern: String, patterns: String*): ResultStreamBuilder[Chunk] = @@ -49,10 +51,18 @@ trait PubSub extends RedisEnvironment { createStreamListBuilder(pattern, patterns.toList, onSubscribe, isPatterned = true) final def unsubscribe(channels: String*): IO[RedisError, Promise[RedisError, Chunk[(String, Long)]]] = - runUnsubscription(PubSubCommand.Unsubscribe(channels.toList)) + for { + ref <- Ref.make(Chunk.empty[(String, Long)]) + callback = createUnsubscriptionCallback(ref) + promise <- runUnsubscription(PubSubCommand.Unsubscribe(channels.toList, callback), ref) + } yield promise final def pUnsubscribe(patterns: String*): IO[RedisError, Promise[RedisError, Chunk[(String, Long)]]] = - runUnsubscription(PubSubCommand.PUnsubscribe(patterns.toList)) + for { + ref <- Ref.make(Chunk.empty[(String, Long)]) + callback = createUnsubscriptionCallback(ref) + promise <- runUnsubscription(PubSubCommand.PUnsubscribe(patterns.toList, callback), ref) + } yield promise final def publish[A: Schema](channel: String, message: A): IO[RedisError, Long] = { val command = RedisCommand(Publish, Tuple2(StringInput, ArbitraryInput[A]()), LongOutput, codec, executor) @@ -81,43 +91,50 @@ trait PubSub extends RedisEnvironment { isPatterned: Boolean ): ResultStreamBuilder[Chunk] = new ResultStreamBuilder[Chunk] { - def returning[R: Schema]: IO[RedisError, Chunk[Stream[RedisError, R]]] = - if (isPatterned) runSubscription(PubSubCommand.PSubscribe(key, keys), callback) - else runSubscription(PubSubCommand.Subscribe(key, keys), callback) + def returning[R: Schema]: IO[RedisError, Chunk[Stream[RedisError, R]]] = { + val subscriptionCallback = createSubscriptionCallback(callback); + if (isPatterned) runSubscription(PubSubCommand.PSubscribe(key, keys, subscriptionCallback)) + else runSubscription(PubSubCommand.Subscribe(key, keys, subscriptionCallback)) + } } private def runUnsubscription( - command: PubSubCommand + command: PubSubCommand, + resultRef: Ref[Chunk[(String, Long)]] ): IO[RedisError, Promise[RedisError, Chunk[(String, Long)]]] = for { promise <- Promise.make[RedisError, Chunk[(String, Long)]] - ref <- Ref.make(Chunk.empty[(String, Long)]) - streams <- RedisPubSubCommand(command, codec, pubSub).run[Unit] { - case PushProtocol.Unsubscribe(channel, numOfSubscription) => - ref.update(_ appended (channel, numOfSubscription)) - case PushProtocol.PUnsubscribe(pattern, numOfSubscription) => - ref.update(_ appended (pattern, numOfSubscription)) - case _ => ZIO.unit - } + streams <- RedisPubSubCommand(command, codec, pubSub).run[Unit]() _ <- streams .fold(ZStream.empty)(_ merge _) .runDrain .onDone( e => promise.fail(e), - _ => ref.get.flatMap(promise.succeed(_)) + _ => resultRef.get.flatMap(promise.succeed(_)) ) .fork } yield promise private def runSubscription[R: Schema]( - command: PubSubCommand, - onSubscribe: PubSubCallback + command: PubSubCommand ): IO[RedisError, Chunk[Stream[RedisError, R]]] = - RedisPubSubCommand(command, codec, pubSub).run[R] { - case PushProtocol.Subscribe(key, numOfSubscription) => onSubscribe(key, numOfSubscription) - case PushProtocol.PSubscribe(key, numOfSubscription) => onSubscribe(key, numOfSubscription) - case _ => ZIO.unit - } + RedisPubSubCommand(command, codec, pubSub).run[R]() + + private def createSubscriptionCallback( + onSubscribe: PubSubCallback + ): PushProtocol => IO[RedisError, Unit] = { + case PushProtocol.Subscribe(key, numOfSubscription) => onSubscribe(key, numOfSubscription) + case PushProtocol.PSubscribe(key, numOfSubscription) => onSubscribe(key, numOfSubscription) + case _ => ZIO.unit + } + + private def createUnsubscriptionCallback(resultRef: Ref[Chunk[(String, Long)]]): PushProtocol => UIO[Unit] = { + case PushProtocol.Unsubscribe(channel, numOfSubscription) => + resultRef.update(_ appended (channel, numOfSubscription)) + case PushProtocol.PUnsubscribe(pattern, numOfSubscription) => + resultRef.update(_ appended (pattern, numOfSubscription)) + case _ => ZIO.unit + } private def extractOne[A](key: String, elements: Chunk[A]) = ZIO.fromOption(elements.headOption).orElseFail(RedisError.NoPubSubStream(key)) From 8b6ce730e4820aa9c7c82945fcf16ec1454fbe94 Mon Sep 17 00:00:00 2001 From: 0pg Date: Wed, 15 Feb 2023 02:01:08 +0900 Subject: [PATCH 21/51] Apply changes to test --- redis/src/test/scala/zio/redis/ApiSpec.scala | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/redis/src/test/scala/zio/redis/ApiSpec.scala b/redis/src/test/scala/zio/redis/ApiSpec.scala index e3e3c83e5..7106110e9 100644 --- a/redis/src/test/scala/zio/redis/ApiSpec.scala +++ b/redis/src/test/scala/zio/redis/ApiSpec.scala @@ -1,8 +1,8 @@ package zio.redis -import zio._ import zio.test.TestAspect._ import zio.test._ +import zio.{ZLayer, _} object ApiSpec extends ConnectionSpec @@ -63,7 +63,8 @@ object ApiSpec RedisPubSub.layer, Redis.layer, ZLayer.succeed(codec), - ZLayer.succeed(RedisClusterConfig(Chunk(RedisUri("localhost", 5000)))) + ZLayer.succeed(RedisClusterConfig(Chunk(RedisUri("localhost", 5000)))), + ZLayer.succeed(RedisConfig("localhost", 5000)) ).filterNotTags(_.contains(BaseSpec.ClusterExecutorUnsupported)) .getOrElse(Spec.empty) } From f44c7fc28b2f02a48d4f5a84ca8f4ae51fb7856d Mon Sep 17 00:00:00 2001 From: 0pg Date: Wed, 15 Feb 2023 02:14:31 +0900 Subject: [PATCH 22/51] Add pubsub layer in Main --- example/src/main/scala/example/Main.scala | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/example/src/main/scala/example/Main.scala b/example/src/main/scala/example/Main.scala index 709e90438..5b8d42c7d 100644 --- a/example/src/main/scala/example/Main.scala +++ b/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) ) From 758ea9be2224a3be1cbbf68011d269a4a7a502f8 Mon Sep 17 00:00:00 2001 From: 0pg Date: Mon, 20 Feb 2023 22:47:35 +0900 Subject: [PATCH 23/51] Remove BinaryCodec layer to RedisPubSub interface --- .../main/scala/zio/redis/RedisPubSub.scala | 10 +- .../scala/zio/redis/RedisPubSubCommand.scala | 7 +- .../zio/redis/SingleNodeRedisPubSub.scala | 110 +++++++++--------- 3 files changed, 61 insertions(+), 66 deletions(-) diff --git a/redis/src/main/scala/zio/redis/RedisPubSub.scala b/redis/src/main/scala/zio/redis/RedisPubSub.scala index 244070e51..c27dee8bf 100644 --- a/redis/src/main/scala/zio/redis/RedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/RedisPubSub.scala @@ -2,10 +2,10 @@ package zio.redis import zio.schema.codec.BinaryCodec import zio.stream._ -import zio.{Chunk, ZIO, ZLayer} +import zio.{Chunk, IO, ZIO, ZLayer} trait RedisPubSub { - def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, RespValue]]] + def execute(command: PubSubCommand): IO[RedisError, Chunk[Stream[RedisError, RespValue]]] } object RedisPubSub { @@ -17,6 +17,10 @@ object RedisPubSub { private lazy val pubSublayer: ZLayer[RedisConnection with BinaryCodec, RedisError.IOError, RedisPubSub] = ZLayer.scoped( - ZIO.service[RedisConnection].flatMap(SingleNodeRedisPubSub.create(_)) + for { + conn <- ZIO.service[RedisConnection] + codec <- ZIO.service[BinaryCodec] + pubSub <- SingleNodeRedisPubSub.create(conn, codec) + } yield pubSub ) } diff --git a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala index c297e8cb5..2da0403cc 100644 --- a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala +++ b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala @@ -4,14 +4,12 @@ import zio.redis.Output.ArbitraryOutput import zio.schema.Schema import zio.schema.codec.BinaryCodec import zio.stream._ -import zio.{Chunk, IO, ZIO, ZLayer} +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]]] = { - val codecLayer = ZLayer.succeed(codec) + def run[R: Schema](): IO[RedisError, Chunk[Stream[RedisError, R]]] = executor .execute(command) - .provideLayer(codecLayer) .map( _.map( _.mapZIO(msg => @@ -21,5 +19,4 @@ private[redis] final case class RedisPubSubCommand(command: PubSubCommand, codec ) ) ) - } } diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala index d10ea1f95..d79025261 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -12,10 +12,11 @@ final class SingleNodeRedisPubSub( pubSubHubsRef: Ref[Map[SubscriptionKey, Hub[Take[RedisError, PushProtocol]]]], reqQueue: Queue[Request], resQueue: Queue[Promise[RedisError, PushProtocol]], - connection: RedisConnection + connection: RedisConnection, + implicit val codec: BinaryCodec ) extends RedisPubSub { - def execute(command: PubSubCommand): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, RespValue]]] = { + def execute(command: PubSubCommand): IO[RedisError, Chunk[Stream[RedisError, RespValue]]] = { def applyCallback( streams: Chunk[Stream[RedisError, PushProtocol]], callback: PushProtocol => IO[RedisError, Unit] @@ -47,7 +48,7 @@ final class SingleNodeRedisPubSub( private def subscribe( channel: String, channels: List[String] - ): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] = + ): IO[RedisError, Chunk[Stream[RedisError, PushProtocol]]] = makeSubscriptionStream( PubSub.Subscribe, channelKey(channel), @@ -57,7 +58,7 @@ final class SingleNodeRedisPubSub( private def pSubscribe( pattern: String, patterns: List[String] - ): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] = + ): IO[RedisError, Chunk[Stream[RedisError, PushProtocol]]] = makeSubscriptionStream( PubSub.PSubscribe, patternKey(pattern), @@ -66,7 +67,7 @@ final class SingleNodeRedisPubSub( private def unsubscribe( channels: List[String] - ): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] = + ): IO[RedisError, Chunk[Stream[RedisError, PushProtocol]]] = makeUnsubscriptionStream( PubSub.Unsubscribe, if (channels.nonEmpty) @@ -77,7 +78,7 @@ final class SingleNodeRedisPubSub( private def pUnsubscribe( patterns: List[String] - ): ZIO[BinaryCodec, RedisError, Chunk[Stream[RedisError, PushProtocol]]] = + ): IO[RedisError, Chunk[Stream[RedisError, PushProtocol]]] = makeUnsubscriptionStream( PubSub.PUnsubscribe, if (patterns.nonEmpty) @@ -90,34 +91,30 @@ final class SingleNodeRedisPubSub( command: String, key: SubscriptionKey, keys: List[SubscriptionKey] - ) = - ZIO.serviceWithZIO[BinaryCodec] { implicit codec => - for { - promises <- Promise.make[RedisError, PushProtocol].replicateZIO(keys.size + 1).map(Chunk.fromIterable(_)) - chunk = StringInput.encode(command) ++ NonEmptyList(StringInput).encode((key.value, keys.map(_.value))) - _ <- reqQueue.offer(Request(chunk, promises)) - streams <- ZIO.foreach((key :: keys) zip promises) { case (key, promise) => - for { - hub <- getHub(key) - stream = ZStream.fromHub(hub).flattenTake - } yield ZStream.fromZIO(promise.await) concat stream - } - } yield Chunk.fromIterable(streams) - } + )(implicit codec: BinaryCodec) = + for { + promises <- Promise.make[RedisError, PushProtocol].replicateZIO(keys.size + 1).map(Chunk.fromIterable(_)) + chunk = StringInput.encode(command) ++ NonEmptyList(StringInput).encode((key.value, keys.map(_.value))) + _ <- reqQueue.offer(Request(chunk, promises)) + streams <- ZIO.foreach((key :: keys) zip promises) { case (key, promise) => + for { + hub <- getHub(key) + stream = ZStream.fromHub(hub).flattenTake + } yield ZStream.fromZIO(promise.await) concat stream + } + } yield Chunk.fromIterable(streams) private def makeUnsubscriptionStream( command: String, keys: UIO[List[SubscriptionKey]] - ) = - ZIO.serviceWithZIO[BinaryCodec] { implicit codec => - for { - targets <- keys - chunk = StringInput.encode(command) ++ Varargs(StringInput).encode(targets.map(_.value)) - promises <- Promise.make[RedisError, PushProtocol].replicateZIO(targets.size).map(Chunk.fromIterable(_)) - _ <- reqQueue.offer(Request(chunk, promises)) - streams = promises.map(promise => ZStream.fromZIO(promise.await)) - } yield streams - } + )(implicit codec: BinaryCodec) = + for { + targets <- keys + chunk = StringInput.encode(command) ++ Varargs(StringInput).encode(targets.map(_.value)) + promises <- Promise.make[RedisError, PushProtocol].replicateZIO(targets.size).map(Chunk.fromIterable(_)) + _ <- reqQueue.offer(Request(chunk, promises)) + streams = promises.map(promise => ZStream.fromZIO(promise.await)) + } yield streams private def getHub(key: SubscriptionKey) = { def makeNewHub = @@ -152,7 +149,7 @@ final class SingleNodeRedisPubSub( ) } - private def receive: ZIO[BinaryCodec, RedisError, Unit] = { + private def receive: IO[RedisError, Unit] = { def handlePushProtocolMessage(msg: PushProtocol): UIO[Unit] = { def releasePendingPromise(msg: PushProtocol): UIO[Unit] = resQueue.take.flatMap(_.succeed(msg)).unit @@ -188,31 +185,28 @@ final class SingleNodeRedisPubSub( } } - ZIO.serviceWithZIO[BinaryCodec] { implicit codec => - connection.read - .mapError(RedisError.IOError(_)) - .via(RespValue.decoder) - .collectSome - .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))) - .refineToOrDie[RedisError] - .foreach(handlePushProtocolMessage(_)) - } + connection.read + .mapError(RedisError.IOError(_)) + .via(RespValue.decoder) + .collectSome + .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))) + .refineToOrDie[RedisError] + .foreach(handlePushProtocolMessage(_)) } - private def resubscribe: ZIO[BinaryCodec, RedisError, Unit] = - ZIO.serviceWithZIO[BinaryCodec] { implicit codec => - def makeCommand(name: String, keys: Set[String]) = - RespValue.Array(StringInput.encode(name) ++ Varargs(StringInput).encode(keys)).serialize - - for { - keySet <- pubSubHubsRef.get.map(_.keySet) - (channels, patterns) = keySet.partition(_.isChannel) - _ <- (connection.write(makeCommand(PubSub.Subscribe, channels.map(_.value))).when(channels.nonEmpty) *> - connection.write(makeCommand(PubSub.PSubscribe, patterns.map(_.value))).when(patterns.nonEmpty)) - .mapError(RedisError.IOError(_)) - .retryWhile(True) - } yield () - } + private def resubscribe: IO[RedisError, Unit] = { + def makeCommand(name: String, keys: Set[String]) = + RespValue.Array(StringInput.encode(name) ++ Varargs(StringInput).encode(keys)).serialize + + for { + keySet <- pubSubHubsRef.get.map(_.keySet) + (channels, patterns) = keySet.partition(_.isChannel) + _ <- (connection.write(makeCommand(PubSub.Subscribe, channels.map(_.value))).when(channels.nonEmpty) *> + connection.write(makeCommand(PubSub.PSubscribe, patterns.map(_.value))).when(patterns.nonEmpty)) + .mapError(RedisError.IOError(_)) + .retryWhile(True) + } yield () + } private def patternKey(key: String) = SubscriptionKey(key, true) private def channelKey(key: String) = SubscriptionKey(key, false) @@ -221,9 +215,9 @@ final class SingleNodeRedisPubSub( * 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: ZIO[BinaryCodec, RedisError, AnyVal] = + val run: IO[RedisError, AnyVal] = ZIO.logTrace(s"$this Executable sender and reader has been started") *> - (send.repeat[BinaryCodec, Long](Schedule.forever) race receive) + (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")) @@ -242,12 +236,12 @@ object SingleNodeRedisPubSub { private final val RequestQueueSize = 16 - def create(conn: RedisConnection) = + def create(conn: RedisConnection, codec: BinaryCodec) = for { hubRef <- Ref.make(Map.empty[SubscriptionKey, Hub[Take[RedisError, PushProtocol]]]) reqQueue <- Queue.bounded[Request](RequestQueueSize) resQueue <- Queue.unbounded[Promise[RedisError, PushProtocol]] - pubSub = new SingleNodeRedisPubSub(hubRef, reqQueue, resQueue, conn) + pubSub = new SingleNodeRedisPubSub(hubRef, reqQueue, resQueue, conn, codec) _ <- pubSub.run.forkScoped _ <- logScopeFinalizer(s"$pubSub Node PubSub is closed") } yield pubSub From 9febcf8b7102f127bd1432e60b5996a3799bfe81 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 21 Feb 2023 02:02:11 +0900 Subject: [PATCH 24/51] Modify SingleNodeRedisPubSub failover step --- redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala index d79025261..a311ae2bb 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala @@ -194,6 +194,8 @@ final class SingleNodeRedisPubSub( .foreach(handlePushProtocolMessage(_)) } + private def drainWith(e: RedisError): UIO[Unit] = resQueue.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) + private def resubscribe: IO[RedisError, Unit] = { def makeCommand(name: String, keys: Set[String]) = RespValue.Array(StringInput.encode(name) ++ Varargs(StringInput).encode(keys)).serialize @@ -216,9 +218,9 @@ final class SingleNodeRedisPubSub( * connection. Only exits by interruption or defect. */ val run: IO[RedisError, AnyVal] = - ZIO.logTrace(s"$this Executable sender and reader has been started") *> + 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) + .tapError(e => ZIO.logWarning(s"Reconnecting due to error: $e") *> drainWith(e) *> resubscribe) .retryWhile(True) .tapError(e => ZIO.logError(s"Executor exiting: $e")) } From 3fc77c0b8ebe005126b5ee7fbdbea282be68b45c Mon Sep 17 00:00:00 2001 From: 0pg Date: Fri, 31 Mar 2023 05:09:25 +0900 Subject: [PATCH 25/51] Separate Publish and Subscribe --- .../redis/benchmarks/BenchmarkRuntime.scala | 2 +- example/src/main/scala/example/Main.scala | 4 +- .../scala/zio/redis/ClusterExecutor.scala | 2 +- redis/src/main/scala/zio/redis/Output.scala | 2 +- .../main/scala/zio/redis/PubSubCommand.scala | 26 -- redis/src/main/scala/zio/redis/Redis.scala | 9 +- .../scala/zio/redis/RedisEnvironment.scala | 1 - .../src/main/scala/zio/redis/RedisError.scala | 7 +- .../main/scala/zio/redis/RedisPubSub.scala | 26 -- .../scala/zio/redis/RedisPubSubCommand.scala | 110 +++++++- .../scala/zio/redis/RedisSubscription.scala | 18 ++ .../main/scala/zio/redis/ResultBuilder.scala | 5 - .../zio/redis/SingleNodeRedisPubSub.scala | 250 ------------------ .../SingleNodeSubscriptionExecutor.scala | 146 ++++++++++ .../zio/redis/SubscribeEnvironment.scala | 8 + .../zio/redis/SubscriptionExecutor.scala | 24 ++ .../src/main/scala/zio/redis/api/PubSub.scala | 154 ----------- .../main/scala/zio/redis/api/Publishing.scala | 39 +++ .../main/scala/zio/redis/api/Subscribe.scala | 74 ++++++ .../main/scala/zio/redis/options/PubSub.scala | 10 +- redis/src/main/scala/zio/redis/package.scala | 3 +- redis/src/test/scala/zio/redis/ApiSpec.scala | 5 +- .../scala/zio/redis/ClusterExecutorSpec.scala | 2 - redis/src/test/scala/zio/redis/KeysSpec.scala | 1 - .../src/test/scala/zio/redis/OutputSpec.scala | 1 + .../src/test/scala/zio/redis/PubSubSpec.scala | 166 +++++++----- 26 files changed, 527 insertions(+), 568 deletions(-) delete mode 100644 redis/src/main/scala/zio/redis/PubSubCommand.scala delete mode 100644 redis/src/main/scala/zio/redis/RedisPubSub.scala create mode 100644 redis/src/main/scala/zio/redis/RedisSubscription.scala delete mode 100644 redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala create mode 100644 redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala create mode 100644 redis/src/main/scala/zio/redis/SubscribeEnvironment.scala create mode 100644 redis/src/main/scala/zio/redis/SubscriptionExecutor.scala delete mode 100644 redis/src/main/scala/zio/redis/api/PubSub.scala create mode 100644 redis/src/main/scala/zio/redis/api/Publishing.scala create mode 100644 redis/src/main/scala/zio/redis/api/Subscribe.scala diff --git a/benchmarks/src/main/scala/zio/redis/benchmarks/BenchmarkRuntime.scala b/benchmarks/src/main/scala/zio/redis/benchmarks/BenchmarkRuntime.scala index c8992787f..656dafcc4 100644 --- a/benchmarks/src/main/scala/zio/redis/benchmarks/BenchmarkRuntime.scala +++ b/benchmarks/src/main/scala/zio/redis/benchmarks/BenchmarkRuntime.scala @@ -35,7 +35,7 @@ object BenchmarkRuntime { private final val Layer = ZLayer.make[Redis]( RedisExecutor.local, - RedisPubSub.local, + SubscriptionExecutor.local, ZLayer.succeed[BinaryCodec](ProtobufCodec), Redis.layer ) diff --git a/example/src/main/scala/example/Main.scala b/example/src/main/scala/example/Main.scala index 5b8d42c7d..e2a373709 100644 --- a/example/src/main/scala/example/Main.scala +++ b/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, RedisPubSub} +import zio.redis.{Redis, RedisExecutor, SubscriptionExecutor} import zio.schema.codec.{BinaryCodec, ProtobufCodec} object Main extends ZIOAppDefault { @@ -33,7 +33,7 @@ object Main extends ZIOAppDefault { ContributorsCache.layer, HttpClientZioBackend.layer(), RedisExecutor.layer, - RedisPubSub.layer, + SubscriptionExecutor.layer, Redis.layer, ZLayer.succeed[BinaryCodec](ProtobufCodec) ) diff --git a/redis/src/main/scala/zio/redis/ClusterExecutor.scala b/redis/src/main/scala/zio/redis/ClusterExecutor.scala index 7c2b702ad..577d534b2 100644 --- a/redis/src/main/scala/zio/redis/ClusterExecutor.scala +++ b/redis/src/main/scala/zio/redis/ClusterExecutor.scala @@ -153,7 +153,7 @@ object ClusterExecutor { 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 pubSubLayer = redisConfigLayer ++ codecLayer >>> SubscriptionExecutor.layer val redisLayer = executorLayer ++ pubSubLayer ++ codecLayer >>> Redis.layer for { closableScope <- Scope.make diff --git a/redis/src/main/scala/zio/redis/Output.scala b/redis/src/main/scala/zio/redis/Output.scala index 72224361e..03f42664d 100644 --- a/redis/src/main/scala/zio/redis/Output.scala +++ b/redis/src/main/scala/zio/redis/Output.scala @@ -18,7 +18,7 @@ package zio.redis import zio._ import zio.redis.options.Cluster.{Node, Partition, SlotRange} -import zio.redis.options.PubSub.NumberOfSubscribers +import zio.redis.options.PubSub.{NumberOfSubscribers, PushProtocol} import zio.schema.Schema import zio.schema.codec.BinaryCodec diff --git a/redis/src/main/scala/zio/redis/PubSubCommand.scala b/redis/src/main/scala/zio/redis/PubSubCommand.scala deleted file mode 100644 index c118e6565..000000000 --- a/redis/src/main/scala/zio/redis/PubSubCommand.scala +++ /dev/null @@ -1,26 +0,0 @@ -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 -} diff --git a/redis/src/main/scala/zio/redis/Redis.scala b/redis/src/main/scala/zio/redis/Redis.scala index 3f7ebaed4..f1c82c08d 100644 --- a/redis/src/main/scala/zio/redis/Redis.scala +++ b/redis/src/main/scala/zio/redis/Redis.scala @@ -32,17 +32,16 @@ trait Redis with api.Streams with api.Scripting with api.Cluster - with api.PubSub + with api.Publishing object Redis { - lazy val layer: URLayer[RedisExecutor with RedisPubSub with BinaryCodec, Redis] = + lazy val layer: URLayer[RedisExecutor with BinaryCodec, Redis] = ZLayer { for { executor <- ZIO.service[RedisExecutor] - pubSub <- ZIO.service[RedisPubSub] codec <- ZIO.service[BinaryCodec] - } yield Live(codec, executor, pubSub) + } yield Live(codec, executor) } - private final case class Live(codec: BinaryCodec, executor: RedisExecutor, pubSub: RedisPubSub) extends Redis + private final case class Live(codec: BinaryCodec, executor: RedisExecutor) extends Redis } diff --git a/redis/src/main/scala/zio/redis/RedisEnvironment.scala b/redis/src/main/scala/zio/redis/RedisEnvironment.scala index 7dc01dc64..886ef9a8d 100644 --- a/redis/src/main/scala/zio/redis/RedisEnvironment.scala +++ b/redis/src/main/scala/zio/redis/RedisEnvironment.scala @@ -21,5 +21,4 @@ import zio.schema.codec.BinaryCodec private[redis] trait RedisEnvironment { protected def codec: BinaryCodec protected def executor: RedisExecutor - protected def pubSub: RedisPubSub } diff --git a/redis/src/main/scala/zio/redis/RedisError.scala b/redis/src/main/scala/zio/redis/RedisError.scala index ccca7efc7..c49ae767b 100644 --- a/redis/src/main/scala/zio/redis/RedisError.scala +++ b/redis/src/main/scala/zio/redis/RedisError.scala @@ -46,6 +46,9 @@ object RedisError { object Moved { def apply(slotAndAddress: (Slot, RedisUri)): Moved = Moved(slotAndAddress._1, slotAndAddress._2) } - final case class NoPubSubStream(key: String) extends RedisError - final case class IOError(exception: IOException) extends RedisError + final case class NoPubSubStream(key: String) 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 } diff --git a/redis/src/main/scala/zio/redis/RedisPubSub.scala b/redis/src/main/scala/zio/redis/RedisPubSub.scala deleted file mode 100644 index c27dee8bf..000000000 --- a/redis/src/main/scala/zio/redis/RedisPubSub.scala +++ /dev/null @@ -1,26 +0,0 @@ -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 - ) -} diff --git a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala index 2da0403cc..03ed53178 100644 --- a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala +++ b/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala @@ -1,22 +1,106 @@ package zio.redis -import zio.redis.Output.ArbitraryOutput +import zio.redis.Input._ +import zio.redis.Output.{ArbitraryOutput, PushProtocolOutput} +import zio.redis.options.PubSub.PubSubCallback import zio.schema.Schema import zio.schema.codec.BinaryCodec +import zio.stream.ZStream.RefineToOrDieOps import zio.stream._ -import zio.{Chunk, IO, ZIO} +import zio.{Chunk, IO, Promise, Ref, ZIO} + +private[redis] final case class RedisPubSubCommand(codec: BinaryCodec, executor: SubscriptionExecutor) extends { + import zio.redis.options.PubSub.PushProtocol._ + + def subscribe[A: Schema]( + channels: Chunk[String], + onSubscribe: PubSubCallback, + onUnsubscribe: PubSubCallback + ): IO[RedisError, Stream[RedisError, (String, A)]] = { + val command = CommandNameInput.encode(api.Subscribe.Subscribe)(codec) ++ + Varargs(ArbitraryKeyInput[String]()).encode(channels)(codec) + + val channelSet = channels.toSet + + for { + unsubscribedRef <- Ref.make(channels.map(_ -> false).toMap) + promise <- Promise.make[RedisError, Unit] + stream <- executor + .execute(command) + .map( + _.mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp)(codec))).mapZIO { + case Subscribe(channel, numOfSubscription) if channelSet contains channel => + onSubscribe(channel, numOfSubscription).as(None) + case Message(channel, message) if channelSet contains channel => + ZIO + .attempt(ArbitraryOutput[A]().unsafeDecode(message)(codec)) + .map(msg => Some((channel, msg))) + case Unsubscribe(channel, numOfSubscription) if channelSet contains channel => + for { + _ <- onUnsubscribe(channel, numOfSubscription) + _ <- unsubscribedRef.update(_.updatedWith(channel)(_ => Some(true))) + _ <- promise.succeed(()).whenZIO(unsubscribedRef.get.map(_.values.forall(identity))) + } yield None + case _ => ZIO.none + }.collectSome + .refineToOrDie[RedisError] + .interruptWhen(promise) + ) + } yield stream + } + + def pSubscribe[A: Schema]( + patterns: Chunk[String], + onSubscribe: PubSubCallback, + onUnsubscribe: PubSubCallback + ): IO[RedisError, Stream[RedisError, (String, A)]] = { + val command = CommandNameInput.encode(api.Subscribe.PSubscribe)(codec) ++ + Varargs(ArbitraryKeyInput[String]()).encode(patterns)(codec) + + val patternSet = patterns.toSet + + for { + unsubscribedRef <- Ref.make(patterns.map(_ -> false).toMap) + promise <- Promise.make[RedisError, Unit] + stream <- executor + .execute(command) + .map( + _.mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp)(codec))).mapZIO { + case PSubscribe(pattern, numOfSubscription) if patternSet contains pattern => + onSubscribe(pattern, numOfSubscription).as(None) + case PMessage(pattern, channel, message) if patternSet contains pattern => + ZIO + .attempt(ArbitraryOutput[A]().unsafeDecode(message)(codec)) + .map(msg => Some((channel, msg))) + case PUnsubscribe(pattern, numOfSubscription) if patternSet contains pattern => + for { + _ <- onUnsubscribe(pattern, numOfSubscription) + _ <- unsubscribedRef.update(_.updatedWith(pattern)(_ => Some(true))) + _ <- promise.succeed(()).whenZIO(unsubscribedRef.get.map(_.values.forall(identity))) + } yield None + case _ => ZIO.none + }.collectSome + .refineToOrDie[RedisError] + .interruptWhen(promise) + ) + } yield stream + } + + def unsubscribe(channels: Chunk[String]): IO[RedisError, Unit] = { + val command = CommandNameInput.encode(api.Subscribe.Unsubscribe)(codec) ++ + Varargs(ArbitraryKeyInput[String]()).encode(channels)(codec) + + executor + .execute(command) + .unit + } + + def pUnsubscribe(patterns: Chunk[String]): IO[RedisError, Unit] = { + val command = CommandNameInput.encode(api.Subscribe.PUnsubscribe)(codec) ++ + Varargs(ArbitraryKeyInput[String]()).encode(patterns)(codec) -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] - ) - ) - ) + .unit + } } diff --git a/redis/src/main/scala/zio/redis/RedisSubscription.scala b/redis/src/main/scala/zio/redis/RedisSubscription.scala new file mode 100644 index 000000000..1e061faf1 --- /dev/null +++ b/redis/src/main/scala/zio/redis/RedisSubscription.scala @@ -0,0 +1,18 @@ +package zio.redis + +import zio.{URLayer, ZIO, ZLayer} +import zio.schema.codec.BinaryCodec + +trait RedisSubscription extends api.Subscribe + +object RedisSubscription { + lazy val layer: URLayer[SubscriptionExecutor with BinaryCodec, RedisSubscription] = + ZLayer { + for { + codec <- ZIO.service[BinaryCodec] + executor <- ZIO.service[SubscriptionExecutor] + } yield Live(codec, executor) + } + + private final case class Live(codec: BinaryCodec, executor: SubscriptionExecutor) extends RedisSubscription +} diff --git a/redis/src/main/scala/zio/redis/ResultBuilder.scala b/redis/src/main/scala/zio/redis/ResultBuilder.scala index db22b22ab..7f89f91e0 100644 --- a/redis/src/main/scala/zio/redis/ResultBuilder.scala +++ b/redis/src/main/scala/zio/redis/ResultBuilder.scala @@ -19,7 +19,6 @@ 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] = ??? @@ -47,8 +46,4 @@ 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]]] - } } diff --git a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala b/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala deleted file mode 100644 index a311ae2bb..000000000 --- a/redis/src/main/scala/zio/redis/SingleNodeRedisPubSub.scala +++ /dev/null @@ -1,250 +0,0 @@ -package zio.redis - -import zio.redis.Input.{NonEmptyList, StringInput, Varargs} -import zio.redis.Output.PushProtocolOutput -import zio.redis.SingleNodeRedisPubSub.{Request, RequestQueueSize, SubscriptionKey, True} -import zio.redis.api.PubSub -import zio.schema.codec.BinaryCodec -import zio.stream._ -import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Schedule, UIO, ZIO} - -final class SingleNodeRedisPubSub( - pubSubHubsRef: Ref[Map[SubscriptionKey, Hub[Take[RedisError, PushProtocol]]]], - reqQueue: Queue[Request], - resQueue: Queue[Promise[RedisError, PushProtocol]], - connection: RedisConnection, - implicit val codec: BinaryCodec -) extends RedisPubSub { - - def execute(command: PubSubCommand): IO[RedisError, Chunk[Stream[RedisError, RespValue]]] = { - def applyCallback( - streams: Chunk[Stream[RedisError, PushProtocol]], - callback: PushProtocol => IO[RedisError, Unit] - ) = - streams.map( - _.tap(callback(_)).map { - case PushProtocol.Message(_, msg) => Some(msg) - case PushProtocol.PMessage(_, _, msg) => Some(msg) - case _ => None - }.collectSome - ) - - command match { - case PubSubCommand.Subscribe(channel, channels, callback) => - subscribe(channel, channels) - .map(applyCallback(_, callback)) - case PubSubCommand.PSubscribe(pattern, patterns, callback) => - pSubscribe(pattern, patterns) - .map(applyCallback(_, callback)) - case PubSubCommand.Unsubscribe(channels, callback) => - unsubscribe(channels) - .map(applyCallback(_, callback)) - case PubSubCommand.PUnsubscribe(patterns, callback) => - pUnsubscribe(patterns) - .map(applyCallback(_, callback)) - } - } - - private def subscribe( - channel: String, - channels: List[String] - ): IO[RedisError, Chunk[Stream[RedisError, PushProtocol]]] = - makeSubscriptionStream( - PubSub.Subscribe, - channelKey(channel), - channels.map(channelKey(_)) - ) - - private def pSubscribe( - pattern: String, - patterns: List[String] - ): IO[RedisError, Chunk[Stream[RedisError, PushProtocol]]] = - makeSubscriptionStream( - PubSub.PSubscribe, - patternKey(pattern), - patterns.map(patternKey(_)) - ) - - private def unsubscribe( - channels: List[String] - ): IO[RedisError, Chunk[Stream[RedisError, PushProtocol]]] = - makeUnsubscriptionStream( - PubSub.Unsubscribe, - if (channels.nonEmpty) - ZIO.succeedNow(channels.map(channelKey(_))) - else - pubSubHubsRef.get.map(_.keys.filter(_.isChannel).toList) - ) - - private def pUnsubscribe( - patterns: List[String] - ): IO[RedisError, Chunk[Stream[RedisError, PushProtocol]]] = - makeUnsubscriptionStream( - PubSub.PUnsubscribe, - if (patterns.nonEmpty) - ZIO.succeedNow(patterns.map(patternKey(_))) - else - pubSubHubsRef.get.map(_.keys.filter(_.isPattern).toList) - ) - - private def makeSubscriptionStream( - command: String, - key: SubscriptionKey, - keys: List[SubscriptionKey] - )(implicit codec: BinaryCodec) = - for { - promises <- Promise.make[RedisError, PushProtocol].replicateZIO(keys.size + 1).map(Chunk.fromIterable(_)) - chunk = StringInput.encode(command) ++ NonEmptyList(StringInput).encode((key.value, keys.map(_.value))) - _ <- reqQueue.offer(Request(chunk, promises)) - streams <- ZIO.foreach((key :: keys) zip promises) { case (key, promise) => - for { - hub <- getHub(key) - stream = ZStream.fromHub(hub).flattenTake - } yield ZStream.fromZIO(promise.await) concat stream - } - } yield Chunk.fromIterable(streams) - - private def makeUnsubscriptionStream( - command: String, - keys: UIO[List[SubscriptionKey]] - )(implicit codec: BinaryCodec) = - for { - targets <- keys - chunk = StringInput.encode(command) ++ Varargs(StringInput).encode(targets.map(_.value)) - promises <- Promise.make[RedisError, PushProtocol].replicateZIO(targets.size).map(Chunk.fromIterable(_)) - _ <- reqQueue.offer(Request(chunk, promises)) - streams = promises.map(promise => ZStream.fromZIO(promise.await)) - } yield streams - - private def getHub(key: SubscriptionKey) = { - def makeNewHub = - Hub - .unbounded[Take[RedisError, PushProtocol]] - .tap(hub => pubSubHubsRef.update(_ + (key -> hub))) - - for { - hubs <- pubSubHubsRef.get - hub <- ZIO.fromOption(hubs.get(key)).orElse(makeNewHub) - } yield hub - } - - 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).serialize - } - - val bytes = buffer.result() - - connection - .write(bytes) - .mapError(RedisError.IOError(_)) - .tapBoth( - e => ZIO.foreachDiscard(reqs.flatMap(_.promises))(_.fail(e)), - _ => ZIO.foreachDiscard(reqs.map(_.promises))(resQueue.offerAll(_)) - ) - } - - private def receive: IO[RedisError, Unit] = { - def handlePushProtocolMessage(msg: PushProtocol): UIO[Unit] = { - def releasePendingPromise(msg: PushProtocol): UIO[Unit] = - resQueue.take.flatMap(_.succeed(msg)).unit - - def handleUnsubscription(key: SubscriptionKey, msg: PushProtocol): UIO[Unit] = - for { - _ <- releasePendingPromise(msg) - pubSubs <- pubSubHubsRef.get - hubOpt = pubSubs.get(key) - _ <- ZIO.fromOption(hubOpt).flatMap(_.offer(Take.end)).orElse(ZIO.unit) - _ <- pubSubHubsRef.update(_ - key) - } yield () - - def handleSubscription(key: SubscriptionKey, msg: PushProtocol): UIO[Unit] = - for { - _ <- resQueue.take.flatMap(_.succeed(msg)) - _ <- getHub(key).flatMap(_.offer(Take.single(msg))) - } yield () - - msg match { - case msg @ PushProtocol.Unsubscribe(channel, _) => - handleUnsubscription(channelKey(channel), msg) - case msg @ PushProtocol.PUnsubscribe(pattern, _) => - handleUnsubscription(patternKey(pattern), msg) - case msg @ PushProtocol.Subscribe(channel, _) => - handleSubscription(channelKey(channel), msg) - case msg @ PushProtocol.PSubscribe(pattern, _) => - handleSubscription(patternKey(pattern), msg) - case msg @ PushProtocol.Message(channel, _) => - getHub(channelKey(channel)).flatMap(_.offer(Take.single(msg))).unit - case msg @ PushProtocol.PMessage(pattern, _, _) => - getHub(patternKey(pattern)).flatMap(_.offer(Take.single(msg))).unit - } - } - - connection.read - .mapError(RedisError.IOError(_)) - .via(RespValue.decoder) - .collectSome - .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))) - .refineToOrDie[RedisError] - .foreach(handlePushProtocolMessage(_)) - } - - private def drainWith(e: RedisError): UIO[Unit] = resQueue.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) - - private def resubscribe: IO[RedisError, Unit] = { - def makeCommand(name: String, keys: Set[String]) = - RespValue.Array(StringInput.encode(name) ++ Varargs(StringInput).encode(keys)).serialize - - for { - keySet <- pubSubHubsRef.get.map(_.keySet) - (channels, patterns) = keySet.partition(_.isChannel) - _ <- (connection.write(makeCommand(PubSub.Subscribe, channels.map(_.value))).when(channels.nonEmpty) *> - connection.write(makeCommand(PubSub.PSubscribe, patterns.map(_.value))).when(patterns.nonEmpty)) - .mapError(RedisError.IOError(_)) - .retryWhile(True) - } yield () - } - - private def patternKey(key: String) = SubscriptionKey(key, true) - private def channelKey(key: String) = SubscriptionKey(key, false) - - /** - * 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") *> drainWith(e) *> resubscribe) - .retryWhile(True) - .tapError(e => ZIO.logError(s"Executor exiting: $e")) -} - -object SingleNodeRedisPubSub { - private final case class SubscriptionKey(value: String, isPattern: Boolean) { - def isChannel: Boolean = isPattern == false - } - private final case class Request( - command: Chunk[RespValue.BulkString], - promises: Chunk[Promise[RedisError, PushProtocol]] - ) - - private final val True: Any => Boolean = _ => true - - private final val RequestQueueSize = 16 - - def create(conn: RedisConnection, codec: BinaryCodec) = - for { - hubRef <- Ref.make(Map.empty[SubscriptionKey, Hub[Take[RedisError, PushProtocol]]]) - reqQueue <- Queue.bounded[Request](RequestQueueSize) - resQueue <- Queue.unbounded[Promise[RedisError, PushProtocol]] - pubSub = new SingleNodeRedisPubSub(hubRef, reqQueue, resQueue, conn, codec) - _ <- pubSub.run.forkScoped - _ <- logScopeFinalizer(s"$pubSub Node PubSub is closed") - } yield pubSub -} diff --git a/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala b/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala new file mode 100644 index 000000000..b3c7d681b --- /dev/null +++ b/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala @@ -0,0 +1,146 @@ +package zio.redis + +import zio.redis.SingleNodeSubscriptionExecutor.{Request, RequestQueueSize, True} +import zio.redis.api.Subscribe +import zio.stream._ +import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Schedule, ZIO} + +final class SingleNodeSubscriptionExecutor( + channelSubsRef: Ref[Set[RespArgument.Key]], + patternSubsRef: Ref[Set[RespArgument.Key]], + hub: Hub[RespValue], + reqQueue: Queue[Request], + connection: RedisConnection +) extends SubscriptionExecutor { + def execute(command: RespCommand): IO[RedisError, Stream[RedisError, RespValue]] = + for { + commandName <- ZIO + .fromOption(command.args.collectFirst { case RespArgument.CommandName(name) => name }) + .orElseFail(RedisError.CommandNameNotFound(command.args.toString())) + stream <- commandName match { + case Subscribe.Subscribe => ZIO.succeed(subscribe(channelSubsRef, command)) + case Subscribe.PSubscribe => ZIO.succeed(subscribe(patternSubsRef, command)) + case Subscribe.Unsubscribe => ZIO.succeed(unsubscribe(channelSubsRef, command)) + case Subscribe.PUnsubscribe => ZIO.succeed(unsubscribe(patternSubsRef, command)) + case other => ZIO.fail(RedisError.InvalidPubSubCommand(other)) + } + } yield stream + + private def subscribe( + subscriptionRef: Ref[Set[RespArgument.Key]], + command: RespCommand + ): Stream[RedisError, RespValue] = + ZStream + .fromZIO( + for { + reqPromise <- Promise.make[RedisError, Unit] + _ <- reqQueue.offer(Request(command.args.map(_.value), reqPromise)) + _ <- reqPromise.await + keys = command.args.collect { case key: RespArgument.Key => key } + _ <- subscriptionRef.update(_ ++ keys) + } yield ZStream.fromHub(hub) + ) + .flatten + + private def unsubscribe( + subscriptionRef: Ref[Set[RespArgument.Key]], + command: RespCommand + ): Stream[RedisError, RespValue] = + ZStream + .fromZIO( + for { + reqPromise <- Promise.make[RedisError, Unit] + _ <- reqQueue.offer(Request(command.args.map(_.value), reqPromise)) + _ <- reqPromise.await + keys = command.args.collect { case key: RespArgument.Key => key } + _ <- subscriptionRef.update(subscribedSet => + if (keys.nonEmpty) + subscribedSet -- keys + else + Set.empty + ) + } yield ZStream.fromHub(hub) + ) + .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).serialize + } + + 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] = + connection.read + .mapError(RedisError.IOError(_)) + .via(RespValue.decoder) + .collectSome + .foreach(hub.offer(_)) + + private def resubscribe: IO[RedisError, Unit] = { + def makeCommand(name: String, keys: Chunk[RespArgument.Key]) = + if (keys.isEmpty) + Chunk.empty + else + RespValue.Array((RespCommand(RespArgument.CommandName(name)) ++ RespCommand(keys)).args.map(_.value)).serialize + + for { + channels <- channelSubsRef.get + patterns <- patternSubsRef.get + commands = makeCommand(Subscribe.Subscribe, Chunk.fromIterable(channels)) ++ + makeCommand(Subscribe.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")) +} + +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) = + for { + hub <- Hub.unbounded[RespValue] + reqQueue <- Queue.bounded[Request](RequestQueueSize) + channelRef <- Ref.make(Set.empty[RespArgument.Key]) + patternRef <- Ref.make(Set.empty[RespArgument.Key]) + pubSub = new SingleNodeSubscriptionExecutor(channelRef, patternRef, hub, reqQueue, conn) + _ <- pubSub.run.forkScoped + _ <- logScopeFinalizer(s"$pubSub Node PubSub is closed") + } yield pubSub +} diff --git a/redis/src/main/scala/zio/redis/SubscribeEnvironment.scala b/redis/src/main/scala/zio/redis/SubscribeEnvironment.scala new file mode 100644 index 000000000..44c92b961 --- /dev/null +++ b/redis/src/main/scala/zio/redis/SubscribeEnvironment.scala @@ -0,0 +1,8 @@ +package zio.redis + +import zio.schema.codec.BinaryCodec + +trait SubscribeEnvironment { + protected def codec: BinaryCodec + protected def executor: SubscriptionExecutor +} diff --git a/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala b/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala new file mode 100644 index 000000000..22125c251 --- /dev/null +++ b/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala @@ -0,0 +1,24 @@ +package zio.redis + +import zio.stream._ +import zio.{IO, Layer, ZIO, ZLayer} + +trait SubscriptionExecutor { + def execute(command: RespCommand): IO[RedisError, Stream[RedisError, RespValue]] +} + +object SubscriptionExecutor { + lazy val layer: ZLayer[RedisConfig, RedisError.IOError, SubscriptionExecutor] = + RedisConnectionLive.layer.fresh >>> pubSublayer + + lazy val local: Layer[RedisError.IOError, SubscriptionExecutor] = + RedisConnectionLive.default.fresh >>> pubSublayer + + private lazy val pubSublayer: ZLayer[RedisConnection, RedisError.IOError, SubscriptionExecutor] = + ZLayer.scoped( + for { + conn <- ZIO.service[RedisConnection] + pubSub <- SingleNodeSubscriptionExecutor.create(conn) + } yield pubSub + ) +} diff --git a/redis/src/main/scala/zio/redis/api/PubSub.scala b/redis/src/main/scala/zio/redis/api/PubSub.scala deleted file mode 100644 index dd7c37b3e..000000000 --- a/redis/src/main/scala/zio/redis/api/PubSub.scala +++ /dev/null @@ -1,154 +0,0 @@ -package zio.redis.api - -import zio.redis.Input._ -import zio.redis.Output._ -import zio.redis.ResultBuilder.ResultStreamBuilder -import zio.redis._ -import zio.redis.options.PubSub.NumberOfSubscribers -import zio.schema.Schema -import zio.stream._ -import zio.{Chunk, IO, Promise, Ref, UIO, ZIO} - -trait PubSub extends RedisEnvironment { - import PubSub._ - - final def subscribe(channel: String): ResultStreamBuilder[Id] = - subscribeWithCallback(channel)(emptyCallback) - - final def subscribeWithCallback(channel: String)(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = - new ResultStreamBuilder[Id] { - def returning[R: Schema]: IO[RedisError, Id[Stream[RedisError, R]]] = - runSubscription(PubSubCommand.Subscribe(channel, List.empty, createSubscriptionCallback(onSubscribe))) - .flatMap(extractOne(channel, _)) - } - - final def subscribe(channel: String, channels: String*): ResultStreamBuilder[Chunk] = - createStreamListBuilder(channel, channels.toList, emptyCallback, isPatterned = false) - - final def subscribeWithCallback(channel: String, channels: String*)( - onSubscribe: PubSubCallback - ): ResultStreamBuilder[Chunk] = - createStreamListBuilder(channel, channels.toList, onSubscribe, isPatterned = false) - - final def pSubscribe(pattern: String): ResultStreamBuilder[Id] = - pSubscribeWithCallback(pattern)(emptyCallback) - - final def pSubscribeWithCallback( - pattern: String - )(onSubscribe: PubSubCallback): ResultStreamBuilder[Id] = - new ResultStreamBuilder[Id] { - def returning[R: Schema]: IO[RedisError, Id[Stream[RedisError, R]]] = - runSubscription(PubSubCommand.PSubscribe(pattern, List.empty, createSubscriptionCallback(onSubscribe))) - .flatMap(extractOne(pattern, _)) - } - - final def pSubscribe(pattern: String, patterns: String*): ResultStreamBuilder[Chunk] = - createStreamListBuilder(pattern, patterns.toList, emptyCallback, isPatterned = true) - - final def pSubscribeWithCallback(pattern: String, patterns: String*)( - onSubscribe: PubSubCallback - ): ResultStreamBuilder[Chunk] = - createStreamListBuilder(pattern, patterns.toList, onSubscribe, isPatterned = true) - - final def unsubscribe(channels: String*): IO[RedisError, Promise[RedisError, Chunk[(String, Long)]]] = - for { - ref <- Ref.make(Chunk.empty[(String, Long)]) - callback = createUnsubscriptionCallback(ref) - promise <- runUnsubscription(PubSubCommand.Unsubscribe(channels.toList, callback), ref) - } yield promise - - final def pUnsubscribe(patterns: String*): IO[RedisError, Promise[RedisError, Chunk[(String, Long)]]] = - for { - ref <- Ref.make(Chunk.empty[(String, Long)]) - callback = createUnsubscriptionCallback(ref) - promise <- runUnsubscription(PubSubCommand.PUnsubscribe(patterns.toList, callback), ref) - } yield promise - - final def publish[A: Schema](channel: String, message: A): IO[RedisError, Long] = { - val command = RedisCommand(Publish, Tuple2(StringInput, ArbitraryInput[A]()), LongOutput, codec, executor) - command.run((channel, message)) - } - - final def pubSubChannels(pattern: String): IO[RedisError, Chunk[String]] = { - val command = RedisCommand(PubSubChannels, StringInput, ChunkOutput(MultiStringOutput), codec, executor) - command.run(pattern) - } - - final def pubSubNumPat: IO[RedisError, Long] = { - val command = RedisCommand(PubSubNumPat, NoInput, LongOutput, codec, executor) - command.run(()) - } - - final def pubSubNumSub(channel: String, channels: String*): IO[RedisError, Chunk[NumberOfSubscribers]] = { - val command = RedisCommand(PubSubNumSub, NonEmptyList(StringInput), NumSubResponseOutput, codec, executor) - command.run((channel, channels.toList)) - } - - private def createStreamListBuilder( - key: String, - keys: List[String], - callback: PubSubCallback, - isPatterned: Boolean - ): ResultStreamBuilder[Chunk] = - new ResultStreamBuilder[Chunk] { - def returning[R: Schema]: IO[RedisError, Chunk[Stream[RedisError, R]]] = { - val subscriptionCallback = createSubscriptionCallback(callback); - if (isPatterned) runSubscription(PubSubCommand.PSubscribe(key, keys, subscriptionCallback)) - else runSubscription(PubSubCommand.Subscribe(key, keys, subscriptionCallback)) - } - } - - private def runUnsubscription( - command: PubSubCommand, - resultRef: Ref[Chunk[(String, Long)]] - ): IO[RedisError, Promise[RedisError, Chunk[(String, Long)]]] = - for { - promise <- Promise.make[RedisError, Chunk[(String, Long)]] - streams <- RedisPubSubCommand(command, codec, pubSub).run[Unit]() - _ <- streams - .fold(ZStream.empty)(_ merge _) - .runDrain - .onDone( - e => promise.fail(e), - _ => resultRef.get.flatMap(promise.succeed(_)) - ) - .fork - } yield promise - - private def runSubscription[R: Schema]( - command: PubSubCommand - ): IO[RedisError, Chunk[Stream[RedisError, R]]] = - RedisPubSubCommand(command, codec, pubSub).run[R]() - - private def createSubscriptionCallback( - onSubscribe: PubSubCallback - ): PushProtocol => IO[RedisError, Unit] = { - case PushProtocol.Subscribe(key, numOfSubscription) => onSubscribe(key, numOfSubscription) - case PushProtocol.PSubscribe(key, numOfSubscription) => onSubscribe(key, numOfSubscription) - case _ => ZIO.unit - } - - private def createUnsubscriptionCallback(resultRef: Ref[Chunk[(String, Long)]]): PushProtocol => UIO[Unit] = { - case PushProtocol.Unsubscribe(channel, numOfSubscription) => - resultRef.update(_ appended (channel, numOfSubscription)) - case PushProtocol.PUnsubscribe(pattern, numOfSubscription) => - resultRef.update(_ appended (pattern, numOfSubscription)) - case _ => ZIO.unit - } - - private def extractOne[A](key: String, elements: Chunk[A]) = - ZIO.fromOption(elements.headOption).orElseFail(RedisError.NoPubSubStream(key)) -} - -private[redis] object PubSub { - private lazy val emptyCallback = (_: String, _: Long) => ZIO.unit - - final val Subscribe = "SUBSCRIBE" - final val Unsubscribe = "UNSUBSCRIBE" - final val PSubscribe = "PSUBSCRIBE" - final val PUnsubscribe = "PUNSUBSCRIBE" - final val Publish = "PUBLISH" - final val PubSubChannels = "PUBSUB CHANNELS" - final val PubSubNumPat = "PUBSUB NUMPAT" - final val PubSubNumSub = "PUBSUB NUMSUB" -} diff --git a/redis/src/main/scala/zio/redis/api/Publishing.scala b/redis/src/main/scala/zio/redis/api/Publishing.scala new file mode 100644 index 000000000..1cdab5443 --- /dev/null +++ b/redis/src/main/scala/zio/redis/api/Publishing.scala @@ -0,0 +1,39 @@ +package zio.redis.api + +import zio.redis.Input._ +import zio.redis.Output._ +import zio.redis._ +import zio.redis.options.PubSub.NumberOfSubscribers +import zio.schema.Schema +import zio.{Chunk, IO} + +trait Publishing extends RedisEnvironment { + import Publishing._ + + final def publish[A: Schema](channel: String, message: A): IO[RedisError, Long] = { + val command = RedisCommand(Publish, Tuple2(StringInput, ArbitraryKeyInput[A]()), LongOutput, codec, executor) + command.run((channel, message)) + } + + final def pubSubChannels(pattern: String): IO[RedisError, Chunk[String]] = { + val command = RedisCommand(PubSubChannels, StringInput, ChunkOutput(MultiStringOutput), codec, executor) + command.run(pattern) + } + + final def pubSubNumPat: IO[RedisError, Long] = { + val command = RedisCommand(PubSubNumPat, NoInput, LongOutput, codec, executor) + command.run(()) + } + + final def pubSubNumSub(channel: String, channels: String*): IO[RedisError, Chunk[NumberOfSubscribers]] = { + val command = RedisCommand(PubSubNumSub, NonEmptyList(StringInput), NumSubResponseOutput, codec, executor) + command.run((channel, channels.toList)) + } +} + +private[redis] object Publishing { + final val Publish = "PUBLISH" + final val PubSubChannels = "PUBSUB CHANNELS" + final val PubSubNumPat = "PUBSUB NUMPAT" + final val PubSubNumSub = "PUBSUB NUMSUB" +} diff --git a/redis/src/main/scala/zio/redis/api/Subscribe.scala b/redis/src/main/scala/zio/redis/api/Subscribe.scala new file mode 100644 index 000000000..43f7bcde1 --- /dev/null +++ b/redis/src/main/scala/zio/redis/api/Subscribe.scala @@ -0,0 +1,74 @@ +package zio.redis.api + +import zio.redis.ResultBuilder.ResultBuilder1 +import zio.redis._ +import zio.redis.options.PubSub.PubSubCallback +import zio.schema.Schema +import zio.stream.Stream +import zio.{Chunk, IO, ZIO} + +trait Subscribe extends SubscribeEnvironment { + import Subscribe._ + + final def subscribe(channel: String, channels: String*) = + new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { + def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = + RedisPubSubCommand(codec, executor).subscribe( + Chunk.single(channel) ++ Chunk.fromIterable(channels), + emptyCallback, + emptyCallback + ) + } + + final def subscribeWithCallback(channel: String, channels: String*)( + onSubscribe: PubSubCallback, + onUnsubscribe: PubSubCallback + ) = + new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { + def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = + RedisPubSubCommand(codec, executor).subscribe( + Chunk.single(channel) ++ Chunk.fromIterable(channels), + onSubscribe, + onUnsubscribe + ) + } + + final def pSubscribe(pattern: String, patterns: String*) = + new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { + def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = + RedisPubSubCommand(codec, executor).pSubscribe( + Chunk.single(pattern) ++ Chunk.fromIterable(patterns), + emptyCallback, + emptyCallback + ) + } + + final def pSubscribeWithCallback( + pattern: String, + patterns: String* + )(onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback) = + new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { + def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = + RedisPubSubCommand(codec, executor).pSubscribe( + Chunk.single(pattern) ++ Chunk.fromIterable(patterns), + onSubscribe, + onUnsubscribe + ) + } + + final def unsubscribe(channels: String*): IO[RedisError, Unit] = + RedisPubSubCommand(codec, executor).unsubscribe(Chunk.fromIterable(channels)) + + final def pUnsubscribe(patterns: String*): IO[RedisError, Unit] = + RedisPubSubCommand(codec, executor).unsubscribe(Chunk.fromIterable(patterns)) + +} + +object Subscribe { + private lazy val emptyCallback = (_: String, _: Long) => ZIO.unit + + final val Subscribe = "SUBSCRIBE" + final val Unsubscribe = "UNSUBSCRIBE" + final val PSubscribe = "PSUBSCRIBE" + final val PUnsubscribe = "PUNSUBSCRIBE" +} diff --git a/redis/src/main/scala/zio/redis/options/PubSub.scala b/redis/src/main/scala/zio/redis/options/PubSub.scala index d06ebc88b..9b6d1988f 100644 --- a/redis/src/main/scala/zio/redis/options/PubSub.scala +++ b/redis/src/main/scala/zio/redis/options/PubSub.scala @@ -1,9 +1,9 @@ package zio.redis.options -import zio.IO -import zio.redis.{RedisError, RespValue} -trait PubSub { - type PubSubCallback = (String, Long) => IO[RedisError, Unit] +import zio.UIO +import zio.redis.RespValue +object PubSub { + type PubSubCallback = (String, Long) => UIO[Unit] sealed trait PushProtocol @@ -15,8 +15,6 @@ trait PubSub { case class Message(channel: String, message: RespValue) extends PushProtocol case class PMessage(pattern: String, channel: String, message: RespValue) extends PushProtocol } -} -object PubSub { final case class NumberOfSubscribers(channel: String, subscriberCount: Long) } diff --git a/redis/src/main/scala/zio/redis/package.scala b/redis/src/main/scala/zio/redis/package.scala index 6dd5ef9f6..948ec1623 100644 --- a/redis/src/main/scala/zio/redis/package.scala +++ b/redis/src/main/scala/zio/redis/package.scala @@ -25,8 +25,7 @@ package object redis with options.Strings with options.Lists with options.Streams - with options.Scripting - with options.PubSub { + with options.Scripting { type Id[+A] = A diff --git a/redis/src/test/scala/zio/redis/ApiSpec.scala b/redis/src/test/scala/zio/redis/ApiSpec.scala index 4bb1b2dce..920b07b22 100644 --- a/redis/src/test/scala/zio/redis/ApiSpec.scala +++ b/redis/src/test/scala/zio/redis/ApiSpec.scala @@ -34,10 +34,13 @@ object ApiSpec hyperLogLogSuite, hashSuite, streamsSuite, - scriptingSpec + scriptingSpec, + pubSubSuite ).provideShared( RedisExecutor.local, + SubscriptionExecutor.local, Redis.layer, + RedisSubscription.layer, ZLayer.succeed(codec) ) diff --git a/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala b/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala index 7eb98f0f8..7f891f3e3 100644 --- a/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala +++ b/redis/src/test/scala/zio/redis/ClusterExecutorSpec.scala @@ -69,7 +69,6 @@ object ClusterExecutorSpec extends BaseSpec { ZLayer.make[Redis]( ZLayer.succeed(RedisConfig(uri.host, uri.port)), RedisExecutor.layer, - RedisPubSub.layer, ZLayer.succeed(codec), Redis.layer ) @@ -80,7 +79,6 @@ object ClusterExecutorSpec extends BaseSpec { ZLayer.make[Redis]( ZLayer.succeed(RedisClusterConfig(Chunk(address1, address2))), ClusterExecutor.layer.orDie, - RedisPubSub.local.orDie, ZLayer.succeed(codec), Redis.layer ) diff --git a/redis/src/test/scala/zio/redis/KeysSpec.scala b/redis/src/test/scala/zio/redis/KeysSpec.scala index ea28409cc..d05915ad1 100644 --- a/redis/src/test/scala/zio/redis/KeysSpec.scala +++ b/redis/src/test/scala/zio/redis/KeysSpec.scala @@ -484,7 +484,6 @@ object KeysSpec { ZLayer.succeed(RedisConfig("localhost", 6380)), RedisConnectionLive.layer, SingleNodeExecutor.layer, - RedisPubSub.layer, ZLayer.succeed[BinaryCodec](ProtobufCodec), Redis.layer ) diff --git a/redis/src/test/scala/zio/redis/OutputSpec.scala b/redis/src/test/scala/zio/redis/OutputSpec.scala index eceae9120..abda34564 100644 --- a/redis/src/test/scala/zio/redis/OutputSpec.scala +++ b/redis/src/test/scala/zio/redis/OutputSpec.scala @@ -3,6 +3,7 @@ package zio.redis import zio._ import zio.redis.Output._ import zio.redis.RedisError._ +import zio.redis.options.PubSub.PushProtocol import zio.test.Assertion._ import zio.test._ diff --git a/redis/src/test/scala/zio/redis/PubSubSpec.scala b/redis/src/test/scala/zio/redis/PubSubSpec.scala index 76bfd4166..e410240be 100644 --- a/redis/src/test/scala/zio/redis/PubSubSpec.scala +++ b/redis/src/test/scala/zio/redis/PubSubSpec.scala @@ -8,49 +8,56 @@ import zio.{Chunk, Promise, ZIO} import scala.util.Random trait PubSubSpec extends BaseSpec { - def pubSubSuite: Spec[Redis, RedisError] = + def pubSubSuite: Spec[Redis with RedisSubscription, RedisError] = suite("pubSubs")( suite("subscribe")( test("subscribe response") { for { - redis <- ZIO.service[Redis] - channel <- generateRandomString() - promise <- Promise.make[RedisError, String] - resBuilder = redis.subscribeWithCallback(channel)((key: String, _: Long) => promise.succeed(key).unit) - stream <- resBuilder.returning[String] - _ <- stream.interruptWhen(promise).runDrain.fork - res <- promise.await + subscription <- ZIO.service[RedisSubscription] + channel <- generateRandomString() + promise <- Promise.make[RedisError, String] + resBuilder = + subscription + .subscribeWithCallback(channel)( + (key: String, _: Long) => promise.succeed(key).unit, + (_, _) => ZIO.unit + ) + stream <- resBuilder.returning[String] + _ <- stream.interruptWhen(promise).runDrain.fork + res <- promise.await } yield assertTrue(res == channel) }, test("message response") { for { - redis <- ZIO.service[Redis] - channel <- generateRandomString() - message = "bar" - promise <- Promise.make[RedisError, String] - stream <- redis.subscribe(channel).returning[String] - fiber <- stream.interruptWhen(promise).runHead.fork + redis <- ZIO.service[Redis] + subscription <- ZIO.service[RedisSubscription] + channel <- generateRandomString() + message = "bar" + promise <- Promise.make[RedisError, String] + stream <- subscription.subscribe(channel).returning[String] + fiber <- stream.interruptWhen(promise).runHead.fork _ <- redis .pubSubChannels(channel) .repeatUntil(_ contains channel) _ <- redis.publish(channel, message) res <- fiber.join - } yield assertTrue(res.get == message) + } yield assertTrue(res.get == (channel, message)) }, test("multiple subscribe") { val numOfPublish = 20 for { - redis <- ZIO.service[Redis] - prefix <- generateRandomString(5) - channel1 <- generateRandomString(prefix) - channel2 <- generateRandomString(prefix) - pattern = prefix + '*' - message <- generateRandomString(5) - stream1 <- redis + redis <- ZIO.service[Redis] + subscription <- ZIO.service[RedisSubscription] + prefix <- generateRandomString(5) + channel1 <- generateRandomString(prefix) + channel2 <- generateRandomString(prefix) + pattern = prefix + '*' + message <- generateRandomString(5) + stream1 <- subscription .subscribe(channel1) .returning[String] .fork - stream2 <- redis + stream2 <- subscription .subscribe(channel2) .returning[String] .fork @@ -59,19 +66,21 @@ trait PubSubSpec extends BaseSpec { .repeatUntil(channels => channels.size >= 2) ch1SubsCount <- redis.publish(channel1, message).replicateZIO(numOfPublish).map(_.head) ch2SubsCount <- redis.publish(channel2, message).replicateZIO(numOfPublish).map(_.head) - promises <- redis.unsubscribe() - _ <- promises.await + _ <- subscription.unsubscribe() _ <- stream1.join _ <- stream2.join } yield assertTrue(ch1SubsCount == 1L) && assertTrue(ch2SubsCount == 1L) }, test("psubscribe response") { for { - redis <- ZIO.service[Redis] - pattern <- generateRandomString() - promise <- Promise.make[RedisError, String] - _ <- redis - .pSubscribeWithCallback(pattern)((key: String, _: Long) => promise.succeed(key).unit) + subscription <- ZIO.service[RedisSubscription] + pattern <- generateRandomString() + promise <- Promise.make[RedisError, String] + _ <- subscription + .pSubscribeWithCallback(pattern)( + (key: String, _: Long) => promise.succeed(key).unit, + (_, _) => ZIO.unit + ) .returning[String] .flatMap(_.interruptWhen(promise).runHead) .fork @@ -80,12 +89,13 @@ trait PubSubSpec extends BaseSpec { }, test("pmessage response") { for { - redis <- ZIO.service[Redis] - prefix <- generateRandomString(5) - pattern = prefix + '*' - channel <- generateRandomString(prefix) - message <- generateRandomString(prefix) - stream <- redis + redis <- ZIO.service[Redis] + subscription <- ZIO.service[RedisSubscription] + prefix <- generateRandomString(5) + pattern = prefix + '*' + channel <- generateRandomString(prefix) + message <- generateRandomString(prefix) + stream <- subscription .pSubscribe(pattern) .returning[String] .flatMap(_.runHead) @@ -93,19 +103,20 @@ trait PubSubSpec extends BaseSpec { _ <- redis.pubSubNumPat.repeatUntil(_ > 0) _ <- redis.publish(channel, message) res <- stream.join - } yield assertTrue(res.get == message) + } yield assertTrue(res.get == (channel, message)) } ), suite("publish")(test("publish long type message") { val message = 1L assertZIO( for { - redis <- ZIO.service[Redis] - channel <- generateRandomString() - stream <- redis + redis <- ZIO.service[Redis] + subscription <- ZIO.service[RedisSubscription] + channel <- generateRandomString() + stream <- subscription .subscribe(channel) .returning[Long] - .flatMap(_.runFoldWhile(0L)(_ < 10L) { case (sum, message) => + .flatMap(_.runFoldWhile(0L)(_ < 10L) { case (sum, (_, message)) => sum + message }.fork) _ <- redis.pubSubChannels(channel).repeatUntil(_ contains channel) @@ -118,12 +129,13 @@ trait PubSubSpec extends BaseSpec { test("don't receive message type after unsubscribe") { val numOfPublished = 5 for { - redis <- ZIO.service[Redis] - prefix <- generateRandomString(5) - pattern = prefix + '*' - channel <- generateRandomString(prefix) - message <- generateRandomString() - _ <- redis + redis <- ZIO.service[Redis] + subscription <- ZIO.service[RedisSubscription] + prefix <- generateRandomString(5) + pattern = prefix + '*' + channel <- generateRandomString(prefix) + message <- generateRandomString() + _ <- subscription .subscribe(channel) .returning[String] .flatMap(_.runCollect) @@ -131,44 +143,60 @@ trait PubSubSpec extends BaseSpec { _ <- redis .pubSubChannels(pattern) .repeatUntil(_ contains channel) - promise <- redis.unsubscribe(channel) - _ <- promise.await + _ <- subscription.unsubscribe(channel) receiverCount <- redis.publish(channel, message).replicateZIO(numOfPublished).map(_.head) } yield assertTrue(receiverCount == 0L) }, test("unsubscribe response") { for { - redis <- ZIO.service[Redis] - channel <- generateRandomString() - res <- redis - .unsubscribe(channel) - .flatMap(_.await) - } yield assertTrue(res.head._1 == channel) + subscription <- ZIO.service[RedisSubscription] + channel <- generateRandomString() + promise <- Promise.make[RedisError, String] + _ <- subscription + .subscribeWithCallback(channel)( + (_, _) => ZIO.unit, + (key, _) => promise.succeed(key).unit + ) + .returning[Unit] + .flatMap(_.runDrain) + .fork + _ <- subscription.unsubscribe(channel) + res <- promise.await + } yield assertTrue(res == channel) }, test("punsubscribe response") { for { - redis <- ZIO.service[Redis] - pattern <- generateRandomString() - res <- redis - .pUnsubscribe(pattern) - .flatMap(_.await) - } yield assertTrue(res.head._1 == pattern) + subscription <- ZIO.service[RedisSubscription] + pattern <- generateRandomString() + promise <- Promise.make[RedisError, String] + _ <- subscription + .pSubscribeWithCallback(pattern)( + (_, _) => ZIO.unit, + (key, _) => promise.succeed(key).unit + ) + .returning[Unit] + .flatMap(_.runDrain) + .fork + _ <- subscription.pUnsubscribe(pattern) + res <- promise.await + } yield assertTrue(res == pattern) }, test("unsubscribe with empty param") { for { - redis <- ZIO.service[Redis] - prefix <- generateRandomString(5) - pattern = prefix + '*' - channel1 <- generateRandomString(prefix) - channel2 <- generateRandomString(prefix) + redis <- ZIO.service[Redis] + subscription <- ZIO.service[RedisSubscription] + prefix <- generateRandomString(5) + pattern = prefix + '*' + channel1 <- generateRandomString(prefix) + channel2 <- generateRandomString(prefix) _ <- - redis + subscription .subscribe(channel1) .returning[String] .flatMap(_.runCollect) .fork _ <- - redis + subscription .subscribe(channel2) .returning[String] .flatMap(_.runCollect) @@ -176,7 +204,7 @@ trait PubSubSpec extends BaseSpec { _ <- redis .pubSubChannels(pattern) .repeatUntil(_.size >= 2) - _ <- redis.unsubscribe().flatMap(_.await) + _ <- subscription.unsubscribe() numSubResponses <- redis.pubSubNumSub(channel1, channel2) } yield assertTrue( numSubResponses == Chunk( From 3c8a30780a9edfd5c392d52c5fa589f1b1e92d70 Mon Sep 17 00:00:00 2001 From: 0pg Date: Fri, 31 Mar 2023 05:40:36 +0900 Subject: [PATCH 26/51] Rename classes --- .../scala/zio/redis/RedisSubscription.scala | 2 +- ...d.scala => RedisSubscriptionCommand.scala} | 14 ++++++------- .../SingleNodeSubscriptionExecutor.scala | 16 +++++++-------- .../{Subscribe.scala => Subscription.scala} | 18 ++++++++--------- .../src/test/scala/zio/redis/PubSubSpec.scala | 20 ++++++++++++++++--- 5 files changed, 42 insertions(+), 28 deletions(-) rename redis/src/main/scala/zio/redis/{RedisPubSubCommand.scala => RedisSubscriptionCommand.scala} (89%) rename redis/src/main/scala/zio/redis/api/{Subscribe.scala => Subscription.scala} (80%) diff --git a/redis/src/main/scala/zio/redis/RedisSubscription.scala b/redis/src/main/scala/zio/redis/RedisSubscription.scala index 1e061faf1..ad2448b39 100644 --- a/redis/src/main/scala/zio/redis/RedisSubscription.scala +++ b/redis/src/main/scala/zio/redis/RedisSubscription.scala @@ -3,7 +3,7 @@ package zio.redis import zio.{URLayer, ZIO, ZLayer} import zio.schema.codec.BinaryCodec -trait RedisSubscription extends api.Subscribe +trait RedisSubscription extends api.Subscription object RedisSubscription { lazy val layer: URLayer[SubscriptionExecutor with BinaryCodec, RedisSubscription] = diff --git a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala b/redis/src/main/scala/zio/redis/RedisSubscriptionCommand.scala similarity index 89% rename from redis/src/main/scala/zio/redis/RedisPubSubCommand.scala rename to redis/src/main/scala/zio/redis/RedisSubscriptionCommand.scala index 03ed53178..3163849d6 100644 --- a/redis/src/main/scala/zio/redis/RedisPubSubCommand.scala +++ b/redis/src/main/scala/zio/redis/RedisSubscriptionCommand.scala @@ -9,7 +9,7 @@ import zio.stream.ZStream.RefineToOrDieOps import zio.stream._ import zio.{Chunk, IO, Promise, Ref, ZIO} -private[redis] final case class RedisPubSubCommand(codec: BinaryCodec, executor: SubscriptionExecutor) extends { +private[redis] final case class RedisSubscriptionCommand(codec: BinaryCodec, executor: SubscriptionExecutor) extends { import zio.redis.options.PubSub.PushProtocol._ def subscribe[A: Schema]( @@ -17,7 +17,7 @@ private[redis] final case class RedisPubSubCommand(codec: BinaryCodec, executor: onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback ): IO[RedisError, Stream[RedisError, (String, A)]] = { - val command = CommandNameInput.encode(api.Subscribe.Subscribe)(codec) ++ + val command = CommandNameInput.encode(api.Subscription.Subscribe)(codec) ++ Varargs(ArbitraryKeyInput[String]()).encode(channels)(codec) val channelSet = channels.toSet @@ -54,7 +54,7 @@ private[redis] final case class RedisPubSubCommand(codec: BinaryCodec, executor: onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback ): IO[RedisError, Stream[RedisError, (String, A)]] = { - val command = CommandNameInput.encode(api.Subscribe.PSubscribe)(codec) ++ + val command = CommandNameInput.encode(api.Subscription.PSubscribe)(codec) ++ Varargs(ArbitraryKeyInput[String]()).encode(patterns)(codec) val patternSet = patterns.toSet @@ -87,20 +87,20 @@ private[redis] final case class RedisPubSubCommand(codec: BinaryCodec, executor: } def unsubscribe(channels: Chunk[String]): IO[RedisError, Unit] = { - val command = CommandNameInput.encode(api.Subscribe.Unsubscribe)(codec) ++ + val command = CommandNameInput.encode(api.Subscription.Unsubscribe)(codec) ++ Varargs(ArbitraryKeyInput[String]()).encode(channels)(codec) executor .execute(command) - .unit + .flatMap(_.runDrain) } def pUnsubscribe(patterns: Chunk[String]): IO[RedisError, Unit] = { - val command = CommandNameInput.encode(api.Subscribe.PUnsubscribe)(codec) ++ + val command = CommandNameInput.encode(api.Subscription.PUnsubscribe)(codec) ++ Varargs(ArbitraryKeyInput[String]()).encode(patterns)(codec) executor .execute(command) - .unit + .flatMap(_.runDrain) } } diff --git a/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala b/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala index b3c7d681b..2d020c567 100644 --- a/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala +++ b/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala @@ -1,7 +1,7 @@ package zio.redis import zio.redis.SingleNodeSubscriptionExecutor.{Request, RequestQueueSize, True} -import zio.redis.api.Subscribe +import zio.redis.api.Subscription import zio.stream._ import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Schedule, ZIO} @@ -18,10 +18,10 @@ final class SingleNodeSubscriptionExecutor( .fromOption(command.args.collectFirst { case RespArgument.CommandName(name) => name }) .orElseFail(RedisError.CommandNameNotFound(command.args.toString())) stream <- commandName match { - case Subscribe.Subscribe => ZIO.succeed(subscribe(channelSubsRef, command)) - case Subscribe.PSubscribe => ZIO.succeed(subscribe(patternSubsRef, command)) - case Subscribe.Unsubscribe => ZIO.succeed(unsubscribe(channelSubsRef, command)) - case Subscribe.PUnsubscribe => ZIO.succeed(unsubscribe(patternSubsRef, command)) + case Subscription.Subscribe => ZIO.succeed(subscribe(channelSubsRef, command)) + case Subscription.PSubscribe => ZIO.succeed(subscribe(patternSubsRef, command)) + case Subscription.Unsubscribe => ZIO.succeed(unsubscribe(channelSubsRef, command)) + case Subscription.PUnsubscribe => ZIO.succeed(unsubscribe(patternSubsRef, command)) case other => ZIO.fail(RedisError.InvalidPubSubCommand(other)) } } yield stream @@ -59,7 +59,7 @@ final class SingleNodeSubscriptionExecutor( else Set.empty ) - } yield ZStream.fromHub(hub) + } yield ZStream.empty ) .flatten @@ -101,8 +101,8 @@ final class SingleNodeSubscriptionExecutor( for { channels <- channelSubsRef.get patterns <- patternSubsRef.get - commands = makeCommand(Subscribe.Subscribe, Chunk.fromIterable(channels)) ++ - makeCommand(Subscribe.PSubscribe, Chunk.fromIterable(patterns)) + commands = makeCommand(Subscription.Subscribe, Chunk.fromIterable(channels)) ++ + makeCommand(Subscription.PSubscribe, Chunk.fromIterable(patterns)) _ <- connection .write(commands) .when(commands.nonEmpty) diff --git a/redis/src/main/scala/zio/redis/api/Subscribe.scala b/redis/src/main/scala/zio/redis/api/Subscription.scala similarity index 80% rename from redis/src/main/scala/zio/redis/api/Subscribe.scala rename to redis/src/main/scala/zio/redis/api/Subscription.scala index 43f7bcde1..e1824fd4c 100644 --- a/redis/src/main/scala/zio/redis/api/Subscribe.scala +++ b/redis/src/main/scala/zio/redis/api/Subscription.scala @@ -7,13 +7,13 @@ import zio.schema.Schema import zio.stream.Stream import zio.{Chunk, IO, ZIO} -trait Subscribe extends SubscribeEnvironment { - import Subscribe._ +trait Subscription extends SubscribeEnvironment { + import Subscription._ final def subscribe(channel: String, channels: String*) = new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = - RedisPubSubCommand(codec, executor).subscribe( + RedisSubscriptionCommand(codec, executor).subscribe( Chunk.single(channel) ++ Chunk.fromIterable(channels), emptyCallback, emptyCallback @@ -26,7 +26,7 @@ trait Subscribe extends SubscribeEnvironment { ) = new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = - RedisPubSubCommand(codec, executor).subscribe( + RedisSubscriptionCommand(codec, executor).subscribe( Chunk.single(channel) ++ Chunk.fromIterable(channels), onSubscribe, onUnsubscribe @@ -36,7 +36,7 @@ trait Subscribe extends SubscribeEnvironment { final def pSubscribe(pattern: String, patterns: String*) = new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = - RedisPubSubCommand(codec, executor).pSubscribe( + RedisSubscriptionCommand(codec, executor).pSubscribe( Chunk.single(pattern) ++ Chunk.fromIterable(patterns), emptyCallback, emptyCallback @@ -49,7 +49,7 @@ trait Subscribe extends SubscribeEnvironment { )(onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback) = new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = - RedisPubSubCommand(codec, executor).pSubscribe( + RedisSubscriptionCommand(codec, executor).pSubscribe( Chunk.single(pattern) ++ Chunk.fromIterable(patterns), onSubscribe, onUnsubscribe @@ -57,14 +57,14 @@ trait Subscribe extends SubscribeEnvironment { } final def unsubscribe(channels: String*): IO[RedisError, Unit] = - RedisPubSubCommand(codec, executor).unsubscribe(Chunk.fromIterable(channels)) + RedisSubscriptionCommand(codec, executor).unsubscribe(Chunk.fromIterable(channels)) final def pUnsubscribe(patterns: String*): IO[RedisError, Unit] = - RedisPubSubCommand(codec, executor).unsubscribe(Chunk.fromIterable(patterns)) + RedisSubscriptionCommand(codec, executor).pUnsubscribe(Chunk.fromIterable(patterns)) } -object Subscribe { +object Subscription { private lazy val emptyCallback = (_: String, _: Long) => ZIO.unit final val Subscribe = "SUBSCRIBE" diff --git a/redis/src/test/scala/zio/redis/PubSubSpec.scala b/redis/src/test/scala/zio/redis/PubSubSpec.scala index e410240be..ca313a6fd 100644 --- a/redis/src/test/scala/zio/redis/PubSubSpec.scala +++ b/redis/src/test/scala/zio/redis/PubSubSpec.scala @@ -56,10 +56,12 @@ trait PubSubSpec extends BaseSpec { stream1 <- subscription .subscribe(channel1) .returning[String] + .flatMap(_.runDrain) .fork stream2 <- subscription .subscribe(channel2) .returning[String] + .flatMap(_.runDrain) .fork _ <- redis .pubSubChannels(pattern) @@ -135,8 +137,9 @@ trait PubSubSpec extends BaseSpec { pattern = prefix + '*' channel <- generateRandomString(prefix) message <- generateRandomString() + promise <- Promise.make[Nothing, Unit] _ <- subscription - .subscribe(channel) + .subscribeWithCallback(channel)((_, _) => ZIO.unit, (_, _) => promise.succeed(()).unit) .returning[String] .flatMap(_.runCollect) .fork @@ -144,6 +147,7 @@ trait PubSubSpec extends BaseSpec { .pubSubChannels(pattern) .repeatUntil(_ contains channel) _ <- subscription.unsubscribe(channel) + _ <- promise.await receiverCount <- redis.publish(channel, message).replicateZIO(numOfPublished).map(_.head) } yield assertTrue(receiverCount == 0L) }, @@ -189,15 +193,23 @@ trait PubSubSpec extends BaseSpec { pattern = prefix + '*' channel1 <- generateRandomString(prefix) channel2 <- generateRandomString(prefix) + promise1 <- Promise.make[Nothing, Unit] + promise2 <- Promise.make[Nothing, Unit] _ <- subscription - .subscribe(channel1) + .subscribeWithCallback(channel1)( + (_, _) => ZIO.unit, + (_, _) => promise1.succeed(()).unit + ) .returning[String] .flatMap(_.runCollect) .fork _ <- subscription - .subscribe(channel2) + .subscribeWithCallback(channel2)( + (_, _) => ZIO.unit, + (_, _) => promise2.succeed(()).unit + ) .returning[String] .flatMap(_.runCollect) .fork @@ -205,6 +217,8 @@ trait PubSubSpec extends BaseSpec { .pubSubChannels(pattern) .repeatUntil(_.size >= 2) _ <- subscription.unsubscribe() + _ <- promise1.await + _ <- promise2.await numSubResponses <- redis.pubSubNumSub(channel1, channel2) } yield assertTrue( numSubResponses == Chunk( From 0ea75a1862c812382584c9e2928af47604a0e19c Mon Sep 17 00:00:00 2001 From: 0pg Date: Fri, 31 Mar 2023 12:18:55 +0900 Subject: [PATCH 27/51] Reduce duplication --- .../zio/redis/RedisSubscriptionCommand.scala | 43 +++++++------------ .../main/scala/zio/redis/options/PubSub.scala | 4 +- 2 files changed, 17 insertions(+), 30 deletions(-) diff --git a/redis/src/main/scala/zio/redis/RedisSubscriptionCommand.scala b/redis/src/main/scala/zio/redis/RedisSubscriptionCommand.scala index 3163849d6..7714c92bd 100644 --- a/redis/src/main/scala/zio/redis/RedisSubscriptionCommand.scala +++ b/redis/src/main/scala/zio/redis/RedisSubscriptionCommand.scala @@ -2,6 +2,7 @@ package zio.redis import zio.redis.Input._ import zio.redis.Output.{ArbitraryOutput, PushProtocolOutput} +import zio.redis.api.Subscription import zio.redis.options.PubSub.PubSubCallback import zio.schema.Schema import zio.schema.codec.BinaryCodec @@ -16,17 +17,13 @@ private[redis] final case class RedisSubscriptionCommand(codec: BinaryCodec, exe channels: Chunk[String], onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback - ): IO[RedisError, Stream[RedisError, (String, A)]] = { - val command = CommandNameInput.encode(api.Subscription.Subscribe)(codec) ++ - Varargs(ArbitraryKeyInput[String]()).encode(channels)(codec) - - val channelSet = channels.toSet - + ): IO[RedisError, Stream[RedisError, (String, A)]] = for { unsubscribedRef <- Ref.make(channels.map(_ -> false).toMap) promise <- Promise.make[RedisError, Unit] + channelSet = channels.toSet stream <- executor - .execute(command) + .execute(makeCommand(Subscription.Subscribe, channels)) .map( _.mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp)(codec))).mapZIO { case Subscribe(channel, numOfSubscription) if channelSet contains channel => @@ -47,23 +44,18 @@ private[redis] final case class RedisSubscriptionCommand(codec: BinaryCodec, exe .interruptWhen(promise) ) } yield stream - } def pSubscribe[A: Schema]( patterns: Chunk[String], onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback - ): IO[RedisError, Stream[RedisError, (String, A)]] = { - val command = CommandNameInput.encode(api.Subscription.PSubscribe)(codec) ++ - Varargs(ArbitraryKeyInput[String]()).encode(patterns)(codec) - - val patternSet = patterns.toSet - + ): IO[RedisError, Stream[RedisError, (String, A)]] = for { unsubscribedRef <- Ref.make(patterns.map(_ -> false).toMap) promise <- Promise.make[RedisError, Unit] + patternSet = patterns.toSet stream <- executor - .execute(command) + .execute(makeCommand(Subscription.PSubscribe, patterns)) .map( _.mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp)(codec))).mapZIO { case PSubscribe(pattern, numOfSubscription) if patternSet contains pattern => @@ -84,23 +76,18 @@ private[redis] final case class RedisSubscriptionCommand(codec: BinaryCodec, exe .interruptWhen(promise) ) } yield stream - } - - def unsubscribe(channels: Chunk[String]): IO[RedisError, Unit] = { - val command = CommandNameInput.encode(api.Subscription.Unsubscribe)(codec) ++ - Varargs(ArbitraryKeyInput[String]()).encode(channels)(codec) + def unsubscribe(channels: Chunk[String]): IO[RedisError, Unit] = executor - .execute(command) + .execute(makeCommand(Subscription.Unsubscribe, channels)) .flatMap(_.runDrain) - } - - def pUnsubscribe(patterns: Chunk[String]): IO[RedisError, Unit] = { - val command = CommandNameInput.encode(api.Subscription.PUnsubscribe)(codec) ++ - Varargs(ArbitraryKeyInput[String]()).encode(patterns)(codec) + def pUnsubscribe(patterns: Chunk[String]): IO[RedisError, Unit] = executor - .execute(command) + .execute(makeCommand(Subscription.PUnsubscribe, patterns)) .flatMap(_.runDrain) - } + + private def makeCommand(commandName: String, keys: Chunk[String]) = + CommandNameInput.encode(commandName)(codec) ++ + Varargs(ArbitraryKeyInput[String]()).encode(keys)(codec) } diff --git a/redis/src/main/scala/zio/redis/options/PubSub.scala b/redis/src/main/scala/zio/redis/options/PubSub.scala index 9b6d1988f..21557d481 100644 --- a/redis/src/main/scala/zio/redis/options/PubSub.scala +++ b/redis/src/main/scala/zio/redis/options/PubSub.scala @@ -5,9 +5,9 @@ import zio.redis.RespValue object PubSub { type PubSubCallback = (String, Long) => UIO[Unit] - sealed trait PushProtocol + private[redis] sealed trait PushProtocol - object PushProtocol { + private[redis] object PushProtocol { case class Subscribe(channel: String, numOfSubscription: Long) extends PushProtocol case class PSubscribe(pattern: String, numOfSubscription: Long) extends PushProtocol case class Unsubscribe(channel: String, numOfSubscription: Long) extends PushProtocol From 478577b98f008395d168ae0cffb23c7fd8851937 Mon Sep 17 00:00:00 2001 From: 0pg Date: Fri, 31 Mar 2023 12:55:02 +0900 Subject: [PATCH 28/51] Apply upstream changes --- .../src/main/scala/zio/redis/Output.scala | 4 +- .../scala/zio/redis/RedisSubscription.scala | 17 +++++++ .../SingleNodeSubscriptionExecutor.scala | 36 ++++++++------- .../{internal => }/SubscriptionExecutor.scala | 7 +-- .../redis/internal/RedisSubscription.scala | 18 -------- .../internal/RedisSubscriptionCommand.scala | 29 ++++++------ .../redis/internal/SubscribeEnvironment.scala | 9 ++-- .../zio/redis/internal/api/Publishing.scala | 9 ++-- .../zio/redis/internal/api/Subscription.scala | 13 +++--- .../zio/redis/internal/options/PubSub.scala | 2 +- .../src/test/scala/zio/redis/ApiSpec.scala | 44 +++++++++---------- 11 files changed, 98 insertions(+), 90 deletions(-) create mode 100644 modules/redis/src/main/scala/zio/redis/RedisSubscription.scala rename modules/redis/src/main/scala/zio/redis/{internal => }/SingleNodeSubscriptionExecutor.scala (76%) rename modules/redis/src/main/scala/zio/redis/{internal => }/SubscriptionExecutor.scala (68%) delete mode 100644 modules/redis/src/main/scala/zio/redis/internal/RedisSubscription.scala diff --git a/modules/redis/src/main/scala/zio/redis/Output.scala b/modules/redis/src/main/scala/zio/redis/Output.scala index 07b522538..4ac5bc20b 100644 --- a/modules/redis/src/main/scala/zio/redis/Output.scala +++ b/modules/redis/src/main/scala/zio/redis/Output.scala @@ -828,7 +828,7 @@ object Output { } case object PushProtocolOutput extends Output[PushProtocol] { - protected def tryDecode(respValue: RespValue)(implicit codec: BinaryCodec): PushProtocol = + protected def tryDecode(respValue: RespValue): PushProtocol = respValue match { case RespValue.NullArray => throw ProtocolError(s"Array must not be empty") case RespValue.Array(values) => @@ -859,7 +859,7 @@ object Output { } case object NumSubResponseOutput extends Output[Chunk[NumberOfSubscribers]] { - protected def tryDecode(respValue: RespValue)(implicit codec: BinaryCodec): Chunk[NumberOfSubscribers] = + protected def tryDecode(respValue: RespValue): Chunk[NumberOfSubscribers] = respValue match { case RespValue.Array(values) => Chunk.fromIterator(values.grouped(2).map { chunk => diff --git a/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala b/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala new file mode 100644 index 000000000..1ec223feb --- /dev/null +++ b/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala @@ -0,0 +1,17 @@ +package zio.redis + +import zio.{URLayer, ZIO, ZLayer} + +trait RedisSubscription extends api.Subscription + +object RedisSubscription { + lazy val layer: URLayer[SubscriptionExecutor with CodecSupplier, RedisSubscription] = + ZLayer { + for { + codecSupplier <- ZIO.service[CodecSupplier] + executor <- ZIO.service[SubscriptionExecutor] + } yield Live(codecSupplier, executor) + } + + private final case class Live(codecSupplier: CodecSupplier, executor: SubscriptionExecutor) extends RedisSubscription +} diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala similarity index 76% rename from modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala rename to modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala index 2d020c567..05fd83eb4 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala @@ -2,32 +2,34 @@ package zio.redis import zio.redis.SingleNodeSubscriptionExecutor.{Request, RequestQueueSize, True} import zio.redis.api.Subscription +import zio.redis.internal.{RedisConnection, RespCommand, RespCommandArgument, RespValue} import zio.stream._ -import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Schedule, ZIO} +import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Schedule, ZIO, redis} final class SingleNodeSubscriptionExecutor( - channelSubsRef: Ref[Set[RespArgument.Key]], - patternSubsRef: Ref[Set[RespArgument.Key]], + channelSubsRef: Ref[Set[RespCommandArgument.Key]], + patternSubsRef: Ref[Set[redis.internal.RespCommandArgument.Key]], hub: Hub[RespValue], reqQueue: Queue[Request], connection: RedisConnection ) extends SubscriptionExecutor { def execute(command: RespCommand): IO[RedisError, Stream[RedisError, RespValue]] = for { - commandName <- ZIO - .fromOption(command.args.collectFirst { case RespArgument.CommandName(name) => name }) - .orElseFail(RedisError.CommandNameNotFound(command.args.toString())) + commandName <- + ZIO + .fromOption(command.args.collectFirst { case redis.internal.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(channelSubsRef, command)) case Subscription.PUnsubscribe => ZIO.succeed(unsubscribe(patternSubsRef, command)) - case other => ZIO.fail(RedisError.InvalidPubSubCommand(other)) + case other => ZIO.fail(RedisError.InvalidPubSubCommand(other)) } } yield stream private def subscribe( - subscriptionRef: Ref[Set[RespArgument.Key]], + subscriptionRef: Ref[Set[redis.internal.RespCommandArgument.Key]], command: RespCommand ): Stream[RedisError, RespValue] = ZStream @@ -36,14 +38,14 @@ final class SingleNodeSubscriptionExecutor( reqPromise <- Promise.make[RedisError, Unit] _ <- reqQueue.offer(Request(command.args.map(_.value), reqPromise)) _ <- reqPromise.await - keys = command.args.collect { case key: RespArgument.Key => key } + keys = command.args.collect { case key: redis.internal.RespCommandArgument.Key => key } _ <- subscriptionRef.update(_ ++ keys) } yield ZStream.fromHub(hub) ) .flatten private def unsubscribe( - subscriptionRef: Ref[Set[RespArgument.Key]], + subscriptionRef: Ref[Set[redis.internal.RespCommandArgument.Key]], command: RespCommand ): Stream[RedisError, RespValue] = ZStream @@ -52,7 +54,7 @@ final class SingleNodeSubscriptionExecutor( reqPromise <- Promise.make[RedisError, Unit] _ <- reqQueue.offer(Request(command.args.map(_.value), reqPromise)) _ <- reqPromise.await - keys = command.args.collect { case key: RespArgument.Key => key } + keys = command.args.collect { case key: redis.internal.RespCommandArgument.Key => key } _ <- subscriptionRef.update(subscribedSet => if (keys.nonEmpty) subscribedSet -- keys @@ -87,16 +89,18 @@ final class SingleNodeSubscriptionExecutor( private def receive: IO[RedisError, Unit] = connection.read .mapError(RedisError.IOError(_)) - .via(RespValue.decoder) + .via(RespValue.Decoder) .collectSome .foreach(hub.offer(_)) private def resubscribe: IO[RedisError, Unit] = { - def makeCommand(name: String, keys: Chunk[RespArgument.Key]) = + def makeCommand(name: String, keys: Chunk[redis.internal.RespCommandArgument.Key]) = if (keys.isEmpty) Chunk.empty else - RespValue.Array((RespCommand(RespArgument.CommandName(name)) ++ RespCommand(keys)).args.map(_.value)).serialize + RespValue + .Array((RespCommand(RespCommandArgument.CommandName(name)) ++ RespCommand(keys)).args.map(_.value)) + .serialize for { channels <- channelSubsRef.get @@ -137,8 +141,8 @@ object SingleNodeSubscriptionExecutor { for { hub <- Hub.unbounded[RespValue] reqQueue <- Queue.bounded[Request](RequestQueueSize) - channelRef <- Ref.make(Set.empty[RespArgument.Key]) - patternRef <- Ref.make(Set.empty[RespArgument.Key]) + channelRef <- Ref.make(Set.empty[RespCommandArgument.Key]) + patternRef <- Ref.make(Set.empty[RespCommandArgument.Key]) pubSub = new SingleNodeSubscriptionExecutor(channelRef, patternRef, hub, reqQueue, conn) _ <- pubSub.run.forkScoped _ <- logScopeFinalizer(s"$pubSub Node PubSub is closed") diff --git a/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala similarity index 68% rename from modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala rename to modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala index 22125c251..003b6f7de 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala @@ -1,18 +1,19 @@ package zio.redis +import zio.redis.internal.{RedisConnection, RespCommand, RespValue} import zio.stream._ import zio.{IO, Layer, ZIO, ZLayer} trait SubscriptionExecutor { - def execute(command: RespCommand): IO[RedisError, Stream[RedisError, RespValue]] + private[redis] def execute(command: RespCommand): IO[RedisError, Stream[RedisError, RespValue]] } object SubscriptionExecutor { lazy val layer: ZLayer[RedisConfig, RedisError.IOError, SubscriptionExecutor] = - RedisConnectionLive.layer.fresh >>> pubSublayer + RedisConnection.layer.fresh >>> pubSublayer lazy val local: Layer[RedisError.IOError, SubscriptionExecutor] = - RedisConnectionLive.default.fresh >>> pubSublayer + RedisConnection.local.fresh >>> pubSublayer private lazy val pubSublayer: ZLayer[RedisConnection, RedisError.IOError, SubscriptionExecutor] = ZLayer.scoped( diff --git a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscription.scala b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscription.scala deleted file mode 100644 index ad2448b39..000000000 --- a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscription.scala +++ /dev/null @@ -1,18 +0,0 @@ -package zio.redis - -import zio.{URLayer, ZIO, ZLayer} -import zio.schema.codec.BinaryCodec - -trait RedisSubscription extends api.Subscription - -object RedisSubscription { - lazy val layer: URLayer[SubscriptionExecutor with BinaryCodec, RedisSubscription] = - ZLayer { - for { - codec <- ZIO.service[BinaryCodec] - executor <- ZIO.service[SubscriptionExecutor] - } yield Live(codec, executor) - } - - private final case class Live(codec: BinaryCodec, executor: SubscriptionExecutor) extends RedisSubscription -} diff --git a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala index 7714c92bd..0c4629539 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala @@ -4,20 +4,19 @@ import zio.redis.Input._ import zio.redis.Output.{ArbitraryOutput, PushProtocolOutput} import zio.redis.api.Subscription import zio.redis.options.PubSub.PubSubCallback -import zio.schema.Schema import zio.schema.codec.BinaryCodec import zio.stream.ZStream.RefineToOrDieOps import zio.stream._ import zio.{Chunk, IO, Promise, Ref, ZIO} -private[redis] final case class RedisSubscriptionCommand(codec: BinaryCodec, executor: SubscriptionExecutor) extends { +private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionExecutor) extends { import zio.redis.options.PubSub.PushProtocol._ - def subscribe[A: Schema]( + def subscribe[A: BinaryCodec]( channels: Chunk[String], onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback - ): IO[RedisError, Stream[RedisError, (String, A)]] = + )(implicit codec: BinaryCodec[String]): IO[RedisError, Stream[RedisError, (String, A)]] = for { unsubscribedRef <- Ref.make(channels.map(_ -> false).toMap) promise <- Promise.make[RedisError, Unit] @@ -25,12 +24,12 @@ private[redis] final case class RedisSubscriptionCommand(codec: BinaryCodec, exe stream <- executor .execute(makeCommand(Subscription.Subscribe, channels)) .map( - _.mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp)(codec))).mapZIO { + _.mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))).mapZIO { case Subscribe(channel, numOfSubscription) if channelSet contains channel => onSubscribe(channel, numOfSubscription).as(None) case Message(channel, message) if channelSet contains channel => ZIO - .attempt(ArbitraryOutput[A]().unsafeDecode(message)(codec)) + .attempt(ArbitraryOutput[A]().unsafeDecode(message)) .map(msg => Some((channel, msg))) case Unsubscribe(channel, numOfSubscription) if channelSet contains channel => for { @@ -45,11 +44,11 @@ private[redis] final case class RedisSubscriptionCommand(codec: BinaryCodec, exe ) } yield stream - def pSubscribe[A: Schema]( + def pSubscribe[A: BinaryCodec]( patterns: Chunk[String], onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback - ): IO[RedisError, Stream[RedisError, (String, A)]] = + )(implicit stringCodec: BinaryCodec[String]): IO[RedisError, Stream[RedisError, (String, A)]] = for { unsubscribedRef <- Ref.make(patterns.map(_ -> false).toMap) promise <- Promise.make[RedisError, Unit] @@ -57,12 +56,12 @@ private[redis] final case class RedisSubscriptionCommand(codec: BinaryCodec, exe stream <- executor .execute(makeCommand(Subscription.PSubscribe, patterns)) .map( - _.mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp)(codec))).mapZIO { + _.mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))).mapZIO { case PSubscribe(pattern, numOfSubscription) if patternSet contains pattern => onSubscribe(pattern, numOfSubscription).as(None) case PMessage(pattern, channel, message) if patternSet contains pattern => ZIO - .attempt(ArbitraryOutput[A]().unsafeDecode(message)(codec)) + .attempt(ArbitraryOutput[A]().unsafeDecode(message)) .map(msg => Some((channel, msg))) case PUnsubscribe(pattern, numOfSubscription) if patternSet contains pattern => for { @@ -77,17 +76,17 @@ private[redis] final case class RedisSubscriptionCommand(codec: BinaryCodec, exe ) } yield stream - def unsubscribe(channels: Chunk[String]): IO[RedisError, Unit] = + def unsubscribe(channels: Chunk[String])(implicit codec: BinaryCodec[String]): IO[RedisError, Unit] = executor .execute(makeCommand(Subscription.Unsubscribe, channels)) .flatMap(_.runDrain) - def pUnsubscribe(patterns: Chunk[String]): IO[RedisError, Unit] = + def pUnsubscribe(patterns: Chunk[String])(implicit codec: BinaryCodec[String]): IO[RedisError, Unit] = executor .execute(makeCommand(Subscription.PUnsubscribe, patterns)) .flatMap(_.runDrain) - private def makeCommand(commandName: String, keys: Chunk[String]) = - CommandNameInput.encode(commandName)(codec) ++ - Varargs(ArbitraryKeyInput[String]()).encode(keys)(codec) + private def makeCommand(commandName: String, keys: Chunk[String])(implicit codec: BinaryCodec[String]) = + CommandNameInput.encode(commandName) ++ + Varargs(ArbitraryKeyInput[String]()).encode(keys) } diff --git a/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala b/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala index 44c92b961..38dbfc8f8 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala @@ -1,8 +1,11 @@ -package zio.redis +package zio.redis.internal +import zio.redis.{CodecSupplier, SubscriptionExecutor} +import zio.schema.Schema import zio.schema.codec.BinaryCodec -trait SubscribeEnvironment { - protected def codec: BinaryCodec +private[redis] trait SubscribeEnvironment { + protected def codecSupplier: CodecSupplier protected def executor: SubscriptionExecutor + protected final implicit def codec[A: Schema]: BinaryCodec[A] = codecSupplier.get } diff --git a/modules/redis/src/main/scala/zio/redis/internal/api/Publishing.scala b/modules/redis/src/main/scala/zio/redis/internal/api/Publishing.scala index 1cdab5443..79cf7a4c1 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/api/Publishing.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/api/Publishing.scala @@ -3,6 +3,7 @@ package zio.redis.api import zio.redis.Input._ import zio.redis.Output._ import zio.redis._ +import zio.redis.internal.{RedisCommand, RedisEnvironment} import zio.redis.options.PubSub.NumberOfSubscribers import zio.schema.Schema import zio.{Chunk, IO} @@ -11,22 +12,22 @@ trait Publishing extends RedisEnvironment { import Publishing._ final def publish[A: Schema](channel: String, message: A): IO[RedisError, Long] = { - val command = RedisCommand(Publish, Tuple2(StringInput, ArbitraryKeyInput[A]()), LongOutput, codec, executor) + val command = RedisCommand(Publish, Tuple2(StringInput, ArbitraryKeyInput[A]()), LongOutput, executor) command.run((channel, message)) } final def pubSubChannels(pattern: String): IO[RedisError, Chunk[String]] = { - val command = RedisCommand(PubSubChannels, StringInput, ChunkOutput(MultiStringOutput), codec, executor) + val command = RedisCommand(PubSubChannels, StringInput, ChunkOutput(MultiStringOutput), executor) command.run(pattern) } final def pubSubNumPat: IO[RedisError, Long] = { - val command = RedisCommand(PubSubNumPat, NoInput, LongOutput, codec, executor) + val command = RedisCommand(PubSubNumPat, NoInput, LongOutput, executor) command.run(()) } final def pubSubNumSub(channel: String, channels: String*): IO[RedisError, Chunk[NumberOfSubscribers]] = { - val command = RedisCommand(PubSubNumSub, NonEmptyList(StringInput), NumSubResponseOutput, codec, executor) + val command = RedisCommand(PubSubNumSub, NonEmptyList(StringInput), NumSubResponseOutput, executor) command.run((channel, channels.toList)) } } diff --git a/modules/redis/src/main/scala/zio/redis/internal/api/Subscription.scala b/modules/redis/src/main/scala/zio/redis/internal/api/Subscription.scala index e1824fd4c..73f88b412 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/api/Subscription.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/api/Subscription.scala @@ -2,6 +2,7 @@ package zio.redis.api import zio.redis.ResultBuilder.ResultBuilder1 import zio.redis._ +import zio.redis.internal.SubscribeEnvironment import zio.redis.options.PubSub.PubSubCallback import zio.schema.Schema import zio.stream.Stream @@ -13,7 +14,7 @@ trait Subscription extends SubscribeEnvironment { final def subscribe(channel: String, channels: String*) = new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = - RedisSubscriptionCommand(codec, executor).subscribe( + RedisSubscriptionCommand(executor).subscribe( Chunk.single(channel) ++ Chunk.fromIterable(channels), emptyCallback, emptyCallback @@ -26,7 +27,7 @@ trait Subscription extends SubscribeEnvironment { ) = new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = - RedisSubscriptionCommand(codec, executor).subscribe( + RedisSubscriptionCommand(executor).subscribe( Chunk.single(channel) ++ Chunk.fromIterable(channels), onSubscribe, onUnsubscribe @@ -36,7 +37,7 @@ trait Subscription extends SubscribeEnvironment { final def pSubscribe(pattern: String, patterns: String*) = new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = - RedisSubscriptionCommand(codec, executor).pSubscribe( + RedisSubscriptionCommand(executor).pSubscribe( Chunk.single(pattern) ++ Chunk.fromIterable(patterns), emptyCallback, emptyCallback @@ -49,7 +50,7 @@ trait Subscription extends SubscribeEnvironment { )(onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback) = new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = - RedisSubscriptionCommand(codec, executor).pSubscribe( + RedisSubscriptionCommand(executor).pSubscribe( Chunk.single(pattern) ++ Chunk.fromIterable(patterns), onSubscribe, onUnsubscribe @@ -57,10 +58,10 @@ trait Subscription extends SubscribeEnvironment { } final def unsubscribe(channels: String*): IO[RedisError, Unit] = - RedisSubscriptionCommand(codec, executor).unsubscribe(Chunk.fromIterable(channels)) + RedisSubscriptionCommand(executor).unsubscribe(Chunk.fromIterable(channels)) final def pUnsubscribe(patterns: String*): IO[RedisError, Unit] = - RedisSubscriptionCommand(codec, executor).pUnsubscribe(Chunk.fromIterable(patterns)) + RedisSubscriptionCommand(executor).pUnsubscribe(Chunk.fromIterable(patterns)) } diff --git a/modules/redis/src/main/scala/zio/redis/internal/options/PubSub.scala b/modules/redis/src/main/scala/zio/redis/internal/options/PubSub.scala index 21557d481..62e4e91ce 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/options/PubSub.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/options/PubSub.scala @@ -1,7 +1,7 @@ package zio.redis.options import zio.UIO -import zio.redis.RespValue +import zio.redis.internal.RespValue object PubSub { type PubSubCallback = (String, Long) => UIO[Unit] diff --git a/modules/redis/src/test/scala/zio/redis/ApiSpec.scala b/modules/redis/src/test/scala/zio/redis/ApiSpec.scala index 2a9fb3af0..48839868f 100644 --- a/modules/redis/src/test/scala/zio/redis/ApiSpec.scala +++ b/modules/redis/src/test/scala/zio/redis/ApiSpec.scala @@ -20,7 +20,7 @@ object ApiSpec with PubSubSpec { def spec: Spec[TestEnvironment, Any] = - suite("Redis commands")(clusterSuite, singleNodeSuite) @@ sequential @@ withLiveEnvironment + suite("Redis commands")(singleNodeSuite) @@ sequential @@ withLiveEnvironment private val singleNodeSuite = suite("Single node executor")( @@ -44,25 +44,25 @@ object ApiSpec ZLayer.succeed(ProtobufCodecSupplier) ) - private val clusterSuite = - suite("Cluster executor")( - connectionSuite, - keysSuite, - listSuite, - stringsSuite, - hashSuite, - setsSuite, - sortedSetsSuite, - hyperLogLogSuite, - geoSuite, - streamsSuite, - scriptingSpec, - clusterSpec - ).provideShared( - ClusterExecutor.layer, - Redis.layer, - ZLayer.succeed(ProtobufCodecSupplier), - ZLayer.succeed(RedisClusterConfig(Chunk(RedisUri("localhost", 5000)))) - ).filterNotTags(_.contains(BaseSpec.ClusterExecutorUnsupported)) - .getOrElse(Spec.empty) +// private val clusterSuite = +// suite("Cluster executor")( +// connectionSuite, +// keysSuite, +// listSuite, +// stringsSuite, +// hashSuite, +// setsSuite, +// sortedSetsSuite, +// hyperLogLogSuite, +// geoSuite, +// streamsSuite, +// scriptingSpec, +// clusterSpec +// ).provideShared( +// ClusterExecutor.layer, +// Redis.layer, +// ZLayer.succeed(ProtobufCodecSupplier), +// ZLayer.succeed(RedisClusterConfig(Chunk(RedisUri("localhost", 5000)))) +// ).filterNotTags(_.contains(BaseSpec.ClusterExecutorUnsupported)) +// .getOrElse(Spec.empty) } From c4105ff4e17a1d19012d9f55eb6413d9bfe8c044 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 4 Apr 2023 16:46:02 +0900 Subject: [PATCH 29/51] Replace Hub to Queue --- .../main/scala/zio/redis/ResultBuilder.scala | 5 + .../SingleNodeSubscriptionExecutor.scala | 145 ++++++++++++------ .../zio/redis/SubscriptionExecutor.scala | 7 +- .../redis/{internal => }/api/Publishing.scala | 0 .../{internal => }/api/Subscription.scala | 37 +++-- .../internal/RedisSubscriptionCommand.scala | 98 +++++------- .../redis/{internal => }/options/PubSub.scala | 0 .../src/test/scala/zio/redis/ApiSpec.scala | 44 +++--- .../scala/zio/redis/internal/PubSubSpec.scala | 53 ++++--- 9 files changed, 215 insertions(+), 174 deletions(-) rename modules/redis/src/main/scala/zio/redis/{internal => }/api/Publishing.scala (100%) rename modules/redis/src/main/scala/zio/redis/{internal => }/api/Subscription.scala (60%) rename modules/redis/src/main/scala/zio/redis/{internal => }/options/PubSub.scala (100%) diff --git a/modules/redis/src/main/scala/zio/redis/ResultBuilder.scala b/modules/redis/src/main/scala/zio/redis/ResultBuilder.scala index f04dbf245..ab76e0f12 100644 --- a/modules/redis/src/main/scala/zio/redis/ResultBuilder.scala +++ b/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]] + } } diff --git a/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala index 05fd83eb4..7a6756878 100644 --- a/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala @@ -1,66 +1,81 @@ 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.{RedisConnection, RespCommand, RespCommandArgument, RespValue} +import zio.redis.options.PubSub.PushProtocol import zio.stream._ -import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Schedule, ZIO, redis} +import zio.{Chunk, ChunkBuilder, IO, Promise, Queue, Ref, Schedule, ZIO} final class SingleNodeSubscriptionExecutor( - channelSubsRef: Ref[Set[RespCommandArgument.Key]], - patternSubsRef: Ref[Set[redis.internal.RespCommandArgument.Key]], - hub: Hub[RespValue], + 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): IO[RedisError, Stream[RedisError, RespValue]] = - for { - commandName <- - ZIO - .fromOption(command.args.collectFirst { case redis.internal.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(channelSubsRef, command)) - case Subscription.PUnsubscribe => ZIO.succeed(unsubscribe(patternSubsRef, command)) - case other => ZIO.fail(RedisError.InvalidPubSubCommand(other)) - } - } yield stream + 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[Set[redis.internal.RespCommandArgument.Key]], + subscriptionRef: Ref[Map[String, Chunk[Queue[Take[RedisError, PushProtocol]]]]], command: RespCommand - ): Stream[RedisError, RespValue] = + ): Stream[RedisError, PushProtocol] = ZStream .fromZIO( for { - reqPromise <- Promise.make[RedisError, Unit] - _ <- reqQueue.offer(Request(command.args.map(_.value), reqPromise)) - _ <- reqPromise.await - keys = command.args.collect { case key: redis.internal.RespCommandArgument.Key => key } - _ <- subscriptionRef.update(_ ++ keys) - } yield ZStream.fromHub(hub) + queues <- ZIO.foreach(command.args.collect { case key: RespCommandArgument.Key => key.value.asString })(key => + Queue + .unbounded[Take[RedisError, PushProtocol]] + .tap(queue => + subscriptionRef.update( + _.updatedWith(key)(_.map(_ appended queue).orElse(Some(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(_.updatedWith(key)(_.map(_.filterNot(_ == queue))))) + } + _ <- promise.await.tapError(_ => ZIO.foreachDiscard(queues) { case (_, queue) => queue.shutdown }) + } yield streams.fold(ZStream.empty)(_ merge _) ) .flatten + .flattenTake - private def unsubscribe( - subscriptionRef: Ref[Set[redis.internal.RespCommandArgument.Key]], - command: RespCommand - ): Stream[RedisError, RespValue] = + private def unsubscribe(command: RespCommand): Stream[RedisError, PushProtocol] = ZStream .fromZIO( for { - reqPromise <- Promise.make[RedisError, Unit] - _ <- reqQueue.offer(Request(command.args.map(_.value), reqPromise)) - _ <- reqPromise.await - keys = command.args.collect { case key: redis.internal.RespCommandArgument.Key => key } - _ <- subscriptionRef.update(subscribedSet => - if (keys.nonEmpty) - subscribedSet -- keys - else - Set.empty - ) + promise <- Promise.make[RedisError, Unit] + _ <- reqQueue.offer(Request(command.args.map(_.value), promise)) + _ <- promise.await } yield ZStream.empty ) .flatten @@ -86,25 +101,56 @@ final class SingleNodeSubscriptionExecutor( ) } - private def receive: IO[RedisError, Unit] = + 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 - .foreach(hub.offer(_)) + .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[redis.internal.RespCommandArgument.Key]) = + def makeCommand(name: String, keys: Chunk[String]) = if (keys.isEmpty) Chunk.empty else RespValue - .Array((RespCommand(RespCommandArgument.CommandName(name)) ++ RespCommand(keys)).args.map(_.value)) + .Array((CommandNameInput.encode(name) ++ Input.Varargs(StringInput).encode(keys)).args.map(_.value)) .serialize for { - channels <- channelSubsRef.get - patterns <- patternSubsRef.get + channels <- channelSubsRef.get.map(_.keys) + patterns <- patternSubsRef.get.map(_.keys) commands = makeCommand(Subscription.Subscribe, Chunk.fromIterable(channels)) ++ makeCommand(Subscription.PSubscribe, Chunk.fromIterable(patterns)) _ <- connection @@ -139,11 +185,10 @@ object SingleNodeSubscriptionExecutor { def create(conn: RedisConnection) = for { - hub <- Hub.unbounded[RespValue] reqQueue <- Queue.bounded[Request](RequestQueueSize) - channelRef <- Ref.make(Set.empty[RespCommandArgument.Key]) - patternRef <- Ref.make(Set.empty[RespCommandArgument.Key]) - pubSub = new SingleNodeSubscriptionExecutor(channelRef, patternRef, hub, reqQueue, conn) + 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 Node PubSub is closed") } yield pubSub diff --git a/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala index 003b6f7de..91cd06a78 100644 --- a/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala @@ -1,11 +1,12 @@ package zio.redis -import zio.redis.internal.{RedisConnection, RespCommand, RespValue} +import zio.redis.internal.{RedisConnection, RespCommand} +import zio.redis.options.PubSub.PushProtocol import zio.stream._ -import zio.{IO, Layer, ZIO, ZLayer} +import zio.{Layer, ZIO, ZLayer} trait SubscriptionExecutor { - private[redis] def execute(command: RespCommand): IO[RedisError, Stream[RedisError, RespValue]] + private[redis] def execute(command: RespCommand): Stream[RedisError, PushProtocol] } object SubscriptionExecutor { diff --git a/modules/redis/src/main/scala/zio/redis/internal/api/Publishing.scala b/modules/redis/src/main/scala/zio/redis/api/Publishing.scala similarity index 100% rename from modules/redis/src/main/scala/zio/redis/internal/api/Publishing.scala rename to modules/redis/src/main/scala/zio/redis/api/Publishing.scala diff --git a/modules/redis/src/main/scala/zio/redis/internal/api/Subscription.scala b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala similarity index 60% rename from modules/redis/src/main/scala/zio/redis/internal/api/Subscription.scala rename to modules/redis/src/main/scala/zio/redis/api/Subscription.scala index 73f88b412..32a422467 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/api/Subscription.scala +++ b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala @@ -1,8 +1,8 @@ package zio.redis.api -import zio.redis.ResultBuilder.ResultBuilder1 +import zio.redis.ResultBuilder.ResultStreamBuilder1 import zio.redis._ -import zio.redis.internal.SubscribeEnvironment +import zio.redis.internal._ import zio.redis.options.PubSub.PubSubCallback import zio.schema.Schema import zio.stream.Stream @@ -11,9 +11,12 @@ import zio.{Chunk, IO, ZIO} trait Subscription extends SubscribeEnvironment { import Subscription._ - final def subscribe(channel: String, channels: String*) = - new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { - def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = + final def subscribe( + channel: String, + channels: String* + ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = + new ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] { + def returning[R: Schema]: Stream[RedisError, (String, R)] = RedisSubscriptionCommand(executor).subscribe( Chunk.single(channel) ++ Chunk.fromIterable(channels), emptyCallback, @@ -24,9 +27,9 @@ trait Subscription extends SubscribeEnvironment { final def subscribeWithCallback(channel: String, channels: String*)( onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback - ) = - new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { - def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = + ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = + new ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] { + def returning[R: Schema]: Stream[RedisError, (String, R)] = RedisSubscriptionCommand(executor).subscribe( Chunk.single(channel) ++ Chunk.fromIterable(channels), onSubscribe, @@ -34,9 +37,12 @@ trait Subscription extends SubscribeEnvironment { ) } - final def pSubscribe(pattern: String, patterns: String*) = - new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { - def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = + final def pSubscribe( + pattern: String, + patterns: String* + ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = + new ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] { + def returning[R: Schema]: Stream[RedisError, (String, R)] = RedisSubscriptionCommand(executor).pSubscribe( Chunk.single(pattern) ++ Chunk.fromIterable(patterns), emptyCallback, @@ -47,9 +53,12 @@ trait Subscription extends SubscribeEnvironment { final def pSubscribeWithCallback( pattern: String, patterns: String* - )(onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback) = - new ResultBuilder1[({ type lambda[x] = Stream[RedisError, (String, x)] })#lambda] { - def returning[R: Schema]: IO[RedisError, Stream[RedisError, (String, R)]] = + )( + onSubscribe: PubSubCallback, + onUnsubscribe: PubSubCallback + ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = + new ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] { + def returning[R: Schema]: Stream[RedisError, (String, R)] = RedisSubscriptionCommand(executor).pSubscribe( Chunk.single(pattern) ++ Chunk.fromIterable(patterns), onSubscribe, diff --git a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala index 0c4629539..2a5d0d39e 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala @@ -1,13 +1,13 @@ -package zio.redis +package zio.redis.internal import zio.redis.Input._ -import zio.redis.Output.{ArbitraryOutput, PushProtocolOutput} +import zio.redis.Output.ArbitraryOutput +import zio.redis._ import zio.redis.api.Subscription import zio.redis.options.PubSub.PubSubCallback import zio.schema.codec.BinaryCodec -import zio.stream.ZStream.RefineToOrDieOps import zio.stream._ -import zio.{Chunk, IO, Promise, Ref, ZIO} +import zio.{Chunk, IO, ZIO} private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionExecutor) extends { import zio.redis.options.PubSub.PushProtocol._ @@ -16,75 +16,53 @@ private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionE channels: Chunk[String], onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback - )(implicit codec: BinaryCodec[String]): IO[RedisError, Stream[RedisError, (String, A)]] = - for { - unsubscribedRef <- Ref.make(channels.map(_ -> false).toMap) - promise <- Promise.make[RedisError, Unit] - channelSet = channels.toSet - stream <- executor - .execute(makeCommand(Subscription.Subscribe, channels)) - .map( - _.mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))).mapZIO { - case Subscribe(channel, numOfSubscription) if channelSet contains channel => - onSubscribe(channel, numOfSubscription).as(None) - case Message(channel, message) if channelSet contains channel => - ZIO - .attempt(ArbitraryOutput[A]().unsafeDecode(message)) - .map(msg => Some((channel, msg))) - case Unsubscribe(channel, numOfSubscription) if channelSet contains channel => - for { - _ <- onUnsubscribe(channel, numOfSubscription) - _ <- unsubscribedRef.update(_.updatedWith(channel)(_ => Some(true))) - _ <- promise.succeed(()).whenZIO(unsubscribedRef.get.map(_.values.forall(identity))) - } yield None - case _ => ZIO.none - }.collectSome - .refineToOrDie[RedisError] - .interruptWhen(promise) - ) - } yield stream + )(implicit codec: BinaryCodec[String]): Stream[RedisError, (String, A)] = + executor + .execute(makeCommand(Subscription.Subscribe, channels)) + .mapZIO { + case Subscribe(channel, numOfSubscription) => + onSubscribe(channel, numOfSubscription).as(None) + case Message(channel, message) => + ZIO + .attempt(ArbitraryOutput[A]().unsafeDecode(message)) + .map(msg => Some((channel, msg))) + case Unsubscribe(channel, numOfSubscription) => + onUnsubscribe(channel, numOfSubscription).as(None) + case _ => ZIO.none + } + .collectSome + .refineToOrDie[RedisError] def pSubscribe[A: BinaryCodec]( patterns: Chunk[String], onSubscribe: PubSubCallback, onUnsubscribe: PubSubCallback - )(implicit stringCodec: BinaryCodec[String]): IO[RedisError, Stream[RedisError, (String, A)]] = - for { - unsubscribedRef <- Ref.make(patterns.map(_ -> false).toMap) - promise <- Promise.make[RedisError, Unit] - patternSet = patterns.toSet - stream <- executor - .execute(makeCommand(Subscription.PSubscribe, patterns)) - .map( - _.mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))).mapZIO { - case PSubscribe(pattern, numOfSubscription) if patternSet contains pattern => - onSubscribe(pattern, numOfSubscription).as(None) - case PMessage(pattern, channel, message) if patternSet contains pattern => - ZIO - .attempt(ArbitraryOutput[A]().unsafeDecode(message)) - .map(msg => Some((channel, msg))) - case PUnsubscribe(pattern, numOfSubscription) if patternSet contains pattern => - for { - _ <- onUnsubscribe(pattern, numOfSubscription) - _ <- unsubscribedRef.update(_.updatedWith(pattern)(_ => Some(true))) - _ <- promise.succeed(()).whenZIO(unsubscribedRef.get.map(_.values.forall(identity))) - } yield None - case _ => ZIO.none - }.collectSome - .refineToOrDie[RedisError] - .interruptWhen(promise) - ) - } yield stream + )(implicit stringCodec: BinaryCodec[String]): Stream[RedisError, (String, A)] = + executor + .execute(makeCommand(Subscription.PSubscribe, patterns)) + .mapZIO { + case PSubscribe(pattern, numOfSubscription) => + onSubscribe(pattern, numOfSubscription).as(None) + case PMessage(_, channel, message) => + ZIO + .attempt(ArbitraryOutput[A]().unsafeDecode(message)) + .map(msg => Some((channel, msg))) + case PUnsubscribe(pattern, numOfSubscription) => + onUnsubscribe(pattern, numOfSubscription).as(None) + case _ => ZIO.none + } + .collectSome + .refineToOrDie[RedisError] def unsubscribe(channels: Chunk[String])(implicit codec: BinaryCodec[String]): IO[RedisError, Unit] = executor .execute(makeCommand(Subscription.Unsubscribe, channels)) - .flatMap(_.runDrain) + .runDrain def pUnsubscribe(patterns: Chunk[String])(implicit codec: BinaryCodec[String]): IO[RedisError, Unit] = executor .execute(makeCommand(Subscription.PUnsubscribe, patterns)) - .flatMap(_.runDrain) + .runDrain private def makeCommand(commandName: String, keys: Chunk[String])(implicit codec: BinaryCodec[String]) = CommandNameInput.encode(commandName) ++ diff --git a/modules/redis/src/main/scala/zio/redis/internal/options/PubSub.scala b/modules/redis/src/main/scala/zio/redis/options/PubSub.scala similarity index 100% rename from modules/redis/src/main/scala/zio/redis/internal/options/PubSub.scala rename to modules/redis/src/main/scala/zio/redis/options/PubSub.scala diff --git a/modules/redis/src/test/scala/zio/redis/ApiSpec.scala b/modules/redis/src/test/scala/zio/redis/ApiSpec.scala index 48839868f..d607633db 100644 --- a/modules/redis/src/test/scala/zio/redis/ApiSpec.scala +++ b/modules/redis/src/test/scala/zio/redis/ApiSpec.scala @@ -20,7 +20,7 @@ object ApiSpec with PubSubSpec { def spec: Spec[TestEnvironment, Any] = - suite("Redis commands")(singleNodeSuite) @@ sequential @@ withLiveEnvironment + suite("Redis commands")(singleNodeSuite, clusterSuite) @@ sequential @@ withLiveEnvironment private val singleNodeSuite = suite("Single node executor")( @@ -44,25 +44,25 @@ object ApiSpec ZLayer.succeed(ProtobufCodecSupplier) ) -// private val clusterSuite = -// suite("Cluster executor")( -// connectionSuite, -// keysSuite, -// listSuite, -// stringsSuite, -// hashSuite, -// setsSuite, -// sortedSetsSuite, -// hyperLogLogSuite, -// geoSuite, -// streamsSuite, -// scriptingSpec, -// clusterSpec -// ).provideShared( -// ClusterExecutor.layer, -// Redis.layer, -// ZLayer.succeed(ProtobufCodecSupplier), -// ZLayer.succeed(RedisClusterConfig(Chunk(RedisUri("localhost", 5000)))) -// ).filterNotTags(_.contains(BaseSpec.ClusterExecutorUnsupported)) -// .getOrElse(Spec.empty) + private val clusterSuite = + suite("Cluster executor")( + connectionSuite, + keysSuite, + listSuite, + stringsSuite, + hashSuite, + setsSuite, + sortedSetsSuite, + hyperLogLogSuite, + geoSuite, + streamsSuite, + scriptingSpec, + clusterSpec + ).provideShared( + ClusterExecutor.layer, + Redis.layer, + ZLayer.succeed(ProtobufCodecSupplier), + ZLayer.succeed(RedisClusterConfig(Chunk(RedisUri("localhost", 5000)))) + ).filterNotTags(_.contains(BaseSpec.ClusterExecutorUnsupported)) + .getOrElse(Spec.empty) } diff --git a/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala b/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala index ca313a6fd..ba706c745 100644 --- a/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala +++ b/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala @@ -22,9 +22,12 @@ trait PubSubSpec extends BaseSpec { (key: String, _: Long) => promise.succeed(key).unit, (_, _) => ZIO.unit ) - stream <- resBuilder.returning[String] - _ <- stream.interruptWhen(promise).runDrain.fork - res <- promise.await + stream = resBuilder.returning[String] + _ <- stream + .interruptWhen(promise) + .runDrain + .fork + res <- promise.await } yield assertTrue(res == channel) }, test("message response") { @@ -34,7 +37,7 @@ trait PubSubSpec extends BaseSpec { channel <- generateRandomString() message = "bar" promise <- Promise.make[RedisError, String] - stream <- subscription.subscribe(channel).returning[String] + stream = subscription.subscribe(channel).returning[String] fiber <- stream.interruptWhen(promise).runHead.fork _ <- redis .pubSubChannels(channel) @@ -53,24 +56,22 @@ trait PubSubSpec extends BaseSpec { channel2 <- generateRandomString(prefix) pattern = prefix + '*' message <- generateRandomString(5) - stream1 <- subscription - .subscribe(channel1) - .returning[String] - .flatMap(_.runDrain) - .fork - stream2 <- subscription - .subscribe(channel2) - .returning[String] - .flatMap(_.runDrain) - .fork + stream1 = subscription + .subscribe(channel1) + .returning[String] + stream2 = subscription + .subscribe(channel2) + .returning[String] + fiber1 <- stream1.runDrain.fork + fiber2 <- stream2.runDrain.fork _ <- redis .pubSubChannels(pattern) .repeatUntil(channels => channels.size >= 2) ch1SubsCount <- redis.publish(channel1, message).replicateZIO(numOfPublish).map(_.head) ch2SubsCount <- redis.publish(channel2, message).replicateZIO(numOfPublish).map(_.head) _ <- subscription.unsubscribe() - _ <- stream1.join - _ <- stream2.join + _ <- fiber1.join + _ <- fiber2.join } yield assertTrue(ch1SubsCount == 1L) && assertTrue(ch2SubsCount == 1L) }, test("psubscribe response") { @@ -84,7 +85,8 @@ trait PubSubSpec extends BaseSpec { (_, _) => ZIO.unit ) .returning[String] - .flatMap(_.interruptWhen(promise).runHead) + .interruptWhen(promise) + .runHead .fork res <- promise.await } yield assertTrue(res == pattern) @@ -100,7 +102,7 @@ trait PubSubSpec extends BaseSpec { stream <- subscription .pSubscribe(pattern) .returning[String] - .flatMap(_.runHead) + .runHead .fork _ <- redis.pubSubNumPat.repeatUntil(_ > 0) _ <- redis.publish(channel, message) @@ -118,9 +120,10 @@ trait PubSubSpec extends BaseSpec { stream <- subscription .subscribe(channel) .returning[Long] - .flatMap(_.runFoldWhile(0L)(_ < 10L) { case (sum, (_, message)) => + .runFoldWhile(0L)(_ < 10L) { case (sum, (_, message)) => sum + message - }.fork) + } + .fork _ <- redis.pubSubChannels(channel).repeatUntil(_ contains channel) _ <- ZIO.replicateZIO(10)(redis.publish(channel, message)) res <- stream.join @@ -141,7 +144,7 @@ trait PubSubSpec extends BaseSpec { _ <- subscription .subscribeWithCallback(channel)((_, _) => ZIO.unit, (_, _) => promise.succeed(()).unit) .returning[String] - .flatMap(_.runCollect) + .runCollect .fork _ <- redis .pubSubChannels(pattern) @@ -162,7 +165,7 @@ trait PubSubSpec extends BaseSpec { (key, _) => promise.succeed(key).unit ) .returning[Unit] - .flatMap(_.runDrain) + .runDrain .fork _ <- subscription.unsubscribe(channel) res <- promise.await @@ -179,7 +182,7 @@ trait PubSubSpec extends BaseSpec { (key, _) => promise.succeed(key).unit ) .returning[Unit] - .flatMap(_.runDrain) + .runDrain .fork _ <- subscription.pUnsubscribe(pattern) res <- promise.await @@ -202,7 +205,7 @@ trait PubSubSpec extends BaseSpec { (_, _) => promise1.succeed(()).unit ) .returning[String] - .flatMap(_.runCollect) + .runCollect .fork _ <- subscription @@ -211,7 +214,7 @@ trait PubSubSpec extends BaseSpec { (_, _) => promise2.succeed(()).unit ) .returning[String] - .flatMap(_.runCollect) + .runCollect .fork _ <- redis .pubSubChannels(pattern) From 149ae4bc800dad23d9379c9768f9be967ff6084f Mon Sep 17 00:00:00 2001 From: 0pg Date: Thu, 27 Apr 2023 17:39:18 +0900 Subject: [PATCH 30/51] Remove unused code --- modules/redis/src/main/scala/zio/redis/RedisError.scala | 1 - 1 file changed, 1 deletion(-) diff --git a/modules/redis/src/main/scala/zio/redis/RedisError.scala b/modules/redis/src/main/scala/zio/redis/RedisError.scala index c49ae767b..00faf0e63 100644 --- a/modules/redis/src/main/scala/zio/redis/RedisError.scala +++ b/modules/redis/src/main/scala/zio/redis/RedisError.scala @@ -46,7 +46,6 @@ object RedisError { object Moved { def apply(slotAndAddress: (Slot, RedisUri)): Moved = Moved(slotAndAddress._1, slotAndAddress._2) } - final case class NoPubSubStream(key: String) extends RedisError final case class IOError(exception: IOException) extends RedisError final case class CommandNameNotFound(message: String) extends RedisError sealed trait PubSubError extends RedisError From 67652e6ab64a2bab0d775bc63bee7b6686a5b62c Mon Sep 17 00:00:00 2001 From: 0pg Date: Thu, 27 Apr 2023 18:09:50 +0900 Subject: [PATCH 31/51] Refine subscription api --- .../scala/zio/redis/api/Subscription.scala | 56 +++++++++++++------ .../scala/zio/redis/internal/PubSubSpec.scala | 6 +- 2 files changed, 43 insertions(+), 19 deletions(-) diff --git a/modules/redis/src/main/scala/zio/redis/api/Subscription.scala b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala index 32a422467..5f4d5aed3 100644 --- a/modules/redis/src/main/scala/zio/redis/api/Subscription.scala +++ b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala @@ -11,18 +11,29 @@ import zio.{Chunk, IO, ZIO} trait Subscription extends SubscribeEnvironment { import Subscription._ + final def subscribe(channel: String): ResultStreamBuilder1[Id] = + subscribeWithCallback(channel)(emptyCallback, emptyCallback) + + final def subscribeWithCallback(channel: String)( + onSubscribe: PubSubCallback, + onUnsubscribe: PubSubCallback + ): ResultStreamBuilder1[Id] = + new ResultStreamBuilder1[Id] { + def returning[R: Schema]: Stream[RedisError, R] = + RedisSubscriptionCommand(executor) + .subscribe( + Chunk.single(channel), + onSubscribe, + onUnsubscribe + ) + .map(_._2) + } + final def subscribe( channel: String, channels: String* ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = - new ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] { - def returning[R: Schema]: Stream[RedisError, (String, R)] = - RedisSubscriptionCommand(executor).subscribe( - Chunk.single(channel) ++ Chunk.fromIterable(channels), - emptyCallback, - emptyCallback - ) - } + subscribeWithCallback(channel, channels: _*)(emptyCallback, emptyCallback) final def subscribeWithCallback(channel: String, channels: String*)( onSubscribe: PubSubCallback, @@ -37,18 +48,31 @@ trait Subscription extends SubscribeEnvironment { ) } + final def pSubscribe(pattern: String): ResultStreamBuilder1[Id] = + pSubscribeWithCallback(pattern)(emptyCallback, emptyCallback) + + final def pSubscribeWithCallback( + pattern: String + )( + onSubscribe: PubSubCallback, + onUnsubscribe: PubSubCallback + ): ResultStreamBuilder1[Id] = + new ResultStreamBuilder1[Id] { + def returning[R: Schema]: Stream[RedisError, R] = + RedisSubscriptionCommand(executor) + .pSubscribe( + Chunk.single(pattern), + onSubscribe, + onUnsubscribe + ) + .map(_._2) + } + final def pSubscribe( pattern: String, patterns: String* ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = - new ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] { - def returning[R: Schema]: Stream[RedisError, (String, R)] = - RedisSubscriptionCommand(executor).pSubscribe( - Chunk.single(pattern) ++ Chunk.fromIterable(patterns), - emptyCallback, - emptyCallback - ) - } + pSubscribeWithCallback(pattern, patterns: _*)(emptyCallback, emptyCallback) final def pSubscribeWithCallback( pattern: String, diff --git a/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala b/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala index ba706c745..933985866 100644 --- a/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala +++ b/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala @@ -44,7 +44,7 @@ trait PubSubSpec extends BaseSpec { .repeatUntil(_ contains channel) _ <- redis.publish(channel, message) res <- fiber.join - } yield assertTrue(res.get == (channel, message)) + } yield assertTrue(res.get == message) }, test("multiple subscribe") { val numOfPublish = 20 @@ -107,7 +107,7 @@ trait PubSubSpec extends BaseSpec { _ <- redis.pubSubNumPat.repeatUntil(_ > 0) _ <- redis.publish(channel, message) res <- stream.join - } yield assertTrue(res.get == (channel, message)) + } yield assertTrue(res.get == message) } ), suite("publish")(test("publish long type message") { @@ -120,7 +120,7 @@ trait PubSubSpec extends BaseSpec { stream <- subscription .subscribe(channel) .returning[Long] - .runFoldWhile(0L)(_ < 10L) { case (sum, (_, message)) => + .runFoldWhile(0L)(_ < 10L) { case (sum, message) => sum + message } .fork From 8606409fed27a3ba7dc9a3d0621e24c293ffc72a Mon Sep 17 00:00:00 2001 From: 0pg Date: Wed, 3 May 2023 01:10:12 +0900 Subject: [PATCH 32/51] Rename fields --- .../src/main/scala/zio/redis/Output.scala | 6 ++--- .../src/main/scala/zio/redis/RedisError.scala | 6 ++--- .../internal/RedisSubscriptionCommand.scala | 16 ++++++------- .../main/scala/zio/redis/options/PubSub.scala | 8 +++---- .../src/test/scala/zio/redis/OutputSpec.scala | 24 +++++++++---------- .../main/scala/zio/redis/RedisExecutor.scala | 0 6 files changed, 30 insertions(+), 30 deletions(-) delete mode 100644 redis/src/main/scala/zio/redis/RedisExecutor.scala diff --git a/modules/redis/src/main/scala/zio/redis/Output.scala b/modules/redis/src/main/scala/zio/redis/Output.scala index 1097084df..eaad45320 100644 --- a/modules/redis/src/main/scala/zio/redis/Output.scala +++ b/modules/redis/src/main/scala/zio/redis/Output.scala @@ -676,9 +676,9 @@ object Output { 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) + 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") } diff --git a/modules/redis/src/main/scala/zio/redis/RedisError.scala b/modules/redis/src/main/scala/zio/redis/RedisError.scala index 00faf0e63..f3dda11f6 100644 --- a/modules/redis/src/main/scala/zio/redis/RedisError.scala +++ b/modules/redis/src/main/scala/zio/redis/RedisError.scala @@ -46,8 +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 CommandNameNotFound(message: String) extends RedisError - sealed trait PubSubError 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 } diff --git a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala index 2a5d0d39e..caf3af882 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala @@ -20,14 +20,14 @@ private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionE executor .execute(makeCommand(Subscription.Subscribe, channels)) .mapZIO { - case Subscribe(channel, numOfSubscription) => - onSubscribe(channel, numOfSubscription).as(None) + case Subscribe(channel, numOfSubs) => + onSubscribe(channel, numOfSubs).as(None) case Message(channel, message) => ZIO .attempt(ArbitraryOutput[A]().unsafeDecode(message)) .map(msg => Some((channel, msg))) - case Unsubscribe(channel, numOfSubscription) => - onUnsubscribe(channel, numOfSubscription).as(None) + case Unsubscribe(channel, numOfSubs) => + onUnsubscribe(channel, numOfSubs).as(None) case _ => ZIO.none } .collectSome @@ -41,14 +41,14 @@ private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionE executor .execute(makeCommand(Subscription.PSubscribe, patterns)) .mapZIO { - case PSubscribe(pattern, numOfSubscription) => - onSubscribe(pattern, numOfSubscription).as(None) + case PSubscribe(pattern, numOfSubs) => + onSubscribe(pattern, numOfSubs).as(None) case PMessage(_, channel, message) => ZIO .attempt(ArbitraryOutput[A]().unsafeDecode(message)) .map(msg => Some((channel, msg))) - case PUnsubscribe(pattern, numOfSubscription) => - onUnsubscribe(pattern, numOfSubscription).as(None) + case PUnsubscribe(pattern, numOfSubs) => + onUnsubscribe(pattern, numOfSubs).as(None) case _ => ZIO.none } .collectSome diff --git a/modules/redis/src/main/scala/zio/redis/options/PubSub.scala b/modules/redis/src/main/scala/zio/redis/options/PubSub.scala index 62e4e91ce..7ea3d8ec8 100644 --- a/modules/redis/src/main/scala/zio/redis/options/PubSub.scala +++ b/modules/redis/src/main/scala/zio/redis/options/PubSub.scala @@ -8,10 +8,10 @@ object PubSub { private[redis] sealed trait PushProtocol private[redis] object PushProtocol { - case class Subscribe(channel: String, numOfSubscription: Long) extends PushProtocol - case class PSubscribe(pattern: String, numOfSubscription: Long) extends PushProtocol - case class Unsubscribe(channel: String, numOfSubscription: Long) extends PushProtocol - case class PUnsubscribe(pattern: String, numOfSubscription: Long) extends PushProtocol + case class Subscribe(channel: String, numOfSubs: Long) extends PushProtocol + case class PSubscribe(pattern: String, numOfSubs: Long) extends PushProtocol + case class Unsubscribe(channel: String, numOfSubs: Long) extends PushProtocol + case class PUnsubscribe(pattern: String, numOfSubs: Long) extends PushProtocol case class Message(channel: String, message: RespValue) extends PushProtocol case class PMessage(pattern: String, channel: String, message: RespValue) extends PushProtocol } diff --git a/modules/redis/src/test/scala/zio/redis/OutputSpec.scala b/modules/redis/src/test/scala/zio/redis/OutputSpec.scala index 6483dac3c..076df831b 100644 --- a/modules/redis/src/test/scala/zio/redis/OutputSpec.scala +++ b/modules/redis/src/test/scala/zio/redis/OutputSpec.scala @@ -891,56 +891,56 @@ object OutputSpec extends BaseSpec { suite("PushProtocol")( test("subscribe") { val channel = "foo" - val numOfSubscription = 1L + val numOfSubs = 1L val input = RespValue.array( RespValue.bulkString("subscribe"), RespValue.bulkString(channel), - RespValue.Integer(numOfSubscription) + RespValue.Integer(numOfSubs) ) - val expected = PushProtocol.Subscribe(channel, numOfSubscription) + val expected = PushProtocol.Subscribe(channel, numOfSubs) assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( equalTo(expected) ) }, test("psubscribe") { val pattern = "f*" - val numOfSubscription = 1L + val numOfSubs = 1L val input = RespValue.array( RespValue.bulkString("psubscribe"), RespValue.bulkString(pattern), - RespValue.Integer(numOfSubscription) + RespValue.Integer(numOfSubs) ) - val expected = PushProtocol.PSubscribe(pattern, numOfSubscription) + val expected = PushProtocol.PSubscribe(pattern, numOfSubs) assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( equalTo(expected) ) }, test("unsubscribe") { val channel = "foo" - val numOfSubscription = 1L + val numOfSubs = 1L val input = RespValue.array( RespValue.bulkString("unsubscribe"), RespValue.bulkString(channel), - RespValue.Integer(numOfSubscription) + RespValue.Integer(numOfSubs) ) - val expected = PushProtocol.Unsubscribe(channel, numOfSubscription) + val expected = PushProtocol.Unsubscribe(channel, numOfSubs) assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( equalTo(expected) ) }, test("punsubscribe") { val pattern = "f*" - val numOfSubscription = 1L + val numOfSubs = 1L val input = RespValue.array( RespValue.bulkString("punsubscribe"), RespValue.bulkString(pattern), - RespValue.Integer(numOfSubscription) + RespValue.Integer(numOfSubs) ) - val expected = PushProtocol.PUnsubscribe(pattern, numOfSubscription) + val expected = PushProtocol.PUnsubscribe(pattern, numOfSubs) assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( equalTo(expected) ) diff --git a/redis/src/main/scala/zio/redis/RedisExecutor.scala b/redis/src/main/scala/zio/redis/RedisExecutor.scala deleted file mode 100644 index e69de29bb..000000000 From 6a3f68f1156ec6acb45e72f7a73eee4f88627918 Mon Sep 17 00:00:00 2001 From: opg1 Date: Wed, 3 May 2023 01:11:44 +0900 Subject: [PATCH 33/51] Update modules/redis/src/main/scala/zio/redis/RedisSubscription.scala MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Aleksandar Novaković --- .../redis/src/main/scala/zio/redis/RedisSubscription.scala | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala b/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala index 6c5efaced..91b86550f 100644 --- a/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala +++ b/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala @@ -12,12 +12,7 @@ object RedisSubscription { SubscriptionExecutor.layer >>> makeLayer private def makeLayer: URLayer[CodecSupplier & SubscriptionExecutor, RedisSubscription] = - ZLayer { - for { - codecSupplier <- ZIO.service[CodecSupplier] - executor <- ZIO.service[SubscriptionExecutor] - } yield Live(codecSupplier, executor) - } + ZLayer.fromFunction(Live.apply _) private final case class Live(codecSupplier: CodecSupplier, executor: SubscriptionExecutor) extends RedisSubscription } From 1a0862f7b30f2dc8a611217c98e8f85814276f2b Mon Sep 17 00:00:00 2001 From: 0pg Date: Fri, 5 May 2023 21:28:11 +0900 Subject: [PATCH 34/51] Add newline --- modules/redis/src/main/scala/zio/redis/options/PubSub.scala | 1 + 1 file changed, 1 insertion(+) diff --git a/modules/redis/src/main/scala/zio/redis/options/PubSub.scala b/modules/redis/src/main/scala/zio/redis/options/PubSub.scala index 7ea3d8ec8..837c428a4 100644 --- a/modules/redis/src/main/scala/zio/redis/options/PubSub.scala +++ b/modules/redis/src/main/scala/zio/redis/options/PubSub.scala @@ -2,6 +2,7 @@ package zio.redis.options import zio.UIO import zio.redis.internal.RespValue + object PubSub { type PubSubCallback = (String, Long) => UIO[Unit] From 92457c4f59bfd12b1f51d31fab3bad38aab34fba Mon Sep 17 00:00:00 2001 From: 0pg Date: Sun, 7 May 2023 22:38:21 +0900 Subject: [PATCH 35/51] Fix broken compile --- .../scala/zio/redis/RedisSubscription.scala | 4 +- .../SingleNodeSubscriptionExecutor.scala | 40 +++++++++++-------- .../internal/RedisSubscriptionCommand.scala | 2 +- 3 files changed, 27 insertions(+), 19 deletions(-) diff --git a/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala b/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala index 91b86550f..2b8a11593 100644 --- a/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala +++ b/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala @@ -1,6 +1,6 @@ package zio.redis -import zio.{&, URLayer, ZIO, ZLayer} +import zio._ trait RedisSubscription extends api.Subscription @@ -12,7 +12,7 @@ object RedisSubscription { SubscriptionExecutor.layer >>> makeLayer private def makeLayer: URLayer[CodecSupplier & SubscriptionExecutor, RedisSubscription] = - ZLayer.fromFunction(Live.apply _) + ZLayer.fromFunction(Live.apply _) private final case class Live(codecSupplier: CodecSupplier, executor: SubscriptionExecutor) extends RedisSubscription } diff --git a/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala index 4f2dcac7b..4231c8d6e 100644 --- a/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala @@ -4,12 +4,12 @@ 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.{RedisConnection, RespCommand, RespCommandArgument, RespValue, logScopeFinalizer} +import zio.redis.internal._ import zio.redis.options.PubSub.PushProtocol import zio.stream._ -import zio.{Chunk, ChunkBuilder, IO, Promise, Queue, Ref, Schedule, ZIO} +import zio.{Chunk, ChunkBuilder, IO, Promise, Queue, Ref, Schedule, Scope, URIO, ZIO} -final class SingleNodeSubscriptionExecutor( +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], @@ -41,16 +41,17 @@ final class SingleNodeSubscriptionExecutor( 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( - _.updatedWith(key)(_.map(_ appended queue).orElse(Some(Chunk.single(queue)))) - ) - ) - .map(key -> _) - ) + 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( @@ -61,7 +62,14 @@ final class SingleNodeSubscriptionExecutor( streams = queues.map { case (key, queue) => ZStream .fromQueueWithShutdown(queue) - .ensuring(subscriptionRef.update(_.updatedWith(key)(_.map(_.filterNot(_ == 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 _) @@ -173,7 +181,7 @@ final class SingleNodeSubscriptionExecutor( .tapError(e => ZIO.logError(s"Executor exiting: $e")) } -object SingleNodeSubscriptionExecutor { +private[redis] object SingleNodeSubscriptionExecutor { private final case class Request( command: Chunk[RespValue.BulkString], promise: Promise[RedisError, Unit] @@ -183,7 +191,7 @@ object SingleNodeSubscriptionExecutor { private final val RequestQueueSize = 16 - def create(conn: RedisConnection) = + def create(conn: RedisConnection): URIO[Scope, SubscriptionExecutor] = for { reqQueue <- Queue.bounded[Request](RequestQueueSize) channelRef <- Ref.make(Map.empty[String, Chunk[Queue[Take[RedisError, PushProtocol]]]]) diff --git a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala index caf3af882..af3306c11 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala @@ -9,7 +9,7 @@ import zio.schema.codec.BinaryCodec import zio.stream._ import zio.{Chunk, IO, ZIO} -private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionExecutor) extends { +private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionExecutor) { import zio.redis.options.PubSub.PushProtocol._ def subscribe[A: BinaryCodec]( From bf52e2625bb0b639b27dacd8f8e9e057017a9415 Mon Sep 17 00:00:00 2001 From: 0pg Date: Sun, 7 May 2023 23:25:44 +0900 Subject: [PATCH 36/51] Formatting --- modules/redis/src/test/scala/zio/redis/OutputSpec.scala | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/modules/redis/src/test/scala/zio/redis/OutputSpec.scala b/modules/redis/src/test/scala/zio/redis/OutputSpec.scala index 076df831b..3c9407346 100644 --- a/modules/redis/src/test/scala/zio/redis/OutputSpec.scala +++ b/modules/redis/src/test/scala/zio/redis/OutputSpec.scala @@ -890,7 +890,7 @@ object OutputSpec extends BaseSpec { ), suite("PushProtocol")( test("subscribe") { - val channel = "foo" + val channel = "foo" val numOfSubs = 1L val input = RespValue.array( @@ -904,7 +904,7 @@ object OutputSpec extends BaseSpec { ) }, test("psubscribe") { - val pattern = "f*" + val pattern = "f*" val numOfSubs = 1L val input = RespValue.array( @@ -918,7 +918,7 @@ object OutputSpec extends BaseSpec { ) }, test("unsubscribe") { - val channel = "foo" + val channel = "foo" val numOfSubs = 1L val input = RespValue.array( @@ -932,7 +932,7 @@ object OutputSpec extends BaseSpec { ) }, test("punsubscribe") { - val pattern = "f*" + val pattern = "f*" val numOfSubs = 1L val input = RespValue.array( From c0272953e13191b713713e61e4cd47e8aec3e612 Mon Sep 17 00:00:00 2001 From: 0pg Date: Mon, 8 May 2023 00:10:00 +0900 Subject: [PATCH 37/51] Fix lint --- .../main/scala/zio/redis/RedisSubscription.scala | 16 ++++++++++++++++ .../redis/SingleNodeSubscriptionExecutor.scala | 16 ++++++++++++++++ .../scala/zio/redis/SubscriptionExecutor.scala | 16 ++++++++++++++++ .../main/scala/zio/redis/api/Publishing.scala | 16 ++++++++++++++++ .../main/scala/zio/redis/api/Subscription.scala | 16 ++++++++++++++++ .../internal/RedisSubscriptionCommand.scala | 16 ++++++++++++++++ .../redis/internal/SubscribeEnvironment.scala | 16 ++++++++++++++++ .../main/scala/zio/redis/options/PubSub.scala | 16 ++++++++++++++++ 8 files changed, 128 insertions(+) diff --git a/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala b/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala index 2b8a11593..1efff1008 100644 --- a/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala +++ b/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala @@ -1,3 +1,19 @@ +/* + * 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._ diff --git a/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala index 4231c8d6e..4355f1240 100644 --- a/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala @@ -1,3 +1,19 @@ +/* + * 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} diff --git a/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala index 91cd06a78..3425c4dc1 100644 --- a/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala @@ -1,3 +1,19 @@ +/* + * 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} diff --git a/modules/redis/src/main/scala/zio/redis/api/Publishing.scala b/modules/redis/src/main/scala/zio/redis/api/Publishing.scala index 79cf7a4c1..ac2110b2f 100644 --- a/modules/redis/src/main/scala/zio/redis/api/Publishing.scala +++ b/modules/redis/src/main/scala/zio/redis/api/Publishing.scala @@ -1,3 +1,19 @@ +/* + * 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.api import zio.redis.Input._ diff --git a/modules/redis/src/main/scala/zio/redis/api/Subscription.scala b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala index 5f4d5aed3..04e16d9ae 100644 --- a/modules/redis/src/main/scala/zio/redis/api/Subscription.scala +++ b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala @@ -1,3 +1,19 @@ +/* + * 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.api import zio.redis.ResultBuilder.ResultStreamBuilder1 diff --git a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala index af3306c11..45a673823 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala @@ -1,3 +1,19 @@ +/* + * 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.internal import zio.redis.Input._ diff --git a/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala b/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala index 38dbfc8f8..bdb7adcc0 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala @@ -1,3 +1,19 @@ +/* + * 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.internal import zio.redis.{CodecSupplier, SubscriptionExecutor} diff --git a/modules/redis/src/main/scala/zio/redis/options/PubSub.scala b/modules/redis/src/main/scala/zio/redis/options/PubSub.scala index 837c428a4..d6f212588 100644 --- a/modules/redis/src/main/scala/zio/redis/options/PubSub.scala +++ b/modules/redis/src/main/scala/zio/redis/options/PubSub.scala @@ -1,3 +1,19 @@ +/* + * 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.options import zio.UIO From a1546838ca3ed7a8b88be844069c489ec5943a19 Mon Sep 17 00:00:00 2001 From: 0pg Date: Thu, 11 May 2023 02:14:08 +0900 Subject: [PATCH 38/51] Reduce duplications and refactor package layout --- .../src/main/scala/zio/redis/Output.scala | 20 +-- .../src/main/scala/zio/redis/RedisError.scala | 9 +- .../scala/zio/redis/RedisSubscription.scala | 1 + .../main/scala/zio/redis/api/Publishing.scala | 1 - .../scala/zio/redis/api/Subscription.scala | 35 ++-- .../scala/zio/redis/internal/PubSub.scala | 38 +++++ .../internal/RedisSubscriptionCommand.scala | 71 ++++----- .../SingleNodeSubscriptionExecutor.scala | 150 +++++++++--------- .../redis/internal/SubscribeEnvironment.scala | 2 +- .../{ => internal}/SubscriptionExecutor.scala | 8 +- .../main/scala/zio/redis/options/PubSub.scala | 37 ----- .../scala/zio/redis/options/Publishing.scala | 21 +++ .../src/main/scala/zio/redis/package.scala | 3 +- .../src/test/scala/zio/redis/OutputSpec.scala | 30 ++-- .../scala/zio/redis/internal/PubSubSpec.scala | 27 ++-- 15 files changed, 238 insertions(+), 215 deletions(-) create mode 100644 modules/redis/src/main/scala/zio/redis/internal/PubSub.scala rename modules/redis/src/main/scala/zio/redis/{ => internal}/SingleNodeSubscriptionExecutor.scala (53%) rename modules/redis/src/main/scala/zio/redis/{ => internal}/SubscriptionExecutor.scala (90%) delete mode 100644 modules/redis/src/main/scala/zio/redis/options/PubSub.scala create mode 100644 modules/redis/src/main/scala/zio/redis/options/Publishing.scala diff --git a/modules/redis/src/main/scala/zio/redis/Output.scala b/modules/redis/src/main/scala/zio/redis/Output.scala index eaad45320..ffb121f28 100644 --- a/modules/redis/src/main/scala/zio/redis/Output.scala +++ b/modules/redis/src/main/scala/zio/redis/Output.scala @@ -17,9 +17,9 @@ package zio.redis import zio._ +import zio.redis.internal.PubSub.{PushMessage, SubscriptionKey} 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 @@ -640,8 +640,8 @@ object Output { } } - case object PushProtocolOutput extends Output[PushProtocol] { - protected def tryDecode(respValue: RespValue): PushProtocol = + private[redis] case object PushMessageOutput extends Output[PushMessage] { + protected def tryDecode(respValue: RespValue): PushMessage = respValue match { case RespValue.NullArray => throw ProtocolError(s"Array must not be empty") case RespValue.Array(values) => @@ -650,21 +650,23 @@ object Output { name match { case "subscribe" => val num = LongOutput.unsafeDecode(values(2)) - PushProtocol.Subscribe(key, num) + PushMessage.Subscribed(SubscriptionKey.Channel(key), num) case "psubscribe" => val num = LongOutput.unsafeDecode(values(2)) - PushProtocol.PSubscribe(key, num) + PushMessage.Subscribed(SubscriptionKey.Pattern(key), num) case "unsubscribe" => val num = LongOutput.unsafeDecode(values(2)) - PushProtocol.Unsubscribe(key, num) + PushMessage.Unsubscribed(SubscriptionKey.Channel(key), num) case "punsubscribe" => val num = LongOutput.unsafeDecode(values(2)) - PushProtocol.PUnsubscribe(key, num) + PushMessage.Unsubscribed(SubscriptionKey.Pattern(key), num) case "message" => - PushProtocol.Message(key, values(2)) + val message = values(2) + PushMessage.Message(SubscriptionKey.Channel(key), key, message) case "pmessage" => val channel = MultiStringOutput.unsafeDecode(values(2)) - PushProtocol.PMessage(key, channel, values(3)) + val message = values(3) + PushMessage.Message(SubscriptionKey.Pattern(key), channel, message) case other => throw ProtocolError(s"$other isn't a pushed message") } case other => throw ProtocolError(s"$other isn't an array") diff --git a/modules/redis/src/main/scala/zio/redis/RedisError.scala b/modules/redis/src/main/scala/zio/redis/RedisError.scala index f3dda11f6..db0ed056f 100644 --- a/modules/redis/src/main/scala/zio/redis/RedisError.scala +++ b/modules/redis/src/main/scala/zio/redis/RedisError.scala @@ -16,6 +16,7 @@ package zio.redis +import zio.redis.internal.RespCommand import zio.redis.options.Cluster.Slot import java.io.IOException @@ -46,8 +47,12 @@ 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 CommandNameNotFound(message: String) extends RedisError + final case class IOError(exception: IOException) extends RedisError + final case class CommandNameNotFound(message: String) extends RedisError + object CommandNameNotFound { + def apply(command: RespCommand): CommandNameNotFound = CommandNameNotFound(command.args.toString()) + } + sealed trait PubSubError extends RedisError final case class InvalidPubSubCommand(command: String) extends PubSubError } diff --git a/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala b/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala index 1efff1008..cb8b95eb5 100644 --- a/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala +++ b/modules/redis/src/main/scala/zio/redis/RedisSubscription.scala @@ -17,6 +17,7 @@ package zio.redis import zio._ +import zio.redis.internal.SubscriptionExecutor trait RedisSubscription extends api.Subscription diff --git a/modules/redis/src/main/scala/zio/redis/api/Publishing.scala b/modules/redis/src/main/scala/zio/redis/api/Publishing.scala index ac2110b2f..f82d26165 100644 --- a/modules/redis/src/main/scala/zio/redis/api/Publishing.scala +++ b/modules/redis/src/main/scala/zio/redis/api/Publishing.scala @@ -20,7 +20,6 @@ import zio.redis.Input._ import zio.redis.Output._ import zio.redis._ import zio.redis.internal.{RedisCommand, RedisEnvironment} -import zio.redis.options.PubSub.NumberOfSubscribers import zio.schema.Schema import zio.{Chunk, IO} diff --git a/modules/redis/src/main/scala/zio/redis/api/Subscription.scala b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala index 04e16d9ae..05a732a0e 100644 --- a/modules/redis/src/main/scala/zio/redis/api/Subscription.scala +++ b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala @@ -18,21 +18,20 @@ package zio.redis.api import zio.redis.ResultBuilder.ResultStreamBuilder1 import zio.redis._ +import zio.redis.api.Subscription.PubSubCallback import zio.redis.internal._ -import zio.redis.options.PubSub.PubSubCallback import zio.schema.Schema import zio.stream.Stream -import zio.{Chunk, IO, ZIO} +import zio.{Chunk, IO, UIO} trait Subscription extends SubscribeEnvironment { - import Subscription._ final def subscribe(channel: String): ResultStreamBuilder1[Id] = - subscribeWithCallback(channel)(emptyCallback, emptyCallback) + subscribeWithCallback(channel)(None, None) final def subscribeWithCallback(channel: String)( - onSubscribe: PubSubCallback, - onUnsubscribe: PubSubCallback + onSubscribe: Option[PubSubCallback], + onUnsubscribe: Option[PubSubCallback] ): ResultStreamBuilder1[Id] = new ResultStreamBuilder1[Id] { def returning[R: Schema]: Stream[RedisError, R] = @@ -49,29 +48,29 @@ trait Subscription extends SubscribeEnvironment { channel: String, channels: String* ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = - subscribeWithCallback(channel, channels: _*)(emptyCallback, emptyCallback) + subscribeWithCallback(channel, channels: _*)(None, None) final def subscribeWithCallback(channel: String, channels: String*)( - onSubscribe: PubSubCallback, - onUnsubscribe: PubSubCallback + onSubscribe: Option[PubSubCallback], + onUnsubscribe: Option[PubSubCallback] ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = new ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] { def returning[R: Schema]: Stream[RedisError, (String, R)] = RedisSubscriptionCommand(executor).subscribe( - Chunk.single(channel) ++ Chunk.fromIterable(channels), + Chunk.fromIterable(channel +: channels), onSubscribe, onUnsubscribe ) } final def pSubscribe(pattern: String): ResultStreamBuilder1[Id] = - pSubscribeWithCallback(pattern)(emptyCallback, emptyCallback) + pSubscribeWithCallback(pattern)(None, None) final def pSubscribeWithCallback( pattern: String )( - onSubscribe: PubSubCallback, - onUnsubscribe: PubSubCallback + onSubscribe: Option[PubSubCallback], + onUnsubscribe: Option[PubSubCallback] ): ResultStreamBuilder1[Id] = new ResultStreamBuilder1[Id] { def returning[R: Schema]: Stream[RedisError, R] = @@ -88,19 +87,19 @@ trait Subscription extends SubscribeEnvironment { pattern: String, patterns: String* ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = - pSubscribeWithCallback(pattern, patterns: _*)(emptyCallback, emptyCallback) + pSubscribeWithCallback(pattern, patterns: _*)(None, None) final def pSubscribeWithCallback( pattern: String, patterns: String* )( - onSubscribe: PubSubCallback, - onUnsubscribe: PubSubCallback + onSubscribe: Option[PubSubCallback], + onUnsubscribe: Option[PubSubCallback] ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = new ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] { def returning[R: Schema]: Stream[RedisError, (String, R)] = RedisSubscriptionCommand(executor).pSubscribe( - Chunk.single(pattern) ++ Chunk.fromIterable(patterns), + Chunk.fromIterable(pattern +: patterns), onSubscribe, onUnsubscribe ) @@ -115,7 +114,7 @@ trait Subscription extends SubscribeEnvironment { } object Subscription { - private lazy val emptyCallback = (_: String, _: Long) => ZIO.unit + type PubSubCallback = (String, Long) => UIO[Unit] final val Subscribe = "SUBSCRIBE" final val Unsubscribe = "UNSUBSCRIBE" diff --git a/modules/redis/src/main/scala/zio/redis/internal/PubSub.scala b/modules/redis/src/main/scala/zio/redis/internal/PubSub.scala new file mode 100644 index 000000000..b6ac1921c --- /dev/null +++ b/modules/redis/src/main/scala/zio/redis/internal/PubSub.scala @@ -0,0 +1,38 @@ +/* + * 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.internal + +object PubSub { + private[redis] sealed trait PushMessage { + def key: SubscriptionKey + } + + private[redis] object PushMessage { + final case class Subscribed(key: SubscriptionKey, numOfSubs: Long) extends PushMessage + final case class Unsubscribed(key: SubscriptionKey, numOfSubs: Long) extends PushMessage + final case class Message(key: SubscriptionKey, destChannel: String, message: RespValue) extends PushMessage + } + + private[redis] sealed trait SubscriptionKey { + def value: String + } + + private[redis] object SubscriptionKey { + final case class Channel(value: String) extends SubscriptionKey + final case class Pattern(value: String) extends SubscriptionKey + } +} diff --git a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala index 45a673823..1bace6550 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala @@ -20,67 +20,56 @@ import zio.redis.Input._ import zio.redis.Output.ArbitraryOutput import zio.redis._ import zio.redis.api.Subscription -import zio.redis.options.PubSub.PubSubCallback +import zio.redis.api.Subscription.PubSubCallback +import zio.redis.internal.PubSub.PushMessage._ import zio.schema.codec.BinaryCodec import zio.stream._ import zio.{Chunk, IO, ZIO} private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionExecutor) { - import zio.redis.options.PubSub.PushProtocol._ - def subscribe[A: BinaryCodec]( channels: Chunk[String], - onSubscribe: PubSubCallback, - onUnsubscribe: PubSubCallback + onSubscribe: Option[PubSubCallback], + onUnsubscribe: Option[PubSubCallback] )(implicit codec: BinaryCodec[String]): Stream[RedisError, (String, A)] = - executor - .execute(makeCommand(Subscription.Subscribe, channels)) - .mapZIO { - case Subscribe(channel, numOfSubs) => - onSubscribe(channel, numOfSubs).as(None) - case Message(channel, message) => - ZIO - .attempt(ArbitraryOutput[A]().unsafeDecode(message)) - .map(msg => Some((channel, msg))) - case Unsubscribe(channel, numOfSubs) => - onUnsubscribe(channel, numOfSubs).as(None) - case _ => ZIO.none - } - .collectSome - .refineToOrDie[RedisError] + executeCommand[A](makeCommand(Subscription.Subscribe, channels), onSubscribe, onUnsubscribe) def pSubscribe[A: BinaryCodec]( patterns: Chunk[String], - onSubscribe: PubSubCallback, - onUnsubscribe: PubSubCallback + onSubscribe: Option[PubSubCallback], + onUnsubscribe: Option[PubSubCallback] )(implicit stringCodec: BinaryCodec[String]): Stream[RedisError, (String, A)] = + executeCommand[A](makeCommand(Subscription.PSubscribe, patterns), onSubscribe, onUnsubscribe) + + def unsubscribe(channels: Chunk[String])(implicit codec: BinaryCodec[String]): IO[RedisError, Unit] = + executeCommand(makeCommand(Subscription.Unsubscribe, channels), None, None).runDrain + + def pUnsubscribe(patterns: Chunk[String])(implicit codec: BinaryCodec[String]): IO[RedisError, Unit] = + executeCommand(makeCommand(Subscription.PUnsubscribe, patterns), None, None).runDrain + + private def makeCommand(commandName: String, keys: Chunk[String])(implicit codec: BinaryCodec[String]) = + CommandNameInput.encode(commandName) ++ + Varargs(ArbitraryKeyInput[String]()).encode(keys) + + private def executeCommand[A: BinaryCodec]( + command: RespCommand, + onSubscribe: Option[PubSubCallback], + onUnsubscribe: Option[PubSubCallback] + ): Stream[RedisError, (String, A)] = executor - .execute(makeCommand(Subscription.PSubscribe, patterns)) + .execute(command) .mapZIO { - case PSubscribe(pattern, numOfSubs) => - onSubscribe(pattern, numOfSubs).as(None) - case PMessage(_, channel, message) => + case Subscribed(key, numOfSubs) => + ZIO.foreach(onSubscribe)(_.apply(key.value, numOfSubs)).as(None) + case Unsubscribed(key, numOfSubs) => + ZIO.foreach(onUnsubscribe)(_.apply(key.value, numOfSubs)).as(None) + case Message(_, channel, message) => ZIO .attempt(ArbitraryOutput[A]().unsafeDecode(message)) .map(msg => Some((channel, msg))) - case PUnsubscribe(pattern, numOfSubs) => - onUnsubscribe(pattern, numOfSubs).as(None) case _ => ZIO.none } .collectSome .refineToOrDie[RedisError] - def unsubscribe(channels: Chunk[String])(implicit codec: BinaryCodec[String]): IO[RedisError, Unit] = - executor - .execute(makeCommand(Subscription.Unsubscribe, channels)) - .runDrain - - def pUnsubscribe(patterns: Chunk[String])(implicit codec: BinaryCodec[String]): IO[RedisError, Unit] = - executor - .execute(makeCommand(Subscription.PUnsubscribe, patterns)) - .runDrain - - private def makeCommand(commandName: String, keys: Chunk[String])(implicit codec: BinaryCodec[String]) = - CommandNameInput.encode(commandName) ++ - Varargs(ArbitraryKeyInput[String]()).encode(keys) } diff --git a/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala similarity index 53% rename from modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala rename to modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala index 4355f1240..396e4e02b 100644 --- a/modules/redis/src/main/scala/zio/redis/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala @@ -14,98 +14,110 @@ * limitations under the License. */ -package zio.redis +package zio.redis.internal import zio.redis.Input.{CommandNameInput, StringInput} -import zio.redis.Output.PushProtocolOutput -import zio.redis.SingleNodeSubscriptionExecutor.{Request, RequestQueueSize, True} +import zio.redis.Output.PushMessageOutput import zio.redis.api.Subscription -import zio.redis.internal._ -import zio.redis.options.PubSub.PushProtocol +import zio.redis.internal.PubSub.{PushMessage, SubscriptionKey} +import zio.redis.internal.SingleNodeSubscriptionExecutor.{Request, RequestQueueSize, True} +import zio.redis.{Input, RedisError} import zio.stream._ -import zio.{Chunk, ChunkBuilder, IO, Promise, Queue, Ref, Schedule, Scope, URIO, ZIO} +import zio.{Chunk, ChunkBuilder, IO, Promise, Queue, Ref, Schedule, Scope, UIO, 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], + subsRef: Ref[Map[SubscriptionKey, Chunk[Queue[Take[RedisError, PushMessage]]]]], + requests: Queue[Request], connection: RedisConnection ) extends SubscriptionExecutor { - def execute(command: RespCommand): Stream[RedisError, PushProtocol] = + def execute(command: RespCommand): Stream[RedisError, PushMessage] = ZStream .fromZIO( for { commandName <- ZIO .fromOption(command.args.collectFirst { case RespCommandArgument.CommandName(name) => name }) - .orElseFail(RedisError.CommandNameNotFound(command.args.toString())) + .orElseFail(RedisError.CommandNameNotFound(command)) 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 Subscription.Subscribe => + ZIO.succeed(subscribe(extractKeys(command).map(SubscriptionKey.Channel(_)), command.args)) + case Subscription.PSubscribe => + ZIO.succeed(subscribe(extractKeys(command).map(SubscriptionKey.Pattern(_)), command.args)) + case Subscription.Unsubscribe => ZIO.succeed(unsubscribe(command.args)) + case Subscription.PUnsubscribe => ZIO.succeed(unsubscribe(command.args)) case other => ZIO.fail(RedisError.InvalidPubSubCommand(other)) } } yield stream ) .flatten + private def extractKeys(command: RespCommand): Chunk[String] = command.args.collect { + case key: RespCommandArgument.Key => key.value.asString + } + + private def releaseQueue(key: SubscriptionKey, queue: Queue[Take[RedisError, PushMessage]]): UIO[Unit] = + subsRef.update(subs => + subs.get(key) match { + case Some(queues) => subs.updated(key, queues.filterNot(_ == queue)) + case None => subs + } + ) + private def subscribe( - subscriptionRef: Ref[Map[String, Chunk[Queue[Take[RedisError, PushProtocol]]]]], - command: RespCommand - ): Stream[RedisError, PushProtocol] = + keys: Chunk[SubscriptionKey], + command: Chunk[RespCommandArgument] + ): Stream[RedisError, PushMessage] = 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 -> _) - ) + queues <- ZIO.foreach(keys)(key => + Queue + .unbounded[Take[RedisError, PushMessage]] + .tap(queue => + subsRef.update(subscription => + subscription.updated( + key, + subscription.getOrElse(key, Chunk.empty) ++ Chunk.single(queue) + ) + ) + ) + .map(key -> _) + ) promise <- Promise.make[RedisError, Unit] - _ <- reqQueue.offer( + _ <- requests.offer( Request( - command.args.map(_.value), + command.map(_.value), promise ) ) + _ <- promise.await.tapError(_ => + ZIO.foreachDiscard(queues) { case (key, queue) => + queue.shutdown *> releaseQueue(key, queue) + } + ) 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 - } - ) - ) + .ensuring(releaseQueue(key, queue)) } - _ <- 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] = + private def unsubscribe(command: Chunk[RespCommandArgument]): Stream[RedisError, PushMessage] = ZStream .fromZIO( for { promise <- Promise.make[RedisError, Unit] - _ <- reqQueue.offer(Request(command.args.map(_.value), promise)) + _ <- requests.offer(Request(command.map(_.value), promise)) _ <- promise.await } yield ZStream.empty ) .flatten private def send = - reqQueue.takeBetween(1, RequestQueueSize).flatMap { reqs => + requests.takeBetween(1, RequestQueueSize).flatMap { reqs => val buffer = ChunkBuilder.make[Byte]() val it = reqs.iterator @@ -126,20 +138,17 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( } 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 offerMessage(msg: PushMessage) = + for { + subscription <- subsRef.get + _ <- ZIO.foreachDiscard(subscription.get(msg.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) = + def releaseStream(key: SubscriptionKey) = for { - subscription <- subscriptionRef.getAndUpdate(_ - key) + subscription <- subsRef.getAndUpdate(_ - key) _ <- ZIO.foreachDiscard(subscription.get(key))( ZIO.foreachDiscard(_)(queue => queue.offer(Take.end).unlessZIO(queue.isShutdown)) ) @@ -149,17 +158,12 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( .mapError(RedisError.IOError(_)) .via(RespValue.Decoder) .collectSome - .mapZIO(resp => ZIO.attempt(PushProtocolOutput.unsafeDecode(resp))) + .mapZIO(resp => ZIO.attempt(PushMessageOutput.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) + case msg: PushMessage.Subscribed => offerMessage(msg) + case msg: PushMessage.Unsubscribed => offerMessage(msg) *> releaseStream(msg.key) + case msg: PushMessage.Message => offerMessage(msg) } } @@ -173,8 +177,11 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( .asBytes for { - channels <- channelSubsRef.get.map(_.keys) - patterns <- patternSubsRef.get.map(_.keys) + subsKeys <- subsRef.get.map(_.keys) + (channels, patterns) = subsKeys.partitionMap { + case SubscriptionKey.Channel(value) => Left(value) + case SubscriptionKey.Pattern(value) => Right(value) + } commands = makeCommand(Subscription.Subscribe, Chunk.fromIterable(channels)) ++ makeCommand(Subscription.PSubscribe, Chunk.fromIterable(patterns)) _ <- connection @@ -209,11 +216,10 @@ private[redis] object SingleNodeSubscriptionExecutor { 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") + reqQueue <- Queue.bounded[Request](RequestQueueSize) + subsRef <- Ref.make(Map.empty[SubscriptionKey, Chunk[Queue[Take[RedisError, PushMessage]]]]) + pubSub = new SingleNodeSubscriptionExecutor(subsRef, reqQueue, conn) + _ <- pubSub.run.forkScoped + _ <- logScopeFinalizer(s"$pubSub Subscription Node is closed") } yield pubSub } diff --git a/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala b/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala index bdb7adcc0..aae92094e 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SubscribeEnvironment.scala @@ -16,7 +16,7 @@ package zio.redis.internal -import zio.redis.{CodecSupplier, SubscriptionExecutor} +import zio.redis.CodecSupplier import zio.schema.Schema import zio.schema.codec.BinaryCodec diff --git a/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala similarity index 90% rename from modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala rename to modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala index 3425c4dc1..6f04fe157 100644 --- a/modules/redis/src/main/scala/zio/redis/SubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala @@ -14,15 +14,15 @@ * limitations under the License. */ -package zio.redis +package zio.redis.internal -import zio.redis.internal.{RedisConnection, RespCommand} -import zio.redis.options.PubSub.PushProtocol +import zio.redis.internal.PubSub.PushMessage +import zio.redis.{RedisConfig, RedisError} import zio.stream._ import zio.{Layer, ZIO, ZLayer} trait SubscriptionExecutor { - private[redis] def execute(command: RespCommand): Stream[RedisError, PushProtocol] + private[redis] def execute(command: RespCommand): Stream[RedisError, PushMessage] } object SubscriptionExecutor { diff --git a/modules/redis/src/main/scala/zio/redis/options/PubSub.scala b/modules/redis/src/main/scala/zio/redis/options/PubSub.scala deleted file mode 100644 index d6f212588..000000000 --- a/modules/redis/src/main/scala/zio/redis/options/PubSub.scala +++ /dev/null @@ -1,37 +0,0 @@ -/* - * 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.options - -import zio.UIO -import zio.redis.internal.RespValue - -object PubSub { - type PubSubCallback = (String, Long) => UIO[Unit] - - private[redis] sealed trait PushProtocol - - private[redis] object PushProtocol { - case class Subscribe(channel: String, numOfSubs: Long) extends PushProtocol - case class PSubscribe(pattern: String, numOfSubs: Long) extends PushProtocol - case class Unsubscribe(channel: String, numOfSubs: Long) extends PushProtocol - case class PUnsubscribe(pattern: String, numOfSubs: Long) extends PushProtocol - case class Message(channel: String, message: RespValue) extends PushProtocol - case class PMessage(pattern: String, channel: String, message: RespValue) extends PushProtocol - } - - final case class NumberOfSubscribers(channel: String, subscriberCount: Long) -} diff --git a/modules/redis/src/main/scala/zio/redis/options/Publishing.scala b/modules/redis/src/main/scala/zio/redis/options/Publishing.scala new file mode 100644 index 000000000..7a880dff5 --- /dev/null +++ b/modules/redis/src/main/scala/zio/redis/options/Publishing.scala @@ -0,0 +1,21 @@ +/* + * 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.options + +trait Publishing { + sealed case class NumberOfSubscribers(channel: String, subscriberCount: Long) +} diff --git a/modules/redis/src/main/scala/zio/redis/package.scala b/modules/redis/src/main/scala/zio/redis/package.scala index cb1131fbe..7026485dd 100644 --- a/modules/redis/src/main/scala/zio/redis/package.scala +++ b/modules/redis/src/main/scala/zio/redis/package.scala @@ -25,7 +25,8 @@ package object redis with options.Strings with options.Lists with options.Streams - with options.Scripting { + with options.Scripting + with options.Publishing { type Id[+A] = A } diff --git a/modules/redis/src/test/scala/zio/redis/OutputSpec.scala b/modules/redis/src/test/scala/zio/redis/OutputSpec.scala index 3c9407346..beab980ff 100644 --- a/modules/redis/src/test/scala/zio/redis/OutputSpec.scala +++ b/modules/redis/src/test/scala/zio/redis/OutputSpec.scala @@ -1,10 +1,10 @@ package zio.redis import zio._ -import zio.redis.Output.{PushProtocolOutput, _} +import zio.redis.Output._ import zio.redis.RedisError._ +import zio.redis.internal.PubSub.{PushMessage, SubscriptionKey} import zio.redis.internal.RespValue -import zio.redis.options.PubSub.PushProtocol import zio.test.Assertion._ import zio.test._ @@ -888,7 +888,7 @@ object OutputSpec extends BaseSpec { } ) ), - suite("PushProtocol")( + suite("PushMessage")( test("subscribe") { val channel = "foo" val numOfSubs = 1L @@ -898,8 +898,8 @@ object OutputSpec extends BaseSpec { RespValue.bulkString(channel), RespValue.Integer(numOfSubs) ) - val expected = PushProtocol.Subscribe(channel, numOfSubs) - assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + val expected = PushMessage.Subscribed(SubscriptionKey.Channel(channel), numOfSubs) + assertZIO(ZIO.attempt(PushMessageOutput.unsafeDecode(input)))( equalTo(expected) ) }, @@ -912,8 +912,8 @@ object OutputSpec extends BaseSpec { RespValue.bulkString(pattern), RespValue.Integer(numOfSubs) ) - val expected = PushProtocol.PSubscribe(pattern, numOfSubs) - assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + val expected = PushMessage.Subscribed(SubscriptionKey.Pattern(pattern), numOfSubs) + assertZIO(ZIO.attempt(PushMessageOutput.unsafeDecode(input)))( equalTo(expected) ) }, @@ -926,8 +926,8 @@ object OutputSpec extends BaseSpec { RespValue.bulkString(channel), RespValue.Integer(numOfSubs) ) - val expected = PushProtocol.Unsubscribe(channel, numOfSubs) - assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + val expected = PushMessage.Unsubscribed(SubscriptionKey.Channel(channel), numOfSubs) + assertZIO(ZIO.attempt(PushMessageOutput.unsafeDecode(input)))( equalTo(expected) ) }, @@ -940,8 +940,8 @@ object OutputSpec extends BaseSpec { RespValue.bulkString(pattern), RespValue.Integer(numOfSubs) ) - val expected = PushProtocol.PUnsubscribe(pattern, numOfSubs) - assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + val expected = PushMessage.Unsubscribed(SubscriptionKey.Pattern(pattern), numOfSubs) + assertZIO(ZIO.attempt(PushMessageOutput.unsafeDecode(input)))( equalTo(expected) ) }, @@ -954,8 +954,8 @@ object OutputSpec extends BaseSpec { RespValue.bulkString(channel), message ) - val expected = PushProtocol.Message(channel, message) - assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + val expected = PushMessage.Message(SubscriptionKey.Channel(channel), channel, message) + assertZIO(ZIO.attempt(PushMessageOutput.unsafeDecode(input)))( equalTo(expected) ) }, @@ -970,8 +970,8 @@ object OutputSpec extends BaseSpec { RespValue.bulkString(channel), message ) - val expected = PushProtocol.PMessage(pattern, channel, message) - assertZIO(ZIO.attempt(PushProtocolOutput.unsafeDecode(input)))( + val expected = PushMessage.Message(SubscriptionKey.Pattern(pattern), channel, message) + assertZIO(ZIO.attempt(PushMessageOutput.unsafeDecode(input)))( equalTo(expected) ) } diff --git a/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala b/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala index 933985866..fc5325da6 100644 --- a/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala +++ b/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala @@ -1,6 +1,5 @@ package zio.redis -import zio.redis.options.PubSub.NumberOfSubscribers import zio.test.Assertion._ import zio.test._ import zio.{Chunk, Promise, ZIO} @@ -19,8 +18,8 @@ trait PubSubSpec extends BaseSpec { resBuilder = subscription .subscribeWithCallback(channel)( - (key: String, _: Long) => promise.succeed(key).unit, - (_, _) => ZIO.unit + Some((key: String, _: Long) => promise.succeed(key).unit), + None ) stream = resBuilder.returning[String] _ <- stream @@ -81,8 +80,8 @@ trait PubSubSpec extends BaseSpec { promise <- Promise.make[RedisError, String] _ <- subscription .pSubscribeWithCallback(pattern)( - (key: String, _: Long) => promise.succeed(key).unit, - (_, _) => ZIO.unit + Some((key: String, _: Long) => promise.succeed(key).unit), + None ) .returning[String] .interruptWhen(promise) @@ -142,7 +141,7 @@ trait PubSubSpec extends BaseSpec { message <- generateRandomString() promise <- Promise.make[Nothing, Unit] _ <- subscription - .subscribeWithCallback(channel)((_, _) => ZIO.unit, (_, _) => promise.succeed(()).unit) + .subscribeWithCallback(channel)(None, Some((_, _) => promise.succeed(()).unit)) .returning[String] .runCollect .fork @@ -161,8 +160,8 @@ trait PubSubSpec extends BaseSpec { promise <- Promise.make[RedisError, String] _ <- subscription .subscribeWithCallback(channel)( - (_, _) => ZIO.unit, - (key, _) => promise.succeed(key).unit + None, + Some((key, _) => promise.succeed(key).unit) ) .returning[Unit] .runDrain @@ -178,8 +177,8 @@ trait PubSubSpec extends BaseSpec { promise <- Promise.make[RedisError, String] _ <- subscription .pSubscribeWithCallback(pattern)( - (_, _) => ZIO.unit, - (key, _) => promise.succeed(key).unit + None, + Some((key, _) => promise.succeed(key).unit) ) .returning[Unit] .runDrain @@ -201,8 +200,8 @@ trait PubSubSpec extends BaseSpec { _ <- subscription .subscribeWithCallback(channel1)( - (_, _) => ZIO.unit, - (_, _) => promise1.succeed(()).unit + None, + Some((_, _) => promise1.succeed(()).unit) ) .returning[String] .runCollect @@ -210,8 +209,8 @@ trait PubSubSpec extends BaseSpec { _ <- subscription .subscribeWithCallback(channel2)( - (_, _) => ZIO.unit, - (_, _) => promise2.succeed(()).unit + None, + Some((_, _) => promise2.succeed(()).unit) ) .returning[String] .runCollect From 592d327ba8a55b57e30a05b084f3690fa3b1be3b Mon Sep 17 00:00:00 2001 From: 0pg Date: Thu, 11 May 2023 02:17:21 +0900 Subject: [PATCH 39/51] Commit suggestion --- .../main/scala/zio/redis/internal/SubscriptionExecutor.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala index 6f04fe157..74400c651 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala @@ -33,10 +33,10 @@ object SubscriptionExecutor { RedisConnection.local.fresh >>> pubSublayer private lazy val pubSublayer: ZLayer[RedisConnection, RedisError.IOError, SubscriptionExecutor] = - ZLayer.scoped( + ZLayer.scoped { for { conn <- ZIO.service[RedisConnection] pubSub <- SingleNodeSubscriptionExecutor.create(conn) } yield pubSub - ) + } } From 386fefd1b7a206990313f780427a84b646d777ae Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 23 May 2023 16:45:24 +0900 Subject: [PATCH 40/51] Fix broken compile --- .../zio/redis/internal/RedisSubscriptionCommand.scala | 1 - .../internal/SingleNodeSubscriptionExecutor.scala | 10 +++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala index 1bace6550..45ef71a0e 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala @@ -67,7 +67,6 @@ private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionE ZIO .attempt(ArbitraryOutput[A]().unsafeDecode(message)) .map(msg => Some((channel, msg))) - case _ => ZIO.none } .collectSome .refineToOrDie[RedisError] diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala index 396e4e02b..dfdbbb648 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala @@ -178,12 +178,12 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( for { subsKeys <- subsRef.get.map(_.keys) - (channels, patterns) = subsKeys.partitionMap { - case SubscriptionKey.Channel(value) => Left(value) - case SubscriptionKey.Pattern(value) => Right(value) + (channels, patterns) = subsKeys.partition { + case _: SubscriptionKey.Channel => false + case _: SubscriptionKey.Pattern => true } - commands = makeCommand(Subscription.Subscribe, Chunk.fromIterable(channels)) ++ - makeCommand(Subscription.PSubscribe, Chunk.fromIterable(patterns)) + commands = makeCommand(Subscription.Subscribe, Chunk.fromIterable(channels).map(_.value)) ++ + makeCommand(Subscription.PSubscribe, Chunk.fromIterable(patterns).map(_.value)) _ <- connection .write(commands) .when(commands.nonEmpty) From 16992118130262d46d655f66155c5de8eabd97f3 Mon Sep 17 00:00:00 2001 From: 0pg Date: Wed, 31 May 2023 03:03:57 +0900 Subject: [PATCH 41/51] Use hub instead of chunks of queue --- .../SingleNodeSubscriptionExecutor.scala | 183 +++++++++--------- 1 file changed, 88 insertions(+), 95 deletions(-) diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala index dfdbbb648..c94ee5d43 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala @@ -23,98 +23,90 @@ import zio.redis.internal.PubSub.{PushMessage, SubscriptionKey} import zio.redis.internal.SingleNodeSubscriptionExecutor.{Request, RequestQueueSize, True} import zio.redis.{Input, RedisError} import zio.stream._ -import zio.{Chunk, ChunkBuilder, IO, Promise, Queue, Ref, Schedule, Scope, UIO, URIO, ZIO} +import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Schedule, Scope, UIO, URIO, ZIO} private[redis] final class SingleNodeSubscriptionExecutor private ( - subsRef: Ref[Map[SubscriptionKey, Chunk[Queue[Take[RedisError, PushMessage]]]]], + subsRef: Ref[Map[SubscriptionKey, Hub[Take[RedisError, PushMessage]]]], requests: Queue[Request], + responses: Queue[Promise[RedisError, PushMessage]], connection: RedisConnection ) extends SubscriptionExecutor { - def execute(command: RespCommand): Stream[RedisError, PushMessage] = - ZStream - .fromZIO( - for { - commandName <- - ZIO - .fromOption(command.args.collectFirst { case RespCommandArgument.CommandName(name) => name }) - .orElseFail(RedisError.CommandNameNotFound(command)) - stream <- commandName match { - case Subscription.Subscribe => - ZIO.succeed(subscribe(extractKeys(command).map(SubscriptionKey.Channel(_)), command.args)) - case Subscription.PSubscribe => - ZIO.succeed(subscribe(extractKeys(command).map(SubscriptionKey.Pattern(_)), command.args)) - case Subscription.Unsubscribe => ZIO.succeed(unsubscribe(command.args)) - case Subscription.PUnsubscribe => ZIO.succeed(unsubscribe(command.args)) - case other => ZIO.fail(RedisError.InvalidPubSubCommand(other)) - } - } yield stream - ) - .flatten - - private def extractKeys(command: RespCommand): Chunk[String] = command.args.collect { - case key: RespCommandArgument.Key => key.value.asString + def execute(command: RespCommand): Stream[RedisError, PushMessage] = { + def getStream(commandName: String): IO[RedisError, Stream[RedisError, PushMessage]] = commandName match { + case Subscription.Subscribe => + subscribe(extractChannelKeys(command.args), command) + case Subscription.PSubscribe => + subscribe(extractPatternKeys(command.args), command) + case Subscription.Unsubscribe => + unsubscribe(extractChannelKeys(command.args), command) + case Subscription.PUnsubscribe => + unsubscribe(extractPatternKeys(command.args), command) + case other => ZIO.fail(RedisError.InvalidPubSubCommand(other)) + } + + for { + commandName <- ZStream.fromZIO( + ZIO + .fromOption(command.args.collectFirst { case RespCommandArgument.CommandName(name) => + name + }) + .orElseFail(RedisError.CommandNameNotFound(command)) + ) + pushMessage <- ZStream.fromZIO(getStream(commandName)).flatten + } yield pushMessage } - private def releaseQueue(key: SubscriptionKey, queue: Queue[Take[RedisError, PushMessage]]): UIO[Unit] = - subsRef.update(subs => - subs.get(key) match { - case Some(queues) => subs.updated(key, queues.filterNot(_ == queue)) - case None => subs - } - ) + private def unsubscribe[T <: SubscriptionKey]( + keys: Chunk[T], + command: RespCommand + ): UIO[Stream[RedisError, PushMessage]] = + for { + targetKeys <- if (keys.isEmpty) + subsRef.get + .map(_.keys.collect { case key: SubscriptionKey.Pattern => key }) + .map(Chunk.fromIterable(_)) + else + ZIO.succeed(keys) + keyPromisePairs <- ZIO.foreach(targetKeys)(key => Promise.make[RedisError, PushMessage].map(key -> _)) + _ <- requests.offer(Request(command.args.map(_.value), keyPromisePairs.map(_._2))) + streams = keyPromisePairs.map { case (key, promise) => + ZStream.fromZIO(promise.await <* subsRef.update(_ - key)) + } + } yield streams.fold(ZStream.empty)(_ merge _) + + private def extractChannelKeys(command: Chunk[RespCommandArgument]): Chunk[SubscriptionKey.Channel] = + command.collect { case RespCommandArgument.Key(key) => + key.asString + }.map(SubscriptionKey.Channel(_)) + + private def extractPatternKeys(command: Chunk[RespCommandArgument]): Chunk[SubscriptionKey.Pattern] = + command.collect { case RespCommandArgument.Key(key) => + key.asString + }.map(SubscriptionKey.Pattern(_)) private def subscribe( keys: Chunk[SubscriptionKey], - command: Chunk[RespCommandArgument] - ): Stream[RedisError, PushMessage] = - ZStream - .fromZIO( - for { - queues <- ZIO.foreach(keys)(key => - Queue - .unbounded[Take[RedisError, PushMessage]] - .tap(queue => - subsRef.update(subscription => - subscription.updated( - key, - subscription.getOrElse(key, Chunk.empty) ++ Chunk.single(queue) - ) - ) - ) - .map(key -> _) - ) - promise <- Promise.make[RedisError, Unit] - _ <- requests.offer( - Request( - command.map(_.value), - promise - ) - ) - _ <- promise.await.tapError(_ => - ZIO.foreachDiscard(queues) { case (key, queue) => - queue.shutdown *> releaseQueue(key, queue) - } - ) - streams = queues.map { case (key, queue) => - ZStream - .fromQueueWithShutdown(queue) - .ensuring(releaseQueue(key, queue)) - } - } yield streams.fold(ZStream.empty)(_ merge _) - ) - .flatten - .flattenTake - - private def unsubscribe(command: Chunk[RespCommandArgument]): Stream[RedisError, PushMessage] = - ZStream - .fromZIO( - for { - promise <- Promise.make[RedisError, Unit] - _ <- requests.offer(Request(command.map(_.value), promise)) - _ <- promise.await - } yield ZStream.empty - ) - .flatten + command: RespCommand + ): UIO[Stream[RedisError, PushMessage]] = { + def getHub(key: SubscriptionKey): UIO[Hub[Take[RedisError, PushMessage]]] = + subsRef.get + .map(_.get(key)) + .someOrElseZIO( + Hub.unbounded[Take[RedisError, PushMessage]].tap(hub => subsRef.update(_ + (key -> hub))) + ) + + for { + keyPromisePairs <- ZIO.foreach(keys)(key => Promise.make[RedisError, PushMessage].map(key -> _)) + _ <- requests.offer(Request(command.args.map(_.value), keyPromisePairs.map(_._2))) + streams = keyPromisePairs.map { case (key, promise) => + ZStream.fromZIO(promise.await) ++ + ZStream + .fromZIO(getHub(key)) + .flatMap(ZStream.fromHub(_)) + .flattenTake + } + } yield streams.fold(ZStream.empty)(_ merge _) + } private def send = requests.takeBetween(1, RequestQueueSize).flatMap { reqs => @@ -132,8 +124,8 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( .write(bytes) .mapError(RedisError.IOError(_)) .tapBoth( - e => ZIO.foreachDiscard(reqs.map(_.promise))(_.fail(e)), - _ => ZIO.foreachDiscard(reqs.map(_.promise))(_.succeed(())) + e => ZIO.foreachDiscard(reqs.flatMap(_.promises))(_.fail(e)), + _ => ZIO.foreachDiscard(reqs)(req => responses.offerAll(req.promises)) ) } @@ -141,19 +133,19 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( def offerMessage(msg: PushMessage) = for { subscription <- subsRef.get - _ <- ZIO.foreachDiscard(subscription.get(msg.key))( - ZIO.foreachDiscard(_)(queue => queue.offer(Take.single(msg)).unlessZIO(queue.isShutdown)) - ) + _ <- ZIO.foreachDiscard(subscription.get(msg.key))(_.offer(Take.single(msg))) } yield () def releaseStream(key: SubscriptionKey) = for { subscription <- subsRef.getAndUpdate(_ - key) - _ <- ZIO.foreachDiscard(subscription.get(key))( - ZIO.foreachDiscard(_)(queue => queue.offer(Take.end).unlessZIO(queue.isShutdown)) - ) + _ <- ZIO.foreachDiscard(subscription.get(key))(_.offer(Take.end)) + _ <- subsRef.update(_ - key) } yield () + def releasePromise(msg: PushMessage): UIO[Unit] = + responses.take.flatMap(_.succeed(msg)).unit + connection.read .mapError(RedisError.IOError(_)) .via(RespValue.Decoder) @@ -161,8 +153,8 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( .mapZIO(resp => ZIO.attempt(PushMessageOutput.unsafeDecode(resp))) .refineToOrDie[RedisError] .foreach { - case msg: PushMessage.Subscribed => offerMessage(msg) - case msg: PushMessage.Unsubscribed => offerMessage(msg) *> releaseStream(msg.key) + case msg: PushMessage.Subscribed => releasePromise(msg) + case msg: PushMessage.Unsubscribed => releasePromise(msg) *> releaseStream(msg.key) case msg: PushMessage.Message => offerMessage(msg) } } @@ -207,7 +199,7 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( private[redis] object SingleNodeSubscriptionExecutor { private final case class Request( command: Chunk[RespValue.BulkString], - promise: Promise[RedisError, Unit] + promises: Chunk[Promise[RedisError, PushMessage]] ) private final val True: Any => Boolean = _ => true @@ -217,8 +209,9 @@ private[redis] object SingleNodeSubscriptionExecutor { def create(conn: RedisConnection): URIO[Scope, SubscriptionExecutor] = for { reqQueue <- Queue.bounded[Request](RequestQueueSize) - subsRef <- Ref.make(Map.empty[SubscriptionKey, Chunk[Queue[Take[RedisError, PushMessage]]]]) - pubSub = new SingleNodeSubscriptionExecutor(subsRef, reqQueue, conn) + resQueue <- Queue.unbounded[Promise[RedisError, PushMessage]] + subsRef <- Ref.make(Map.empty[SubscriptionKey, Hub[Take[RedisError, PushMessage]]]) + pubSub = new SingleNodeSubscriptionExecutor(subsRef, reqQueue, resQueue, conn) _ <- pubSub.run.forkScoped _ <- logScopeFinalizer(s"$pubSub Subscription Node is closed") } yield pubSub From be11d670bf2de5d31c94f9ba86b7d5df78ae84da Mon Sep 17 00:00:00 2001 From: 0pg Date: Wed, 31 May 2023 03:40:02 +0900 Subject: [PATCH 42/51] Add release on error --- .../zio/redis/internal/SingleNodeSubscriptionExecutor.scala | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala index c94ee5d43..66d86c6d9 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala @@ -159,6 +159,7 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( } } + private def drainWith(e: RedisError): UIO[Unit] = responses.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) private def resubscribe: IO[RedisError, Unit] = { def makeCommand(name: String, keys: Chunk[String]) = if (keys.isEmpty) @@ -191,7 +192,7 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( 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) + .tapError(e => ZIO.logWarning(s"Reconnecting due to error: $e") *> drainWith(e) *> resubscribe) .retryWhile(True) .tapError(e => ZIO.logError(s"Executor exiting: $e")) } From 552e0021326093b4031fa5a3b7c17a93aa7d587b Mon Sep 17 00:00:00 2001 From: 0pg Date: Wed, 31 May 2023 04:06:05 +0900 Subject: [PATCH 43/51] Extraact common logic --- .../zio/redis/internal/RequestQueue.scala | 24 ++++++++++ .../redis/internal/SingleNodeExecutor.scala | 29 ++++-------- .../zio/redis/internal/SingleNodeRunner.scala | 44 +++++++++++++++++++ .../SingleNodeSubscriptionExecutor.scala | 36 +++++---------- 4 files changed, 88 insertions(+), 45 deletions(-) create mode 100644 modules/redis/src/main/scala/zio/redis/internal/RequestQueue.scala create mode 100644 modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala diff --git a/modules/redis/src/main/scala/zio/redis/internal/RequestQueue.scala b/modules/redis/src/main/scala/zio/redis/internal/RequestQueue.scala new file mode 100644 index 000000000..9f2a6def7 --- /dev/null +++ b/modules/redis/src/main/scala/zio/redis/internal/RequestQueue.scala @@ -0,0 +1,24 @@ +/* + * 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.internal + +import zio.{Queue, UIO} + +private[redis] object RequestQueue { + private final val RequestQueueSize = 16 + def create[A]: UIO[Queue[A]] = Queue.bounded[A](RequestQueueSize) +} diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala index 4b13a07db..0bb043d51 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala @@ -24,7 +24,8 @@ private[redis] final class SingleNodeExecutor private ( connection: RedisConnection, requests: Queue[Request], responses: Queue[Promise[RedisError, RespValue]] -) extends RedisExecutor { +) extends SingleNodeRunner + with RedisExecutor { // TODO NodeExecutor doesn't throw connection errors, timeout errors, it is hanging forever def execute(command: RespCommand): IO[RedisError, RespValue] = @@ -32,21 +33,10 @@ private[redis] final class SingleNodeExecutor private ( .make[RedisError, RespValue] .flatMap(promise => requests.offer(Request(command.args.map(_.value), promise)) *> promise.await) - /** - * Opens a connection to the server and launches send and receive operations. All failures are retried by opening a - * new connection. Only exits by interruption or defect. - */ - private val run: IO[RedisError, AnyVal] = - ZIO.logTrace(s"$this Executable sender and reader has been started") *> - (send.repeat[Any, Long](Schedule.forever) race receive) - .tapError(e => ZIO.logWarning(s"Reconnecting due to error: $e") *> drainWith(e)) - .retryWhile(True) - .tapError(e => ZIO.logError(s"Executor exiting: $e")) + def onError(e: RedisError): UIO[Unit] = responses.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) - private def drainWith(e: RedisError): UIO[Unit] = responses.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) - - private def send: IO[RedisError.IOError, Option[Unit]] = - requests.takeBetween(1, RequestQueueSize).flatMap { requests => + def send: IO[RedisError.IOError, Unit] = + requests.takeAll.flatMap { requests => val bytes = requests .foldLeft(new ChunkBuilder.Byte())((builder, req) => builder ++= RespValue.Array(req.command).asBytes) @@ -59,9 +49,10 @@ private[redis] final class SingleNodeExecutor private ( e => ZIO.foreachDiscard(requests)(_.promise.fail(e)), _ => ZIO.foreachDiscard(requests)(req => responses.offer(req.promise)) ) + .unit } - private def receive: IO[RedisError, Unit] = + def receive: IO[RedisError, Unit] = connection.read .mapError(RedisError.IOError(_)) .via(RespValue.Decoder) @@ -79,7 +70,7 @@ private[redis] object SingleNodeExecutor { def create(connection: RedisConnection): URIO[Scope, SingleNodeExecutor] = for { - requests <- Queue.bounded[Request](RequestQueueSize) + requests <- RequestQueue.create[Request] responses <- Queue.unbounded[Promise[RedisError, RespValue]] executor = new SingleNodeExecutor(connection, requests, responses) _ <- executor.run.forkScoped @@ -88,10 +79,6 @@ private[redis] object SingleNodeExecutor { private final case class Request(command: Chunk[RespValue.BulkString], promise: Promise[RedisError, RespValue]) - private final val True: Any => Boolean = _ => true - - private final val RequestQueueSize = 16 - private def makeLayer: ZLayer[RedisConnection, RedisError.IOError, RedisExecutor] = ZLayer.scoped(ZIO.serviceWithZIO[RedisConnection](create)) } diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala new file mode 100644 index 000000000..d60040eb4 --- /dev/null +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala @@ -0,0 +1,44 @@ +/* + * 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.internal + +import zio.redis.RedisError +import zio.redis.internal.SingleNodeRunner.True +import zio.{IO, Schedule, ZIO} + +private[redis] trait SingleNodeRunner { + def send: IO[RedisError.IOError, Unit] + + def receive: IO[RedisError, Unit] + + def onError(e: RedisError): IO[RedisError, Unit] + + /** + * 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 sender and reader has been started") *> + (send.repeat(Schedule.forever) race receive) + .tapError(e => ZIO.logWarning(s"Reconnecting due to error: $e") *> onError(e)) + .retryWhile(True) + .tapError(e => ZIO.logError(s"Executor exiting: $e")) +} + +private[redis] object SingleNodeRunner { + final val True: Any => Boolean = _ => true +} diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala index 66d86c6d9..0d83cb383 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala @@ -20,17 +20,19 @@ import zio.redis.Input.{CommandNameInput, StringInput} import zio.redis.Output.PushMessageOutput import zio.redis.api.Subscription import zio.redis.internal.PubSub.{PushMessage, SubscriptionKey} -import zio.redis.internal.SingleNodeSubscriptionExecutor.{Request, RequestQueueSize, True} +import zio.redis.internal.SingleNodeRunner.True +import zio.redis.internal.SingleNodeSubscriptionExecutor.Request import zio.redis.{Input, RedisError} import zio.stream._ -import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Schedule, Scope, UIO, URIO, ZIO} +import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Scope, UIO, URIO, ZIO} private[redis] final class SingleNodeSubscriptionExecutor private ( subsRef: Ref[Map[SubscriptionKey, Hub[Take[RedisError, PushMessage]]]], requests: Queue[Request], responses: Queue[Promise[RedisError, PushMessage]], connection: RedisConnection -) extends SubscriptionExecutor { +) extends SingleNodeRunner + with SubscriptionExecutor { def execute(command: RespCommand): Stream[RedisError, PushMessage] = { def getStream(commandName: String): IO[RedisError, Stream[RedisError, PushMessage]] = commandName match { case Subscription.Subscribe => @@ -108,8 +110,8 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( } yield streams.fold(ZStream.empty)(_ merge _) } - private def send = - requests.takeBetween(1, RequestQueueSize).flatMap { reqs => + def send: IO[RedisError.IOError, Unit] = + requests.takeAll.flatMap { reqs => val buffer = ChunkBuilder.make[Byte]() val it = reqs.iterator @@ -127,9 +129,10 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( e => ZIO.foreachDiscard(reqs.flatMap(_.promises))(_.fail(e)), _ => ZIO.foreachDiscard(reqs)(req => responses.offerAll(req.promises)) ) + .unit } - private def receive: IO[RedisError, Unit] = { + def receive: IO[RedisError, Unit] = { def offerMessage(msg: PushMessage) = for { subscription <- subsRef.get @@ -159,8 +162,7 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( } } - private def drainWith(e: RedisError): UIO[Unit] = responses.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) - private def resubscribe: IO[RedisError, Unit] = { + def onError(e: RedisError): IO[RedisError, Unit] = { def makeCommand(name: String, keys: Chunk[String]) = if (keys.isEmpty) Chunk.empty @@ -170,6 +172,7 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( .asBytes for { + _ <- responses.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) subsKeys <- subsRef.get.map(_.keys) (channels, patterns) = subsKeys.partition { case _: SubscriptionKey.Channel => false @@ -184,17 +187,6 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( .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") *> drainWith(e) *> resubscribe) - .retryWhile(True) - .tapError(e => ZIO.logError(s"Executor exiting: $e")) } private[redis] object SingleNodeSubscriptionExecutor { @@ -203,13 +195,9 @@ private[redis] object SingleNodeSubscriptionExecutor { promises: Chunk[Promise[RedisError, PushMessage]] ) - 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) + reqQueue <- RequestQueue.create[Request] resQueue <- Queue.unbounded[Promise[RedisError, PushMessage]] subsRef <- Ref.make(Map.empty[SubscriptionKey, Hub[Take[RedisError, PushMessage]]]) pubSub = new SingleNodeSubscriptionExecutor(subsRef, reqQueue, resQueue, conn) From e218b08ecec03981d354b9091e2c1e844918ff5b Mon Sep 17 00:00:00 2001 From: 0pg Date: Wed, 31 May 2023 15:48:39 +0900 Subject: [PATCH 44/51] Ensure order of subs/unsubs --- .../src/main/scala/zio/redis/RedisError.scala | 6 +- .../SingleNodeSubscriptionExecutor.scala | 119 ++++++++++-------- .../zio/redis/{internal => }/PubSubSpec.scala | 58 +++++---- 3 files changed, 101 insertions(+), 82 deletions(-) rename modules/redis/src/test/scala/zio/redis/{internal => }/PubSubSpec.scala (85%) diff --git a/modules/redis/src/main/scala/zio/redis/RedisError.scala b/modules/redis/src/main/scala/zio/redis/RedisError.scala index db0ed056f..d0213227b 100644 --- a/modules/redis/src/main/scala/zio/redis/RedisError.scala +++ b/modules/redis/src/main/scala/zio/redis/RedisError.scala @@ -16,6 +16,7 @@ package zio.redis +import zio.redis.internal.PubSub.SubscriptionKey import zio.redis.internal.RespCommand import zio.redis.options.Cluster.Slot @@ -53,6 +54,7 @@ object RedisError { def apply(command: RespCommand): CommandNameNotFound = CommandNameNotFound(command.args.toString()) } - sealed trait PubSubError extends RedisError - final case class InvalidPubSubCommand(command: String) extends PubSubError + sealed trait PubSubError extends RedisError + final case class InvalidPubSubCommand(command: String) extends PubSubError + final case class SubscriptionStreamAlreadyClosed(key: SubscriptionKey) extends PubSubError } diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala index 0d83cb383..b1606357b 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala @@ -29,7 +29,7 @@ import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Scope, UIO, URIO, private[redis] final class SingleNodeSubscriptionExecutor private ( subsRef: Ref[Map[SubscriptionKey, Hub[Take[RedisError, PushMessage]]]], requests: Queue[Request], - responses: Queue[Promise[RedisError, PushMessage]], + subsResponses: Queue[Promise[RedisError, PushMessage]], connection: RedisConnection ) extends SingleNodeRunner with SubscriptionExecutor { @@ -40,9 +40,9 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( case Subscription.PSubscribe => subscribe(extractPatternKeys(command.args), command) case Subscription.Unsubscribe => - unsubscribe(extractChannelKeys(command.args), command) + unsubscribe(command).as(ZStream.empty) case Subscription.PUnsubscribe => - unsubscribe(extractPatternKeys(command.args), command) + unsubscribe(command).as(ZStream.empty) case other => ZIO.fail(RedisError.InvalidPubSubCommand(other)) } @@ -58,23 +58,45 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( } yield pushMessage } - private def unsubscribe[T <: SubscriptionKey]( - keys: Chunk[T], + private def subscribe( + keys: Chunk[SubscriptionKey], command: RespCommand - ): UIO[Stream[RedisError, PushMessage]] = + ): IO[RedisError, Stream[RedisError, PushMessage]] = { + def getHub(key: SubscriptionKey): IO[RedisError, Hub[Take[RedisError, PushMessage]]] = + subsRef.get + .map(_.get(key)) + .flatMap(ZIO.fromOption(_)) + .orElseFail(RedisError.SubscriptionStreamAlreadyClosed(key)) + + def getStream(promise: Promise[RedisError, PushMessage]) = + for { + subscribed <- ZStream.fromZIO(promise.await) + head = ZStream(subscribed) + tail = ZStream.fromZIO(getHub(subscribed.key)).flatMap(ZStream.fromHub(_)).flattenTake + message <- head ++ tail + } yield message + + def makeHub(key: SubscriptionKey): UIO[Unit] = + Hub + .unbounded[Take[RedisError, PushMessage]] + .tap(hub => subsRef.update(_ + (key -> hub))) + .unlessZIO(subsRef.get.map(_.contains(key))) + .unit + for { - targetKeys <- if (keys.isEmpty) - subsRef.get - .map(_.keys.collect { case key: SubscriptionKey.Pattern => key }) - .map(Chunk.fromIterable(_)) - else - ZIO.succeed(keys) - keyPromisePairs <- ZIO.foreach(targetKeys)(key => Promise.make[RedisError, PushMessage].map(key -> _)) - _ <- requests.offer(Request(command.args.map(_.value), keyPromisePairs.map(_._2))) - streams = keyPromisePairs.map { case (key, promise) => - ZStream.fromZIO(promise.await <* subsRef.update(_ - key)) - } + _ <- ZIO.foreachDiscard(keys)(makeHub(_)) + promises <- Promise.make[RedisError, PushMessage].replicateZIO(keys.size).map(Chunk.fromIterable(_)) + _ <- requests.offer(Request.Subscribe(command.args.map(_.value), promises)) + streams = promises.map(getStream(_)) } yield streams.fold(ZStream.empty)(_ merge _) + } + + private def unsubscribe(command: RespCommand): IO[RedisError, Unit] = + for { + promise <- Promise.make[RedisError, Unit] + _ <- requests.offer(Request.Unsubscribe(command.args.map(_.value), promise)) + _ <- promise.await + } yield () private def extractChannelKeys(command: Chunk[RespCommandArgument]): Chunk[SubscriptionKey.Channel] = command.collect { case RespCommandArgument.Key(key) => @@ -86,30 +108,6 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( key.asString }.map(SubscriptionKey.Pattern(_)) - private def subscribe( - keys: Chunk[SubscriptionKey], - command: RespCommand - ): UIO[Stream[RedisError, PushMessage]] = { - def getHub(key: SubscriptionKey): UIO[Hub[Take[RedisError, PushMessage]]] = - subsRef.get - .map(_.get(key)) - .someOrElseZIO( - Hub.unbounded[Take[RedisError, PushMessage]].tap(hub => subsRef.update(_ + (key -> hub))) - ) - - for { - keyPromisePairs <- ZIO.foreach(keys)(key => Promise.make[RedisError, PushMessage].map(key -> _)) - _ <- requests.offer(Request(command.args.map(_.value), keyPromisePairs.map(_._2))) - streams = keyPromisePairs.map { case (key, promise) => - ZStream.fromZIO(promise.await) ++ - ZStream - .fromZIO(getHub(key)) - .flatMap(ZStream.fromHub(_)) - .flattenTake - } - } yield streams.fold(ZStream.empty)(_ merge _) - } - def send: IO[RedisError.IOError, Unit] = requests.takeAll.flatMap { reqs => val buffer = ChunkBuilder.make[Byte]() @@ -126,8 +124,16 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( .write(bytes) .mapError(RedisError.IOError(_)) .tapBoth( - e => ZIO.foreachDiscard(reqs.flatMap(_.promises))(_.fail(e)), - _ => ZIO.foreachDiscard(reqs)(req => responses.offerAll(req.promises)) + e => + ZIO.foreachDiscard(reqs) { + case Request.Subscribe(_, promises) => ZIO.foreachDiscard(promises)(_.fail(e)) + case Request.Unsubscribe(_, promise) => promise.fail(e) + }, + _ => + ZIO.foreachDiscard(reqs) { + case Request.Subscribe(_, promises) => subsResponses.offerAll(promises) + case Request.Unsubscribe(_, promise) => promise.succeed(()) + } ) .unit } @@ -139,16 +145,15 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( _ <- ZIO.foreachDiscard(subscription.get(msg.key))(_.offer(Take.single(msg))) } yield () - def releaseStream(key: SubscriptionKey) = + def releasePromise(msg: PushMessage): UIO[Unit] = + subsResponses.take.flatMap(_.succeed(msg)).unit + + def releaseHub(key: SubscriptionKey): UIO[Unit] = for { subscription <- subsRef.getAndUpdate(_ - key) _ <- ZIO.foreachDiscard(subscription.get(key))(_.offer(Take.end)) - _ <- subsRef.update(_ - key) } yield () - def releasePromise(msg: PushMessage): UIO[Unit] = - responses.take.flatMap(_.succeed(msg)).unit - connection.read .mapError(RedisError.IOError(_)) .via(RespValue.Decoder) @@ -157,7 +162,7 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( .refineToOrDie[RedisError] .foreach { case msg: PushMessage.Subscribed => releasePromise(msg) - case msg: PushMessage.Unsubscribed => releasePromise(msg) *> releaseStream(msg.key) + case msg: PushMessage.Unsubscribed => offerMessage(msg) *> releaseHub(msg.key) case msg: PushMessage.Message => offerMessage(msg) } } @@ -172,7 +177,7 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( .asBytes for { - _ <- responses.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) + _ <- subsResponses.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) subsKeys <- subsRef.get.map(_.keys) (channels, patterns) = subsKeys.partition { case _: SubscriptionKey.Channel => false @@ -190,10 +195,16 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( } private[redis] object SingleNodeSubscriptionExecutor { - private final case class Request( - command: Chunk[RespValue.BulkString], - promises: Chunk[Promise[RedisError, PushMessage]] - ) + private sealed trait Request { + def command: Chunk[RespValue.BulkString] + } + + private object Request { + final case class Subscribe(command: Chunk[RespValue.BulkString], promises: Chunk[Promise[RedisError, PushMessage]]) + extends Request + final case class Unsubscribe(command: Chunk[RespValue.BulkString], promise: Promise[RedisError, Unit]) + extends Request + } def create(conn: RedisConnection): URIO[Scope, SubscriptionExecutor] = for { diff --git a/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala b/modules/redis/src/test/scala/zio/redis/PubSubSpec.scala similarity index 85% rename from modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala rename to modules/redis/src/test/scala/zio/redis/PubSubSpec.scala index fc5325da6..200b35ae2 100644 --- a/modules/redis/src/test/scala/zio/redis/internal/PubSubSpec.scala +++ b/modules/redis/src/test/scala/zio/redis/PubSubSpec.scala @@ -41,7 +41,7 @@ trait PubSubSpec extends BaseSpec { _ <- redis .pubSubChannels(channel) .repeatUntil(_ contains channel) - _ <- redis.publish(channel, message) + _ <- redis.publish(channel, message).replicateZIO(10) res <- fiber.join } yield assertTrue(res.get == message) }, @@ -84,7 +84,6 @@ trait PubSubSpec extends BaseSpec { None ) .returning[String] - .interruptWhen(promise) .runHead .fork res <- promise.await @@ -104,7 +103,7 @@ trait PubSubSpec extends BaseSpec { .runHead .fork _ <- redis.pubSubNumPat.repeatUntil(_ > 0) - _ <- redis.publish(channel, message) + _ <- redis.publish(channel, message).replicateZIO(10) res <- stream.join } yield assertTrue(res.get == message) } @@ -124,7 +123,7 @@ trait PubSubSpec extends BaseSpec { } .fork _ <- redis.pubSubChannels(channel).repeatUntil(_ contains channel) - _ <- ZIO.replicateZIO(10)(redis.publish(channel, message)) + _ <- ZIO.replicateZIO(50)(redis.publish(channel, message)) res <- stream.join } yield res )(equalTo(10L)) @@ -136,18 +135,19 @@ trait PubSubSpec extends BaseSpec { redis <- ZIO.service[Redis] subscription <- ZIO.service[RedisSubscription] prefix <- generateRandomString(5) - pattern = prefix + '*' channel <- generateRandomString(prefix) message <- generateRandomString() + subsPromise <- Promise.make[Nothing, Unit] promise <- Promise.make[Nothing, Unit] _ <- subscription - .subscribeWithCallback(channel)(None, Some((_, _) => promise.succeed(()).unit)) + .subscribeWithCallback(channel)( + Some((_, _) => subsPromise.succeed(()).unit), + Some((_, _) => promise.succeed(()).unit) + ) .returning[String] - .runCollect + .runDrain .fork - _ <- redis - .pubSubChannels(pattern) - .repeatUntil(_ contains channel) + _ <- subsPromise.await _ <- subscription.unsubscribe(channel) _ <- promise.await receiverCount <- redis.publish(channel, message).replicateZIO(numOfPublished).map(_.head) @@ -157,16 +157,19 @@ trait PubSubSpec extends BaseSpec { for { subscription <- ZIO.service[RedisSubscription] channel <- generateRandomString() + subsPromise <- Promise.make[RedisError, Unit] promise <- Promise.make[RedisError, String] - _ <- subscription - .subscribeWithCallback(channel)( - None, - Some((key, _) => promise.succeed(key).unit) - ) - .returning[Unit] - .runDrain - .fork + fiber <- subscription + .subscribeWithCallback(channel)( + Some((_, _) => subsPromise.succeed(()).unit), + Some((key, _) => promise.succeed(key).unit) + ) + .returning[Unit] + .runDrain + .fork + _ <- subsPromise.await _ <- subscription.unsubscribe(channel) + _ <- fiber.join res <- promise.await } yield assertTrue(res == channel) }, @@ -174,16 +177,19 @@ trait PubSubSpec extends BaseSpec { for { subscription <- ZIO.service[RedisSubscription] pattern <- generateRandomString() + subsPromise <- Promise.make[RedisError, Unit] promise <- Promise.make[RedisError, String] - _ <- subscription - .pSubscribeWithCallback(pattern)( - None, - Some((key, _) => promise.succeed(key).unit) - ) - .returning[Unit] - .runDrain - .fork + fiber <- subscription + .pSubscribeWithCallback(pattern)( + Some((_, _) => subsPromise.succeed(()).unit), + Some((key, _) => promise.succeed(key).unit) + ) + .returning[Unit] + .runDrain + .fork + _ <- subsPromise.await _ <- subscription.pUnsubscribe(pattern) + _ <- fiber.join res <- promise.await } yield assertTrue(res == pattern) }, From ac85cb9feea4c141f38c4ccf4649370d764c6abf Mon Sep 17 00:00:00 2001 From: opg1 Date: Wed, 7 Jun 2023 22:59:41 +0900 Subject: [PATCH 45/51] Update modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Dejan Mijić --- .../zio/redis/internal/SingleNodeSubscriptionExecutor.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala index b1606357b..81d338144 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala @@ -24,7 +24,7 @@ import zio.redis.internal.SingleNodeRunner.True import zio.redis.internal.SingleNodeSubscriptionExecutor.Request import zio.redis.{Input, RedisError} import zio.stream._ -import zio.{Chunk, ChunkBuilder, Hub, IO, Promise, Queue, Ref, Scope, UIO, URIO, ZIO} +import zio._ private[redis] final class SingleNodeSubscriptionExecutor private ( subsRef: Ref[Map[SubscriptionKey, Hub[Take[RedisError, PushMessage]]]], From dadb5e795b3c39105f106255e37c9d0bf7a04c7f Mon Sep 17 00:00:00 2001 From: 0pg Date: Wed, 7 Jun 2023 23:36:04 +0900 Subject: [PATCH 46/51] Use ConcurrentMap instead of Map of Ref --- build.sbt | 1 + .../scala/zio/redis/internal/PubSub.scala | 2 +- .../internal/RedisSubscriptionCommand.scala | 1 - .../zio/redis/internal/RequestQueue.scala | 24 --------- .../redis/internal/SingleNodeExecutor.scala | 2 +- .../zio/redis/internal/SingleNodeRunner.scala | 2 +- .../SingleNodeSubscriptionExecutor.scala | 51 ++++++++++--------- .../scala/zio/redis/internal/package.scala | 1 + 8 files changed, 32 insertions(+), 52 deletions(-) delete mode 100644 modules/redis/src/main/scala/zio/redis/internal/RequestQueue.scala diff --git a/build.sbt b/build.sbt index e90e2d55c..a9b0d9ca4 100644 --- a/build.sbt +++ b/build.sbt @@ -45,6 +45,7 @@ lazy val redis = .settings( libraryDependencies ++= List( "dev.zio" %% "zio-streams" % zioVersion, + "dev.zio" %% "zio-concurrent" % zioVersion, "dev.zio" %% "zio-schema" % zioSchemaVersion, "dev.zio" %% "zio-schema-protobuf" % zioSchemaVersion % Test, "dev.zio" %% "zio-test" % zioVersion % Test, diff --git a/modules/redis/src/main/scala/zio/redis/internal/PubSub.scala b/modules/redis/src/main/scala/zio/redis/internal/PubSub.scala index b6ac1921c..606eba9df 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/PubSub.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/PubSub.scala @@ -16,7 +16,7 @@ package zio.redis.internal -object PubSub { +private[redis] object PubSub { private[redis] sealed trait PushMessage { def key: SubscriptionKey } diff --git a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala index 45ef71a0e..8e1b65341 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala @@ -70,5 +70,4 @@ private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionE } .collectSome .refineToOrDie[RedisError] - } diff --git a/modules/redis/src/main/scala/zio/redis/internal/RequestQueue.scala b/modules/redis/src/main/scala/zio/redis/internal/RequestQueue.scala deleted file mode 100644 index 9f2a6def7..000000000 --- a/modules/redis/src/main/scala/zio/redis/internal/RequestQueue.scala +++ /dev/null @@ -1,24 +0,0 @@ -/* - * 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.internal - -import zio.{Queue, UIO} - -private[redis] object RequestQueue { - private final val RequestQueueSize = 16 - def create[A]: UIO[Queue[A]] = Queue.bounded[A](RequestQueueSize) -} diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala index 0bb043d51..dc8bd1f62 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala @@ -70,7 +70,7 @@ private[redis] object SingleNodeExecutor { def create(connection: RedisConnection): URIO[Scope, SingleNodeExecutor] = for { - requests <- RequestQueue.create[Request] + requests <- Queue.bounded[Request](RequestQueueSize) responses <- Queue.unbounded[Promise[RedisError, RespValue]] executor = new SingleNodeExecutor(connection, requests, responses) _ <- executor.run.forkScoped diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala index d60040eb4..eb69eaa93 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala @@ -31,7 +31,7 @@ private[redis] trait SingleNodeRunner { * 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] = + private[redis] final val run: IO[RedisError, AnyVal] = ZIO.logTrace(s"$this sender and reader has been started") *> (send.repeat(Schedule.forever) race receive) .tapError(e => ZIO.logWarning(s"Reconnecting due to error: $e") *> onError(e)) diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala index 81d338144..aa86709e9 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala @@ -16,6 +16,8 @@ package zio.redis.internal +import zio._ +import zio.concurrent.ConcurrentMap import zio.redis.Input.{CommandNameInput, StringInput} import zio.redis.Output.PushMessageOutput import zio.redis.api.Subscription @@ -24,10 +26,9 @@ import zio.redis.internal.SingleNodeRunner.True import zio.redis.internal.SingleNodeSubscriptionExecutor.Request import zio.redis.{Input, RedisError} import zio.stream._ -import zio._ private[redis] final class SingleNodeSubscriptionExecutor private ( - subsRef: Ref[Map[SubscriptionKey, Hub[Take[RedisError, PushMessage]]]], + subscriptionMap: ConcurrentMap[SubscriptionKey, Hub[Take[RedisError, PushMessage]]], requests: Queue[Request], subsResponses: Queue[Promise[RedisError, PushMessage]], connection: RedisConnection @@ -63,8 +64,8 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( command: RespCommand ): IO[RedisError, Stream[RedisError, PushMessage]] = { def getHub(key: SubscriptionKey): IO[RedisError, Hub[Take[RedisError, PushMessage]]] = - subsRef.get - .map(_.get(key)) + subscriptionMap + .get(key) .flatMap(ZIO.fromOption(_)) .orElseFail(RedisError.SubscriptionStreamAlreadyClosed(key)) @@ -79,15 +80,14 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( def makeHub(key: SubscriptionKey): UIO[Unit] = Hub .unbounded[Take[RedisError, PushMessage]] - .tap(hub => subsRef.update(_ + (key -> hub))) - .unlessZIO(subsRef.get.map(_.contains(key))) + .tap(subscriptionMap.putIfAbsent(key, _)) .unit for { - _ <- ZIO.foreachDiscard(keys)(makeHub(_)) - promises <- Promise.make[RedisError, PushMessage].replicateZIO(keys.size).map(Chunk.fromIterable(_)) + _ <- ZIO.foreachDiscard(keys)(makeHub) + promises <- Promise.make[RedisError, PushMessage].replicateZIO(keys.size).map(Chunk.fromIterable) _ <- requests.offer(Request.Subscribe(command.args.map(_.value), promises)) - streams = promises.map(getStream(_)) + streams = promises.map(getStream) } yield streams.fold(ZStream.empty)(_ merge _) } @@ -101,12 +101,12 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( private def extractChannelKeys(command: Chunk[RespCommandArgument]): Chunk[SubscriptionKey.Channel] = command.collect { case RespCommandArgument.Key(key) => key.asString - }.map(SubscriptionKey.Channel(_)) + }.map(SubscriptionKey.Channel.apply) private def extractPatternKeys(command: Chunk[RespCommandArgument]): Chunk[SubscriptionKey.Pattern] = command.collect { case RespCommandArgument.Key(key) => key.asString - }.map(SubscriptionKey.Pattern(_)) + }.map(SubscriptionKey.Pattern.apply) def send: IO[RedisError.IOError, Unit] = requests.takeAll.flatMap { reqs => @@ -122,7 +122,7 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( connection .write(bytes) - .mapError(RedisError.IOError(_)) + .mapError(RedisError.IOError.apply) .tapBoth( e => ZIO.foreachDiscard(reqs) { @@ -141,21 +141,24 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( def receive: IO[RedisError, Unit] = { def offerMessage(msg: PushMessage) = for { - subscription <- subsRef.get - _ <- ZIO.foreachDiscard(subscription.get(msg.key))(_.offer(Take.single(msg))) + hub <- subscriptionMap.get(msg.key) + _ <- ZIO.foreachDiscard(hub)(_.offer(Take.single(msg))) } yield () def releasePromise(msg: PushMessage): UIO[Unit] = - subsResponses.take.flatMap(_.succeed(msg)).unit + for { + promise <- subsResponses.take + _ <- promise.succeed(msg) + } yield () def releaseHub(key: SubscriptionKey): UIO[Unit] = for { - subscription <- subsRef.getAndUpdate(_ - key) - _ <- ZIO.foreachDiscard(subscription.get(key))(_.offer(Take.end)) + hub <- subscriptionMap.remove(key) + _ <- ZIO.foreachDiscard(hub)(_.offer(Take.end)) } yield () connection.read - .mapError(RedisError.IOError(_)) + .mapError(RedisError.IOError.apply) .via(RespValue.Decoder) .collectSome .mapZIO(resp => ZIO.attempt(PushMessageOutput.unsafeDecode(resp))) @@ -177,9 +180,9 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( .asBytes for { - _ <- subsResponses.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) - subsKeys <- subsRef.get.map(_.keys) - (channels, patterns) = subsKeys.partition { + _ <- subsResponses.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) + subscriptions <- subscriptionMap.toChunk + (channels, patterns) = subscriptions.map(_._1).partition { case _: SubscriptionKey.Channel => false case _: SubscriptionKey.Pattern => true } @@ -188,7 +191,7 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( _ <- connection .write(commands) .when(commands.nonEmpty) - .mapError(RedisError.IOError(_)) + .mapError(RedisError.IOError.apply) .retryWhile(True) } yield () } @@ -208,9 +211,9 @@ private[redis] object SingleNodeSubscriptionExecutor { def create(conn: RedisConnection): URIO[Scope, SubscriptionExecutor] = for { - reqQueue <- RequestQueue.create[Request] + reqQueue <- Queue.bounded[Request](RequestQueueSize) resQueue <- Queue.unbounded[Promise[RedisError, PushMessage]] - subsRef <- Ref.make(Map.empty[SubscriptionKey, Hub[Take[RedisError, PushMessage]]]) + subsRef <- ConcurrentMap.empty[SubscriptionKey, Hub[Take[RedisError, PushMessage]]] pubSub = new SingleNodeSubscriptionExecutor(subsRef, reqQueue, resQueue, conn) _ <- pubSub.run.forkScoped _ <- logScopeFinalizer(s"$pubSub Subscription Node is closed") diff --git a/modules/redis/src/main/scala/zio/redis/internal/package.scala b/modules/redis/src/main/scala/zio/redis/internal/package.scala index 16aad1fc7..06e3656a6 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/package.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/package.scala @@ -19,6 +19,7 @@ package zio.redis import zio._ package object internal { + private[redis] final val RequestQueueSize = 16 private[redis] def logScopeFinalizer(msg: String): URIO[Scope, Unit] = for { scope <- ZIO.scope From 88c3f4ed4e9fb43c8b2dcd42fdf856c825f6719b Mon Sep 17 00:00:00 2001 From: 0pg Date: Wed, 7 Jun 2023 23:40:55 +0900 Subject: [PATCH 47/51] Add private accessor --- .../main/scala/zio/redis/internal/SubscriptionExecutor.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala index 74400c651..fbe7f1157 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SubscriptionExecutor.scala @@ -21,11 +21,11 @@ import zio.redis.{RedisConfig, RedisError} import zio.stream._ import zio.{Layer, ZIO, ZLayer} -trait SubscriptionExecutor { +private[redis] trait SubscriptionExecutor { private[redis] def execute(command: RespCommand): Stream[RedisError, PushMessage] } -object SubscriptionExecutor { +private[redis] object SubscriptionExecutor { lazy val layer: ZLayer[RedisConfig, RedisError.IOError, SubscriptionExecutor] = RedisConnection.layer.fresh >>> pubSublayer From 6e50d4d471b4d82980e619a2d890cbecd544ce10 Mon Sep 17 00:00:00 2001 From: 0pg Date: Wed, 7 Jun 2023 23:44:48 +0900 Subject: [PATCH 48/51] Fix accessor --- .../src/main/scala/zio/redis/internal/SingleNodeRunner.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala index eb69eaa93..258cea146 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala @@ -31,7 +31,7 @@ private[redis] trait SingleNodeRunner { * 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. */ - private[redis] final val run: IO[RedisError, AnyVal] = + protected final val run: IO[RedisError, AnyVal] = ZIO.logTrace(s"$this sender and reader has been started") *> (send.repeat(Schedule.forever) race receive) .tapError(e => ZIO.logWarning(s"Reconnecting due to error: $e") *> onError(e)) From 5ca98e7d763e1b85998623aeabf4935f8fe69405 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 27 Jun 2023 04:08:02 +0900 Subject: [PATCH 49/51] Add doc --- .../src/main/scala/zio/redis/Output.scala | 12 +- .../main/scala/zio/redis/api/Publishing.scala | 36 ++++- .../scala/zio/redis/api/Subscription.scala | 143 +++++++++++++----- .../internal/RedisSubscriptionCommand.scala | 20 +-- .../redis/internal/SingleNodeExecutor.scala | 2 +- .../zio/redis/internal/SingleNodeRunner.scala | 2 +- .../SingleNodeSubscriptionExecutor.scala | 2 +- .../scala/zio/redis/options/Publishing.scala | 21 --- .../src/main/scala/zio/redis/package.scala | 3 +- .../src/test/scala/zio/redis/PubSubSpec.scala | 56 +++---- 10 files changed, 182 insertions(+), 115 deletions(-) delete mode 100644 modules/redis/src/main/scala/zio/redis/options/Publishing.scala diff --git a/modules/redis/src/main/scala/zio/redis/Output.scala b/modules/redis/src/main/scala/zio/redis/Output.scala index ffb121f28..ba71b4419 100644 --- a/modules/redis/src/main/scala/zio/redis/Output.scala +++ b/modules/redis/src/main/scala/zio/redis/Output.scala @@ -673,15 +673,17 @@ object Output { } } - case object NumSubResponseOutput extends Output[Chunk[NumberOfSubscribers]] { - protected def tryDecode(respValue: RespValue): Chunk[NumberOfSubscribers] = + case object NumSubResponseOutput extends Output[Map[String, Long]] { + protected def tryDecode(respValue: RespValue): Map[String, Long] = respValue match { case RespValue.Array(values) => - Chunk.fromIterator(values.grouped(2).map { chunk => + val builder = Map.newBuilder[String, Long] + values.grouped(2).foreach { chunk => val channel = MultiStringOutput.unsafeDecode(chunk(0)) val numOfSubs = LongOutput.unsafeDecode(chunk(1)) - NumberOfSubscribers(channel, numOfSubs) - }) + builder += channel -> numOfSubs + } + builder.result() case other => throw ProtocolError(s"$other isn't an array") } } diff --git a/modules/redis/src/main/scala/zio/redis/api/Publishing.scala b/modules/redis/src/main/scala/zio/redis/api/Publishing.scala index f82d26165..112a171a5 100644 --- a/modules/redis/src/main/scala/zio/redis/api/Publishing.scala +++ b/modules/redis/src/main/scala/zio/redis/api/Publishing.scala @@ -26,22 +26,56 @@ import zio.{Chunk, IO} trait Publishing extends RedisEnvironment { import Publishing._ + /** + * Posts a message to the given channel. + * + * @param channel + * Target channel name for publishing messages. + * @param message + * The value of the message to be published to the channel. + * @return + * Returns the number of clients that received the message. + */ final def publish[A: Schema](channel: String, message: A): IO[RedisError, Long] = { val command = RedisCommand(Publish, Tuple2(StringInput, ArbitraryKeyInput[A]()), LongOutput, executor) command.run((channel, message)) } + /** + * Lists the currently active channel that has one or more subscribers (excluding clients subscribed to patterns). + * + * @param pattern + * Pattern to get matching channels. + * @return + * Returns a list of active channels matching the specified pattern. + */ final def pubSubChannels(pattern: String): IO[RedisError, Chunk[String]] = { val command = RedisCommand(PubSubChannels, StringInput, ChunkOutput(MultiStringOutput), executor) command.run(pattern) } + /** + * The number of unique patterns that are subscribed to by clients. + * + * @return + * Returns the number of patterns all the clients are subscribed to. + */ final def pubSubNumPat: IO[RedisError, Long] = { val command = RedisCommand(PubSubNumPat, NoInput, LongOutput, executor) command.run(()) } - final def pubSubNumSub(channel: String, channels: String*): IO[RedisError, Chunk[NumberOfSubscribers]] = { + /** + * The number of subscribers (exclusive of clients subscribed to patterns) for the specified channels. + * + * @param channel + * Channel name to get the number of subscribers. + * @param channels + * Channel names to get the number of subscribers. + * @return + * Returns a map of channel and number of subscribers for channel. + */ + final def pubSubNumSub(channel: String, channels: String*): IO[RedisError, Map[String, Long]] = { val command = RedisCommand(PubSubNumSub, NonEmptyList(StringInput), NumSubResponseOutput, executor) command.run((channel, channels.toList)) } diff --git a/modules/redis/src/main/scala/zio/redis/api/Subscription.scala b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala index 05a732a0e..c3d0aa3c7 100644 --- a/modules/redis/src/main/scala/zio/redis/api/Subscription.scala +++ b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala @@ -18,20 +18,39 @@ package zio.redis.api import zio.redis.ResultBuilder.ResultStreamBuilder1 import zio.redis._ -import zio.redis.api.Subscription.PubSubCallback +import zio.redis.api.Subscription.{NoopCallback, PubSubCallback} import zio.redis.internal._ import zio.schema.Schema import zio.stream.Stream -import zio.{Chunk, IO, UIO} +import zio.{Chunk, IO, UIO, ZIO} trait Subscription extends SubscribeEnvironment { - final def subscribe(channel: String): ResultStreamBuilder1[Id] = - subscribeWithCallback(channel)(None, None) - - final def subscribeWithCallback(channel: String)( - onSubscribe: Option[PubSubCallback], - onUnsubscribe: Option[PubSubCallback] + /** + * Subscribes the client to the specified channel. + * + * @param channel + * Channel name to subscribe to. + * @return + * Returns stream that only emits published value for the given channel. + */ + final def subscribeSingle(channel: String): ResultStreamBuilder1[Id] = + subscribeSingleWith(channel)() + + /** + * Subscribes the client to the specified channel with callbacks for Subscribe and Unsubscribe messages. + * + * @param channel + * Channel name to subscribe to. + * @param onSubscribe + * Callback for given channel name and the number of subscribers of channel from upstream when subscribed. + * @param onUnsubscribe + * Callback for given channel name and the number of subscribers of channel from upstream when unsubscribed. + * @return + */ + final def subscribeSingleWith(channel: String)( + onSubscribe: PubSubCallback = NoopCallback, + onUnsubscribe: PubSubCallback = NoopCallback ): ResultStreamBuilder1[Id] = new ResultStreamBuilder1[Id] { def returning[R: Schema]: Stream[RedisError, R] = @@ -44,15 +63,37 @@ trait Subscription extends SubscribeEnvironment { .map(_._2) } + /** + * Subscribes the client to the specified channels. + * @param channel + * Channel name to subscribe to. + * @param channels + * Channel names to subscribe to consecutively. + * @return + * Returns stream that emits published value keyed by that channel name for the given channels. + */ final def subscribe( channel: String, channels: String* ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = - subscribeWithCallback(channel, channels: _*)(None, None) - - final def subscribeWithCallback(channel: String, channels: String*)( - onSubscribe: Option[PubSubCallback], - onUnsubscribe: Option[PubSubCallback] + subscribeWith(channel, channels: _*)(NoopCallback, NoopCallback) + + /** + * Subscribes the client to the specified channels with callbacks for Subscribe and Unsubscribe messages. + * @param channel + * Channel name to subscribe to. + * @param channels + * Channel names to subscribe to consecutively. + * @param onSubscribe + * Callback for given channel name and the number of subscribers of channel from upstream when subscribed. + * @param onUnsubscribe + * Callback for given channel name and the number of subscribers of channel from upstream when unsubscribed. + * @return + * Returns stream that emits published value keyed by that channel name for the given channels. + */ + final def subscribeWith(channel: String, channels: String*)( + onSubscribe: PubSubCallback = NoopCallback, + onUnsubscribe: PubSubCallback = NoopCallback ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = new ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] { def returning[R: Schema]: Stream[RedisError, (String, R)] = @@ -63,38 +104,40 @@ trait Subscription extends SubscribeEnvironment { ) } - final def pSubscribe(pattern: String): ResultStreamBuilder1[Id] = - pSubscribeWithCallback(pattern)(None, None) - - final def pSubscribeWithCallback( - pattern: String - )( - onSubscribe: Option[PubSubCallback], - onUnsubscribe: Option[PubSubCallback] - ): ResultStreamBuilder1[Id] = - new ResultStreamBuilder1[Id] { - def returning[R: Schema]: Stream[RedisError, R] = - RedisSubscriptionCommand(executor) - .pSubscribe( - Chunk.single(pattern), - onSubscribe, - onUnsubscribe - ) - .map(_._2) - } - + /** + * Subscribe to messages published to all channels that match the given patterns. + * + * @param pattern + * Pattern to subscribing to matching channels. (for the more details about pattern format, see + * https://redis.io/commands/psubscribe) + * @param patterns + * Patterns to subscribing to matching channels consecutively. + * @return + * Returns stream that emits published value keyed by the matching channel name. + */ final def pSubscribe( pattern: String, patterns: String* ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = - pSubscribeWithCallback(pattern, patterns: _*)(None, None) - - final def pSubscribeWithCallback( - pattern: String, - patterns: String* - )( - onSubscribe: Option[PubSubCallback], - onUnsubscribe: Option[PubSubCallback] + pSubscribeWith(pattern, patterns: _*)(NoopCallback, NoopCallback) + + /** + * Subscribe to messages published to all channels that match the given pattern with callbacks for PSubscribe and + * PUnsubscribe messages. + * @param pattern + * Pattern to subscribing to matching channels. (for the more details about pattern format, see + * https://redis.io/commands/psubscribe) + * @param patterns + * Patterns to subscribing to matching channels consecutively. + * @param onSubscribe + * Callback for given pattern and the number of subscribers of pattern from upstream when subscribed. + * @param onUnsubscribe + * Callback for given pattern and the number of subscribers of pattern from upstream when unsubscribed. + * @return + */ + final def pSubscribeWith(pattern: String, patterns: String*)( + onSubscribe: PubSubCallback = NoopCallback, + onUnsubscribe: PubSubCallback = NoopCallback ): ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] = new ResultStreamBuilder1[({ type lambda[x] = (String, x) })#lambda] { def returning[R: Schema]: Stream[RedisError, (String, R)] = @@ -105,9 +148,25 @@ trait Subscription extends SubscribeEnvironment { ) } + /** + * Unsubscribes the client from the given channels, or from all of them if none is given. + * @param channels + * Channels to unsubscribe from. if this is empty, all channels that this client is subscribed to will be + * unsubscribed. + * @return + * Returns unit if successful, error otherwise. + */ final def unsubscribe(channels: String*): IO[RedisError, Unit] = RedisSubscriptionCommand(executor).unsubscribe(Chunk.fromIterable(channels)) + /** + * Unsubscribes the client from the given patterns, or from all of them if none is given. + * @param patterns + * Patterns to unsubscribe from. if this is empty, all patterns that this client is subscribed to will be + * unsubscribed. + * @return + * Returns unit if successful, error otherwise. + */ final def pUnsubscribe(patterns: String*): IO[RedisError, Unit] = RedisSubscriptionCommand(executor).pUnsubscribe(Chunk.fromIterable(patterns)) @@ -120,4 +179,6 @@ object Subscription { final val Unsubscribe = "UNSUBSCRIBE" final val PSubscribe = "PSUBSCRIBE" final val PUnsubscribe = "PUNSUBSCRIBE" + + private final val NoopCallback: PubSubCallback = (_, _) => ZIO.unit } diff --git a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala index 8e1b65341..83b78e8f7 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/RedisSubscriptionCommand.scala @@ -29,23 +29,23 @@ import zio.{Chunk, IO, ZIO} private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionExecutor) { def subscribe[A: BinaryCodec]( channels: Chunk[String], - onSubscribe: Option[PubSubCallback], - onUnsubscribe: Option[PubSubCallback] + onSubscribe: PubSubCallback, + onUnsubscribe: PubSubCallback )(implicit codec: BinaryCodec[String]): Stream[RedisError, (String, A)] = executeCommand[A](makeCommand(Subscription.Subscribe, channels), onSubscribe, onUnsubscribe) def pSubscribe[A: BinaryCodec]( patterns: Chunk[String], - onSubscribe: Option[PubSubCallback], - onUnsubscribe: Option[PubSubCallback] + onSubscribe: PubSubCallback, + onUnsubscribe: PubSubCallback )(implicit stringCodec: BinaryCodec[String]): Stream[RedisError, (String, A)] = executeCommand[A](makeCommand(Subscription.PSubscribe, patterns), onSubscribe, onUnsubscribe) def unsubscribe(channels: Chunk[String])(implicit codec: BinaryCodec[String]): IO[RedisError, Unit] = - executeCommand(makeCommand(Subscription.Unsubscribe, channels), None, None).runDrain + executor.execute(makeCommand(Subscription.Unsubscribe, channels)).runDrain def pUnsubscribe(patterns: Chunk[String])(implicit codec: BinaryCodec[String]): IO[RedisError, Unit] = - executeCommand(makeCommand(Subscription.PUnsubscribe, patterns), None, None).runDrain + executor.execute(makeCommand(Subscription.PUnsubscribe, patterns)).runDrain private def makeCommand(commandName: String, keys: Chunk[String])(implicit codec: BinaryCodec[String]) = CommandNameInput.encode(commandName) ++ @@ -53,16 +53,16 @@ private[redis] final case class RedisSubscriptionCommand(executor: SubscriptionE private def executeCommand[A: BinaryCodec]( command: RespCommand, - onSubscribe: Option[PubSubCallback], - onUnsubscribe: Option[PubSubCallback] + onSubscribe: PubSubCallback, + onUnsubscribe: PubSubCallback ): Stream[RedisError, (String, A)] = executor .execute(command) .mapZIO { case Subscribed(key, numOfSubs) => - ZIO.foreach(onSubscribe)(_.apply(key.value, numOfSubs)).as(None) + onSubscribe(key.value, numOfSubs).as(None) case Unsubscribed(key, numOfSubs) => - ZIO.foreach(onUnsubscribe)(_.apply(key.value, numOfSubs)).as(None) + onUnsubscribe(key.value, numOfSubs).as(None) case Message(_, channel, message) => ZIO .attempt(ArbitraryOutput[A]().unsafeDecode(message)) diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala index dc8bd1f62..5febe70c4 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeExecutor.scala @@ -36,7 +36,7 @@ private[redis] final class SingleNodeExecutor private ( def onError(e: RedisError): UIO[Unit] = responses.takeAll.flatMap(ZIO.foreachDiscard(_)(_.fail(e))) def send: IO[RedisError.IOError, Unit] = - requests.takeAll.flatMap { requests => + requests.takeBetween(1, RequestQueueSize).flatMap { requests => val bytes = requests .foldLeft(new ChunkBuilder.Byte())((builder, req) => builder ++= RespValue.Array(req.command).asBytes) diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala index 258cea146..b2f67689e 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeRunner.scala @@ -31,7 +31,7 @@ private[redis] trait SingleNodeRunner { * 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. */ - protected final val run: IO[RedisError, AnyVal] = + private[internal] final val run: IO[RedisError, AnyVal] = ZIO.logTrace(s"$this sender and reader has been started") *> (send.repeat(Schedule.forever) race receive) .tapError(e => ZIO.logWarning(s"Reconnecting due to error: $e") *> onError(e)) diff --git a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala index aa86709e9..7fb4adc1d 100644 --- a/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala +++ b/modules/redis/src/main/scala/zio/redis/internal/SingleNodeSubscriptionExecutor.scala @@ -109,7 +109,7 @@ private[redis] final class SingleNodeSubscriptionExecutor private ( }.map(SubscriptionKey.Pattern.apply) def send: IO[RedisError.IOError, Unit] = - requests.takeAll.flatMap { reqs => + requests.takeBetween(1, RequestQueueSize).flatMap { reqs => val buffer = ChunkBuilder.make[Byte]() val it = reqs.iterator diff --git a/modules/redis/src/main/scala/zio/redis/options/Publishing.scala b/modules/redis/src/main/scala/zio/redis/options/Publishing.scala deleted file mode 100644 index 7a880dff5..000000000 --- a/modules/redis/src/main/scala/zio/redis/options/Publishing.scala +++ /dev/null @@ -1,21 +0,0 @@ -/* - * 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.options - -trait Publishing { - sealed case class NumberOfSubscribers(channel: String, subscriberCount: Long) -} diff --git a/modules/redis/src/main/scala/zio/redis/package.scala b/modules/redis/src/main/scala/zio/redis/package.scala index 7026485dd..cb1131fbe 100644 --- a/modules/redis/src/main/scala/zio/redis/package.scala +++ b/modules/redis/src/main/scala/zio/redis/package.scala @@ -25,8 +25,7 @@ package object redis with options.Strings with options.Lists with options.Streams - with options.Scripting - with options.Publishing { + with options.Scripting { type Id[+A] = A } diff --git a/modules/redis/src/test/scala/zio/redis/PubSubSpec.scala b/modules/redis/src/test/scala/zio/redis/PubSubSpec.scala index 200b35ae2..da3e46709 100644 --- a/modules/redis/src/test/scala/zio/redis/PubSubSpec.scala +++ b/modules/redis/src/test/scala/zio/redis/PubSubSpec.scala @@ -2,7 +2,7 @@ package zio.redis import zio.test.Assertion._ import zio.test._ -import zio.{Chunk, Promise, ZIO} +import zio.{Promise, ZIO} import scala.util.Random @@ -17,10 +17,7 @@ trait PubSubSpec extends BaseSpec { promise <- Promise.make[RedisError, String] resBuilder = subscription - .subscribeWithCallback(channel)( - Some((key: String, _: Long) => promise.succeed(key).unit), - None - ) + .subscribeSingleWith(channel)(onSubscribe = (key: String, _: Long) => promise.succeed(key).unit) stream = resBuilder.returning[String] _ <- stream .interruptWhen(promise) @@ -36,7 +33,7 @@ trait PubSubSpec extends BaseSpec { channel <- generateRandomString() message = "bar" promise <- Promise.make[RedisError, String] - stream = subscription.subscribe(channel).returning[String] + stream = subscription.subscribeSingle(channel).returning[String] fiber <- stream.interruptWhen(promise).runHead.fork _ <- redis .pubSubChannels(channel) @@ -56,10 +53,10 @@ trait PubSubSpec extends BaseSpec { pattern = prefix + '*' message <- generateRandomString(5) stream1 = subscription - .subscribe(channel1) + .subscribeSingle(channel1) .returning[String] stream2 = subscription - .subscribe(channel2) + .subscribeSingle(channel2) .returning[String] fiber1 <- stream1.runDrain.fork fiber2 <- stream2.runDrain.fork @@ -79,10 +76,7 @@ trait PubSubSpec extends BaseSpec { pattern <- generateRandomString() promise <- Promise.make[RedisError, String] _ <- subscription - .pSubscribeWithCallback(pattern)( - Some((key: String, _: Long) => promise.succeed(key).unit), - None - ) + .pSubscribeWith(pattern)(onSubscribe = (key: String, _: Long) => promise.succeed(key).unit) .returning[String] .runHead .fork @@ -105,7 +99,7 @@ trait PubSubSpec extends BaseSpec { _ <- redis.pubSubNumPat.repeatUntil(_ > 0) _ <- redis.publish(channel, message).replicateZIO(10) res <- stream.join - } yield assertTrue(res.get == message) + } yield assertTrue(res.map(_._2).get == message) } ), suite("publish")(test("publish long type message") { @@ -116,7 +110,7 @@ trait PubSubSpec extends BaseSpec { subscription <- ZIO.service[RedisSubscription] channel <- generateRandomString() stream <- subscription - .subscribe(channel) + .subscribeSingle(channel) .returning[Long] .runFoldWhile(0L)(_ < 10L) { case (sum, message) => sum + message @@ -140,9 +134,9 @@ trait PubSubSpec extends BaseSpec { subsPromise <- Promise.make[Nothing, Unit] promise <- Promise.make[Nothing, Unit] _ <- subscription - .subscribeWithCallback(channel)( - Some((_, _) => subsPromise.succeed(()).unit), - Some((_, _) => promise.succeed(()).unit) + .subscribeSingleWith(channel)( + onSubscribe = (_, _) => subsPromise.succeed(()).unit, + onUnsubscribe = (_, _) => promise.succeed(()).unit ) .returning[String] .runDrain @@ -160,9 +154,9 @@ trait PubSubSpec extends BaseSpec { subsPromise <- Promise.make[RedisError, Unit] promise <- Promise.make[RedisError, String] fiber <- subscription - .subscribeWithCallback(channel)( - Some((_, _) => subsPromise.succeed(()).unit), - Some((key, _) => promise.succeed(key).unit) + .subscribeSingleWith(channel)( + onSubscribe = (_, _) => subsPromise.succeed(()).unit, + onUnsubscribe = (key, _) => promise.succeed(key).unit ) .returning[Unit] .runDrain @@ -180,9 +174,9 @@ trait PubSubSpec extends BaseSpec { subsPromise <- Promise.make[RedisError, Unit] promise <- Promise.make[RedisError, String] fiber <- subscription - .pSubscribeWithCallback(pattern)( - Some((_, _) => subsPromise.succeed(()).unit), - Some((key, _) => promise.succeed(key).unit) + .pSubscribeWith(pattern)( + onSubscribe = (_, _) => subsPromise.succeed(()).unit, + onUnsubscribe = (key, _) => promise.succeed(key).unit ) .returning[Unit] .runDrain @@ -205,18 +199,16 @@ trait PubSubSpec extends BaseSpec { promise2 <- Promise.make[Nothing, Unit] _ <- subscription - .subscribeWithCallback(channel1)( - None, - Some((_, _) => promise1.succeed(()).unit) + .subscribeSingleWith(channel1)( + onUnsubscribe = (_, _) => promise1.succeed(()).unit ) .returning[String] .runCollect .fork _ <- subscription - .subscribeWithCallback(channel2)( - None, - Some((_, _) => promise2.succeed(()).unit) + .subscribeSingleWith(channel2)( + onUnsubscribe = (_, _) => promise2.succeed(()).unit ) .returning[String] .runCollect @@ -229,9 +221,9 @@ trait PubSubSpec extends BaseSpec { _ <- promise2.await numSubResponses <- redis.pubSubNumSub(channel1, channel2) } yield assertTrue( - numSubResponses == Chunk( - NumberOfSubscribers(channel1, 0L), - NumberOfSubscribers(channel2, 0L) + numSubResponses == Map( + channel1 -> 0L, + channel2 -> 0L ) ) } From 539ef3fcd9884af7c627e18909d275e0474de860 Mon Sep 17 00:00:00 2001 From: 0pg Date: Tue, 27 Jun 2023 04:19:38 +0900 Subject: [PATCH 50/51] Add missing params --- modules/redis/src/main/scala/zio/redis/api/Subscription.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/redis/src/main/scala/zio/redis/api/Subscription.scala b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala index c3d0aa3c7..d5162978e 100644 --- a/modules/redis/src/main/scala/zio/redis/api/Subscription.scala +++ b/modules/redis/src/main/scala/zio/redis/api/Subscription.scala @@ -35,7 +35,7 @@ trait Subscription extends SubscribeEnvironment { * Returns stream that only emits published value for the given channel. */ final def subscribeSingle(channel: String): ResultStreamBuilder1[Id] = - subscribeSingleWith(channel)() + subscribeSingleWith(channel)(NoopCallback, NoopCallback) /** * Subscribes the client to the specified channel with callbacks for Subscribe and Unsubscribe messages. From 63cc263a7cc9d98f2fb12f1924bff8d5bec63bcf Mon Sep 17 00:00:00 2001 From: 0pg Date: Fri, 7 Jul 2023 14:06:14 +0900 Subject: [PATCH 51/51] Fix test codes --- .../src/test/scala/zio/redis/PubSubSpec.scala | 77 +++++++++++-------- 1 file changed, 43 insertions(+), 34 deletions(-) diff --git a/modules/redis/src/test/scala/zio/redis/PubSubSpec.scala b/modules/redis/src/test/scala/zio/redis/PubSubSpec.scala index da3e46709..045ccc1bb 100644 --- a/modules/redis/src/test/scala/zio/redis/PubSubSpec.scala +++ b/modules/redis/src/test/scala/zio/redis/PubSubSpec.scala @@ -32,43 +32,48 @@ trait PubSubSpec extends BaseSpec { subscription <- ZIO.service[RedisSubscription] channel <- generateRandomString() message = "bar" - promise <- Promise.make[RedisError, String] - stream = subscription.subscribeSingle(channel).returning[String] - fiber <- stream.interruptWhen(promise).runHead.fork - _ <- redis - .pubSubChannels(channel) - .repeatUntil(_ contains channel) - _ <- redis.publish(channel, message).replicateZIO(10) + promise <- Promise.make[RedisError, Unit] + fiber <- subscription + .subscribeSingleWith(channel)( + onSubscribe = (_, _) => promise.succeed(()).unit + ) + .returning[String] + .runHead + .fork + _ <- promise.await + _ <- redis.publish(channel, message) res <- fiber.join } yield assertTrue(res.get == message) }, test("multiple subscribe") { - val numOfPublish = 20 for { redis <- ZIO.service[Redis] subscription <- ZIO.service[RedisSubscription] prefix <- generateRandomString(5) channel1 <- generateRandomString(prefix) channel2 <- generateRandomString(prefix) - pattern = prefix + '*' message <- generateRandomString(5) + subsPromise1 <- Promise.make[RedisError, Unit] + subsPromise2 <- Promise.make[RedisError, Unit] stream1 = subscription - .subscribeSingle(channel1) + .subscribeSingleWith(channel1)( + onSubscribe = (_, _) => subsPromise1.succeed(()).unit + ) .returning[String] stream2 = subscription - .subscribeSingle(channel2) + .subscribeSingleWith(channel2)( + onSubscribe = (_, _) => subsPromise2.succeed(()).unit + ) .returning[String] - fiber1 <- stream1.runDrain.fork - fiber2 <- stream2.runDrain.fork - _ <- redis - .pubSubChannels(pattern) - .repeatUntil(channels => channels.size >= 2) - ch1SubsCount <- redis.publish(channel1, message).replicateZIO(numOfPublish).map(_.head) - ch2SubsCount <- redis.publish(channel2, message).replicateZIO(numOfPublish).map(_.head) + fiber1 <- stream1.runDrain.fork + fiber2 <- stream2.runDrain.fork + _ <- subsPromise1.await *> subsPromise2.await + ch1SubsCount <- redis.publish(channel1, message) + ch2SubsCount <- redis.publish(channel2, message) _ <- subscription.unsubscribe() _ <- fiber1.join _ <- fiber2.join - } yield assertTrue(ch1SubsCount == 1L) && assertTrue(ch2SubsCount == 1L) + } yield assertTrue(ch1SubsCount == 1L, ch2SubsCount == 1L) }, test("psubscribe response") { for { @@ -91,13 +96,16 @@ trait PubSubSpec extends BaseSpec { pattern = prefix + '*' channel <- generateRandomString(prefix) message <- generateRandomString(prefix) + promise <- Promise.make[RedisError, Unit] stream <- subscription - .pSubscribe(pattern) + .pSubscribeWith(pattern)( + onSubscribe = (_, _) => promise.succeed(()).unit + ) .returning[String] .runHead .fork - _ <- redis.pubSubNumPat.repeatUntil(_ > 0) - _ <- redis.publish(channel, message).replicateZIO(10) + _ <- promise.await + _ <- redis.publish(channel, message) res <- stream.join } yield assertTrue(res.map(_._2).get == message) } @@ -109,28 +117,29 @@ trait PubSubSpec extends BaseSpec { redis <- ZIO.service[Redis] subscription <- ZIO.service[RedisSubscription] channel <- generateRandomString() + promise <- Promise.make[RedisError, Unit] stream <- subscription - .subscribeSingle(channel) + .subscribeSingleWith(channel)( + onSubscribe = (_, _) => promise.succeed(()).unit + ) .returning[Long] .runFoldWhile(0L)(_ < 10L) { case (sum, message) => sum + message } .fork - _ <- redis.pubSubChannels(channel).repeatUntil(_ contains channel) - _ <- ZIO.replicateZIO(50)(redis.publish(channel, message)) + _ <- promise.await + _ <- ZIO.replicateZIO(10)(redis.publish(channel, message)) res <- stream.join } yield res )(equalTo(10L)) }), suite("unsubscribe")( test("don't receive message type after unsubscribe") { - val numOfPublished = 5 for { redis <- ZIO.service[Redis] subscription <- ZIO.service[RedisSubscription] prefix <- generateRandomString(5) channel <- generateRandomString(prefix) - message <- generateRandomString() subsPromise <- Promise.make[Nothing, Unit] promise <- Promise.make[Nothing, Unit] _ <- subscription @@ -144,7 +153,7 @@ trait PubSubSpec extends BaseSpec { _ <- subsPromise.await _ <- subscription.unsubscribe(channel) _ <- promise.await - receiverCount <- redis.publish(channel, message).replicateZIO(numOfPublished).map(_.head) + receiverCount <- redis.pubSubNumSub(channel).map(_.getOrElse(channel, 0L)) } yield assertTrue(receiverCount == 0L) }, test("unsubscribe response") { @@ -192,14 +201,16 @@ trait PubSubSpec extends BaseSpec { redis <- ZIO.service[Redis] subscription <- ZIO.service[RedisSubscription] prefix <- generateRandomString(5) - pattern = prefix + '*' channel1 <- generateRandomString(prefix) channel2 <- generateRandomString(prefix) + subsPromise1 <- Promise.make[Nothing, Unit] + subsPromise2 <- Promise.make[Nothing, Unit] promise1 <- Promise.make[Nothing, Unit] promise2 <- Promise.make[Nothing, Unit] _ <- subscription .subscribeSingleWith(channel1)( + onSubscribe = (_, _) => subsPromise1.succeed(()).unit, onUnsubscribe = (_, _) => promise1.succeed(()).unit ) .returning[String] @@ -208,17 +219,15 @@ trait PubSubSpec extends BaseSpec { _ <- subscription .subscribeSingleWith(channel2)( + onSubscribe = (_, _) => subsPromise2.succeed(()).unit, onUnsubscribe = (_, _) => promise2.succeed(()).unit ) .returning[String] .runCollect .fork - _ <- redis - .pubSubChannels(pattern) - .repeatUntil(_.size >= 2) + _ <- subsPromise1.await *> subsPromise2.await _ <- subscription.unsubscribe() - _ <- promise1.await - _ <- promise2.await + _ <- promise1.await *> promise2.await numSubResponses <- redis.pubSubNumSub(channel1, channel2) } yield assertTrue( numSubResponses == Map(