Implement Hashicorps's recommendations when watching consul key (#1713)

* Implemented recommendations from Consul docs

https://www.consul.io/api/features/blocking.html#implementation-details

Fixes issue #1708

Signed-off-by: Peter Štibraný <[email protected]>

* Added tests for rate limiting, and index resets.

Signed-off-by: Peter Štibraný <[email protected]>

* Call cancel function, even though we don't use it.
This makes lint happy.

Signed-off-by: Peter Štibraný <[email protected]>

* Use rate limiter from golang.org/x/time/rate

Removed copy of modified jaegers' rate-limiter

Signed-off-by: Peter Štibraný <[email protected]>

* Rate limit configuration

Signed-off-by: Peter Štibraný <[email protected]>

* WatchPrefix should use same index checks as WatchKey

Signed-off-by: Peter Štibraný <[email protected]>

* Added comment for NewInMemoryClientWithConfig

Signed-off-by: Peter Štibraný <[email protected]>

* Use rate.Inf as rate limit, when it is not set.

Signed-off-by: Peter Štibraný <[email protected]>

* TestReset doesn't use rate limit now.

Signed-off-by: Peter Štibraný <[email protected]>

* Remove unused variable.

Signed-off-by: Peter Štibraný <[email protected]>

* Rename watch-rate -> watch-rate-limit, watch-burst -> watch-burst-size.

Signed-off-by: Peter Štibraný <[email protected]>

* Use 0 instead of "Zero"

Signed-off-by: Peter Štibraný <[email protected]>

* Updated changelog with consul client enhancements.

Signed-off-by: Peter Štibraný <[email protected]>

* Moved changelog entries to the unreleased version.

Signed-off-by: Peter Štibraný <[email protected]>

Author: pstibrany

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 * [CHANGE] Removed `Delta` encoding. Any old chunks with `Delta` encoding cannot be read anymore. If `ingester.chunk-encoding` is set to `Delta` the ingester will fail to start. #1706
 * [ENHANCEMENT] Allocation improvements in adding samples to Chunk. #1706
+* [ENHANCEMENT] Consul client now follows recommended practices for blocking queries wrt returned Index value. #1708
+* [ENHANCEMENT] Consul client can optionally rate-limit itself during Watch (used e.g. by ring watchers) and WatchPrefix (used by HA feature) operations. Rate limiting is disabled by default. New flags added: `--consul.watch-rate-limit`, and `--consul.watch-burst-size`. #1708
 
 ## 0.3.0 / 2019-10-11
 

pkg/ring/kv/consul/client.go

