From 6b1cd1e43bbfa7c7b41b19216993befc011f8732 Mon Sep 17 00:00:00 2001
From: Seebs <seebs@seebs.net>
Date: Fri, 26 Mar 2021 14:28:33 -0500
Subject: [PATCH 1/5] drop etcd-with-cache option

The new etcd implementation has internal caching-like behavior which is
much more reliable (it doesn't use a TTL, it just updates when there's updates
to process) so we don't need this cache.
---
 etcd/cache.go    | 57 ------------------------------------------------
 server/server.go |  2 +-
 2 files changed, 1 insertion(+), 58 deletions(-)
 delete mode 100644 etcd/cache.go

diff --git a/etcd/cache.go b/etcd/cache.go
deleted file mode 100644
index 4cd5bc684..000000000
--- a/etcd/cache.go
+++ /dev/null
@@ -1,57 +0,0 @@
-// Copyright 2017 Pilosa Corp.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package etcd
-
-import (
-	"sync"
-	"time"
-
-	"github.com/pilosa/pilosa/v2/topology"
-)
-
-// EtcdWithCache is a wrapper around the Etcd type which will return a
-// cached value when the number of requests come in below a configured
-// frequency. It also breaks the cache after a configured TTL.
-type EtcdWithCache struct {
-	*Etcd
-
-	peersMu sync.Mutex // peer-list cache updates
-
-	nodes            []*topology.Node // unmarshalled Node data
-	nodesTTL         int              // seconds
-	nodesLastRequest time.Time        // last time requested
-}
-
-// NewEtcdWithCache returns a new instance of Cache.
-func NewEtcdWithCache(opt Options, replicas int) *EtcdWithCache {
-	return &EtcdWithCache{
-		Etcd: NewEtcd(opt, replicas),
-
-		nodesTTL: 6,
-	}
-}
-
-// Nodes caches the result of the underlying implementation's node list.
-func (c *EtcdWithCache) Nodes() []*topology.Node {
-	c.peersMu.Lock()
-	defer c.peersMu.Unlock()
-
-	now := time.Now()
-	if now.Sub(c.nodesLastRequest) > (time.Duration(c.nodesTTL) * time.Second) {
-		c.nodes = c.Etcd.Nodes()
-		c.nodesLastRequest = now
-	}
-	return c.nodes
-}
diff --git a/server/server.go b/server/server.go
index 429dfc39c..d9c52cc99 100644
--- a/server/server.go
+++ b/server/server.go
@@ -396,7 +396,7 @@ func (m *Command) SetupServer() error {
 		m.Config.Etcd.Dir = filepath.Join(path, pilosa.DiscoDir)
 	}
 
