Memberlist integration test fix

go.mod

@@ -61,12 +61,12 @@ require (
 	github.com/thanos-io/thanos v0.8.1-0.20200109203923-552ffa4c1a0d
 	github.com/tinylib/msgp v0.0.0-20161221055906-38a6f61a768d // indirect
 	github.com/tmc/grpc-websocket-proxy v0.0.0-20190109142713-0ad062ec5ee5 // indirect
-	github.com/uber-go/atomic v1.4.0
 	github.com/uber/jaeger-client-go v2.20.1+incompatible
 	github.com/weaveworks/billing-client v0.0.0-20171006123215-be0d55e547b1
 	github.com/weaveworks/common v0.0.0-20200201141823-27e183090ab1
 	go.etcd.io/bbolt v1.3.3
 	go.etcd.io/etcd v0.0.0-20190709142735-eb7dd97135a5
+	go.uber.org/atomic v1.5.0
 	golang.org/x/net v0.0.0-20191112182307-2180aed22343
 	golang.org/x/time v0.0.0-20191024005414-555d28b269f0
 	google.golang.org/api v0.14.0

go.sum

@@ -763,8 +763,6 @@ github.com/tmc/grpc-websocket-proxy v0.0.0-20170815181823-89b8d40f7ca8/go.mod h1
 github.com/tmc/grpc-websocket-proxy v0.0.0-20190109142713-0ad062ec5ee5 h1:LnC5Kc/wtumK+WB441p7ynQJzVuNRJiqddSIE3IlSEQ=
 github.com/tmc/grpc-websocket-proxy v0.0.0-20190109142713-0ad062ec5ee5/go.mod h1:ncp9v5uamzpCO7NfCPTXjqaC+bZgJeR0sMTm6dMHP7U=
 github.com/tv42/httpunix v0.0.0-20150427012821-b75d8614f926/go.mod h1:9ESjWnEqriFuLhtthL60Sar/7RFoluCcXsuvEwTV5KM=
-github.com/uber-go/atomic v1.4.0 h1:yOuPqEq4ovnhEjpHmfFwsqBXDYbQeT6Nb0bwD6XnD5o=
-github.com/uber-go/atomic v1.4.0/go.mod h1:/Ct5t2lcmbJ4OSe/waGBoaVvVqtO0bmtfVNex1PFV8g=
 github.com/uber/jaeger-client-go v2.15.0+incompatible/go.mod h1:WVhlPFC8FDjOFMMWRy2pZqQJSXxYSwNYOkTr/Z6d3Kk=
 github.com/uber/jaeger-client-go v2.20.1+incompatible h1:HgqpYBng0n7tLJIlyT4kPCIv5XgCsF+kai1NnnrJzEU=
 github.com/uber/jaeger-client-go v2.20.1+incompatible/go.mod h1:WVhlPFC8FDjOFMMWRy2pZqQJSXxYSwNYOkTr/Z6d3Kk=

integration/integration_memberlist_single_binary_test.go

@@ -62,6 +62,7 @@ func startSingleBinary(s *framework.Scenario, name string, join string) error {
 		"-ingester.observe-period":       "5s", // to avoid conflicts in tokens
 		"-ring.store":                    "memberlist",
 		"-memberlist.bind-port":          "8000",
+		"-memberlist.pullpush-interval":  "3s", // speed up state convergence to make test faster and avoid flakiness
 	}
 
 	if join != "" {

pkg/ring/kv/memberlist/memberlist_client.go

@@ -15,6 +15,7 @@ import (
 	"github.com/go-kit/kit/log/level"
 	"github.com/hashicorp/memberlist"
 	"github.com/prometheus/client_golang/prometheus"
+	"go.uber.org/atomic"
 
 	"github.com/cortexproject/cortex/pkg/ring/kv/codec"
 	"github.com/cortexproject/cortex/pkg/util"
@@ -120,6 +121,9 @@ type KV struct {
 	memberlist *memberlist.Memberlist
 	broadcasts *memberlist.TransmitLimitedQueue
 
+	// Disabled on Stop()
+	casBroadcastsEnabled *atomic.Bool
+
 	// KV Store.
 	storeMu sync.Mutex
 	store   map[string]valueDesc
@@ -224,42 +228,43 @@ func NewKV(cfg KVConfig) (*KV, error) {
 	// As we don't use UDP for sending packets, we can use higher value here.
 	mlCfg.UDPBufferSize = 10 * 1024 * 1024
 
-	memberlistClient := &KV{
-		cfg:            cfg,
-		store:          make(map[string]valueDesc),
-		codecs:         make(map[string]codec.Codec),
-		watchers:       make(map[string][]chan string),
-		prefixWatchers: make(map[string][]chan string),
-		shutdown:       make(chan struct{}),
-		maxCasRetries:  maxCasRetries,
+	mlkv := &KV{
+		cfg:                  cfg,
+		store:                make(map[string]valueDesc),
+		codecs:               make(map[string]codec.Codec),
+		watchers:             make(map[string][]chan string),
+		prefixWatchers:       make(map[string][]chan string),
+		shutdown:             make(chan struct{}),
+		maxCasRetries:        maxCasRetries,
+		casBroadcastsEnabled: atomic.NewBool(true),
 	}
 
-	mlCfg.Delegate = memberlistClient
+	mlCfg.Delegate = mlkv
 
 	list, err := memberlist.Create(mlCfg)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create memberlist: %v", err)
 	}
 
 	// finish delegate initialization
-	memberlistClient.memberlist = list
-	memberlistClient.broadcasts = &memberlist.TransmitLimitedQueue{
+	mlkv.memberlist = list
+	mlkv.broadcasts = &memberlist.TransmitLimitedQueue{
 		NumNodes:       list.NumMembers,
 		RetransmitMult: cfg.RetransmitMult,
 	}
 
 	// Almost ready...
-	memberlistClient.createAndRegisterMetrics()
+	mlkv.createAndRegisterMetrics()
 
 	for _, c := range cfg.Codecs {
-		memberlistClient.codecs[c.CodecID()] = c
+		mlkv.codecs[c.CodecID()] = c
 	}
 
 	// Join the cluster
 	if len(cfg.JoinMembers) > 0 {
-		reached, err := memberlistClient.JoinMembers(cfg.JoinMembers)
+		reached, err := mlkv.JoinMembers(cfg.JoinMembers)
 		if err != nil && cfg.AbortIfJoinFails {
-			_ = memberlistClient.memberlist.Shutdown()
+			_ = mlkv.memberlist.Shutdown()
 			return nil, err
 		}
 
@@ -270,7 +275,7 @@ func NewKV(cfg KVConfig) (*KV, error) {
 		}
 	}
 
-	return memberlistClient, nil
+	return mlkv, nil
 }
 
 // GetCodec returns codec for given ID or nil.
@@ -294,7 +299,23 @@ func (m *KV) JoinMembers(members []string) (int, error) {
 func (m *KV) Stop() {
 	level.Info(util.Logger).Log("msg", "leaving memberlist cluster")
 
-	// TODO: should we empty our broadcast queue before leaving? That would make sure that we have sent out everything we wanted.
+	m.casBroadcastsEnabled.Store(false)
+
+	// Wait until broadcast queue is empty, but don't wait for too long.
+	// Also don't wait if there is just one node left.
+	// Problem is that broadcast queue is also filled up by state changes received from other nodes,
+	// so it may never be empty in a busy cluster. However, we generally only care about messages
+	// generated on this node via CAS, and those are disabled now (via casBroadcastsEnabled), and should be able
+	// to get out in this timeout.
+
+	waitTimeout := time.Now().Add(10 * time.Second)
+	for m.broadcasts.NumQueued() > 0 && m.memberlist.NumMembers() > 1 && time.Now().Before(waitTimeout) {
+		time.Sleep(250 * time.Millisecond)
+	}
+
+	if cnt := m.broadcasts.NumQueued(); cnt > 0 {
+		level.Warn(util.Logger).Log("msg", "broadcast messages left in queue", "count", cnt, "nodes", m.memberlist.NumMembers())
+	}
 
 	err := m.memberlist.Leave(m.cfg.LeaveTimeout)
 	if err != nil {
@@ -524,7 +545,12 @@ outer:
 		if change != nil {
 			m.casSuccesses.Inc()
 			m.notifyWatchers(key)
-			m.broadcastNewValue(key, change, newver, codec)
+
+			if m.casBroadcastsEnabled.Load() {
+				m.broadcastNewValue(key, change, newver, codec)
+			} else {
+				level.Warn(util.Logger).Log("msg", "skipped broadcasting CAS update because memberlist KV is shutting down", "key", key)
+			}
 		}
 
 		return nil

pkg/ring/kv/multi.go

@@ -9,11 +9,11 @@ import (
 
 	"github.com/go-kit/kit/log"
 	"github.com/prometheus/client_golang/prometheus"
+	"go.uber.org/atomic"
 
 	"github.com/cortexproject/cortex/pkg/util"
 
 	"github.com/go-kit/kit/log/level"
-	"github.com/uber-go/atomic"
 )
 
 var (

pkg/util/runtimeconfig/manager_test.go

@@ -7,7 +7,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/require"
-	"github.com/uber-go/atomic"
+	"go.uber.org/atomic"
 	"gopkg.in/yaml.v2"
 )