@@ -10,8 +10,9 @@ import (
 
 	"github.com/go-kit/kit/log/level"
 	consul "github.com/hashicorp/consul/api"
-	cleanhttp "github.com/hashicorp/go-cleanhttp"
+	"github.com/hashicorp/go-cleanhttp"
 	"github.com/weaveworks/common/instrument"
+	"golang.org/x/time/rate"
 
 	"github.com/cortexproject/cortex/pkg/ring/kv/codec"
 	"github.com/cortexproject/cortex/pkg/util"
@@ -39,6 +40,8 @@ type Config struct {
 	ACLToken          string
 	HTTPClientTimeout time.Duration
 	ConsistentReads   bool
+	WatchKeyRateLimit float64 // Zero disables rate limit
+	WatchKeyBurstSize int     // Burst when doing rate-limit, defaults to 1
 }
 
 type kv interface {
@@ -62,6 +65,8 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet, prefix string) {
 	f.StringVar(&cfg.ACLToken, prefix+"consul.acltoken", "", "ACL Token used to interact with Consul.")
 	f.DurationVar(&cfg.HTTPClientTimeout, prefix+"consul.client-timeout", 2*longPollDuration, "HTTP timeout when talking to Consul")
 	f.BoolVar(&cfg.ConsistentReads, prefix+"consul.consistent-reads", true, "Enable consistent reads to Consul.")
+	f.Float64Var(&cfg.WatchKeyRateLimit, prefix+"consul.watch-rate-limit", 0, "Rate limit when watching key or prefix in Consul, in requests per second. 0 disables the rate limit.")
+	f.IntVar(&cfg.WatchKeyBurstSize, prefix+"consul.watch-burst-size", 1, "Burst size used in rate limit. Values less than 1 are treated as 1.")
 }
 
 // NewClient returns a new Client.
@@ -170,8 +175,17 @@ func (c *Client) WatchKey(ctx context.Context, key string, f func(interface{}) b
 	var (
 		backoff = util.NewBackoff(ctx, backoffConfig)
 		index   = uint64(0)
+		limiter = c.createRateLimiter()
 	)
+
 	for backoff.Ongoing() {
+		err := limiter.Wait(ctx)
+		if err != nil {
+			level.Error(util.Logger).Log("msg", "error while rate-limiting", "key", key, "err", err)
+			backoff.Wait()
+			continue
+		}
+
 		queryOptions := &consul.QueryOptions{
 			AllowStale:        !c.cfg.ConsistentReads,
 			RequireConsistent: c.cfg.ConsistentReads,
@@ -187,12 +201,11 @@ func (c *Client) WatchKey(ctx context.Context, key string, f func(interface{}) b
 		}
 		backoff.Reset()
 
-		// Skip if the index is the same as last time, because the key value is
-		// guaranteed to be the same as last time
-		if index == meta.LastIndex {
+		skip := false
+		index, skip = checkLastIndex(index, meta.LastIndex)
+		if skip {
 			continue
 		}
-		index = meta.LastIndex
 
 		out, err := c.codec.Decode(kvp.Value)
 		if err != nil {
@@ -212,8 +225,16 @@ func (c *Client) WatchPrefix(ctx context.Context, prefix string, f func(string,
 	var (
 		backoff = util.NewBackoff(ctx, backoffConfig)
 		index   = uint64(0)
+		limiter = c.createRateLimiter()
 	)
 	for backoff.Ongoing() {
+		err := limiter.Wait(ctx)
+		if err != nil {
+			level.Error(util.Logger).Log("msg", "error while rate-limiting", "prefix", prefix, "err", err)
+			backoff.Wait()
+			continue
+		}
+
 		queryOptions := &consul.QueryOptions{
 			AllowStale:        !c.cfg.ConsistentReads,
 			RequireConsistent: c.cfg.ConsistentReads,
@@ -228,13 +249,13 @@ func (c *Client) WatchPrefix(ctx context.Context, prefix string, f func(string,
 			continue
 		}
 		backoff.Reset()
-		// Skip if the index is the same as last time, because the key value is
-		// guaranteed to be the same as last time
-		if index == meta.LastIndex {
+
+		skip := false
+		index, skip = checkLastIndex(index, meta.LastIndex)
+		if skip {
 			continue
 		}
 
-		index = meta.LastIndex
 		for _, kvp := range kvps {
 			out, err := c.codec.Decode(kvp.Value)
 			if err != nil {
@@ -264,3 +285,33 @@ func (c *Client) Get(ctx context.Context, key string) (interface{}, error) {
 	}
 	return c.codec.Decode(kvp.Value)
 }
+
+func checkLastIndex(index, metaLastIndex uint64) (newIndex uint64, skip bool) {
+	// See https://www.consul.io/api/features/blocking.html#implementation-details for logic behind these checks.
+	if metaLastIndex == 0 {
+		// Don't just keep using index=0.
+		// After blocking request, returned index must be at least 1.
+		return 1, false
+	} else if metaLastIndex < index {
+		// Index reset.
+		return 0, false
+	} else if index == metaLastIndex {
+		// Skip if the index is the same as last time, because the key value is
+		// guaranteed to be the same as last time
+		return metaLastIndex, true
+	} else {
+		return metaLastIndex, false
+	}
+}
+
+func (c *Client) createRateLimiter() *rate.Limiter {
+	if c.cfg.WatchKeyRateLimit <= 0 {
+		// burst is ignored when limit = rate.Inf
+		return rate.NewLimiter(rate.Inf, 0)
+	}
+	burst := c.cfg.WatchKeyBurstSize
+	if burst < 1 {
+		burst = 1
+	}
+	return rate.NewLimiter(rate.Limit(c.cfg.WatchKeyRateLimit), burst)
+}

pkg/ring/kv/consul/client_test.go

@@ -0,0 +1,161 @@
+package consul
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/go-kit/kit/log/level"
+	consul "github.com/hashicorp/consul/api"
+
+	"github.com/cortexproject/cortex/pkg/util"
+)
+
+type stringCodec struct{}
+
+func (c stringCodec) Encode(d interface{}) ([]byte, error) {
+	if d == nil {
+		return nil, fmt.Errorf("nil")
+	}
+	s, ok := d.(string)
+	if !ok {
+		return nil, fmt.Errorf("not string: %T", d)
+	}
+
+	return []byte(s), nil
+}
+
+func (c stringCodec) Decode(d []byte) (interface{}, error) {
+	return string(d), nil
+}
+
+func writeValuesToKV(client *Client, key string, start, end int, sleep time.Duration) <-chan struct{} {
+	ch := make(chan struct{})
+	go func() {
+		defer close(ch)
+		for i := start; i <= end; i++ {
+			level.Debug(util.Logger).Log("ts", time.Now(), "msg", "writing value", "val", i)
+			_, _ = client.Put(&consul.KVPair{Key: key, Value: []byte(fmt.Sprintf("%d", i))}, nil)
+			time.Sleep(sleep)
+		}
+	}()
+	return ch
+}
+
+func TestWatchKey(t *testing.T) {
+	c := NewInMemoryClientWithConfig(&stringCodec{}, Config{
+		WatchKeyRateLimit: 5.0,
+		WatchKeyBurstSize: 1,
+	})
+
+	const key = "test"
+	const max = 100
+
+	// Make sure to start with non-empty value, otherwise WatchKey will bail
+	_, _ = c.Put(&consul.KVPair{Key: key, Value: []byte("start")}, nil)
+
+	ch := writeValuesToKV(c, key, 0, max, 10*time.Millisecond)
+
+	observed := observeValueForSomeTime(c, key, 1200*time.Millisecond) // little over 1 second
+
+	// wait until updater finishes
+	<-ch
+
+	if testing.Verbose() {
+		t.Log(observed)
+	}
+	// Let's see how many updates we have observed. Given the rate limit and our observing time, it should be 6
+	// We should also have seen one of the later values, as we're observing for longer than a second, so rate limit should allow
+	// us to see it.
+	if len(observed) < 5 || len(observed) > 10 {
+		t.Error("Expected ~6 observed values, got", observed)
+	}
+	last := observed[len(observed)-1]
+	n, _ := strconv.Atoi(last)
+	if n < max/2 {
+		t.Error("Expected to see high last observed value, got", observed)
+	}
+}
+
+func TestWatchKeyNoRateLimit(t *testing.T) {
+	c := NewInMemoryClientWithConfig(&stringCodec{}, Config{
+		WatchKeyRateLimit: 0,
+	})
+
+	const key = "test"
+	const max = 100
+
+	// Make sure to start with non-empty value, otherwise WatchKey will bail
+	_, _ = c.Put(&consul.KVPair{Key: key, Value: []byte("start")}, nil)
+
+	ch := writeValuesToKV(c, key, 0, max, time.Millisecond)
+	observed := observeValueForSomeTime(c, key, 500*time.Millisecond)
+
+	// wait until updater finishes
+	<-ch
+
+	// With no limit, we should see most written values (we can lose some values if watching
+	// code is busy while multiple new values are written)
+	if len(observed) < 3*max/4 {
+		t.Error("Expected at least 3/4 of all values, got", observed)
+	}
+}
+
+func TestReset(t *testing.T) {
+	c := NewInMemoryClient(&stringCodec{})
+
+	const key = "test"
+	const max = 5
+
+	// Make sure to start with non-empty value, otherwise WatchKey will bail
+	_, _ = c.Put(&consul.KVPair{Key: key, Value: []byte("start")}, nil)
+
+	ch := make(chan error)
+	go func() {
+		defer close(ch)
+		for i := 0; i <= max; i++ {
+			level.Debug(util.Logger).Log("ts", time.Now(), "msg", "writing value", "val", i)
+			_, _ = c.Put(&consul.KVPair{Key: key, Value: []byte(fmt.Sprintf("%d", i))}, nil)
+			if i == 1 {
+				c.kv.(*mockKV).ResetIndex()
+			}
+			if i == 2 {
+				c.kv.(*mockKV).ResetIndexForKey(key)
+			}
+			time.Sleep(10 * time.Millisecond)
+		}
+	}()
+
+	observed := observeValueForSomeTime(c, key, 25*max*time.Millisecond)
+
+	// wait until updater finishes
+	<-ch
+
+	// Let's see how many updates we have observed. Given the rate limit and our observing time, we should see all numeric values
+	if testing.Verbose() {
+		t.Log(observed)
+	}
+	if len(observed) < max {
+		t.Error("Expected all values, got", observed)
+	} else if observed[len(observed)-1] != fmt.Sprintf("%d", max) {
+		t.Error("Expected to see last written value, got", observed)
+	}
+}
+
+func observeValueForSomeTime(client *Client, key string, timeout time.Duration) []string {
+	observed := []string(nil)
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	client.WatchKey(ctx, key, func(i interface{}) bool {
+		s, ok := i.(string)
+		if !ok {
+			return false
+		}
+		level.Debug(util.Logger).Log("ts", time.Now(), "msg", "observed value", "val", s)
+		observed = append(observed, s)
+		return true
+	})
+	return observed
+}

pkg/ring/kv/consul/mock.go

@@ -1,7 +1,7 @@
 package consul
 
 import (
-	fmt "fmt"
+	"fmt"
 	"sync"
 	"time"
 
@@ -21,6 +21,11 @@ type mockKV struct {
 
 // NewInMemoryClient makes a new mock consul client.
 func NewInMemoryClient(codec codec.Codec) *Client {
+	return NewInMemoryClientWithConfig(codec, Config{})
+}
+
+// NewInMemoryClientWithConfig makes a new mock consul client with supplied Config.
+func NewInMemoryClientWithConfig(codec codec.Codec, cfg Config) *Client {
 	m := mockKV{
 		kvps: map[string]*consul.KVPair{},
 	}
@@ -29,6 +34,7 @@ func NewInMemoryClient(codec codec.Codec) *Client {
 	return &Client{
 		kv:    &m,
 		codec: codec,
+		cfg:   cfg,
 	}
 }
 
@@ -67,6 +73,8 @@ func (m *mockKV) Put(p *consul.KVPair, q *consul.WriteOptions) (*consul.WriteMet
 	}
 
 	m.cond.Broadcast()
+
+	level.Debug(util.Logger).Log("msg", "Put", "key", p.Key, "value", fmt.Sprintf("%.40q", p.Value), "modify_index", m.current)
 	return nil, nil
 }
 
@@ -109,9 +117,16 @@ func (m *mockKV) Get(key string, q *consul.QueryOptions) (*consul.KVPair, *consu
 		return nil, &consul.QueryMeta{LastIndex: m.current}, nil
 	}
 
-	if q.WaitTime > 0 {
+	if q.WaitIndex >= value.ModifyIndex && q.WaitTime > 0 {
 		deadline := time.Now().Add(q.WaitTime)
-		for q.WaitIndex >= value.ModifyIndex && time.Now().Before(deadline) {
+		if ctxDeadline, ok := q.Context().Deadline(); ok && ctxDeadline.Before(deadline) {
+			// respect deadline from context, if set.
+			deadline = ctxDeadline
+		}
+
+		// simply wait until value.ModifyIndex changes. This allows us to test reporting old index values by resetting them.
+		startModify := value.ModifyIndex
+		for startModify == value.ModifyIndex && time.Now().Before(deadline) {
 			m.cond.Wait()
 		}
 		if time.Now().After(deadline) {
@@ -144,3 +159,25 @@ func (m *mockKV) List(prefix string, q *consul.QueryOptions) (consul.KVPairs, *c
 	}
 	return result, &consul.QueryMeta{LastIndex: m.current}, nil
 }
+
+func (m *mockKV) ResetIndex() {
+	m.mtx.Lock()
+	defer m.mtx.Unlock()
+
+	m.current = 0
+	m.cond.Broadcast()
+
+	level.Debug(util.Logger).Log("msg", "Reset")
+}
+
+func (m *mockKV) ResetIndexForKey(key string) {
+	m.mtx.Lock()
+	defer m.mtx.Unlock()
+
+	if value, ok := m.kvps[key]; ok {
+		value.ModifyIndex = 0
+	}
+
+	m.cond.Broadcast()
+	level.Debug(util.Logger).Log("msg", "ResetIndexForKey", "key", key)
+}