diff --git a/.scalafmt.conf b/.scalafmt.conf index 201abc09..fd7c2805 100644 --- a/.scalafmt.conf +++ b/.scalafmt.conf @@ -16,3 +16,4 @@ spaces { optIn.annotationNewlines = true rewrite.rules = [SortImports, RedundantBraces] + diff --git a/benchmarks/src/main/scala/com/devsisters/shardcake/Client.scala b/benchmarks/src/main/scala/com/devsisters/shardcake/Client.scala index fe432b0d..a143c17d 100644 --- a/benchmarks/src/main/scala/com/devsisters/shardcake/Client.scala +++ b/benchmarks/src/main/scala/com/devsisters/shardcake/Client.scala @@ -1,6 +1,6 @@ package com.devsisters.shardcake -import com.devsisters.shardcake.Server.Message.Ping +import com.devsisters.shardcake.Server.Message.{ Ping, StreamPing } import com.devsisters.shardcake.Server.PingPongEntity import zio.{ Config => _, _ } @@ -17,4 +17,20 @@ object Client { } yield () ) .provide(config, Server.sharding) + + def sendStream(streams: Int, messagesPerStream: Int, parallelism: Int): Task[Unit] = + ZIO + .scoped[Sharding]( + for { + ping <- Sharding.messenger(PingPongEntity) + _ <- ZIO + .foreachParDiscard(1 to streams) { _ => + ping + .sendAndReceiveStream("ping")(StreamPing("ping", messagesPerStream, _)) + .flatMap(_.runDrain) + } + .withParallelism(parallelism) + } yield () + ) + .provide(config, Server.sharding) } diff --git a/benchmarks/src/main/scala/com/devsisters/shardcake/SendBenchmark.scala b/benchmarks/src/main/scala/com/devsisters/shardcake/SendBenchmark.scala index 1385e6bf..494d7055 100644 --- a/benchmarks/src/main/scala/com/devsisters/shardcake/SendBenchmark.scala +++ b/benchmarks/src/main/scala/com/devsisters/shardcake/SendBenchmark.scala @@ -8,8 +8,8 @@ import java.util.concurrent.TimeUnit @State(Scope.Thread) @BenchmarkMode(Array(Mode.Throughput)) @OutputTimeUnit(TimeUnit.SECONDS) -@Warmup(iterations = 5, time = 5, timeUnit = TimeUnit.SECONDS) -@Measurement(iterations = 5, time = 5, timeUnit = TimeUnit.SECONDS) +@Warmup(iterations = 3, time = 3, timeUnit = TimeUnit.SECONDS) +@Measurement(iterations = 3, time = 3, timeUnit = TimeUnit.SECONDS) @Fork(1) class SendBenchmark { private var fiber: Fiber[Any, Any] = _ diff --git a/benchmarks/src/main/scala/com/devsisters/shardcake/SendStreamBenchmark.scala b/benchmarks/src/main/scala/com/devsisters/shardcake/SendStreamBenchmark.scala new file mode 100644 index 00000000..8948b815 --- /dev/null +++ b/benchmarks/src/main/scala/com/devsisters/shardcake/SendStreamBenchmark.scala @@ -0,0 +1,31 @@ +package com.devsisters.shardcake + +import org.openjdk.jmh.annotations._ +import zio.{ durationInt, Fiber, Runtime, Unsafe, ZIO } + +import java.util.concurrent.TimeUnit + +@State(Scope.Thread) +@BenchmarkMode(Array(Mode.Throughput)) +@OutputTimeUnit(TimeUnit.SECONDS) +@Warmup(iterations = 3, time = 3, timeUnit = TimeUnit.SECONDS) +@Measurement(iterations = 3, time = 3, timeUnit = TimeUnit.SECONDS) +@Fork(1) +class SendStreamBenchmark { + private var fiber: Fiber[Any, Any] = _ + + @Setup + def setup(): Unit = + fiber = Unsafe.unsafe(implicit unsafe => + Runtime.default.unsafe.run(Server.run.forkDaemon <* ZIO.sleep(3.seconds)).getOrThrow() + ) + + @TearDown + def tearDown(): Unit = + Unsafe.unsafe(implicit unsafe => Runtime.default.unsafe.run(fiber.interrupt)) + + // 8 parallel server-streams, each receiving 100 messages → 800 messages per op + @Benchmark + def serverStream(): Unit = + Unsafe.unsafe(implicit unsafe => Runtime.default.unsafe.run(Client.sendStream(8, 100, 8))) +} diff --git a/benchmarks/src/main/scala/com/devsisters/shardcake/Server.scala b/benchmarks/src/main/scala/com/devsisters/shardcake/Server.scala index 3c985f34..78ec9d0d 100644 --- a/benchmarks/src/main/scala/com/devsisters/shardcake/Server.scala +++ b/benchmarks/src/main/scala/com/devsisters/shardcake/Server.scala @@ -10,13 +10,18 @@ object Server { sealed trait Message object Message { - case class Ping(msg: String, replier: Replier[String]) extends Message + case class Ping(msg: String, replier: Replier[String]) extends Message + case class StreamPing(msg: String, count: Int, replier: StreamReplier[String]) extends Message } object PingPongEntity extends EntityType[Message]("ping-pong") private def behavior(entityId: String, messages: Dequeue[Message]): RIO[Sharding, Nothing] = - messages.take.flatMap { case Message.Ping(msg, replier) => replier.reply(msg) }.forever + messages.take.flatMap { + case Message.Ping(msg, replier) => replier.reply(msg) + case Message.StreamPing(msg, count, replier) => + replier.replyStream(zio.stream.ZStream.repeat(msg).take(count.toLong)) + }.forever private val shardManagerClient: ZLayer[Config, Nothing, ShardManagerClient] = ZLayer { diff --git a/build.sbt b/build.sbt index a4e5dc0c..cec52626 100644 --- a/build.sbt +++ b/build.sbt @@ -1,7 +1,7 @@ val scala3 = "3.3.7" val zioVersion = "2.1.24" -val zioGrpcVersion = "0.6.3" +val proteusVersion = "0.4.1" val grpcNettyVersion = "1.71.0" val zioK8sVersion = "3.2.0" val zioK8sSttpVersion = "3.11.0" @@ -41,7 +41,7 @@ inThisBuild( name := "shardcake" addCommandAlias("fmt", "all scalafmtSbt scalafmt test:scalafmt") -addCommandAlias("check", "all scalafmtSbtCheck scalafmtCheck test:scalafmtCheck") +addCommandAlias("check", "all scalafmtSbtCheck scalafmtCheck test:scalafmtCheck grpcProtocol/checkProto") lazy val root = project .in(file(".")) @@ -151,25 +151,47 @@ lazy val serializationKryo = project ) ) +lazy val generateProto = taskKey[Unit]("Regenerate sharding.proto from the Scala protocol definition.") +lazy val checkProto = taskKey[Unit]("Fail if sharding.proto is out of sync with the Scala protocol definition.") + lazy val grpcProtocol = project .in(file("protocol-grpc")) .settings(name := "shardcake-protocol-grpc") .settings(commonSettings) - .settings(protobuf: _*) - .settings( - Compile / PB.targets := Seq( - scalapb.gen(grpc = true) -> (Compile / sourceManaged).value, - scalapb.zio_grpc.ZioCodeGenerator -> (Compile / sourceManaged).value - ) - ) .dependsOn(core, entities) .settings( libraryDependencies ++= Seq( - "com.thesamet.scalapb" %% "scalapb-runtime" % scalapb.compiler.Version.scalapbVersion % "protobuf", - "com.thesamet.scalapb" %% "scalapb-runtime-grpc" % scalapb.compiler.Version.scalapbVersion, - "com.thesamet.scalapb.zio-grpc" %% "zio-grpc-core" % zioGrpcVersion, - "io.grpc" % "grpc-netty" % grpcNettyVersion - ) + "com.github.ghostdogpr" %% "proteus-grpc" % proteusVersion, + "com.github.ghostdogpr" %% "proteus-grpc-zio" % proteusVersion, + "io.grpc" % "grpc-netty" % grpcNettyVersion, + "io.grpc" % "grpc-services" % grpcNettyVersion + ), + generateProto := { + val cp = (Compile / fullClasspath).value + val log = streams.value.log + val output = (Compile / sourceDirectory).value / "protobuf" + runner.value + .run( + "com.devsisters.shardcake.protocol.GenerateProto", + cp.files, + Seq(output.getAbsolutePath), + log + ) + .get + log.info(s"Regenerated $output/sharding.proto") + }, + checkProto := { + val log = streams.value.log + val proto = (Compile / sourceDirectory).value / "protobuf" / "sharding.proto" + val _ = generateProto.value + import scala.sys.process._ + val diff = s"git diff --exit-code -- ${proto.getAbsolutePath}".! + if (diff != 0) { + sys.error( + "sharding.proto is out of sync with the Scala protocol definition. Run `sbt grpcProtocol/generateProto` and commit the result." + ) + } else log.info("sharding.proto is in sync.") + } ) lazy val examples = project @@ -194,10 +216,6 @@ lazy val benchmarks = project .enablePlugins(JmhPlugin) .dependsOn(grpcProtocol, serializationKryo) -lazy val protobuf = Seq( - PB.protocVersion := "3.19.2" -) ++ Project.inConfig(Test)(sbtprotoc.ProtocPlugin.protobufConfigSettings) - lazy val commonSettings = Def.settings( testFrameworks := Seq(new TestFramework("zio.test.sbt.ZTestFramework")), libraryDependencies ++= diff --git a/examples/src/test/scala/example/GrpcAuthExampleSpec.scala b/examples/src/test/scala/example/GrpcAuthExampleSpec.scala index fd286644..9d25520d 100644 --- a/examples/src/test/scala/example/GrpcAuthExampleSpec.scala +++ b/examples/src/test/scala/example/GrpcAuthExampleSpec.scala @@ -3,7 +3,6 @@ package example import com.devsisters.shardcake._ import com.devsisters.shardcake.interfaces.{ Pods, Storage } import io.grpc.{ Metadata, Status } -import scalapb.zio_grpc.{ ZClientInterceptor, ZTransform } import zio.test._ import zio.{ Config => _, _ } @@ -11,7 +10,7 @@ object GrpcAuthExampleSpec extends ZIOSpecDefault { private val validAuthenticationKey = "validAuthenticationKey" - private val authKey = Metadata.Key.of("authentication-key", io.grpc.Metadata.ASCII_STRING_MARSHALLER) + private val authKey = Metadata.Key.of("authentication-key", Metadata.ASCII_STRING_MARSHALLER) private val config = ZLayer.succeed(Config.default.copy(simulateRemotePods = true)) @@ -19,14 +18,12 @@ object GrpcAuthExampleSpec extends ZIOSpecDefault { ZLayer.succeed( GrpcConfig.default.copy( clientInterceptors = Seq( - ZClientInterceptor.headersUpdater((_, _, md) => md.put(authKey, clientAuthKey).unit) + ShardingClientInterceptor.headersUpdater(_.put(authKey, clientAuthKey)) ), serverInterceptors = Seq( - ZTransform { requestContext => - for { - authenticated <- requestContext.metadata.get(authKey).map(_.contains(validAuthenticationKey)) - _ <- ZIO.when(!authenticated)(ZIO.fail(Status.UNAUTHENTICATED.asException)) - } yield requestContext + ShardingServerInterceptor.beforeEach { ctx => + val authenticated = Option(ctx.requestMetadata.get(authKey)).contains(validAuthenticationKey) + ZIO.unless(authenticated)(ZIO.fail(Status.UNAUTHENTICATED.asException())).unit } ) ) diff --git a/project/plugins.sbt b/project/plugins.sbt index 3fa05fc4..eb982250 100644 --- a/project/plugins.sbt +++ b/project/plugins.sbt @@ -1,7 +1,3 @@ addSbtPlugin("org.scalameta" % "sbt-scalafmt" % "2.4.6") addSbtPlugin("com.github.sbt" % "sbt-ci-release" % "1.11.1") -addSbtPlugin("com.thesamet" % "sbt-protoc" % "1.0.7") addSbtPlugin("pl.project13.scala" % "sbt-jmh" % "0.4.7") - -libraryDependencies += "com.thesamet.scalapb" %% "compilerplugin" % "0.11.17" -libraryDependencies += "com.thesamet.scalapb.zio-grpc" %% "zio-grpc-codegen" % "0.6.3" diff --git a/protocol-grpc/src/main/protobuf/sharding.proto b/protocol-grpc/src/main/protobuf/sharding.proto index 2037cbb1..bb991567 100644 --- a/protocol-grpc/src/main/protobuf/sharding.proto +++ b/protocol-grpc/src/main/protobuf/sharding.proto @@ -1,7 +1,5 @@ syntax = "proto3"; -option java_package = "com.devsisters.shardcake.protobuf"; - service ShardingService { rpc AssignShards (AssignShardsRequest) returns (AssignShardsResponse) {} rpc UnassignShards (UnassignShardsRequest) returns (UnassignShardsResponse) {} diff --git a/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcConfig.scala b/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcConfig.scala index 8b0bccb0..3d448c34 100644 --- a/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcConfig.scala +++ b/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcConfig.scala @@ -1,30 +1,37 @@ package com.devsisters.shardcake +import io.grpc.ClientInterceptor import zio._ -import scalapb.zio_grpc.RequestContext -import scalapb.zio_grpc.ZClientInterceptor -import scalapb.zio_grpc.ZTransform import java.util.concurrent.Executor /** - * The configuration for the gRPC client. + * The configuration for the gRPC client and server. * - * @param maxInboundMessageSize the maximum message size allowed to be received by the grpc client + * @param maxInboundMessageSize the maximum message size allowed to be received by the grpc client and server * @param executor a custom executor to pass to grpc-java when creating gRPC clients and servers * @param shutdownTimeout the timeout to wait for the gRPC server to shutdown before forcefully shutting it down - * @param clientInterceptors the interceptors to be used by the gRPC client, e.g for adding tracing or logging - * @param serverInterceptors the interceptors to be used by the gRPC Server, e.g for adding tracing or logging + * @param streamingPrefetch the in-flight window for streaming RPCs (request/response messages fetched ahead of the consumer) + * @param clientInterceptors the interceptors to be used by the gRPC client, e.g. for adding tracing or logging + * @param serverInterceptors the interceptors to be used by the gRPC server, e.g. for adding tracing or logging */ case class GrpcConfig( maxInboundMessageSize: Int, executor: Option[Executor], shutdownTimeout: Duration, - clientInterceptors: Seq[ZClientInterceptor], - serverInterceptors: Seq[ZTransform[RequestContext, RequestContext]] + streamingPrefetch: Int, + clientInterceptors: Seq[ClientInterceptor], + serverInterceptors: Seq[ShardingServerInterceptor] ) object GrpcConfig { val default: GrpcConfig = - GrpcConfig(maxInboundMessageSize = 32 * 1024 * 1024, None, 3.seconds, Seq.empty, Seq.empty) + GrpcConfig( + maxInboundMessageSize = 32 * 1024 * 1024, + executor = None, + shutdownTimeout = 3.seconds, + streamingPrefetch = 16, + clientInterceptors = Seq.empty, + serverInterceptors = Seq.empty + ) } diff --git a/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcPods.scala b/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcPods.scala index 1f6eca84..ebdb0c58 100644 --- a/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcPods.scala +++ b/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcPods.scala @@ -3,28 +3,29 @@ package com.devsisters.shardcake import com.devsisters.shardcake.errors._ import com.devsisters.shardcake.interfaces.Pods import com.devsisters.shardcake.interfaces.Pods.BinaryMessage -import com.devsisters.shardcake.protobuf.sharding._ -import com.devsisters.shardcake.protobuf.sharding.ZioSharding.ShardingServiceClient -import com.google.protobuf.ByteString -import io.grpc.{ ManagedChannelBuilder, Status } -import scalapb.zio_grpc.ZManagedChannel +import com.devsisters.shardcake.protocol.Sharding._ +import io.grpc.{ ClientInterceptors, ManagedChannel, ManagedChannelBuilder, Status, StatusException } +import proteus.client.ZioClientBackend import zio._ import zio.stream.ZStream +import java.util.concurrent.TimeUnit + class GrpcPods( config: GrpcConfig, - connections: Ref.Synchronized[Map[PodAddress, (ShardingServiceClient, Fiber[Throwable, Nothing])]] + connections: Ref.Synchronized[Map[PodAddress, (PodClient, Fiber[Throwable, Nothing])]] ) extends Pods { - private def getConnection(pod: PodAddress): Task[ShardingServiceClient] = + + private def getConnection(pod: PodAddress): Task[PodClient] = // optimize happy path and only get first connections.get.flatMap(_.get(pod) match { - case Some((channel, _)) => ZIO.succeed(channel) - case None => + case Some((client, _)) => ZIO.succeed(client) + case None => // then do modify in the case it doesn't already exist connections.modifyZIO { map => map.get(pod) match { - case Some((channel, _)) => ZIO.succeed((channel, map)) - case None => + case Some((client, _)) => ZIO.succeed((client, map)) + case None => val builder = { config.executor match { case Some(executor) => @@ -41,94 +42,94 @@ class GrpcPods( } } - val channel = ZManagedChannel(builder, config.clientInterceptors) + val acquireChannel: RIO[Scope, ManagedChannel] = + ZIO.acquireRelease(ZIO.attempt(builder.build())) { channel => + ZIO.attemptBlocking { + channel.shutdown() + channel.awaitTermination(config.shutdownTimeout.toMillis, TimeUnit.MILLISECONDS) + channel.shutdownNow(): Unit + }.ignore + } + // create a fiber that never ends and keeps the connection alive for { - _ <- ZIO.logDebug(s"Opening connection to pod $pod") - promise <- Promise.make[Nothing, ShardingServiceClient] - fiber <- - ZIO - .scoped( - ShardingServiceClient - .scoped(channel) - .flatMap(promise.succeed(_) *> ZIO.never) - .ensuring(connections.update(_ - pod) *> ZIO.logDebug(s"Closed connection to pod $pod")) - ) - .forkDaemon - connection <- promise.await - } yield (connection, map.updated(pod, (connection, fiber))) + _ <- ZIO.logDebug(s"Opening connection to pod $pod") + promise <- Promise.make[Throwable, PodClient] + fiber <- ZIO + .scoped[Any] { + acquireChannel.flatMap { rawChannel => + val channel = + if (config.clientInterceptors.isEmpty) rawChannel + else ClientInterceptors.intercept(rawChannel, config.clientInterceptors*) + val backend = ZioClientBackend(channel, config.streamingPrefetch) + val client = PodClient.from(backend) + promise.succeed(client) *> ZIO.never + } + } + .onError(promise.failCause(_)) + .ensuring( + connections.update(_ - pod) *> ZIO.logDebug(s"Closed connection to pod $pod") + ) + .forkDaemon + client <- promise.await + } yield (client, map.updated(pod, (client, fiber))) } } }) + private def mapClientError(pod: PodAddress, entityId: String, isStream: Boolean)(t: Throwable): Throwable = + t match { + case ex: StatusException if ex.getStatus.getCode == Status.Code.RESOURCE_EXHAUSTED => + // entity is not managed by this pod, wait and retry (assignments will be updated) + EntityNotManagedByThisPod(entityId) + case ex: StatusException if ex.getStatus.getCode == Status.Code.UNAVAILABLE => PodUnavailable(pod) + case ex: StatusException if ex.getStatus.getCode == Status.Code.CANCELLED => + if (isStream) StreamCancelled else PodUnavailable(pod) + case other => other + } + + private def toSendRequest(message: BinaryMessage): SendRequest = + SendRequest(message.entityId, message.entityType, message.body, message.replyId) + def assignShards(pod: PodAddress, shards: Set[ShardId]): Task[Unit] = - getConnection(pod).flatMap(_.assignShards(AssignShardsRequest(shards.toSeq)).unit) + getConnection(pod).flatMap(_.assignShards(AssignShardsRequest(shards.toList)).unit) def unassignShards(pod: PodAddress, shards: Set[ShardId]): Task[Unit] = - getConnection(pod).flatMap(_.unassignShards(UnassignShardsRequest(shards.toSeq)).unit) + getConnection(pod).flatMap(_.unassignShards(UnassignShardsRequest(shards.toList)).unit) def ping(pod: PodAddress): Task[Unit] = getConnection(pod).flatMap(_.pingShards(PingShardsRequest()).unit) def sendMessage(pod: PodAddress, message: BinaryMessage): Task[Option[Array[Byte]]] = - getConnection(pod) - .flatMap( - _.send(SendRequest(message.entityId, message.entityType, ByteString.copyFrom(message.body), message.replyId)) - .mapBoth( - ex => - if (ex.getStatus.getCode == Status.Code.RESOURCE_EXHAUSTED) { - // entity is not managed by this pod, wait and retry (assignments will be updated) - EntityNotManagedByThisPod(message.entityId) - } else if ( - ex.getStatus.getCode == Status.Code.UNAVAILABLE || ex.getStatus.getCode == Status.Code.CANCELLED - ) - PodUnavailable(pod) - else ex, - res => if (res.body.isEmpty) None else Some(res.body.toByteArray) - ) - ) + getConnection(pod).flatMap { conn => + conn + .send(toSendRequest(message)) + .mapBoth( + mapClientError(pod, message.entityId, isStream = false), + res => if (res.body.isEmpty) None else Some(res.body) + ) + } def sendStream( pod: PodAddress, entityId: String, messages: ZStream[Any, Throwable, BinaryMessage] ): Task[Option[Array[Byte]]] = - getConnection(pod) - .flatMap( - _.sendStream( - messages.mapBoth( - Status.INTERNAL.withCause(_).asException(), - message => - SendRequest(message.entityId, message.entityType, ByteString.copyFrom(message.body), message.replyId) - ) - ).mapBoth( - ex => - if (ex.getStatus.getCode == Status.Code.RESOURCE_EXHAUSTED) { - // entity is not managed by this pod, wait and retry (assignments will be updated) - EntityNotManagedByThisPod(entityId) - } else if (ex.getStatus.getCode == Status.Code.UNAVAILABLE) PodUnavailable(pod) - else if (ex.getStatus.getCode == Status.Code.CANCELLED) StreamCancelled - else ex, - res => if (res.body.isEmpty) None else Some(res.body.toByteArray) + getConnection(pod).flatMap { conn => + conn + .sendStream(messages.mapBoth(Status.INTERNAL.withCause(_).asException(), toSendRequest)) + .mapBoth( + mapClientError(pod, entityId, isStream = true), + res => if (res.body.isEmpty) None else Some(res.body) ) - ) + } def sendMessageAndReceiveStream(pod: PodAddress, message: BinaryMessage): ZStream[Any, Throwable, Array[Byte]] = ZStream .fromZIO(getConnection(pod)) .flatMap( - _.sendAndReceiveStream( - SendRequest(message.entityId, message.entityType, ByteString.copyFrom(message.body), message.replyId) - ).mapBoth( - ex => - if (ex.getStatus.getCode == Status.Code.RESOURCE_EXHAUSTED) { - // entity is not managed by this pod, wait and retry (assignments will be updated) - EntityNotManagedByThisPod(message.entityId) - } else if (ex.getStatus.getCode == Status.Code.UNAVAILABLE) PodUnavailable(pod) - else if (ex.getStatus.getCode == Status.Code.CANCELLED) StreamCancelled - else ex, - _.body.toByteArray - ) + _.sendAndReceiveStream(toSendRequest(message)) + .mapBoth(mapClientError(pod, message.entityId, isStream = true), _.body) ) def sendStreamAndReceiveStream( @@ -140,21 +141,8 @@ class GrpcPods( .fromZIO(getConnection(pod)) .flatMap( _.sendStreamAndReceiveStream( - messages.mapBoth( - Status.INTERNAL.withCause(_).asException(), - message => - SendRequest(message.entityId, message.entityType, ByteString.copyFrom(message.body), message.replyId) - ) - ).mapBoth( - ex => - if (ex.getStatus.getCode == Status.Code.RESOURCE_EXHAUSTED) { - // entity is not managed by this pod, wait and retry (assignments will be updated) - EntityNotManagedByThisPod(entityId) - } else if (ex.getStatus.getCode == Status.Code.UNAVAILABLE) PodUnavailable(pod) - else if (ex.getStatus.getCode == Status.Code.CANCELLED) StreamCancelled - else ex, - _.body.toByteArray - ) + messages.mapBoth(Status.INTERNAL.withCause(_).asException(), toSendRequest) + ).mapBoth(mapClientError(pod, entityId, isStream = true), _.body) ) } @@ -167,13 +155,12 @@ object GrpcPods { ZLayer.scoped { for { config <- ZIO.service[GrpcConfig] - connections <- - Ref.Synchronized - .make(Map.empty[PodAddress, (ShardingServiceClient, Fiber[Throwable, Nothing])]) - .withFinalizer( - // stop all connection fibers on release - _.get.flatMap(connections => ZIO.foreachDiscard(connections) { case (_, (_, fiber)) => fiber.interrupt }) - ) + connections <- Ref.Synchronized + .make(Map.empty[PodAddress, (PodClient, Fiber[Throwable, Nothing])]) + .withFinalizer( + // stop all connection fibers on release + _.get.flatMap(conns => ZIO.foreachDiscard(conns) { case (_, (_, fiber)) => fiber.interrupt }) + ) } yield new GrpcPods(config, connections) } } diff --git a/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcShardingService.scala b/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcShardingService.scala index e95bec9d..9f8724f9 100644 --- a/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcShardingService.scala +++ b/protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcShardingService.scala @@ -2,18 +2,17 @@ package com.devsisters.shardcake import com.devsisters.shardcake.errors.EntityNotManagedByThisPod import com.devsisters.shardcake.interfaces.Pods.BinaryMessage -import com.devsisters.shardcake.protobuf.sharding.ZioSharding.ShardingService -import com.devsisters.shardcake.protobuf.sharding._ -import com.google.protobuf.ByteString +import com.devsisters.shardcake.protocol.Sharding +import com.devsisters.shardcake.protocol.Sharding._ import io.grpc._ -import io.grpc.protobuf.services.ProtoReflectionService -import scalapb.zio_grpc.ServiceList +import io.grpc.protobuf.services.ProtoReflectionServiceV1 +import proteus.server.{ GrpcContext, ServerService, ZioServerBackend } import zio.stream.ZStream import zio.{ Config => _, _ } import java.util.concurrent.TimeUnit -abstract class GrpcShardingService(sharding: Sharding, timeout: Duration) extends ShardingService { +abstract class GrpcShardingService(sharding: Sharding, timeout: Duration) { def assignShards(request: AssignShardsRequest): ZIO[Any, StatusException, AssignShardsResponse] = sharding.assign(request.shards.toSet).as(AssignShardsResponse()) @@ -22,61 +21,70 @@ abstract class GrpcShardingService(sharding: Sharding, timeout: Duration) extend def send(request: SendRequest): ZIO[Any, StatusException, SendResponse] = sharding - .sendToLocalEntity( - BinaryMessage(request.entityId, request.entityType, request.body.toByteArray, request.replyId) - ) - .map { - case None => ByteString.EMPTY - case Some(res) => ByteString.copyFrom(res) - } - .mapBoth(mapErrorToStatusWithInternalDetails, SendResponse(_)) - .timeoutFail(Status.ABORTED.withDescription("Timeout while handling sharding send grpc").asException())(timeout) - - def sendAndReceiveStream(request: SendRequest): ZStream[Any, StatusException, SendResponse] = - sharding - .sendToLocalEntityAndReceiveStream( - BinaryMessage(request.entityId, request.entityType, request.body.toByteArray, request.replyId) - ) - .mapBoth(mapErrorToStatusWithInternalDetails, bytes => SendResponse(ByteString.copyFrom(bytes))) + .sendToLocalEntity(GrpcShardingService.toBinary(request)) + .map(GrpcShardingService.toSendResponse) + .mapError(GrpcShardingService.mapErrorToStatus) + .timeoutFail(GrpcShardingService.timeoutException)(timeout) def sendStream( requests: ZStream[Any, StatusException, SendRequest] ): ZIO[Any, StatusException, SendResponse] = sharding - .sendStreamToLocalEntity( - requests.map(request => - BinaryMessage(request.entityId, request.entityType, request.body.toByteArray, request.replyId) - ) - ) - .map { - case None => ByteString.EMPTY - case Some(res) => ByteString.copyFrom(res) - } - .mapBoth(mapErrorToStatusWithInternalDetails, SendResponse(_)) + .sendStreamToLocalEntity(requests.map(GrpcShardingService.toBinary)) + .map(GrpcShardingService.toSendResponse) + .mapError(GrpcShardingService.mapErrorToStatus) + + def sendAndReceiveStream(request: SendRequest): ZStream[Any, StatusException, SendResponse] = + sharding + .sendToLocalEntityAndReceiveStream(GrpcShardingService.toBinary(request)) + .map(SendResponse(_)) + .mapError(GrpcShardingService.mapErrorToStatus) def sendStreamAndReceiveStream( requests: ZStream[Any, StatusException, SendRequest] ): ZStream[Any, StatusException, SendResponse] = sharding - .sendStreamToLocalEntityAndReceiveStream( - requests.map(request => - BinaryMessage(request.entityId, request.entityType, request.body.toByteArray, request.replyId) - ) - ) - .mapBoth(mapErrorToStatusWithInternalDetails, bytes => SendResponse(ByteString.copyFrom(bytes))) + .sendStreamToLocalEntityAndReceiveStream(requests.map(GrpcShardingService.toBinary)) + .map(SendResponse(_)) + .mapError(GrpcShardingService.mapErrorToStatus) def pingShards(request: PingShardsRequest): ZIO[Any, StatusException, PingShardsResponse] = ZIO.succeed(PingShardsResponse()) +} - private def mapErrorToStatusWithInternalDetails: Function[Throwable, StatusException] = { +object GrpcShardingService { + + private[shardcake] val timeoutException: StatusException = + Status.ABORTED.withDescription("Timeout while handling sharding send grpc").asException() + + private val emptySendResponse: SendResponse = SendResponse(Array.emptyByteArray) + + private[shardcake] def toBinary(req: SendRequest): BinaryMessage = + BinaryMessage(req.entityId, req.entityType, req.body, req.replyId) + + private[shardcake] def toSendResponse(body: Option[Array[Byte]]): SendResponse = + body.fold(emptySendResponse)(SendResponse(_)) + + private[shardcake] val mapErrorToStatus: Throwable => StatusException = { case e: StatusException => e case e: StatusRuntimeException => e.getStatus.asException() case e: EntityNotManagedByThisPod => Status.RESOURCE_EXHAUSTED.withCause(e).asException() case e => Status.INTERNAL.withCause(e).withDescription(e.getMessage).asException() } -} -object GrpcShardingService { + private def buildServiceDefinition( + service: GrpcShardingService, + backend: ZioServerBackend[Any, StatusException, GrpcContext] + ): ServerServiceDefinition = + ServerService(using backend) + .rpc(Sharding.assignShards, service.assignShards) + .rpc(Sharding.unassignShards, service.unassignShards) + .rpc(Sharding.send, service.send) + .rpc(Sharding.sendStream, service.sendStream) + .rpc(Sharding.sendAndReceiveStream, service.sendAndReceiveStream) + .rpc(Sharding.sendStreamAndReceiveStream, service.sendStreamAndReceiveStream) + .rpc(Sharding.pingShards, service.pingShards) + .build(Sharding.service) /** * A layer that creates a gRPC server that exposes the Pods API. @@ -84,38 +92,34 @@ object GrpcShardingService { val live: ZLayer[Config with Sharding with GrpcConfig, Throwable, Unit] = ZLayer.scoped[Config with Sharding with GrpcConfig] { for { - config <- ZIO.service[Config] - grpcConfig <- ZIO.service[GrpcConfig] - sharding <- ZIO.service[Sharding] - builder = grpcConfig.executor match { - case Some(executor) => - ServerBuilder - .forPort(config.shardingPort) - .executor(executor) - case None => - ServerBuilder.forPort(config.shardingPort) - } - grpcShardingService = new GrpcShardingService(sharding, config.sendTimeout) {} - services <- - ServiceList - .add( - grpcConfig.serverInterceptors - .reduceOption(_.andThen(_)) - .map(t => grpcShardingService.transform(t)) - .getOrElse(grpcShardingService.asGeneric) - ) - .bindAll - server: Server = services - .foldLeft(builder) { case (builder0, service) => builder0.addService(service) } - .addService(ProtoReflectionService.newInstance()) - .build() - _ <- ZIO.acquireRelease(ZIO.attempt(server.start()))(server => - ZIO.attemptBlocking { - server.shutdown() - server.awaitTermination(grpcConfig.shutdownTimeout.toMillis, TimeUnit.MILLISECONDS) - server.shutdownNow() - }.ignore - ) + config <- ZIO.service[Config] + grpcConfig <- ZIO.service[GrpcConfig] + sharding <- ZIO.service[Sharding] + runtime <- ZIO.runtime[Any] + interceptor = ShardingServerInterceptor.compose(grpcConfig.serverInterceptors) + backend = ZioServerBackend(interceptor, runtime, grpcConfig.streamingPrefetch) + service = new GrpcShardingService(sharding, config.sendTimeout) {} + serviceDefinition = buildServiceDefinition(service, backend) + baseBuilder = grpcConfig.executor match { + case Some(executor) => + ServerBuilder + .forPort(config.shardingPort) + .executor(executor) + case None => + ServerBuilder.forPort(config.shardingPort) + } + builder = baseBuilder + .maxInboundMessageSize(grpcConfig.maxInboundMessageSize) + .addService(serviceDefinition) + .addService(ProtoReflectionServiceV1.newInstance()) + server: Server = builder.build() + _ <- ZIO.acquireRelease(ZIO.attempt(server.start()))(server => + ZIO.attemptBlocking { + server.shutdown() + server.awaitTermination(grpcConfig.shutdownTimeout.toMillis, TimeUnit.MILLISECONDS) + server.shutdownNow() + }.ignore + ) } yield () } } diff --git a/protocol-grpc/src/main/scala/com/devsisters/shardcake/PodClient.scala b/protocol-grpc/src/main/scala/com/devsisters/shardcake/PodClient.scala new file mode 100644 index 00000000..533bd416 --- /dev/null +++ b/protocol-grpc/src/main/scala/com/devsisters/shardcake/PodClient.scala @@ -0,0 +1,48 @@ +package com.devsisters.shardcake + +import com.devsisters.shardcake.protocol.Sharding +import com.devsisters.shardcake.protocol.Sharding._ +import io.grpc.StatusException +import proteus.client.ZioClientBackend +import zio._ +import zio.stream.ZStream + +/** + * Bundles the per-RPC client functions for a single pod. The functions are thin + * wrappers around the proteus backend; the underlying channel's lifecycle is + * managed by the scope that builds the backend. + */ +private[shardcake] trait PodClient { + def assignShards(request: AssignShardsRequest): IO[StatusException, AssignShardsResponse] + def unassignShards(request: UnassignShardsRequest): IO[StatusException, UnassignShardsResponse] + def pingShards(request: PingShardsRequest): IO[StatusException, PingShardsResponse] + def send(request: SendRequest): IO[StatusException, SendResponse] + def sendStream(requests: ZStream[Any, StatusException, SendRequest]): IO[StatusException, SendResponse] + def sendAndReceiveStream(request: SendRequest): ZStream[Any, StatusException, SendResponse] + def sendStreamAndReceiveStream( + requests: ZStream[Any, StatusException, SendRequest] + ): ZStream[Any, StatusException, SendResponse] +} + +private[shardcake] object PodClient { + def from(backend: ZioClientBackend): PodClient = + new PodClient { + private val assignShardsCall = backend.client(Sharding.assignShards, Sharding.service) + private val unassignShardsCall = backend.client(Sharding.unassignShards, Sharding.service) + private val pingShardsCall = backend.client(Sharding.pingShards, Sharding.service) + private val sendCall = backend.client(Sharding.send, Sharding.service) + private val sendStreamCall = backend.client(Sharding.sendStream, Sharding.service) + private val sendAndReceiveStreamCall = backend.client(Sharding.sendAndReceiveStream, Sharding.service) + private val sendStreamAndReceiveStreamCall = + backend.client(Sharding.sendStreamAndReceiveStream, Sharding.service) + + def assignShards(request: AssignShardsRequest) = assignShardsCall(request) + def unassignShards(request: UnassignShardsRequest) = unassignShardsCall(request) + def pingShards(request: PingShardsRequest) = pingShardsCall(request) + def send(request: SendRequest) = sendCall(request) + def sendStream(requests: ZStream[Any, StatusException, SendRequest]) = sendStreamCall(requests) + def sendAndReceiveStream(request: SendRequest) = sendAndReceiveStreamCall(request) + def sendStreamAndReceiveStream(requests: ZStream[Any, StatusException, SendRequest]) = + sendStreamAndReceiveStreamCall(requests) + } +} diff --git a/protocol-grpc/src/main/scala/com/devsisters/shardcake/ShardingClientInterceptor.scala b/protocol-grpc/src/main/scala/com/devsisters/shardcake/ShardingClientInterceptor.scala new file mode 100644 index 00000000..af6765e0 --- /dev/null +++ b/protocol-grpc/src/main/scala/com/devsisters/shardcake/ShardingClientInterceptor.scala @@ -0,0 +1,31 @@ +package com.devsisters.shardcake + +import io.grpc.ForwardingClientCall.SimpleForwardingClientCall +import io.grpc.{ CallOptions, Channel, ClientCall, ClientInterceptor, Metadata, MethodDescriptor } + +/** + * Helpers for building common [[io.grpc.ClientInterceptor]] instances to put in + * [[GrpcConfig.clientInterceptors]]. + */ +object ShardingClientInterceptor { + + /** + * Builds a client interceptor that lets you mutate the outgoing request headers before every + * call. The `update` callback receives the [[Metadata]] for the current call and may put, + * remove or read entries on it. + */ + def headersUpdater(update: Metadata => Unit): ClientInterceptor = + new ClientInterceptor { + def interceptCall[ReqT, RespT]( + method: MethodDescriptor[ReqT, RespT], + callOptions: CallOptions, + next: Channel + ): ClientCall[ReqT, RespT] = + new SimpleForwardingClientCall[ReqT, RespT](next.newCall(method, callOptions)) { + override def start(responseListener: ClientCall.Listener[RespT], headers: Metadata): Unit = { + update(headers) + super.start(responseListener, headers) + } + } + } +} diff --git a/protocol-grpc/src/main/scala/com/devsisters/shardcake/ShardingServerInterceptor.scala b/protocol-grpc/src/main/scala/com/devsisters/shardcake/ShardingServerInterceptor.scala new file mode 100644 index 00000000..b7f02d36 --- /dev/null +++ b/protocol-grpc/src/main/scala/com/devsisters/shardcake/ShardingServerInterceptor.scala @@ -0,0 +1,108 @@ +package com.devsisters.shardcake + +import io.grpc.StatusException +import proteus.ProtobufCodec +import proteus.server.{ GrpcContext, ServerInterceptor } +import zio._ +import zio.stream.ZStream + +/** + * Type alias for a proteus [[ServerInterceptor]] scoped to shardcake's effect types + * ([[IO]] and [[ZStream]] with [[StatusException]] errors) and the default + * [[GrpcContext]]. Use the smart constructors in the companion object — or implement + * the trait directly for advanced pre/post/error wrapping. + */ +type ShardingServerInterceptor = ServerInterceptor[ + [A] =>> IO[StatusException, A], + [A] =>> IO[StatusException, A], + [A] =>> ZStream[Any, StatusException, A], + [A] =>> ZStream[Any, StatusException, A], + GrpcContext, + GrpcContext +] + +object ShardingServerInterceptor { + private type UnaryEffect[A] = IO[StatusException, A] + private type StreamEffect[A] = ZStream[Any, StatusException, A] + + /** + * An interceptor that does nothing — useful as a starting point or for composition. + */ + val identity: ShardingServerInterceptor = + ServerInterceptor.empty[UnaryEffect, StreamEffect, GrpcContext] + + /** + * Builds an interceptor that runs `effect` before each handler. If the effect fails, + * the RPC fails with that error and the handler is never called. + */ + def beforeEach(effect: GrpcContext => IO[StatusException, Unit]): ShardingServerInterceptor = + new ShardingServerInterceptor { + def unary[Req: ProtobufCodec, Resp: ProtobufCodec]( + io: GrpcContext => UnaryEffect[Resp] + ): Req => GrpcContext => UnaryEffect[Resp] = + _ => ctx => effect(ctx) *> io(ctx) + + def clientStreaming[Req: ProtobufCodec, Resp: ProtobufCodec]( + io: StreamEffect[Req] => GrpcContext => UnaryEffect[Resp] + ): StreamEffect[Req] => GrpcContext => UnaryEffect[Resp] = + stream => ctx => effect(ctx) *> io(stream)(ctx) + + def serverStreaming[Req: ProtobufCodec, Resp: ProtobufCodec]( + io: GrpcContext => StreamEffect[Resp] + ): Req => GrpcContext => StreamEffect[Resp] = + _ => ctx => ZStream.fromZIO(effect(ctx)) *> io(ctx) + + def bidiStreaming[Req: ProtobufCodec, Resp: ProtobufCodec]( + io: StreamEffect[Req] => GrpcContext => StreamEffect[Resp] + ): StreamEffect[Req] => GrpcContext => StreamEffect[Resp] = + stream => ctx => ZStream.fromZIO(effect(ctx)) *> io(stream)(ctx) + } + + /** + * Composes a sequence of interceptors into a single one. The first interceptor in the + * sequence runs outermost (i.e. its pre-logic runs first, its post-logic runs last). + */ + def compose(interceptors: Seq[ShardingServerInterceptor]): ShardingServerInterceptor = + interceptors match { + case Seq() => identity + case Seq(only) => only + case _ => + new ShardingServerInterceptor { + def unary[Req: ProtobufCodec, Resp: ProtobufCodec]( + io: GrpcContext => UnaryEffect[Resp] + ): Req => GrpcContext => UnaryEffect[Resp] = + interceptors.foldRight(((_: Req) => io): Req => GrpcContext => UnaryEffect[Resp]) { (i, acc) => req => + { + val applied: Req => GrpcContext => UnaryEffect[Resp] = i.unary[Req, Resp](acc(req)) + applied(req) + } + } + + def clientStreaming[Req: ProtobufCodec, Resp: ProtobufCodec]( + io: StreamEffect[Req] => GrpcContext => UnaryEffect[Resp] + ): StreamEffect[Req] => GrpcContext => UnaryEffect[Resp] = + interceptors.foldRight(io) { (i, acc) => + val step: StreamEffect[Req] => GrpcContext => UnaryEffect[Resp] = i.clientStreaming[Req, Resp](acc) + step + } + + def serverStreaming[Req: ProtobufCodec, Resp: ProtobufCodec]( + io: GrpcContext => StreamEffect[Resp] + ): Req => GrpcContext => StreamEffect[Resp] = + interceptors.foldRight(((_: Req) => io): Req => GrpcContext => StreamEffect[Resp]) { (i, acc) => req => + { + val applied: Req => GrpcContext => StreamEffect[Resp] = i.serverStreaming[Req, Resp](acc(req)) + applied(req) + } + } + + def bidiStreaming[Req: ProtobufCodec, Resp: ProtobufCodec]( + io: StreamEffect[Req] => GrpcContext => StreamEffect[Resp] + ): StreamEffect[Req] => GrpcContext => StreamEffect[Resp] = + interceptors.foldRight(io) { (i, acc) => + val step: StreamEffect[Req] => GrpcContext => StreamEffect[Resp] = i.bidiStreaming[Req, Resp](acc) + step + } + } + } +} diff --git a/protocol-grpc/src/main/scala/com/devsisters/shardcake/protocol/Sharding.scala b/protocol-grpc/src/main/scala/com/devsisters/shardcake/protocol/Sharding.scala new file mode 100644 index 00000000..152e21a0 --- /dev/null +++ b/protocol-grpc/src/main/scala/com/devsisters/shardcake/protocol/Sharding.scala @@ -0,0 +1,48 @@ +package com.devsisters.shardcake.protocol + +import proteus._ + +object Sharding { + given ProtobufDeriver = ProtobufDeriver + + case class AssignShardsRequest(shards: List[Int]) derives ProtobufCodec + case class AssignShardsResponse() derives ProtobufCodec + + case class UnassignShardsRequest(shards: List[Int]) derives ProtobufCodec + case class UnassignShardsResponse() derives ProtobufCodec + + case class SendRequest(entityId: String, entityType: String, body: Array[Byte], replyId: Option[String]) + derives ProtobufCodec + case class SendResponse(body: Array[Byte]) derives ProtobufCodec + + case class PingShardsRequest() derives ProtobufCodec + case class PingShardsResponse() derives ProtobufCodec + + val assignShards = Rpc.unary[AssignShardsRequest, AssignShardsResponse]("AssignShards") + val unassignShards = Rpc.unary[UnassignShardsRequest, UnassignShardsResponse]("UnassignShards") + val send = Rpc.unary[SendRequest, SendResponse]("Send") + val sendStream = Rpc.clientStreaming[SendRequest, SendResponse]("SendStream") + val sendAndReceiveStream = Rpc.serverStreaming[SendRequest, SendResponse]("SendAndReceiveStream") + val sendStreamAndReceiveStream = Rpc.bidiStreaming[SendRequest, SendResponse]("SendStreamAndReceiveStream") + val pingShards = Rpc.unary[PingShardsRequest, PingShardsResponse]("PingShards") + + val service = + Service("ShardingService") + .rpc(assignShards) + .rpc(unassignShards) + .rpc(send) + .rpc(sendStream) + .rpc(sendAndReceiveStream) + .rpc(sendStreamAndReceiveStream) + .rpc(pingShards) +} + +object GenerateProto { + def main(args: Array[String]): Unit = { + val outputDir = args.headOption.getOrElse { + sys.error("Usage: GenerateProto ") + } + Sharding.service.renderToFile(List.empty, outputDir, "sharding") + println(s"Wrote $outputDir/sharding.proto") + } +} diff --git a/vuepress/docs/docs/customization.md b/vuepress/docs/docs/customization.md index 7a5186b5..8a69f758 100644 --- a/vuepress/docs/docs/customization.md +++ b/vuepress/docs/docs/customization.md @@ -77,9 +77,19 @@ trait Pods { def assignShards(pod: PodAddress, shards: Set[ShardId]): Task[Unit] def unassignShards(pod: PodAddress, shards: Set[ShardId]): Task[Unit] def ping(pod: PodAddress): Task[Unit] - + def sendMessage(pod: PodAddress, message: BinaryMessage): Task[Option[Array[Byte]]] - def sendMessageStreaming(pod: PodAddress, message: BinaryMessage): ZStream[Any, Throwable, Array[Byte]] + def sendStream( + pod: PodAddress, + entityId: String, + messages: ZStream[Any, Throwable, BinaryMessage] + ): Task[Option[Array[Byte]]] + def sendMessageAndReceiveStream(pod: PodAddress, message: BinaryMessage): ZStream[Any, Throwable, Array[Byte]] + def sendStreamAndReceiveStream( + pod: PodAddress, + entityId: String, + messages: ZStream[Any, Throwable, BinaryMessage] + ): ZStream[Any, Throwable, Array[Byte]] } ``` For testing, you can use the `Pods.noop` layer that does nothing.