Merge pull request #92 from weaveworks/readiness

Add readiness probe

Author: tomwilkie

cmd/cortex/main.go

@@ -107,17 +107,18 @@ func main() {
 		log.Fatalf("Error initializing Consul client: %v", err)
 	}
 	consul = ring.PrefixClient(consul, cfg.consulPrefix)
+	r := ring.New(consul, cfg.distributorConfig.HeartbeatTimeout)
+	defer r.Stop()
 
 	switch cfg.mode {
 	case modeDistributor:
-		ring := ring.New(consul, cfg.distributorConfig.HeartbeatTimeout)
-		cfg.distributorConfig.Ring = ring
+		cfg.distributorConfig.Ring = r
 		cfg.distributorConfig.ClientFactory = func(address string) (*cortex.IngesterClient, error) {
 			return cortex.NewIngesterClient(address, cfg.remoteTimeout)
 		}
-		defer ring.Stop()
 		setupDistributor(cfg.distributorConfig, chunkStore, cfg.logSuccess)
 	case modeIngester:
+		cfg.ingesterConfig.Ring = r
 		registration, err := ring.RegisterIngester(consul, cfg.listenPort, cfg.numTokens)
 		if err != nil {
 			// This only happens for errors in configuration & set-up, not for
@@ -267,6 +268,7 @@ func setupIngester(
 	http.Handle("/query", instrument(logSuccess, cortex.QueryHandler(ingester)))
 	http.Handle("/label_values", instrument(logSuccess, cortex.LabelValuesHandler(ingester)))
 	http.Handle("/user_stats", instrument(logSuccess, cortex.IngesterUserStatsHandler(ingester.UserStats)))
+	http.Handle("/ready", instrument(logSuccess, cortex.IngesterReadinessHandler(ingester)))
 	return ingester
 }
 

ingester/ingester.go

@@ -10,11 +10,12 @@ import (
 	"github.com/prometheus/common/log"
 	"github.com/prometheus/common/model"
 	prom_chunk "github.com/prometheus/prometheus/storage/local/chunk"
-	cortex "github.com/weaveworks/cortex/chunk"
-	"github.com/weaveworks/cortex/user"
+	"github.com/prometheus/prometheus/storage/metric"
 	"golang.org/x/net/context"
 
-	"github.com/prometheus/prometheus/storage/metric"
+	cortex "github.com/weaveworks/cortex/chunk"
+	"github.com/weaveworks/cortex/ring"
+	"github.com/weaveworks/cortex/user"
 )
 
 const (
@@ -78,6 +79,7 @@ type Config struct {
 	FlushCheckPeriod time.Duration
 	MaxChunkAge      time.Duration
 	RateUpdatePeriod time.Duration
+	Ring             *ring.Ring
 }
 
 // UserStats models ingestion statistics for one user.
@@ -154,6 +156,12 @@ func New(cfg Config, chunkStore cortex.Store) (*Ingester, error) {
 	return i, nil
 }
 
+// Ready is used to indicate to k8s when the ingesters are ready for
+// the addition / removal of another ingester.
+func (i *Ingester) Ready() bool {
+	return i.cfg.Ring.Ready()
+}
+
 func (i *Ingester) getStateFor(ctx context.Context) (*userState, error) {
 	userID, err := user.GetID(ctx)
 	if err != nil {

ring/ring.go

@@ -184,6 +184,26 @@ func (r *Ring) GetAll() []IngesterDesc {
 	return ingesters
 }
 
+// Ready is true when all ingesters are active and healthy.
+func (r *Ring) Ready() bool {
+	r.mtx.RLock()
+	defer r.mtx.RUnlock()
+
+	for _, ingester := range r.ringDesc.Ingesters {
+		if time.Now().Sub(ingester.Timestamp) > r.heartbeatTimeout {
+			return false
+		}
+	}
+
+	for _, token := range r.ringDesc.Tokens {
+		if token.State != Active {
+			return false
+		}
+	}
+
+	return len(r.ringDesc.Tokens) > 0
+}
+
 func (r *Ring) search(key uint32) int {
 	i := sort.Search(len(r.ringDesc.Tokens), func(x int) bool {
 		return r.ringDesc.Tokens[x].Token > key

server.go

@@ -260,6 +260,19 @@ func IngesterUserStatsHandler(statsFn func(context.Context) (*ingester.UserStats
 	})
 }
 
+// IngesterReadinessHandler returns 204 when the ingester is ready,
+// 500 otherwise.  Its use by kubernetes to indicate if the ingester
+// pool is ready to have ingesters added / removed.
+func IngesterReadinessHandler(i *ingester.Ingester) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if i.Ready() {
+			w.WriteHeader(http.StatusNoContent)
+		} else {
+			w.WriteHeader(http.StatusInternalServerError)
+		}
+	})
+}
+
 // DistributorUserStatsHandler handles user stats to the Distributor.
 func DistributorUserStatsHandler(statsFn func(context.Context) (*ingester.UserStats, error)) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {