Make grpc server shutdown timeout configurable

project/build.properties

@@ -1 +1 @@
-sbt.version = 1.7.1
+sbt.version = 1.9.3

protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcConfig.scala

@@ -1,16 +1,19 @@
 package com.devsisters.shardcake
 
+import zio._
+
 import java.util.concurrent.Executor
 
 /**
  * The configuration for the gRPC client.
  *
  * @param maxInboundMessageSize the maximum message size allowed to be received by the grpc client
  * @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
  */
-case class GrpcConfig(maxInboundMessageSize: Int, executor: Option[Executor])
+case class GrpcConfig(maxInboundMessageSize: Int, executor: Option[Executor], shutdownTimeout: Duration)
 
 object GrpcConfig {
   val default: GrpcConfig =
-    GrpcConfig(maxInboundMessageSize = 32 * 1024 * 1024, None)
+    GrpcConfig(maxInboundMessageSize = 32 * 1024 * 1024, None, 3.seconds)
 }

protocol-grpc/src/main/scala/com/devsisters/shardcake/GrpcShardingService.scala

@@ -5,11 +5,12 @@ 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 io.grpc.protobuf.services.ProtoReflectionService
-import io.grpc.{ ServerBuilder, Status, StatusException, StatusRuntimeException }
-import scalapb.zio_grpc.{ ScopedServer, ServiceList }
-import zio.{ Config => _, _ }
+import io.grpc._
+import scalapb.zio_grpc.ServiceList
 import zio.stream.ZStream
+import zio.{ Config => _, _ }
+
+import java.util.concurrent.TimeUnit
 
 abstract class GrpcShardingService(sharding: Sharding, timeout: Duration) extends ShardingService {
   def assignShards(request: AssignShardsRequest): ZIO[Any, StatusException, AssignShardsResponse] =
@@ -56,19 +57,26 @@ 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)
-                      }
-        services    = ServiceList.add(new GrpcShardingService(sharding, config.sendTimeout) {})
-        _          <- ScopedServer.fromServiceList(builder.addService(ProtoReflectionService.newInstance()), services)
+        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)
+                         }
+        services      <- ServiceList.add(new GrpcShardingService(sharding, config.sendTimeout) {}).bindAll
+        server: Server = services.foldLeft(builder) { case (builder0, service) => builder0.addService(service) }.build()
+        _             <- ZIO.acquireRelease(ZIO.attempt(server.start()))(server =>
+                           ZIO.attemptBlocking {
+                             server.shutdown()
+                             server.awaitTermination(grpcConfig.shutdownTimeout.toMillis, TimeUnit.MILLISECONDS)
+                             server.shutdownNow()
+                           }.ignore
+                         )
       } yield ()
     }
 }