Merge pull request #19 from scalableminds/grpc-health

support standard grpc health check

Author: fm3

build.sbt

@@ -23,8 +23,9 @@ libraryDependencies ++= Seq(
   "ch.qos.logback" % "logback-classic" % "1.2.3",
   "com.typesafe.scala-logging" %% "scala-logging" % "3.7.2",
   "org.scalatest" % "scalatest_2.12" % "3.0.4" % "test",
-  "io.grpc" % "grpc-netty" % com.trueaccord.scalapb.compiler.Version.grpcJavaVersion,
-  "com.trueaccord.scalapb" %% "scalapb-runtime-grpc" % com.trueaccord.scalapb.compiler.Version.scalapbVersion,
+  "io.grpc" % "grpc-netty" % scalapb.compiler.Version.grpcJavaVersion,
+  "io.grpc" % "grpc-services" % scalapb.compiler.Version.grpcJavaVersion,
+  "com.thesamet.scalapb" %% "scalapb-runtime-grpc" % scalapb.compiler.Version.scalapbVersion,
   "org.rocksdb" % "rocksdbjni" % "5.1.2",
   "com.github.scopt" %% "scopt" % "3.7.0"
 )

project/scalapb.sbt

@@ -1,3 +1,3 @@
-addSbtPlugin("com.thesamet" % "sbt-protoc" % "0.99.12")
+addSbtPlugin("com.thesamet" % "sbt-protoc" % "0.99.18")
 
-libraryDependencies += "com.trueaccord.scalapb" %% "compilerplugin" % "0.6.6"
+libraryDependencies += "com.thesamet.scalapb" %% "compilerplugin" % "0.7.4"

src/main/scala/com/scalableminds/fossildb/FossilDBGrpcImpl.scala

@@ -13,7 +13,9 @@ import com.typesafe.scalalogging.LazyLogging
 
 import scala.concurrent.Future
 
-class FossilDBGrpcImpl(storeManager: StoreManager) extends FossilDBGrpc.FossilDB with LazyLogging {
+class FossilDBGrpcImpl(storeManager: StoreManager)
+  extends FossilDBGrpc.FossilDB
+  with LazyLogging {
 
   override def health(req: HealthRequest) = withExceptionHandler(req) {
     HealthReply(true)

src/main/scala/com/scalableminds/fossildb/FossilDBServer.scala

@@ -5,18 +5,26 @@ package com.scalableminds.fossildb
 
 import com.scalableminds.fossildb.db.StoreManager
 import com.scalableminds.fossildb.proto.fossildbapi.FossilDBGrpc
+import io.grpc.health.v1.HealthCheckResponse
 import com.typesafe.scalalogging.LazyLogging
 import io.grpc.Server
 import io.grpc.netty.NettyServerBuilder
+import io.grpc.services.HealthStatusManager
 
 import scala.concurrent.ExecutionContext
 
 class FossilDBServer(storeManager: StoreManager, port: Int, executionContext: ExecutionContext) extends LazyLogging
 { self =>
   private[this] var server: Server = null
+  private[this] var healthStatusManager: HealthStatusManager = null
 
   def start(): Unit = {
-    server = NettyServerBuilder.forPort(port).maxMessageSize(Int.MaxValue).addService(FossilDBGrpc.bindService(new FossilDBGrpcImpl(storeManager), executionContext)).build.start
+    healthStatusManager = new HealthStatusManager()
+    server = NettyServerBuilder.forPort(port).maxMessageSize(Int.MaxValue)
+      .addService(FossilDBGrpc.bindService(new FossilDBGrpcImpl(storeManager), executionContext))
+      .addService(healthStatusManager.getHealthService())
+      .build.start
+    healthStatusManager.setStatus("", HealthCheckResponse.ServingStatus.SERVING)
     logger.info("Server started, listening on " + port)
     sys.addShutdownHook {
       logger.info("Shutting down gRPC server since JVM is shutting down")
@@ -29,6 +37,7 @@ class FossilDBServer(storeManager: StoreManager, port: Int, executionContext: Ex
     if (server != null) {
       server.shutdown()
       storeManager.close
+      healthStatusManager.setStatus("", HealthCheckResponse.ServingStatus.NOT_SERVING)
     }
   }
 

src/test/scala/com/scalableminds/fossildb/FossilDBSuite.scala

@@ -9,6 +9,7 @@ import java.nio.file.Paths
 import com.google.protobuf.ByteString
 import com.scalableminds.fossildb.db.StoreManager
 import com.scalableminds.fossildb.proto.fossildbapi._
+import io.grpc.health.v1._
 import io.grpc.netty.NettyChannelBuilder
 import org.scalatest.{BeforeAndAfterEach, FlatSpec}
 
@@ -21,7 +22,9 @@ class FossilDBSuite extends FlatSpec with BeforeAndAfterEach {
 
   val port = 21505
   var serverOpt: Option[FossilDBServer] = None
-  val client = FossilDBGrpc.blockingStub(NettyChannelBuilder.forAddress("127.0.0.1", port).maxInboundMessageSize(Int.MaxValue).usePlaintext(true).build)
+  val channel = NettyChannelBuilder.forAddress("127.0.0.1", port).maxInboundMessageSize(Int.MaxValue).usePlaintext().build
+  val client = FossilDBGrpc.blockingStub(channel)
+  val healthClient = HealthGrpc.newBlockingStub(channel)
 
   val collectionA = "collectionA"
   val collectionB = "collectionB"
@@ -65,6 +68,11 @@ class FossilDBSuite extends FlatSpec with BeforeAndAfterEach {
     assert(reply.success)
   }
 
+  "GRPC Standard Health Check" should "report SERVING" in {
+    val reply = healthClient.check(HealthCheckRequest.getDefaultInstance())
+    assert(reply.getStatus.toString == "SERVING")
+  }
+
   "Put" should "overwrite old value" in {
     client.put(PutRequest(collectionA, aKey, Some(0), testData1))
     client.put(PutRequest(collectionA, aKey, Some(0), testData2))