-	e := petcd.NewEtcdWithCache(m.Config.Etcd, m.Config.Cluster.ReplicaN)
+	e := petcd.NewEtcd(m.Config.Etcd, m.Config.Cluster.ReplicaN)
 	discoOpt := pilosa.OptServerDisCo(e, e, e, e, e, e, e)
 
 	serverOptions := []pilosa.ServerOption{

From fd8a19278c936810e67d31f0532d8063e0a0c97a Mon Sep 17 00:00:00 2001
From: Seebs <seebs@seebs.net>
Date: Fri, 26 Mar 2021 12:40:43 -0500
Subject: [PATCH 2/5] add TestMain wrapper in ctl

The TestMain wrapper gets us the fancy testhook stuff tracking whether we're
deallocating things as expected, and we probably want that.
---
 ctl/main_test.go | 43 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100644 ctl/main_test.go

diff --git a/ctl/main_test.go b/ctl/main_test.go
new file mode 100644
index 000000000..3a724eafd
--- /dev/null
+++ b/ctl/main_test.go
@@ -0,0 +1,43 @@
+// Copyright 2020 Pilosa Corp.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ctl_test
+
+import (
+	"fmt"
+	"net"
+	"net/http"
+	"testing"
+
+	_ "net/http/pprof"
+
+	"github.com/pilosa/pilosa/v2/testhook"
+)
+
+func TestMain(m *testing.M) {
+	l, err := net.Listen("tcp", ":0")
+	if err != nil {
+		panic(err)
+	}
+	port := l.Addr().(*net.TCPAddr).Port
+	fmt.Printf("pilosa/ctl TestMain: online stack-traces: curl http://localhost:%v/debug/pprof/goroutine?debug=2\n", port)
+	go func() {
+		err := http.Serve(l, nil)
+		if err != nil {
+			panic(err)
+		}
+	}()
+	testhook.RunTestsWithHooks(m)
+
+}

From d1752a7af76dd72705c1c42fd15b0b55258788a2 Mon Sep 17 00:00:00 2001
From: Seebs <seebs@seebs.net>
Date: Fri, 19 Mar 2021 15:15:37 -0500
Subject: [PATCH 3/5] switch to using a watcher to watch etcd changes

This is a significant overhaul! Quite a lot of things changed here.

Basically: Prior to this, every request for data from etcd implies
requesting the current live data from etcd, and then unpacking it or
extracting it in some way. This is expensive, which is why we have
a cache in front of it.

We don't need to do that! We can use a Watch, which notifies us
of changes as changes happen. However, there's some challenges and
difficulties along the way, and there's a couple of other changes
which are included here because it's a pain to try to separate them
out.

1. We require a logger to be provided to create our internal Etcd
wrapper. We then use that logger, instead of `fmt.Printf`. This makes
debugging messages work better, and also diagnostics, and so on.

2. The internal client that we are reusing can enter a failed state
after a leader election, in which case we have to recreate the client
to have a working client. We add a new internal-use method,
`retryClient`, which wraps a function which takes an etcd client
and returns an error, and checks for leader-election type errors
and retries creating the client when they happen. That last bit
has not been successfully tested because it's actually really hard
to trigger this now. (Because it was related in part to the
amount of etcd traffic we were producing, which is reduced.)

3. The general swap over from looking things up to unpacking things
as they come in, then returning those already-unpacked things when
we get requests.

With this change, *many tests will fail*. That is addressed by
a separate commit which addresses the secondary problem, which is
that some of our test harness code was relying on the assumption
that if any node in a cluster thinks the cluster is up, every node
will. That was usually true when we were doing everything as
expensive fully-synchronized cluster checks, but becomes significantly
less reliably true in real-world cases where nodes are also
going down sometimes, or nodes are going up and down unexpectedly.
---
 cluster.go            |   8 -
 etcd/embed.go         | 598 +++++++++++++++++++++++++++++++-----------
 etcd/leasedkv.go      |  54 ++--
 etcd/leasedkv_test.go |   4 +-
 server/server.go      |   2 +-
 5 files changed, 481 insertions(+), 185 deletions(-)

diff --git a/cluster.go b/cluster.go
index d7d54edc9..a869e8d88 100644
--- a/cluster.go
+++ b/cluster.go
@@ -702,7 +702,6 @@ func (c *cluster) Nodes() []*topology.Node {
 	copiedNodes := make([]topology.Node, len(nodes))
 	result := make([]*topology.Node, len(nodes))
 
-	// Create a snapshot of the cluster to use for node/partition calculations.
 	primary := topology.PrimaryNode(nodes, c.Hasher)
 
 	// Set node states and IsPrimary.
@@ -712,13 +711,6 @@ func (c *cluster) Nodes() []*topology.Node {
 		if node == primary {
 			copiedNodes[i].IsPrimary = true
 		}
-		s, err := c.stator.NodeState(context.Background(), node.ID)
-		if err != nil {
-			// TODO should we delete this?
-			copiedNodes[i].State = disco.NodeStateUnknown
-			continue
-		}
-		copiedNodes[i].State = s
 	}
 	return result
 }
diff --git a/etcd/embed.go b/etcd/embed.go
index 09afb7b22..59304c922 100644
--- a/etcd/embed.go
+++ b/etcd/embed.go
@@ -18,13 +18,17 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"fmt"
 	"log"
 	"net"
 	"path"
 	"sort"
 	"strings"
+	"sync"
+	"time"
 
 	"github.com/pilosa/pilosa/v2/disco"
+	"github.com/pilosa/pilosa/v2/logger"
 	"github.com/pilosa/pilosa/v2/roaring"
 	"github.com/pilosa/pilosa/v2/topology"
 	"github.com/pkg/errors"
@@ -32,6 +36,7 @@ import (
 	"go.etcd.io/etcd/clientv3/clientv3util"
 	"go.etcd.io/etcd/clientv3/concurrency"
 	"go.etcd.io/etcd/embed"
+	"go.etcd.io/etcd/etcdserver"
 	"go.etcd.io/etcd/etcdserver/api/v3client"
 	"go.etcd.io/etcd/mvcc/mvccpb"
 	"go.etcd.io/etcd/pkg/transport"
@@ -70,28 +75,85 @@ var (
 )
 
 const (
-	heartbeatPrefix = "/heartbeat/"
+	// We put all the things the node-watcher watches in /node so we
+	// can use a single watcher for them.
+	nodePrefix      = "/node/"
+	heartbeatPrefix = "/node/heartbeat/"
 	schemaPrefix    = "/schema/"
-	resizePrefix    = "/resize/"
-	metadataPrefix  = "/metadata/"
+	resizePrefix    = "/node/resize/"
+	metadataPrefix  = "/node/metadata/"
 	shardPrefix     = "/shard/"
 	lockPrefix      = "/lock/"
 )
 
+var (
+	etcdLeaderChanged = etcdserver.ErrLeaderChanged.Error()
+)
+
+// nodeData is an internal tracker of the data we're keeping about
+// nodes in etcd, which we update from data collected either directly
+// from the KV, or via heartbeats.
+//
+// Any change to the topology.Node should create a new Node rather
+// than reusing the old one, so we can return the structure and not worry
+// about data races.
+//
+// We have to track the revisions of individual components so we can
+// discard updates which are genuinely out-of-order for a given field,
+// but still handle cases where we get updates to several fields that
+// reach us out of order.
+type nodeData struct {
+	heartbeatState string
+	resizeState    string
+	metadata       []byte
+	topologyNode   *topology.Node
+}
+
+func (n *nodeData) computedState() disco.NodeState {
+	if n.resizeState != "" {
+		return disco.NodeStateResizing
+	}
+	if n.heartbeatState != "" {
+		return disco.NodeState(n.heartbeatState)
+	}
+	return disco.NodeStateUnknown
+}
+
 type Etcd struct {
 	options  Options
 	replicas int
 
-	e   *embed.Etcd
-	cli *clientv3.Client
+	e     *embed.Etcd
+	cli   *clientv3.Client
+	cliMu sync.Mutex
 
 	heartbeatLeasedKV, resizeLeasedKV *leasedKV
+
+	// We have a watcher running. watchCancel() cancels its context.
+	watchCancel func()
+
+	// knownNodes and sortedNodes get updated by data coming in from
+	// watchers. Any change to the contents of a *topology.Node here
+	// should be implemented by making a new one and replacing the pointer,
+	// so the old pointer stays valid and can be used.
+	nodeMu          sync.Mutex
+	nodeRev         int64
+	knownNodes      map[string]*nodeData
+	sortedNodes     []*topology.Node
+	nodeStates      map[string]disco.NodeState
+	nodeStatesDirty bool // do we need to remake the nodeStates map to use it?
+
+	// we want to inherit parent's logging functionality
+	logger logger.Logger
 }
 
-func NewEtcd(opt Options, replicas int) *Etcd {
+func NewEtcd(opt Options, logger logger.Logger, replicas int) *Etcd {
 	e := &Etcd{
-		options:  opt,
-		replicas: replicas,
+		options:    opt,
+		logger:     logger,
+		replicas:   replicas,
+		knownNodes: make(map[string]*nodeData),
+		nodeStates: make(map[string]disco.NodeState),
 	}
 
 	if e.options.HeartbeatTTL == 0 {
@@ -102,6 +164,9 @@ func NewEtcd(opt Options, replicas int) *Etcd {
 
 // Close implements io.Closer
 func (e *Etcd) Close() error {
+	if e.watchCancel != nil {
+		e.watchCancel()
+	}
 	if e.e != nil {
 		if e.resizeLeasedKV != nil {
 			e.resizeLeasedKV.Stop()
@@ -122,6 +187,27 @@ func (e *Etcd) Close() error {
 	return nil
 }
 
+// retryClient attempts to do a thing, but also tries to handle the
+// specific case where the client fails because of a leader election,
+// in which case we need to restart the client and retry the thing.
+func (e *Etcd) retryClient(fn func(cli *clientv3.Client) error) (err error) {
+	e.cliMu.Lock()
+	cli := e.cli
+	e.cliMu.Unlock()
+	if err = fn(cli); err == nil || err.Error() != etcdLeaderChanged {
+		// either it's nil or it's an error we don't try to handle here
+		return err
+	}
+	// we can't do much with an error from closing e.cli at this point, so
+	// we try again.
+	_ = cli.Close()
+	cli = v3client.New(e.e.Server)
+	e.cliMu.Lock()
+	e.cli = cli
+	e.cliMu.Unlock()
+	return fn(cli)
+}
+
 func parseOptions(opt Options) *embed.Config {
 	cfg := embed.NewConfig()
 	cfg.Debug = false // true gives data races on grpc.EnableTracing in etcd
@@ -194,38 +280,61 @@ func parseOptions(opt Options) *embed.Config {
 }
 
 // Start starts etcd and hearbeat
-func (e *Etcd) Start(ctx context.Context) (disco.InitialClusterState, error) {
+func (e *Etcd) Start(ctx context.Context) (_ disco.InitialClusterState, err error) {
 	opts := parseOptions(e.options)
 	state := disco.InitialClusterState(opts.ClusterState)
 
-	etcd, err := embed.StartEtcd(opts)
+	e.e, err = embed.StartEtcd(opts)
 	if err != nil {
 		return state, errors.Wrap(err, "starting etcd")
 	}
-	e.e = etcd
+	// If we are returning an error, the caller won't be shutting us down
+	// later, so we have to stop the server ourselves.
+	defer func() {
+		if err != nil {
+			e.e.Server.Stop()
+		}
+	}()
 	e.cli = v3client.New(e.e.Server)
 
 	select {
 	case <-ctx.Done():
-		e.e.Server.Stop()
 		return state, ctx.Err()
 
 	case err := <-e.e.Err():
 		return state, err
 
 	case <-e.e.Server.ReadyNotify():
-		return state, e.startHeartbeat(ctx)
+		members := e.e.Server.Cluster().Members()
+		e.nodeMu.Lock()
+		defer e.nodeMu.Unlock()
+		// mark everything unknown so we show a state for nodes we haven't
+		// heard from yet.
+		for _, member := range members {
+			peerID := member.ID.String()
+			e.knownNodes[peerID] = &nodeData{}
+			e.nodeStates[peerID] = disco.NodeStateUnknown
+		}
+		e.nodeStatesDirty = true
+		return state, e.startHeartbeatAndWatcher(ctx)
 	}
 }
 
-func (e *Etcd) startHeartbeat(ctx context.Context) error {
+// startHeartbeatAndWatcher spins up the heartbeat, and also a background
+// watcher that watches for changes to events we care about.
+func (e *Etcd) startHeartbeatAndWatcher(ctx context.Context) error {
 	key := heartbeatPrefix + e.e.Server.ID().String()
-	e.heartbeatLeasedKV = newLeasedKV(e.cli, key, e.options.HeartbeatTTL)
+	// e.logger.Printf("starting heartbeat: %q => %q", e.options.Name, e.e.Server.ID())
+	e.heartbeatLeasedKV = newLeasedKV(e, key, e.options.HeartbeatTTL)
 
 	if err := e.heartbeatLeasedKV.Start(string(disco.NodeStateStarting)); err != nil {
 		return errors.Wrap(err, "startHeartbeat: starting a new heartbeat")
 	}
-
+	// WatchNodes does not check for an error, and will need to be shut
+	// down later. We only get this far at a point where we're returning
+	// a nil error, and thus, the caller is expected to cleanly shut down
+	// the server later.
+	go e.WatchNodes()
 	return nil
 }
 
@@ -234,70 +343,17 @@ func (e *Etcd) NodeState(ctx context.Context, peerID string) (disco.NodeState, e
 }
 
 func (e *Etcd) nodeState(ctx context.Context, peerID string) (disco.NodeState, error) {
-	resp, err := e.cli.Txn(ctx).
-		If(clientv3util.KeyMissing(path.Join(resizePrefix, peerID))).
-		Then(clientv3.OpGet(path.Join(heartbeatPrefix, peerID))).
-		Commit()
-	if err != nil {
-		return disco.NodeStateUnknown, err
-	}
-
-	if !resp.Succeeded {
-		return disco.NodeStateResizing, nil
-	}
-
-	if len(resp.Responses) == 0 {
-		return disco.NodeStateUnknown, disco.ErrNoResults
-	}
-
-	kvs := resp.Responses[0].GetResponseRange().Kvs
-	if len(kvs) == 0 {
-		return disco.NodeStateUnknown, disco.ErrNoResults
-	}
-	if len(kvs) > 1 {
-		return disco.NodeStateUnknown, disco.ErrTooManyResults
-	}
-
-	return disco.NodeState(kvs[0].Value), nil
+	e.nodeMu.Lock()
+	defer e.nodeMu.Unlock()
+	err := e.populateNodeStates(ctx)
+	return e.nodeStates[peerID], err
 }
 
 func (e *Etcd) NodeStates(ctx context.Context) (map[string]disco.NodeState, error) {
-	members := e.e.Server.Cluster().Members()
-	ops := make([]clientv3.Op, 2*(len(members)))
-	for i, member := range members {
-		peerID := member.ID.String()
-		ops[2*i] = clientv3.OpGet(path.Join(resizePrefix, peerID), clientv3.WithCountOnly())
-		ops[2*i+1] = clientv3.OpGet(path.Join(heartbeatPrefix, peerID))
-	}
-
-	resp, err := e.cli.Txn(ctx).Then(ops...).Commit()
-	if err != nil {
-		return nil, err
-	}
-
-	out := make(map[string]disco.NodeState, len(members))
-	for i, member := range members {
-		peerID := member.ID.String()
-		if resp.Responses[2*i].GetResponseRange().Count > 0 {
-			// This node is processing a resize operation.
-			out[peerID] = disco.NodeStateResizing
-			continue
-		}
-
-		kvs := resp.Responses[2*i+1].GetResponseRange().Kvs
-		switch len(kvs) {
-		case 0:
-			// The node has not reported a state.
-			out[peerID] = disco.NodeStateUnknown
-		case 1:
-			// The node has reported its state.
-			out[peerID] = disco.NodeState(kvs[0].Value)
-		default:
-			return nil, disco.ErrTooManyResults
-		}
-	}
-
-	return out, nil
+	e.nodeMu.Lock()
+	defer e.nodeMu.Unlock()
+	err := e.populateNodeStates(ctx)
+	return e.nodeStates, err
 }
 
 func (e *Etcd) Started(ctx context.Context) (err error) {
@@ -337,7 +393,7 @@ func (e *Etcd) Leader() *disco.Peer {
 	return peer
 }
 
-func (e *Etcd) ClusterState(ctx context.Context) (disco.ClusterState, error) {
+func (e *Etcd) ClusterState(ctx context.Context) (out disco.ClusterState, err error) {
 	if e.e == nil {
 		return disco.ClusterStateUnknown, nil
 	}
@@ -347,11 +403,17 @@ func (e *Etcd) ClusterState(ctx context.Context) (disco.ClusterState, error) {
 		resize     bool
 		starting   bool
 	)
-	states, err := e.NodeStates(ctx)
+	e.nodeMu.Lock()
+	err = e.populateNodeStates(ctx)
+	states := e.nodeStates
+	e.nodeMu.Unlock()
+	// defer func() {
+	// 	e.logger.Printf("ClusterState %q: out %s, states %#v", e.options.Name, out, states)
+	// }()
 	if err != nil {
+		e.logger.Printf("ClusterState %q: getting node states: %v", e.options.Name, states)
 		return disco.ClusterStateUnknown, err
 	}
-
 	for _, state := range states {
 		switch state {
 		case disco.NodeStateStarting:
@@ -387,7 +449,7 @@ func (e *Etcd) ClusterState(ctx context.Context) (disco.ClusterState, error) {
 func (e *Etcd) Resize(ctx context.Context) (func([]byte) error, error) {
 	key := path.Join(resizePrefix, e.e.Server.ID().String())
 	if e.resizeLeasedKV == nil {
-		e.resizeLeasedKV = newLeasedKV(e.cli, key, e.options.HeartbeatTTL)
+		e.resizeLeasedKV = newLeasedKV(e, key, e.options.HeartbeatTTL)
 	}
 
 	if err := e.resizeLeasedKV.Start(""); err != nil {
@@ -435,6 +497,217 @@ func (e *Etcd) Watch(ctx context.Context, peerID string, onUpdate func([]byte) e
 	return nil
 }
 
+// parseNodeKey reads "/node/heartbeat/23" and yields ("heartbeat", "23", nil).
+func parseNodeKey(key []byte) (prefix string, peerID string, err error) {
+	// we're looking for "/node/"
+	if !bytes.HasPrefix(key, []byte("/node/")) {
+		return "", "", fmt.Errorf("not a node key: %q", key)
+	}
+	vals := bytes.SplitN(key[6:], []byte("/"), 2)
+	if len(vals) != 2 {
+		return "", "", fmt.Errorf("not a valid node key: %q", key)
+	}
+	return string(vals[0]), string(vals[1]), nil
+}
+
+// deleteNodeData is like putNodeData, but handles deletes rather than cases
+// where a value exists. you should call it with the node mutex locked.
+func (e *Etcd) deleteNodeData(key []byte, revision int64) error {
+	prefix, peerID, err := parseNodeKey(key)
+	if err != nil {
+		return err
+	}
+	if revision > e.nodeRev {
+		e.nodeRev = revision
+	}
+	switch prefix {
+	case "heartbeat":
+		if e.knownNodes[peerID] == nil {
+			e.knownNodes[peerID] = &nodeData{}
+		}
+		e.knownNodes[peerID].heartbeatState = ""
+		e.nodeStatesDirty = true
+	case "metadata":
+		if e.knownNodes[peerID] == nil {
+			e.knownNodes[peerID] = &nodeData{}
+		}
+		e.knownNodes[peerID].metadata = nil
+		e.knownNodes[peerID].topologyNode = &topology.Node{}
+		e.nodeStatesDirty = true
+	case "resize":
+		if e.knownNodes[peerID] == nil {
+			e.knownNodes[peerID] = &nodeData{}
+		}
+		e.knownNodes[peerID].resizeState = ""
+		e.nodeStatesDirty = true
+	default:
+		return fmt.Errorf("node watch: invalid prefix %q\n", prefix)
+	}
+	return nil
+}
+
+// putNodeData does the actual updating of the node state maps, etc,
+// given an incoming heartbeat, metadata, or resizing change. It requires
+// that you already hold the node mutex.
+func (e *Etcd) putNodeData(key []byte, value []byte, revision int64) (err error) {
+	// defer func() {
+	// 	e.logger.Printf("putNodeData %q: key %q, rev %d, dirty %t, err %v", e.options.Name, key, e.nodeRev, e.nodeStatesDirty, err)
+	// }()
+	prefix, peerID, err := parseNodeKey(key)
+	if err != nil {
+		return err
+	}
+	if revision > e.nodeRev {
+		e.nodeRev = revision
+	}
+	switch prefix {
+	case "heartbeat":
+		if e.knownNodes[peerID] == nil {
+			e.knownNodes[peerID] = &nodeData{}
+		}
+		e.knownNodes[peerID].heartbeatState = string(value)
+		e.nodeStatesDirty = true
+	case "metadata":
+		if e.knownNodes[peerID] == nil {
+			e.knownNodes[peerID] = &nodeData{}
+		}
+		e.knownNodes[peerID].metadata = value
+		var newNode topology.Node
+		err := json.Unmarshal(value, &newNode)
+		if err != nil {
+			return fmt.Errorf("json unmarshal of node metadata: %v\n", err)
+		}
+		// e.logger.Printf("unmarshalling topologyNode for %q", peerID)
+		e.knownNodes[peerID].topologyNode = &newNode
+		// This saves us one remake of the node later, probably.
+		e.knownNodes[peerID].topologyNode.State = e.knownNodes[peerID].computedState()
+		e.nodeStatesDirty = true
+	case "resize":
+		if e.knownNodes[peerID] == nil {
+			e.knownNodes[peerID] = &nodeData{}
+		}
+		e.knownNodes[peerID].resizeState = string(value)
+		e.nodeStatesDirty = true
+	default:
+		return fmt.Errorf("node watch: invalid prefix %q\n", prefix)
+	}
+	return nil
+}
+
+// compute the states of all the nodes. we compute all of them because
+// we might have returned the old map in response to a query, so we want to
+// make a new one. You should have the node state lock held when you call this.
+func (e *Etcd) populateNodeStates(ctx context.Context) error {
+	if !e.nodeStatesDirty {
+		return nil
+	}
+	updated := false
+	for peerID, data := range e.knownNodes {
+		newState := data.computedState()
+		if data.topologyNode == nil || newState != data.topologyNode.State || newState != e.nodeStates[peerID] {
+			updated = true
+			break
+		}
+	}
+	// no changes
+	if !updated {
+		e.nodeStatesDirty = false
+		return nil
+	}
+	// make a new node state map, because something changed.
+	e.nodeStates = make(map[string]disco.NodeState, len(e.knownNodes))
+	e.sortedNodes = make([]*topology.Node, 0, len(e.knownNodes))
+	for peerID, data := range e.knownNodes {
+		newState := data.computedState()
+		e.nodeStates[peerID] = newState
+		// update the state with the current state, so we can
+		// reuse these nodes later. sortedNodes may end up shorter
+		// than the whole node list if we don't have all the nodes
+		// yet!
+		if data.topologyNode != nil {
+			if data.topologyNode.State != newState {
+				newNode := *data.topologyNode
+				newNode.State = newState
+				data.topologyNode = &newNode
+			}
+			e.sortedNodes = append(e.sortedNodes, data.topologyNode)
+		} else {
+			e.logger.Printf("no topologyNode for peer %q, metadata %q", peerID, data.metadata)
+		}
+	}
+	// sort list by ID. list now contains sorted nodes which have their
+	// current states.
+	sort.Sort(topology.ByID(e.sortedNodes))
+	e.nodeStatesDirty = false
+	return nil
+}
+
+// watchNodesOnce is a helper function to use with the retry logic
+// to let us restart the client if we need to.
+func (e *Etcd) watchNodesOnce(ctx context.Context, cli *clientv3.Client) (err error) {
+	e.nodeMu.Lock()
+	// we are looking for revisions HIGHER than the highest revision we've
+	// currently seen, we don't want one equal to it.
+	minRev := e.nodeRev + 1
+	e.nodeMu.Unlock()
+	// e.logger.Printf("watchNodesOnce %q: minRev %d", e.options.Name, minRev)
+	for resp := range cli.Watch(ctx, nodePrefix, clientv3.WithPrefix(), clientv3.WithRev(minRev)) {
+		// e.logger.Printf("watchp %q: resp rev %d with err %v, %d events\n", e.options.Name, resp.Header.Revision, resp.Err(), len(resp.Events))
+		if err := resp.Err(); err != nil {
+			e.logger.Printf("node watch %q: resp err %v", e.options.Name, err)
+			return err
+		}
+		// lock the node mutex for this whole process of updating so
+		// we never see partial updates; everything that comes into the
+		// watcher as a single message will be processed atomically.
+		e.nodeMu.Lock()
+		for _, ev := range resp.Events {
+			switch ev.Type {
+			case mvccpb.PUT:
+				// e.logger.Printf("watchp %q: PUT rev %d value %q=%q\n", e.options.Name, ev.Kv.ModRevision, ev.Kv.Key, ev.Kv.Value)
+				err := e.putNodeData(ev.Kv.Key, ev.Kv.Value, ev.Kv.ModRevision)
+				if err != nil {
+					e.logger.Printf("put event: %v", err)
+				}
+			case mvccpb.DELETE:
+				// e.logger.Printf("watchp %q: DELETE %q\n", e.options.Name, ev.Kv.Key)
+				err := e.deleteNodeData(ev.Kv.Key, ev.Kv.ModRevision)
+				if err != nil {
+					e.logger.Printf("delete event: %v", err)
+				}
+			default:
+				e.logger.Printf("watchp %q: unknown event %#v", e.options.Name, ev)
+			}
+		}
+		e.nodeMu.Unlock()
+
+	}
+	return nil
+}
+
+// WatchNodes monitors changes to /heartbeat/, /resizing/, and /metadata/;
+// basically, it catches changes to cluster state, but ignores the schema.
+func (e *Etcd) WatchNodes() {
+	ctx, cancel := context.WithCancel(context.Background())
+	e.watchCancel = cancel
+	// e.logger.Printf("WatchNodes %q: starting with rev %d", e.options.Name, e.nodeRev)
+	watchInContext := func(cli *clientv3.Client) error {
+		return e.watchNodesOnce(ctx, cli)
+	}
+	// retryClient will retry on leader failure, but not for other failures
+	// such as ErrCompacted which can terminate a watch. But we want to resume
+	// watching again as long as our context isn't cancelled. The context
+	// should get cancelled when this Etcd gets shut down.
+	for ctx.Err() == nil {
+		err := e.retryClient(watchInContext)
+		if err != nil {
+			e.logger.Printf("WatchNodes: error from watch client: %v", err)
+		}
+		// delay slightly on error so we don't go completely crazy
+		time.Sleep(1 * time.Second)
+	}
+}
+
 func (e *Etcd) DeleteNode(ctx context.Context, nodeID string) error {
 	id, err := types.IDFromString(nodeID)
 	if err != nil {
@@ -505,29 +778,17 @@ func (e *Etcd) Schema(ctx context.Context) (disco.Schema, error) {
 }
 
 func (e *Etcd) Metadata(ctx context.Context, peerID string) ([]byte, error) {
-	key := path.Join(metadataPrefix, peerID)
-
-	resp, err := e.cli.Txn(ctx).
-		If(clientv3util.KeyExists(key)).
-		Then(clientv3.OpGet(key)).
-		Commit()
+	e.nodeMu.Lock()
+	defer e.nodeMu.Unlock()
+	err := e.populateNodeStates(ctx)
 	if err != nil {
-		return nil, errors.Wrapf(err, "Metadata(%s)", key)
+		return nil, err
 	}
-
-	if !resp.Succeeded || len(resp.Responses) == 0 {
-		return nil, disco.ErrNoResults
+	data, ok := e.knownNodes[peerID]
+	if !ok {
+		return nil, errors.New("node not found")
 	}
-
-	kvs := resp.Responses[0].GetResponseRange().Kvs
-	if len(kvs) == 0 {
-		return nil, disco.ErrNoResults
-	}
-	if len(kvs) > 1 {
-		return nil, disco.ErrTooManyResults
-	}
-
-	return kvs[0].Value, nil
+	return data.metadata, nil
 }
 
 func (e *Etcd) SetMetadata(ctx context.Context, metadata []byte) error {
@@ -550,10 +811,14 @@ func (e *Etcd) CreateIndex(ctx context.Context, name string, val []byte) error {
 	op.WithValueBytes(val)
 
 	// Check for key existence, and execute Op within a transaction.
-	resp, err := e.cli.Txn(ctx).
-		If(clientv3util.KeyMissing(key)).
-		Then(op).
-		Commit()
+	var resp *clientv3.TxnResponse
+	err := e.retryClient(func(cli *clientv3.Client) (err error) {
+		resp, err = cli.Txn(ctx).
+			If(clientv3util.KeyMissing(key)).
+			Then(op).
+			Commit()
+		return err
+	})
 	if err != nil {
 		return errors.Wrap(err, "executing transaction")
 	}
@@ -572,12 +837,16 @@ func (e *Etcd) Index(ctx context.Context, name string) ([]byte, error) {
 func (e *Etcd) DeleteIndex(ctx context.Context, name string) (err error) {
 	key := schemaPrefix + name
 	// Deleting index and fields in one transaction.
-	_, err = e.cli.Txn(ctx).
-		If(clientv3.Compare(clientv3.Version(key), ">", -1)).
-		Then(
-			clientv3.OpDelete(key+"/", clientv3.WithPrefix()), // deleting index fields
-			clientv3.OpDelete(key),                            // deleting index
-		).Commit()
+
+	err = e.retryClient(func(cli *clientv3.Client) error {
+		_, err = cli.Txn(ctx).
+			If(clientv3.Compare(clientv3.Version(key), ">", -1)).
+			Then(
+				clientv3.OpDelete(key+"/", clientv3.WithPrefix()), // deleting index fields
+				clientv3.OpDelete(key),                            // deleting index
+			).Commit()
+		return err
+	})
 
 	return errors.Wrap(err, "DeleteIndex")
 }
@@ -595,10 +864,15 @@ func (e *Etcd) CreateField(ctx context.Context, indexName string, name string, v
 	op.WithValueBytes(val)
 
 	// Check for key existence, and execute Op within a transaction.
-	resp, err := e.cli.Txn(ctx).
-		If(clientv3util.KeyMissing(key)).
-		Then(op).
-		Commit()
+	var resp *clientv3.TxnResponse
+
+	err := e.retryClient(func(cli *clientv3.Client) (err error) {
+		resp, err = cli.Txn(ctx).
+			If(clientv3util.KeyMissing(key)).
+			Then(op).
+			Commit()
+		return err
+	})
 	if err != nil {
 		return errors.Wrap(err, "executing transaction")
 	}
@@ -613,12 +887,15 @@ func (e *Etcd) CreateField(ctx context.Context, indexName string, name string, v
 func (e *Etcd) DeleteField(ctx context.Context, indexname string, name string) (err error) {
 	key := schemaPrefix + indexname + "/" + name
 	// Deleting field and views in one transaction.
-	_, err = e.cli.Txn(ctx).
-		If(clientv3.Compare(clientv3.Version(key), ">", -1)).
-		Then(
-			clientv3.OpDelete(key+"/", clientv3.WithPrefix()), // deleting field views
-			clientv3.OpDelete(key),                            // deleting field
-		).Commit()
+	err = e.retryClient(func(cli *clientv3.Client) (err error) {
+		_, err = cli.Txn(ctx).
+			If(clientv3.Compare(clientv3.Version(key), ">", -1)).
+			Then(
+				clientv3.OpDelete(key+"/", clientv3.WithPrefix()), // deleting field views
+				clientv3.OpDelete(key),                            // deleting field
+			).Commit()
+		return err
+	})
 
 	return errors.Wrap(err, "DeleteField")
 }
@@ -635,10 +912,13 @@ func (e *Etcd) CreateView(ctx context.Context, indexName, fieldName, name string
 	key := schemaPrefix + indexName + "/" + fieldName + "/" + name
 
 	// Check for key existence, and execute Op within a transaction.
-	_, err = e.cli.Txn(ctx).
-		If(clientv3util.KeyMissing(key)).
-		Then(clientv3.OpPut(key, "")).
-		Commit()
+	err = e.retryClient(func(cli *clientv3.Client) (err error) {
+		_, err = cli.Txn(ctx).
+			If(clientv3util.KeyMissing(key)).
+			Then(clientv3.OpPut(key, "")).
+			Commit()
+		return err
+	})
 	if err != nil {
 		return errors.Wrap(err, "executing transaction")
 	}
@@ -651,19 +931,23 @@ func (e *Etcd) DeleteView(ctx context.Context, indexName, fieldName, name string
 }
 
 func (e *Etcd) putKey(ctx context.Context, key, val string, opts ...clientv3.OpOption) error {
-	if _, err := e.cli.Txn(ctx).
-		Then(clientv3.OpPut(key, val, opts...)).
-		Commit(); err != nil {
-		return errors.Wrapf(err, "putKey: Put(%s, %s)", key, val)
-	}
-
-	return nil
+	err := e.retryClient(func(cli *clientv3.Client) (err error) {
+		_, err = cli.Txn(ctx).
+			Then(clientv3.OpPut(key, val, opts...)).
+			Commit()
+		return err
+	})
+	return errors.Wrapf(err, "putKey: Put(%s, %s)", key, val)
 }
 
 func (e *Etcd) getKeyBytes(ctx context.Context, key string) ([]byte, error) {
 	// Get the current value for the key.
 	op := clientv3.OpGet(key)
-	resp, err := e.cli.Txn(ctx).Then(op).Commit()
+	var resp *clientv3.TxnResponse
+	err := e.retryClient(func(cli *clientv3.Client) (err error) {
+		resp, err = cli.Txn(ctx).Then(op).Commit()
+		return err
+	})
 	if err != nil {
 		return nil, err
 	}
@@ -681,9 +965,12 @@ func (e *Etcd) getKeyBytes(ctx context.Context, key string) ([]byte, error) {
 }
 
 func (e *Etcd) getKeyWithPrefix(ctx context.Context, key string) (keys []string, values [][]byte, err error) {
-	resp, err := e.cli.Txn(ctx).
-		Then(clientv3.OpGet(key, clientv3.WithPrefix())).
-		Commit()
+	op := clientv3.OpGet(key, clientv3.WithPrefix())
+	var resp *clientv3.TxnResponse
+	err = e.retryClient(func(cli *clientv3.Client) (err error) {
+		resp, err = cli.Txn(ctx).Then(op).Commit()
+		return err
+	})
 	if err != nil {
 		return nil, nil, errors.Wrapf(err, "getKeyWithPrefix(%s)", key)
 	}
@@ -708,10 +995,14 @@ func (e *Etcd) getKeyWithPrefix(ctx context.Context, key string) (keys []string,
 }
 
 func (e *Etcd) keyExists(ctx context.Context, key string) (bool, error) {
-	resp, err := e.cli.Txn(ctx).
-		If(clientv3util.KeyExists(key)).
-		Then(clientv3.OpGet(key, clientv3.WithCountOnly())).
-		Commit()
+	var resp *clientv3.TxnResponse
+	err := e.retryClient(func(cli *clientv3.Client) (err error) {
+		resp, err = cli.Txn(ctx).
+			If(clientv3util.KeyExists(key)).
+			Then(clientv3.OpGet(key, clientv3.WithCountOnly())).
+			Commit()
+		return err
+	})
 	if err != nil {
 		return false, err
 	}
@@ -980,26 +1271,15 @@ func (e *Etcd) RemoveShard(ctx context.Context, index, field string, shard uint6
 // Nodes implements the Noder interface. It returns the sorted list of nodes
 // based on the etcd peers.
 func (e *Etcd) Nodes() []*topology.Node {
-	peers := e.Peers()
-	nodes := make([]*topology.Node, 0, len(peers))
-	for _, peer := range peers {
-		var node *topology.Node
-
-		if meta, err := e.Metadata(context.Background(), peer.ID); err != nil {
-			log.Println(err, "getting metadata") // TODO: handle this with a logger
-		} else if err := json.Unmarshal(meta, &node); err != nil {
-			log.Println(err, "unmarshaling json metadata")
-		}
-
-		if node != nil && peer != nil {
-			node.ID = peer.ID
-			nodes = append(nodes, node)
-		}
+	e.nodeMu.Lock()
+	defer e.nodeMu.Unlock()
+	err := e.populateNodeStates(context.TODO())
+	if err != nil {
+		// e.logger.Printf("requesting node list: %v", err)
+		return nil
 	}
-
-	// Nodes must be sorted.
-	sort.Sort(topology.ByID(nodes))
-	return nodes
+	// e.logger.Printf("Nodes(): returning %d nodes", len(e.sortedNodes))
+	return e.sortedNodes
 }
 
 // PrimaryNodeID implements the Noder interface.
diff --git a/etcd/leasedkv.go b/etcd/leasedkv.go
index 17a35ed0b..4356821c3 100644
--- a/etcd/leasedkv.go
+++ b/etcd/leasedkv.go
@@ -30,7 +30,7 @@ import (
 // It will try to renew the lease at any cost after losing it.
 // It will recreate the previous existing value for the key again.
 type leasedKV struct {
-	cli    *clientv3.Client
+	e      *Etcd
 	cancel context.CancelFunc
 
 	key        string
@@ -41,9 +41,9 @@ type leasedKV struct {
 	stopped bool   // protected by mu
 }
 
-func newLeasedKV(cli *clientv3.Client, key string, ttlSeconds int64) *leasedKV {
+func newLeasedKV(e *Etcd, key string, ttlSeconds int64) *leasedKV {
 	return &leasedKV{
-		cli:        cli,
+		e:          e,
 		key:        key,
 		ttlSeconds: ttlSeconds,
 	}
@@ -72,26 +72,38 @@ func (l *leasedKV) create(initValue string) (<-chan *clientv3.LeaseKeepAliveResp
 	}
 
 	l.cancel = cancel
+	var leaseResp *clientv3.LeaseGrantResponse
 
-	leaseResp, err := l.cli.Grant(ctx, l.ttlSeconds)
+	err := l.e.retryClient(func(cli *clientv3.Client) (err error) {
+		leaseResp, err = cli.Grant(ctx, l.ttlSeconds)
+		return err
+	})
 	if err != nil {
 		return nil, errors.Wrap(err, "creating a lease")
 	}
 
-	if _, err := l.cli.Txn(ctx).
-		Then(clientv3.OpPut(l.key, initValue, clientv3.WithLease(leaseResp.ID))).
-		Commit(); err != nil {
+	err = l.e.retryClient(func(cli *clientv3.Client) (err error) {
+		_, err = cli.Txn(ctx).
+			Then(clientv3.OpPut(l.key, initValue, clientv3.WithLease(leaseResp.ID))).
+			Commit()
+		return err
+	})
+	if err != nil {
 		return nil, errors.Wrapf(err, "creating key %s with value [%s]", l.key, initValue)
 	}
 
-	kaChann, err := l.cli.KeepAlive(ctx, leaseResp.ID)
+	var kaChan <-chan *clientv3.LeaseKeepAliveResponse
+	err = l.e.retryClient(func(cli *clientv3.Client) (err error) {
+		kaChan, err = cli.KeepAlive(ctx, leaseResp.ID)
+		return err
+	})
 	if err != nil {
 		return nil, errors.Wrapf(err, "keeping alive the lease for the key %s with value %s", l.key, l.value)
 	}
 
 	l.value = initValue
 
-	return kaChann, nil
+	return kaChan, nil
 }
 
 func (l *leasedKV) consumeLease(ch <-chan *clientv3.LeaseKeepAliveResponse) {
@@ -144,9 +156,15 @@ func (l *leasedKV) Set(ctx context.Context, value string) error {
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
-	if _, err := l.cli.Txn(ctx).
-		Then(clientv3.OpPut(l.key, value, clientv3.WithIgnoreLease())).
-		Commit(); err != nil {
+	err := l.e.retryClient(func(cli *clientv3.Client) (err error) {
+		_, err = cli.Txn(ctx).
+			Then(clientv3.OpPut(l.key, value, clientv3.WithIgnoreLease())).
+			Commit()
+		return err
+	})
+	// l.e.logger.Printf("set key %q on %q value %q: err %v", l.key, l.e.options.Name, value, err)
+
+	if err != nil {
 		return errors.Wrapf(err, "creating key %s with value [%s]", l.key, l.value)
 	}
 
@@ -160,10 +178,14 @@ func (l *leasedKV) Get(ctx context.Context) (string, error) {
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
-	getResp, err := l.cli.Txn(ctx).
-		If(clientv3util.KeyExists(l.key)).
-		Then(clientv3.OpGet(l.key, clientv3.WithIgnoreLease())).
-		Commit()
+	var getResp *clientv3.TxnResponse
+	err := l.e.retryClient(func(cli *clientv3.Client) (err error) {
+		getResp, err = cli.Txn(ctx).
+			If(clientv3util.KeyExists(l.key)).
+			Then(clientv3.OpGet(l.key, clientv3.WithIgnoreLease())).
+			Commit()
+		return err
+	})
 	if err != nil {
 		return "", errors.Wrapf(err, "getting key %s", l.key)
 	}
diff --git a/etcd/leasedkv_test.go b/etcd/leasedkv_test.go
index fb54fab93..8a8077764 100644
--- a/etcd/leasedkv_test.go
+++ b/etcd/leasedkv_test.go
@@ -22,6 +22,7 @@ import (
 	"time"
 
 	"github.com/pilosa/pilosa/v2/disco"
+	"github.com/pilosa/pilosa/v2/logger"
 	"github.com/pilosa/pilosa/v2/testhook"
 	"go.etcd.io/etcd/embed"
 	"go.etcd.io/etcd/etcdserver/api/v3client"
@@ -50,8 +51,9 @@ func TestLeasedKv(t *testing.T) {
 		<-etcd.Server.StopNotify()
 		cli.Close()
 	}()
+	wrapper := &Etcd{e: etcd, cli: cli, logger: logger.NewLogfLogger(t)}
 
-	lkv := newLeasedKV(cli, "/test", 1)
+	lkv := newLeasedKV(wrapper, "/test", 1)
 
 	ctx := context.Background()
 
diff --git a/server/server.go b/server/server.go
index d9c52cc99..1f70d29d0 100644
--- a/server/server.go
+++ b/server/server.go
@@ -396,7 +396,7 @@ func (m *Command) SetupServer() error {
 		m.Config.Etcd.Dir = filepath.Join(path, pilosa.DiscoDir)
 	}
 
-	e := petcd.NewEtcd(m.Config.Etcd, m.Config.Cluster.ReplicaN)
+	e := petcd.NewEtcd(m.Config.Etcd, m.logger, m.Config.Cluster.ReplicaN)
 	discoOpt := pilosa.OptServerDisCo(e, e, e, e, e, e, e)
 
 	serverOptions := []pilosa.ServerOption{

From 0638101d2a2096a995d72654ac5a6e59a307f73c Mon Sep 17 00:00:00 2001
From: Seebs <seebs@seebs.net>
Date: Wed, 24 Mar 2021 11:36:14 -0500
Subject: [PATCH 4/5] cluster state checking cleanups and fixes

A while back we started just polling the reported cluster state of one node
when starting a cluster for tests. This works fine if we're doing fresh
new etcd queries for every single operation -- but that's insanely
expensive, it turns out.

When we use the watcher, some nodes will report stale data for "a
while", where "a while" appears to be easily a couple dozen milliseconds.
This is probably irrelevant in most real-world cases, because the common
case (detecting a node going down) means that we have at least five
seconds after a node goes down before etcd notices the lease expiring,
and a few milliseconds more or less won't matter.

But we have tests that assume either that node 0 is always the
coordinator (wrong) or that waiting for node 0 to think the cluster
is up means that every node in the cluster thinks the cluster is up,
or at least that it means that the coordinator thinks the cluster is
up. We retried later operations but not the initial ones against
the coordinator.

In fact, we probably want to wait for the entire cluster to think
it's up before we start trying things on clusters.

We also replace the "CheckClusterState" function with the existing
AwaitState call, or a new AssertState which errors out since that's
the way we usually use AwaitState anyway.

In the AwaitPrimaryState function, which used to be
AwaitCoordinatorState in a different long-lost revision, we have
to delay until a primary node is available, or fail if one does
not become available, to avoid a panic. This probably shouldn't
happen anymore, because of the last change:

Also, rovide dummy topology.Node entries before metadata is read.

During initial startup, we want to be able to do things like determine
which node is the primary, even before we've read metadata from them.

To do this, we populate the node list with dummy entries that just have
the ID (the only part we need to sort our list), and a node state of
UNKNOWN.

This breaks the fancy logic for determining whether or not to update
the node data, because the initial status of UNKNOWN matches what we
get from SetMetadata giving us new data so we end up not realizing
that this was actually a meaningful change. But actually, that's
a pretty niche optimization; we usually only get state changes when
there's an actual change in state. The updates here are cheap
and only happen after a write (or on the first query) so it's not
worth making the logic a lot fancier to make it work, when we can
just do the simple thing and update any time the dirty flag is set.

We also standardize on a 50ms delay, because 1ms delays were
really expensive when each check was hitting etcd multiple times,
and 50ms is Usually Long Enough.
---
 etcd/embed.go          | 82 +++++++++++++++++-------------------------
 executor_test.go       |  4 +--
 server/cluster_test.go | 82 ++++++++++--------------------------------
 server/server_test.go  | 22 ++++++------
 test/cluster.go        | 80 +++++++++++++++++++++++++++++++----------
 test/pilosa.go         | 19 +++++++++-
 6 files changed, 145 insertions(+), 144 deletions(-)

diff --git a/etcd/embed.go b/etcd/embed.go
index 59304c922..6afe9f547 100644
--- a/etcd/embed.go
+++ b/etcd/embed.go
@@ -78,10 +78,10 @@ const (
 	// We put all the things the node-watcher watches in /node so we
 	// can use a single watcher for them.
 	nodePrefix      = "/node/"
-	heartbeatPrefix = "/node/heartbeat/"
+	heartbeatPrefix = nodePrefix + "heartbeat/"
 	schemaPrefix    = "/schema/"
-	resizePrefix    = "/node/resize/"
-	metadataPrefix  = "/node/metadata/"
+	resizePrefix    = nodePrefix + "resize/"
+	metadataPrefix  = nodePrefix + "metadata/"
 	shardPrefix     = "/shard/"
 	lockPrefix      = "/lock/"
 )
@@ -190,6 +190,12 @@ func (e *Etcd) Close() error {
 // retryClient attempts to do a thing, but also tries to handle the
 // specific case where the client fails because of a leader election,
 // in which case we need to restart the client and retry the thing.
+//
+// We have to let go of the lock while calling `fn` because some fn are
+// long-lasting ones, like watchNodesOnce. So we grab a local copy of
+// the client object, then call things on that object. This should error
+// out sanely instead of panicing if we close the client while something
+// is running on it.
 func (e *Etcd) retryClient(fn func(cli *clientv3.Client) error) (err error) {
 	e.cliMu.Lock()
 	cli := e.cli
@@ -200,9 +206,14 @@ func (e *Etcd) retryClient(fn func(cli *clientv3.Client) error) (err error) {
 	}
 	// we can't do much with an error from closing e.cli at this point, so
 	// we try again.
+	e.cliMu.Lock()
+	if cli != e.cli {
+		cli = e.cli
+		e.cliMu.Unlock()
+		return fn(cli)
+	}
 	_ = cli.Close()
 	cli = v3client.New(e.e.Server)
-	e.cliMu.Lock()
 	e.cli = cli
 	e.cliMu.Unlock()
 	return fn(cli)
@@ -312,7 +323,12 @@ func (e *Etcd) Start(ctx context.Context) (_ disco.InitialClusterState, err erro
 		// heard from yet.
 		for _, member := range members {
 			peerID := member.ID.String()
-			e.knownNodes[peerID] = &nodeData{}
+			e.knownNodes[peerID] = &nodeData{
+				topologyNode: &topology.Node{
+					ID:    peerID,
+					State: disco.NodeStateUnknown,
+				},
+			}
 			e.nodeStates[peerID] = disco.NodeStateUnknown
 		}
 		e.nodeStatesDirty = true
@@ -324,7 +340,6 @@ func (e *Etcd) Start(ctx context.Context) (_ disco.InitialClusterState, err erro
 // watcher that watches for changes to events we care about.
 func (e *Etcd) startHeartbeatAndWatcher(ctx context.Context) error {
 	key := heartbeatPrefix + e.e.Server.ID().String()
-	// e.logger.Printf("starting heartbeat: %q => %q", e.options.Name, e.e.Server.ID())
 	e.heartbeatLeasedKV = newLeasedKV(e, key, e.options.HeartbeatTTL)
 
 	if err := e.heartbeatLeasedKV.Start(string(disco.NodeStateStarting)); err != nil {
@@ -407,9 +422,6 @@ func (e *Etcd) ClusterState(ctx context.Context) (out disco.ClusterState, err er
 	err = e.populateNodeStates(ctx)
 	states := e.nodeStates
 	e.nodeMu.Unlock()
-	// defer func() {
-	// 	e.logger.Printf("ClusterState %q: out %s, states %#v", e.options.Name, out, states)
-	// }()
 	if err != nil {
 		e.logger.Printf("ClusterState %q: getting node states: %v", e.options.Name, states)
 		return disco.ClusterStateUnknown, err
@@ -497,17 +509,17 @@ func (e *Etcd) Watch(ctx context.Context, peerID string, onUpdate func([]byte) e
 	return nil
 }
 
-// parseNodeKey reads "/node/heartbeat/23" and yields ("heartbeat", "23", nil).
+// parseNodeKey reads heartbeatPrefix + "23" and yields (heartbeatPrefix, "23", nil).
 func parseNodeKey(key []byte) (prefix string, peerID string, err error) {
-	// we're looking for "/node/"
-	if !bytes.HasPrefix(key, []byte("/node/")) {
+	// we're looking for things starting with nodePrefix
+	if !bytes.HasPrefix(key, []byte(nodePrefix)) {
 		return "", "", fmt.Errorf("not a node key: %q", key)
 	}
-	vals := bytes.SplitN(key[6:], []byte("/"), 2)
-	if len(vals) != 2 {
+	peerIndex := bytes.LastIndex(key, []byte("/"))
+	if peerIndex < 6 {
 		return "", "", fmt.Errorf("not a valid node key: %q", key)
 	}
-	return string(vals[0]), string(vals[1]), nil
+	return string(key[:peerIndex+1]), string(key[peerIndex+1:]), nil
 }
 
 // deleteNodeData is like putNodeData, but handles deletes rather than cases
@@ -521,20 +533,20 @@ func (e *Etcd) deleteNodeData(key []byte, revision int64) error {
 		e.nodeRev = revision
 	}
 	switch prefix {
-	case "heartbeat":
+	case heartbeatPrefix:
 		if e.knownNodes[peerID] == nil {
 			e.knownNodes[peerID] = &nodeData{}
 		}
 		e.knownNodes[peerID].heartbeatState = ""
 		e.nodeStatesDirty = true
-	case "metadata":
+	case metadataPrefix:
 		if e.knownNodes[peerID] == nil {
 			e.knownNodes[peerID] = &nodeData{}
 		}
 		e.knownNodes[peerID].metadata = nil
 		e.knownNodes[peerID].topologyNode = &topology.Node{}
 		e.nodeStatesDirty = true
-	case "resize":
+	case resizePrefix:
 		if e.knownNodes[peerID] == nil {
 			e.knownNodes[peerID] = &nodeData{}
 		}
@@ -550,9 +562,6 @@ func (e *Etcd) deleteNodeData(key []byte, revision int64) error {
 // given an incoming heartbeat, metadata, or resizing change. It requires
 // that you already hold the node mutex.
 func (e *Etcd) putNodeData(key []byte, value []byte, revision int64) (err error) {
-	// defer func() {
-	// 	e.logger.Printf("putNodeData %q: key %q, rev %d, dirty %t, err %v", e.options.Name, key, e.nodeRev, e.nodeStatesDirty, err)
-	// }()
 	prefix, peerID, err := parseNodeKey(key)
 	if err != nil {
 		return err
@@ -561,13 +570,13 @@ func (e *Etcd) putNodeData(key []byte, value []byte, revision int64) (err error)
 		e.nodeRev = revision
 	}
 	switch prefix {
-	case "heartbeat":
+	case heartbeatPrefix:
 		if e.knownNodes[peerID] == nil {
 			e.knownNodes[peerID] = &nodeData{}
 		}
 		e.knownNodes[peerID].heartbeatState = string(value)
 		e.nodeStatesDirty = true
-	case "metadata":
+	case metadataPrefix:
 		if e.knownNodes[peerID] == nil {
 			e.knownNodes[peerID] = &nodeData{}
 		}
@@ -577,12 +586,11 @@ func (e *Etcd) putNodeData(key []byte, value []byte, revision int64) (err error)
 		if err != nil {
 			return fmt.Errorf("json unmarshal of node metadata: %v\n", err)
 		}
-		// e.logger.Printf("unmarshalling topologyNode for %q", peerID)
 		e.knownNodes[peerID].topologyNode = &newNode
 		// This saves us one remake of the node later, probably.
 		e.knownNodes[peerID].topologyNode.State = e.knownNodes[peerID].computedState()
 		e.nodeStatesDirty = true
-	case "resize":
+	case resizePrefix:
 		if e.knownNodes[peerID] == nil {
 			e.knownNodes[peerID] = &nodeData{}
 		}
@@ -601,20 +609,6 @@ func (e *Etcd) populateNodeStates(ctx context.Context) error {
 	if !e.nodeStatesDirty {
 		return nil
 	}
-	updated := false
-	for peerID, data := range e.knownNodes {
-		newState := data.computedState()
-		if data.topologyNode == nil || newState != data.topologyNode.State || newState != e.nodeStates[peerID] {
-			updated = true
-			break
-		}
-	}
-	// no changes
-	if !updated {
-		e.nodeStatesDirty = false
-		return nil
-	}
-	// make a new node state map, because something changed.
 	e.nodeStates = make(map[string]disco.NodeState, len(e.knownNodes))
 	e.sortedNodes = make([]*topology.Node, 0, len(e.knownNodes))
 	for peerID, data := range e.knownNodes {
@@ -631,8 +625,6 @@ func (e *Etcd) populateNodeStates(ctx context.Context) error {
 				data.topologyNode = &newNode
 			}
 			e.sortedNodes = append(e.sortedNodes, data.topologyNode)
-		} else {
-			e.logger.Printf("no topologyNode for peer %q, metadata %q", peerID, data.metadata)
 		}
 	}
 	// sort list by ID. list now contains sorted nodes which have their
@@ -650,11 +642,8 @@ func (e *Etcd) watchNodesOnce(ctx context.Context, cli *clientv3.Client) (err er
 	// currently seen, we don't want one equal to it.
 	minRev := e.nodeRev + 1
 	e.nodeMu.Unlock()
-	// e.logger.Printf("watchNodesOnce %q: minRev %d", e.options.Name, minRev)
 	for resp := range cli.Watch(ctx, nodePrefix, clientv3.WithPrefix(), clientv3.WithRev(minRev)) {
-		// e.logger.Printf("watchp %q: resp rev %d with err %v, %d events\n", e.options.Name, resp.Header.Revision, resp.Err(), len(resp.Events))
 		if err := resp.Err(); err != nil {
-			e.logger.Printf("node watch %q: resp err %v", e.options.Name, err)
 			return err
 		}
 		// lock the node mutex for this whole process of updating so
@@ -664,13 +653,11 @@ func (e *Etcd) watchNodesOnce(ctx context.Context, cli *clientv3.Client) (err er
 		for _, ev := range resp.Events {
 			switch ev.Type {
 			case mvccpb.PUT:
-				// e.logger.Printf("watchp %q: PUT rev %d value %q=%q\n", e.options.Name, ev.Kv.ModRevision, ev.Kv.Key, ev.Kv.Value)
 				err := e.putNodeData(ev.Kv.Key, ev.Kv.Value, ev.Kv.ModRevision)
 				if err != nil {
 					e.logger.Printf("put event: %v", err)
 				}
 			case mvccpb.DELETE:
-				// e.logger.Printf("watchp %q: DELETE %q\n", e.options.Name, ev.Kv.Key)
 				err := e.deleteNodeData(ev.Kv.Key, ev.Kv.ModRevision)
 				if err != nil {
 					e.logger.Printf("delete event: %v", err)
@@ -690,7 +677,6 @@ func (e *Etcd) watchNodesOnce(ctx context.Context, cli *clientv3.Client) (err er
 func (e *Etcd) WatchNodes() {
 	ctx, cancel := context.WithCancel(context.Background())
 	e.watchCancel = cancel
-	// e.logger.Printf("WatchNodes %q: starting with rev %d", e.options.Name, e.nodeRev)
 	watchInContext := func(cli *clientv3.Client) error {
 		return e.watchNodesOnce(ctx, cli)
 	}
@@ -1275,10 +1261,8 @@ func (e *Etcd) Nodes() []*topology.Node {
 	defer e.nodeMu.Unlock()
 	err := e.populateNodeStates(context.TODO())
 	if err != nil {
-		// e.logger.Printf("requesting node list: %v", err)
 		return nil
 	}
-	// e.logger.Printf("Nodes(): returning %d nodes", len(e.sortedNodes))
 	return e.sortedNodes
 }
 
diff --git a/executor_test.go b/executor_test.go
index 425a028a7..330ef1fd6 100644
--- a/executor_test.go
+++ b/executor_test.go
@@ -3764,7 +3764,7 @@ func TestReopenCluster(t *testing.T) {
 	if err := node0.Reopen(); err != nil {
 		t.Fatal(err)
 	}
-	if err := node0.AwaitState(disco.ClusterStateNormal, 10*time.Second); err != nil {
+	if err := c.AwaitState(disco.ClusterStateNormal, 10*time.Second); err != nil {
 		t.Fatalf("restarting cluster: %v", err)
 	}
 
@@ -3828,7 +3828,7 @@ func TestExecutor_Execute_Existence(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		if err := node0.AwaitState(disco.ClusterStateNormal, 10*time.Second); err != nil {
+		if err := c.AwaitState(disco.ClusterStateNormal, 10*time.Second); err != nil {
 			t.Fatalf("restarting cluster: %v", err)
 		}
 
diff --git a/server/cluster_test.go b/server/cluster_test.go
index 97b592abf..16ae4f7a5 100644
--- a/server/cluster_test.go
+++ b/server/cluster_test.go
@@ -159,13 +159,8 @@ func TestClusterResize_AddNode(t *testing.T) {
 		clus := test.MustRunCluster(t, 3)
 		defer clus.Close()
 
-		state0, err0 := clus.GetNode(0).API.State()
-		state1, err1 := clus.GetNode(1).API.State()
-		if err0 != nil || !test.CheckClusterState(clus.GetNode(0), disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node0 cluster state: %s, error: %v", state0, err0)
-		} else if err1 != nil || !test.CheckClusterState(clus.GetNode(1), disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node1 cluster state: %s, error: %v", state1, err1)
-		}
+		clus.GetNode(0).AssertState(t, disco.ClusterStateNormal, 1*time.Second)
+		clus.GetNode(1).AssertState(t, disco.ClusterStateNormal, 1*time.Second)
 	})
 	t.Run("WithIndex", func(t *testing.T) {
 		// Configure node0
@@ -207,13 +202,8 @@ func TestClusterResize_AddNode(t *testing.T) {
 
 		defer m1.Close()
 
-		state0, err0 := m0.API.State()
-		state1, err1 := m1.API.State()
-		if err0 != nil || !test.CheckClusterState(m0, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node0 cluster state: %s, error: %v", state0, err0)
-		} else if err1 != nil || !test.CheckClusterState(m1, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node1 cluster state: %s, error; %v", state1, err1)
-		}
+		m0.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
+		m1.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
 	})
 	t.Run("ContinuousShards", func(t *testing.T) {
 
@@ -253,13 +243,8 @@ func TestClusterResize_AddNode(t *testing.T) {
 		m1 := c.GetNode(1)
 		defer m1.Close()
 
-		state0, err0 := m0.API.State()
-		state1, err1 := m1.API.State()
-		if err0 != nil || !test.CheckClusterState(m0, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node0 cluster state: %s, error: %v", state0, err0)
-		} else if err1 != nil || !test.CheckClusterState(m1, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node1 cluster state: %s, error: %v", state1, err1)
-		}
+		m0.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
+		m1.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
 
 		// Verify the data exists on both nodes.
 		m0.QueryExpect(t, "i", "", `Row(f=1)`, exp)
@@ -301,13 +286,8 @@ func TestClusterResize_AddNode(t *testing.T) {
 		m1 := c.GetNode(1)
 		defer m1.Close()
 
-		state0, err0 := m0.API.State()
-		state1, err1 := m1.API.State()
-		if err0 != nil || !test.CheckClusterState(m0, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node0 cluster state: %s, error: %v", state0, err0)
-		} else if err1 != nil || !test.CheckClusterState(m1, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node1 cluster state: %s, error: %v", state1, err1)
-		}
+		m0.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
+		m1.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
 
 		// Verify the data exists on both nodes.
 		m0.QueryExpect(t, "i", "", `Row(f=1)`, exp)
@@ -353,13 +333,8 @@ func TestClusterResize_AddNode(t *testing.T) {
 		m1 := c.GetNode(1)
 		defer m1.Close()
 
-		state0, err0 := m0.API.State()
-		state1, err1 := m1.API.State()
-		if err0 != nil || !test.CheckClusterState(m0, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node0 cluster state: %s, error: %v", state0, err0)
-		} else if err1 != nil || !test.CheckClusterState(m1, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node1 cluster state: %s, error: %v", state1, err1)
-		}
+		m0.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
+		m1.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
 
 		// Verify the data exists on both nodes.
 		m0.QueryExpect(t, "i", "", `Row(f=1)`, exp)
@@ -399,13 +374,8 @@ func TestClusterResize_AddNodeConcurrentIndex(t *testing.T) {
 		m1 := c.GetNode(1)
 		defer m1.Close()
 
-		state0, err0 := m0.API.State()
-		state1, err1 := m1.API.State()
-		if err0 != nil || !test.CheckClusterState(m0, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node0 cluster state: %s, error: %v", state0, err0)
-		} else if err1 != nil || !test.CheckClusterState(m1, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node1 cluster state: %s, error: %v", state1, err1)
-		}
+		m0.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
+		m1.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
 
 		if err := <-errc; err != nil {
 			t.Fatalf("error from index creation: %v", err)
@@ -450,13 +420,8 @@ func TestClusterResize_AddNodeConcurrentIndex(t *testing.T) {
 		m1 := c.GetNode(1)
 		defer m1.Close()
 
-		state0, err0 := m0.API.State()
-		state1, err1 := m1.API.State()
-		if err0 != nil || !test.CheckClusterState(m0, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node0 cluster state: %s, error: %v", state0, err0)
-		} else if err1 != nil || !test.CheckClusterState(m1, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node1 cluster state: %s, error: %v", state1, err1)
-		}
+		m0.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
+		m1.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
 
 		// Verify the data exists on both nodes.
 		m0.QueryExpect(t, "i", "", `Row(f=1)`, exp)
@@ -501,13 +466,8 @@ func TestClusterResize_AddNodeConcurrentIndex(t *testing.T) {
 		m1 := c.GetNode(1)
 		defer m1.Close()
 
-		state0, err0 := m0.API.State()
-		state1, err1 := m1.API.State()
-		if err0 != nil || !test.CheckClusterState(m0, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node0 cluster state: %s, error: %v", state0, err0)
-		} else if err1 != nil || !test.CheckClusterState(m1, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node1 cluster state: %s, error: %v", state1, err1)
-		}
+		m0.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
+		m1.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
 
 		// Verify the data exists on both nodes.
 		m0.QueryExpect(t, "i", "", `Row(f=1)`, exp)
@@ -550,13 +510,9 @@ func TestClusterResize_AddNodeConcurrentIndex(t *testing.T) {
 		m1 := c.GetNode(1)
 		defer m1.Close()
 
-		state0, err0 := m0.API.State()
-		state1, err1 := m1.API.State()
-		if err0 != nil || !test.CheckClusterState(m0, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node0 cluster state: %s, error: %v", state0, err0)
-		} else if err1 != nil || !test.CheckClusterState(m1, disco.ClusterStateNormal, 1000) {
-			t.Fatalf("unexpected node1 cluster state: %s, error: %v", state1, err1)
-		}
+		m0.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
+		m1.AssertState(t, disco.ClusterStateNormal, 1*time.Second)
+
 		m0.QueryExpect(t, "i", "", `Row(f=1)`, exp)
 		m1.QueryExpect(t, "i", "", `Row(f=1)`, exp)
 	})
diff --git a/server/server_test.go b/server/server_test.go
index d2381f8de..771735f8b 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -371,7 +371,7 @@ func TestConcurrentFieldCreation(t *testing.T) {
 	defer cluster.Close()
 
 	node0 := cluster.GetNode(0)
-	err := node0.AwaitState(disco.ClusterStateNormal, 100*time.Millisecond)
+	err := cluster.AwaitState(disco.ClusterStateNormal, 100*time.Millisecond)
 	if err != nil {
 		t.Fatalf("starting cluster: %v", err)
 	}
@@ -643,7 +643,7 @@ func TestClusteringNodesReplica1(t *testing.T) {
 	cluster := test.MustRunCluster(t, 3)
 	defer cluster.Close()
 
-	if err := cluster.GetNode(0).AwaitState(disco.ClusterStateNormal, 100*time.Millisecond); err != nil {
+	if err := cluster.AwaitState(disco.ClusterStateNormal, 100*time.Millisecond); err != nil {
 		t.Fatalf("starting cluster: %v", err)
 	}
 
@@ -651,7 +651,7 @@ func TestClusteringNodesReplica1(t *testing.T) {
 		t.Fatalf("closing third node: %v", err)
 	}
 
-	if err := cluster.GetPrimary().AwaitState(disco.ClusterStateDown, 30*time.Second); err != nil {
+	if err := cluster.AwaitPrimaryState(disco.ClusterStateDown, 30*time.Second); err != nil {
 		t.Fatalf("starting cluster: %v", err)
 	}
 
@@ -681,7 +681,7 @@ func TestClusteringNodesReplica2(t *testing.T) {
 		t.Fatalf("closing third node: %v", err)
 	}
 
-	err = coord.AwaitState(disco.ClusterStateDegraded, 30*time.Second)
+	err = cluster.AwaitPrimaryState(disco.ClusterStateDegraded, 30*time.Second)
 	if err != nil {
 		t.Fatalf("after closing first server: %v", err)
 	}
@@ -699,7 +699,7 @@ func TestClusteringNodesReplica2(t *testing.T) {
 		t.Fatalf("closing 2nd node: %v", err)
 	}
 
-	err = coord.AwaitState(disco.ClusterStateDown, 30*time.Second)
+	err = cluster.AwaitPrimaryState(disco.ClusterStateDown, 30*time.Second)
 	if err != nil {
 		t.Fatalf("after closing second server: %v", err)
 	}
@@ -730,7 +730,7 @@ func TestRemoveNodeAfterItDies(t *testing.T) {
 
 	coord, others := cluster.GetPrimary(), cluster.GetNonPrimaries()
 
-	err = coord.AwaitState(disco.ClusterStateNormal, 100*time.Millisecond)
+	err = cluster.AwaitState(disco.ClusterStateNormal, 100*time.Millisecond)
 	if err != nil {
 		t.Fatalf("starting cluster: %v", err)
 	}
@@ -741,16 +741,16 @@ func TestRemoveNodeAfterItDies(t *testing.T) {
 		t.Fatalf("closing third node: %v", err)
 	}
 
-	err = coord.AwaitState(disco.ClusterStateDegraded, 30*time.Second)
+	err = cluster.AwaitPrimaryState(disco.ClusterStateDegraded, 30*time.Second)
 	if err != nil {
-		t.Fatalf("starting cluster: %v", err)
+		t.Fatalf("degrading cluster: %v", err)
 	}
 
 	if _, err := coord.API.RemoveNode(disabled.API.Node().ID); err != nil {
 		t.Fatalf("removing failed node: %v", err)
 	}
 
-	err = coord.AwaitState(disco.ClusterStateNormal, 30*time.Second)
+	err = cluster.AwaitPrimaryState(disco.ClusterStateNormal, 30*time.Second)
 	if err != nil {
 		t.Fatalf("removing disabled node: %v", err)
 	}
@@ -774,7 +774,7 @@ func TestRemoveConcurrentIndexCreation(t *testing.T) {
 	defer cluster.Close()
 
 	node0 := cluster.GetNode(0)
-	err = node0.AwaitState(disco.ClusterStateNormal, 100*time.Millisecond)
+	err = cluster.AwaitState(disco.ClusterStateNormal, 100*time.Millisecond)
 	if err != nil {
 		t.Fatalf("starting cluster: %v", err)
 	}
@@ -789,7 +789,7 @@ func TestRemoveConcurrentIndexCreation(t *testing.T) {
 		t.Fatalf("removing node: %v", err)
 	}
 
-	err = cluster.GetPrimary().AwaitState(disco.ClusterStateNormal, 100*time.Millisecond)
+	err = cluster.AwaitPrimaryState(disco.ClusterStateNormal, 100*time.Millisecond)
 	if err != nil {
 		t.Fatalf("starting cluster: %v", err)
 	}
diff --git a/test/cluster.go b/test/cluster.go
index ec98f3e0f..22489bc9a 100644
--- a/test/cluster.go
+++ b/test/cluster.go
@@ -415,7 +415,7 @@ func (c *Cluster) Start() error {
 	if err != nil {
 		return errors.Wrap(err, "starting cluster")
 	}
-	return c.GetNode(0).AwaitState(disco.ClusterStateNormal, 30*time.Second)
+	return c.AwaitState(disco.ClusterStateNormal, 30*time.Second)
 }
 
 // Close stops a Cluster
@@ -447,6 +447,67 @@ func (c *Cluster) CloseAndRemove(n int) error {
 	return err
 }
 
+// AwaitPrimaryState waits for the cluster primary to reach a specified cluster state.
+// When this happens, we know etcd reached a combination of node states that
+// would imply this cluster state, but some nodes may not have caught up yet;
+// we just test that the coordinator thought the cluster was in the given state.
+func (c *Cluster) AwaitPrimaryState(expectedState disco.ClusterState, timeout time.Duration) error {
+	if len(c.Nodes) < 1 {
+		return errors.New("can't await coordinator state on an empty cluster")
+	}
+	primary := c.GetPrimary()
+	if primary == nil {
+		startTime := time.Now()
+		var elapsed time.Duration
+		for elapsed = 0; elapsed <= timeout; elapsed = time.Since(startTime) {
+			time.Sleep(50 * time.Millisecond)
+			primary = c.GetPrimary()
+			if primary != nil {
+				break
+			}
+		}
+		if primary == nil {
+			return errors.New("timed out waiting for cluster to have valid topology")
+		}
+		// we used up some of our timeout waiting for this
+		c.tb.Logf("had to wait %v for cluster topology", elapsed)
+		timeout -= elapsed
+	}
+	onlyCoordinator := &Cluster{Nodes: []*Command{primary}}
+	return onlyCoordinator.AwaitState(expectedState, timeout)
+}
+
+// ExceptionalState returns an error if any node in the cluster is not
+// in the expected state.
+func (c *Cluster) ExceptionalState(expectedState disco.ClusterState) error {
+	for _, node := range c.Nodes {
+		state, err := node.API.State()
+		if err != nil || state != expectedState {
+			return fmt.Errorf("node %q: state %s: err %v", node.ID(), state, err)
+		}
+	}
+	return nil
+}
+
+// AwaitState waits for the whole cluster to reach a specified state.
+func (c *Cluster) AwaitState(expectedState disco.ClusterState, timeout time.Duration) (err error) {
+	if len(c.Nodes) < 1 {
+		return errors.New("can't await state of an empty cluster")
+	}
+	startTime := time.Now()
+	var elapsed time.Duration
+	for elapsed = 0; elapsed <= timeout; elapsed = time.Since(startTime) {
+		// Counterintuitive: We're returning if the err *is* nil,
+		// meaning we've reached the expected state.
+		if err = c.ExceptionalState(expectedState); err == nil {
+			return err
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+	return fmt.Errorf("waited %v for cluster to reach state %q: %v",
+		elapsed, expectedState, err)
+}
+
 // MustNewCluster creates a new cluster. If opts contains only one
 // slice of command options, those options are used with every node.
 // If it is empty, default options are used. Otherwise, it must contain size
@@ -466,23 +527,6 @@ func MustNewCluster(tb testing.TB, size int, opts ...[]server.CommandOption) *Cl
 	return c
 }
 
-// CheckClusterState polls a given cluster for its state until it
-// receives a matching state. It polls up to n times before returning.
-func CheckClusterState(m *Command, state disco.ClusterState, n int) bool {
-	for i := 0; i < n; i++ {
-
-		apiState, err := m.API.State()
-		if err != nil {
-			return false
-		}
-		if apiState == state {
-			return true
-		}
-		time.Sleep(100 * time.Millisecond)
-	}
-	return false
-}
-
 // newCluster creates a new cluster
 func newCluster(tb testing.TB, size int, opts ...[]server.CommandOption) (*Cluster, error) {
 	if size == 0 {
diff --git a/test/pilosa.go b/test/pilosa.go
index e9c19f552..66b17ea4e 100644
--- a/test/pilosa.go
+++ b/test/pilosa.go
@@ -384,12 +384,29 @@ func (m *Command) AwaitState(expectedState disco.ClusterState, timeout time.Dura
 		if err = m.exceptionalState(expectedState); err == nil {
 			return err
 		}
-		time.Sleep(1 * time.Millisecond)
+		time.Sleep(50 * time.Millisecond)
 	}
 	return fmt.Errorf("waited %v for command to reach state %q: %v",
 		elapsed, expectedState, err)
 }
 
+// AssertState waits for the whole cluster to reach a specified state, or
+// fails the calling test if it can't.
+func (m *Command) AssertState(t testing.TB, expectedState disco.ClusterState, timeout time.Duration) {
+	startTime := time.Now()
+	var elapsed time.Duration
+	var err error
+	for elapsed = 0; elapsed <= timeout; elapsed = time.Since(startTime) {
+		// Counterintuitive: We're returning if the err *is* nil,
+		// meaning we've reached the expected state.
+		if err = m.exceptionalState(expectedState); err == nil {
+			return
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+	t.Fatalf("waited %v for command to reach state %q: %v", elapsed, expectedState, err)
+}
+
 // exceptionalState returns an error if the node is not in the expected state.
 func (m *Command) exceptionalState(expectedState disco.ClusterState) error {
 	state, err := m.API.State()

From 7497d5fbe2d3e88a665b025c32397855e46b4c1f Mon Sep 17 00:00:00 2001
From: Seebs <seebs@seebs.net>
Date: Wed, 31 Mar 2021 12:27:34 -0500
Subject: [PATCH 5/5] artificially increase heartbeat TTL for tests only

CI systems sometimes hiccup for five seconds, which causes heartbeat leases
to fail and breaks all sorts of things. As a workaround, update heartbeat
TTL for tests only. This might in turn cause different failures to do
with leader elections, but in theory those should be handled now?
---
 test/disco.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/disco.go b/test/disco.go
index ade6a2b2a..fad3783b9 100644
--- a/test/disco.go
+++ b/test/disco.go
@@ -106,7 +106,7 @@ func GetPortsGenConfigs(tb testing.TB, nodes []*Command) error {
 			AClientURL:    clientURL,
 			LPeerURL:      peerURL,
 			APeerURL:      peerURL,
-			HeartbeatTTL:  5,
+			HeartbeatTTL:  12,
 			LPeerSocket:   []*net.TCPListener{peerListener},
 			LClientSocket: []*net.TCPListener{clientListener},
 		}