From b5d8c77b99cd84ba38c69e1fecadc92b3eaf6cb8 Mon Sep 17 00:00:00 2001
From: Aditya C S <aditya.gnu@gmail.com>
Date: Tue, 6 Aug 2019 20:58:44 +0530
Subject: [PATCH] Readiness probe for querier (#854)

---
 pkg/distributor/distributor.go            |  3 ++-
 pkg/loki/modules.go                       |  3 ++-
 pkg/querier/querier.go                    | 19 +++++++++++++++++++
 production/ksonnet/loki/querier.libsonnet |  4 ++++
 4 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/pkg/distributor/distributor.go b/pkg/distributor/distributor.go
index 79f417a5de..1234ebe108 100644
--- a/pkg/distributor/distributor.go
+++ b/pkg/distributor/distributor.go
@@ -104,7 +104,8 @@ type pushTracker struct {
 	err            chan error
 }
 
-// ReadinessHandler is handler for Distributor
+// ReadinessHandler is used to indicate to k8s when the distributor is ready.
+// Returns 200 when the distributor is ready, 500 otherwise.
 func (d *Distributor) ReadinessHandler(w http.ResponseWriter, r *http.Request) {
 	_, err := d.ring.GetAll()
 	if err != nil {
diff --git a/pkg/loki/modules.go b/pkg/loki/modules.go
index c57d4b0ca1..4e611eeecd 100644
--- a/pkg/loki/modules.go
+++ b/pkg/loki/modules.go
@@ -125,11 +125,11 @@ func (t *Loki) initDistributor() (err error) {
 	if err != nil {
 		return
 	}
+
 	t.server.HTTP.Path("/ready").Handler(http.HandlerFunc(t.distributor.ReadinessHandler))
 	t.server.HTTP.Handle("/api/prom/push", middleware.Merge(
 		t.httpAuthMiddleware,
 	).Wrap(http.HandlerFunc(t.distributor.PushHandler)))
-
 	return
 }
 
@@ -142,6 +142,7 @@ func (t *Loki) initQuerier() (err error) {
 	httpMiddleware := middleware.Merge(
 		t.httpAuthMiddleware,
 	)
+	t.server.HTTP.Path("/ready").Handler(http.HandlerFunc(t.querier.ReadinessHandler))
 	t.server.HTTP.Handle("/api/prom/query", httpMiddleware.Wrap(http.HandlerFunc(t.querier.QueryHandler)))
 	t.server.HTTP.Handle("/api/prom/label", httpMiddleware.Wrap(http.HandlerFunc(t.querier.LabelHandler)))
 	t.server.HTTP.Handle("/api/prom/label/{name}/values", httpMiddleware.Wrap(http.HandlerFunc(t.querier.LabelHandler)))
diff --git a/pkg/querier/querier.go b/pkg/querier/querier.go
index 10cd97c1cf..d3ec24a3ec 100644
--- a/pkg/querier/querier.go
+++ b/pkg/querier/querier.go
@@ -3,11 +3,13 @@ package querier
 import (
 	"context"
 	"flag"
+	"net/http"
 	"time"
 
 	cortex_client "github.com/cortexproject/cortex/pkg/ingester/client"
 	"github.com/cortexproject/cortex/pkg/ring"
 	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/go-kit/kit/log/level"
 	"github.com/prometheus/common/model"
 	"google.golang.org/grpc/health/grpc_health_v1"
 
@@ -18,6 +20,8 @@ import (
 	"github.com/grafana/loki/pkg/storage"
 )
 
+var readinessProbeSuccess = []byte("Ready")
+
 // Config for a querier.
 type Config struct {
 	TailMaxDuration time.Duration `yaml:"tail_max_duration"`
@@ -63,6 +67,21 @@ type responseFromIngesters struct {
 	response interface{}
 }
 
+// ReadinessHandler is used to indicate to k8s when the querier is ready.
+// Returns 200 when the querier is ready, 500 otherwise.
+func (q *Querier) ReadinessHandler(w http.ResponseWriter, r *http.Request) {
+	_, err := q.ring.GetAll()
+	if err != nil {
+		http.Error(w, "Not ready: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	w.WriteHeader(http.StatusOK)
+	if _, err := w.Write(readinessProbeSuccess); err != nil {
+		level.Error(util.Logger).Log("msg", "error writing success message", "error", err)
+	}
+}
+
 // forAllIngesters runs f, in parallel, for all ingesters
 // TODO taken from Cortex, see if we can refactor out an usable interface.
 func (q *Querier) forAllIngesters(f func(logproto.QuerierClient) (interface{}, error)) ([]responseFromIngesters, error) {
diff --git a/production/ksonnet/loki/querier.libsonnet b/production/ksonnet/loki/querier.libsonnet
index 40329ab34c..b10af76287 100644
--- a/production/ksonnet/loki/querier.libsonnet
+++ b/production/ksonnet/loki/querier.libsonnet
@@ -10,6 +10,10 @@
     container.new('querier', $._images.querier) +
     container.withPorts($.util.defaultPorts) +
     container.withArgsMixin($.util.mapToFlags($.querier_args)),
+    container.mixin.readinessProbe.httpGet.withPath('/ready') +
+    container.mixin.readinessProbe.httpGet.withPort(80) +
+    container.mixin.readinessProbe.withInitialDelaySeconds(15) +
+    container.mixin.readinessProbe.withTimeoutSeconds(1) +
 
   local deployment = $.apps.v1beta1.deployment,