Bugfix for sensuctl cluster health (#3469)

Signed-off-by: Eric Chlebek <[email protected]>

backend/apid/routers/health.go

@@ -40,9 +40,9 @@ func (r *HealthRouter) health(w http.ResponseWriter, req *http.Request) {
 	}
 	ctx := req.Context()
 	if timeout > 0 {
-		tctx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
-		defer cancel()
-		ctx = tctx
+		// We're storing the timeout as a value so it can be used by several
+		// contexts in GetClusterHealth, which is a concurrent gatherer.
+		ctx = context.WithValue(ctx, "timeout", time.Duration(timeout)*time.Second)
 	}
 	clusterHealth := r.controller.GetClusterHealth(ctx)
 	_ = json.NewEncoder(w).Encode(clusterHealth)

backend/store/etcd/health_store.go

@@ -69,10 +69,22 @@ func (s *Store) getHealth(ctx context.Context, id uint64, name string, urls []st
 func (s *Store) GetClusterHealth(ctx context.Context, cluster clientv3.Cluster, etcdClientTLSConfig *tls.Config) *corev2.HealthResponse {
 	healthResponse := &corev2.HealthResponse{}
 
+	var timeout time.Duration
+	if val := ctx.Value("timeout"); val != nil {
+		timeout, _ = val.(time.Duration)
+	}
+
 	// Do a get op against every cluster member. Collect the  memberIDs and
 	// op errors into a response map, and return this map as etcd health
 	// information.
-	mList, err := cluster.MemberList(ctx)
+	tctx := ctx
+	if timeout > 0 {
+		var cancel context.CancelFunc
+		tctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
+	}
+
+	mList, err := cluster.MemberList(tctx)
 	if err != nil {
 		logger.WithError(err).Error("could not get the cluster member list")
 		healthResponse.ClusterHealth = []*corev2.ClusterHealth{&corev2.ClusterHealth{
@@ -96,10 +108,13 @@ func (s *Store) GetClusterHealth(ctx context.Context, cluster clientv3.Cluster,
 	for _, member := range mList.Members {
 		go func(id uint64, name string, urls []string) {
 			defer wg.Done()
-			select {
-			case healths <- s.getHealth(ctx, id, name, urls, etcdClientTLSConfig):
-			case <-ctx.Done():
+			tctx := ctx
+			if timeout > 0 {
+				var cancel context.CancelFunc
+				tctx, cancel = context.WithTimeout(ctx, timeout)
+				defer cancel()
 			}
+			healths <- s.getHealth(tctx, id, name, urls, etcdClientTLSConfig)
 		}(member.ID, member.Name, member.ClientURLs)
 	}
 
@@ -112,7 +127,13 @@ func (s *Store) GetClusterHealth(ctx context.Context, cluster clientv3.Cluster,
 		return healthResponse.ClusterHealth[i].Name < healthResponse.ClusterHealth[j].Name
 	})
 
-	alarmResponse, err := s.client.Maintenance.AlarmList(ctx)
+	if timeout > 0 {
+		var cancel context.CancelFunc
+		tctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
+	}
+
+	alarmResponse, err := s.client.Maintenance.AlarmList(tctx)
 	if err != nil {
 		logger.WithError(err).Error("failed to fetch etcd alarm list")
 	} else {

backend/store/etcd/health_store_test.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"crypto/tls"
 	"testing"
+	"time"
 
 	"github.com/coreos/etcd/clientv3"
 	"github.com/sensu/sensu-go/backend/store"
@@ -18,3 +19,10 @@ func TestGetClusterHealth(t *testing.T) {
 		assert.Empty(t, healthResult.ClusterHealth[0].Err)
 	})
 }
+
+func TestGetClusterHealthTimeout(t *testing.T) {
+	testWithEtcdClient(t, func(store store.Store, client *clientv3.Client) {
+		result := store.GetClusterHealth(context.WithValue(context.Background(), "timeout", time.Nanosecond), client.Cluster, (*tls.Config)(nil))
+		assert.NotEmpty(t, result.ClusterHealth[0].Err)
+	})
+}