From 01972a8fe7789f4ed55a3e160f1430ba6b5430b3 Mon Sep 17 00:00:00 2001
From: Adrien Fillon <adrien.fillon@cdiscount.com>
Date: Tue, 12 Mar 2019 13:33:07 +0100
Subject: [PATCH] query: cleanup store statuses as they come and go

Signed-off-by: Adrien Fillon <adrien.fillon@cdiscount.com>
---
 pkg/query/storeset.go | 61 ++++++++++++++++++++++++++++++++-----------
 1 file changed, 46 insertions(+), 15 deletions(-)

diff --git a/pkg/query/storeset.go b/pkg/query/storeset.go
index bb9c7698dff..6e8a04c60d4 100644
--- a/pkg/query/storeset.go
+++ b/pkg/query/storeset.go
@@ -22,6 +22,7 @@ import (
 const (
 	unhealthyStoreMessage = "removing store because it's unhealthy or does not exist"
 	droppingStoreMessage  = "dropping store, external labels are not unique"
+	storeUnhealthyTimeout = time.Minute
 )
 
 type StoreSpec interface {
@@ -36,13 +37,14 @@ type StoreSpec interface {
 }
 
 type StoreStatus struct {
-	Name      string
-	LastCheck time.Time
-	LastError error
-	Labels    []storepb.Label
-	StoreType component.StoreAPI
-	MinTime   int64
-	MaxTime   int64
+	Name             string
+	LastCheck        time.Time
+	LastHealthyCheck *time.Time
+	LastError        error
+	Labels           []storepb.Label
+	StoreType        component.StoreAPI
+	MinTime          int64
+	MaxTime          int64
 }
 
 type grpcStoreSpec struct {
@@ -251,6 +253,7 @@ func (s *StoreSet) Update(ctx context.Context) {
 	}
 	s.externalLabelStores = externalLabelStores
 	s.storeNodeConnections.Set(float64(len(s.stores)))
+	s.cleanUpStoreStatuses()
 }
 
 func (s *StoreSet) getHealthyStores(ctx context.Context) map[string]*storeRef {
@@ -341,16 +344,30 @@ func (s *StoreSet) updateStoreStatus(store *storeRef, err error) {
 	s.storesStatusesMtx.Lock()
 	defer s.storesStatusesMtx.Unlock()
 
+	var status StoreStatus
 	now := time.Now()
-	s.storeStatuses[store.addr] = &StoreStatus{
-		Name:      store.addr,
-		LastError: err,
-		LastCheck: now,
-		Labels:    store.labels,
-		StoreType: store.storeType,
-		MinTime:   store.minTime,
-		MaxTime:   store.maxTime,
+
+	if previousStatus, ok := s.storeStatuses[store.addr]; ok {
+		status = *previousStatus
+		status.LastError = err
+		status.LastCheck = now
+	} else {
+		status = StoreStatus{
+			Name:      store.addr,
+			LastError: err,
+			LastCheck: now,
+		}
+	}
+
+	if err == nil {
+		status.LastHealthyCheck = &now
+		status.Labels = store.labels
+		status.StoreType = store.storeType
+		status.MinTime = store.minTime
+		status.MaxTime = store.maxTime
 	}
+
+	s.storeStatuses[store.addr] = &status
 }
 
 func (s *StoreSet) GetStoreStatus() []StoreStatus {
@@ -397,3 +414,17 @@ func (s *StoreSet) Close() {
 		st.close()
 	}
 }
+
+func (s *StoreSet) cleanUpStoreStatuses() {
+	s.storesStatusesMtx.RLock()
+	defer s.storesStatusesMtx.RUnlock()
+
+	now := time.Now()
+	for k, status := range s.storeStatuses {
+		if _, ok := s.stores[k]; !ok {
+			if status.LastHealthyCheck != nil && now.Sub(*status.LastHealthyCheck) >= storeUnhealthyTimeout {
+				delete(s.storeStatuses, k)
+			}
+		}
+	}
+}