From 9fdd2802a1b2696c489387f9cf8b5b6ad9d0a489 Mon Sep 17 00:00:00 2001
From: Cyril Tovena <cyril.tovena@gmail.com>
Date: Tue, 31 Jan 2023 11:41:01 +0100
Subject: [PATCH] Correctly implement connect health checker service (#491)

Connect healthcheck service is now correctly checking for services health and only registered on the http port.

Fixes #472
---
 pkg/phlare/health.go  | 27 +++++++++++++++++++++++++++
 pkg/phlare/modules.go |  3 ---
 pkg/phlare/phlare.go  |  3 +--
 3 files changed, 28 insertions(+), 5 deletions(-)
 create mode 100644 pkg/phlare/health.go

diff --git a/pkg/phlare/health.go b/pkg/phlare/health.go
new file mode 100644
index 000000000..a82f146f4
--- /dev/null
+++ b/pkg/phlare/health.go
@@ -0,0 +1,27 @@
+package phlare
+
+import (
+	"context"
+
+	grpchealth "github.com/bufbuild/connect-grpchealth-go"
+	"github.com/gorilla/mux"
+	"github.com/grafana/dskit/grpcutil"
+)
+
+type checker struct {
+	checks []grpcutil.Check
+}
+
+func RegisterHealthServer(mux *mux.Router, checks ...grpcutil.Check) {
+	prefix, handler := grpchealth.NewHandler(&checker{checks: checks})
+	mux.NewRoute().PathPrefix(prefix).Handler(handler)
+}
+
+func (c *checker) Check(ctx context.Context, req *grpchealth.CheckRequest) (*grpchealth.CheckResponse, error) {
+	for _, check := range c.checks {
+		if !check(ctx) {
+			return &grpchealth.CheckResponse{Status: grpchealth.StatusNotServing}, nil
+		}
+	}
+	return &grpchealth.CheckResponse{Status: grpchealth.StatusServing}, nil
+}
diff --git a/pkg/phlare/modules.go b/pkg/phlare/modules.go
index 20a95afed..901b5f738 100644
--- a/pkg/phlare/modules.go
+++ b/pkg/phlare/modules.go
@@ -6,7 +6,6 @@ import (
 	"net/http"
 	"os"
 
-	grpchealth "github.com/bufbuild/connect-grpchealth-go"
 	"github.com/felixge/fgprof"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
@@ -206,8 +205,6 @@ func (f *Phlare) initIngester() (_ services.Service, err error) {
 	if err != nil {
 		return nil, err
 	}
-	prefix, handler := grpchealth.NewHandler(grpchealth.NewStaticChecker(ingesterv1connect.IngesterServiceName))
-	f.Server.HTTP.NewRoute().PathPrefix(prefix).Handler(handler)
 	ingesterv1connect.RegisterIngesterServiceHandler(f.Server.HTTP, ingester, f.auth)
 	return ingester, nil
 }
diff --git a/pkg/phlare/phlare.go b/pkg/phlare/phlare.go
index 1a6a3d417..2415c54ed 100644
--- a/pkg/phlare/phlare.go
+++ b/pkg/phlare/phlare.go
@@ -28,7 +28,6 @@ import (
 	"github.com/weaveworks/common/server"
 	"github.com/weaveworks/common/signals"
 	wwtracing "github.com/weaveworks/common/tracing"
-	"google.golang.org/grpc/health/grpc_health_v1"
 
 	pushv1connect "github.com/grafana/phlare/api/gen/proto/go/push/v1/pushv1connect"
 	"github.com/grafana/phlare/pkg/agent"
@@ -306,7 +305,7 @@ func (f *Phlare) Run() error {
 	}
 	f.Server.HTTP.Path("/ready").Methods("GET").Handler(f.readyHandler(sm))
 
-	grpc_health_v1.RegisterHealthServer(f.Server.GRPC, grpcutil.NewHealthCheck(sm))
+	RegisterHealthServer(f.Server.HTTP, grpcutil.WithManager(sm))
 	healthy := func() { level.Info(f.logger).Log("msg", "Phlare started", "version", version.Info()) }
 
 	serviceFailed := func(service services.Service) {