mvcc: avoid double decrement of watcher gauge on close/cancel race

This occurs specifically when the watch is for a compacted revision, as there's
a possible interleaving of cancel/close that invokes the `cancelWatch` function
twice.

It's fairly difficult to provoke the race condition but it is possible to
observe on `main` the racing test can fail with a negative gauge:

```
$ go test ./...  -run TestNewWatcherCountGauge/compacted_watch,_close/cancel_race
--- FAIL: TestNewWatcherCountGauge (0.34s)
    watchable_store_test.go:86:  # HELP etcd_debugging_mvcc_watcher_total Total number of watchers.
         # TYPE etcd_debugging_mvcc_watcher_total gauge
        -etcd_debugging_mvcc_watcher_total -1
        +etcd_debugging_mvcc_watcher_total 0

FAIL
FAIL    go.etcd.io/etcd/server/v3/storage/mvcc  0.830s
?       go.etcd.io/etcd/server/v3/storage/mvcc/testutil [no test files]
FAIL
```

It seems as though it is partially expected for the cancel function to be
invoked multiple times and to handle that safely (i.e., the existing `ch == nil`
check) - the bug here is that in the `if/else if` branches it comes "too late",
and multiple invocations where `wa.compacted` is true will both decrement the
counter. Shifting the case up one ensures that we can't follow that decrement
branch multiple times.

In fact, it seems logically more sensible to put this `wa.ch == nil` case
_first_, as a guard for the function being invoked multiple times, but moving i
before the sync/unsynced watch set delete functions could have a greater
inadvertent functional impact (i.e., if we never deleted cancelled watches from
these sets it would presumably introduce a leak), so from an abundance of
caution I've made the smallest change I think will fix my issue.

Signed-off-by: Kieran Gorman <[email protected]>

Author: kjgorman

server/storage/mvcc/watchable_store.go

@@ -162,12 +162,12 @@ func (s *watchableStore) cancelWatcher(wa *watcher) {
 		} else if s.synced.delete(wa) {
 			watcherGauge.Dec()
 			break
-		} else if wa.compacted {
-			watcherGauge.Dec()
-			break
 		} else if wa.ch == nil {
 			// already canceled (e.g., cancel/close race)
 			break
+		} else if wa.compacted {
+			watcherGauge.Dec()
+			break
 		}
 
 		if !wa.victim {

server/storage/mvcc/watchable_store_test.go

@@ -17,10 +17,12 @@ package mvcc
 import (
 	"fmt"
 	"reflect"
+	"strings"
 	"sync"
 	"testing"
 	"time"
 
+	"github.com/prometheus/client_golang/prometheus/testutil"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap/zaptest"
@@ -73,6 +75,175 @@ func TestNewWatcherCancel(t *testing.T) {
 	}
 }
 
+func TestNewWatcherCountGauge(t *testing.T) {
+	expectWatchGauge := func(watchers int) {
+		expected := fmt.Sprintf(`# HELP etcd_debugging_mvcc_watcher_total Total number of watchers.
+# TYPE etcd_debugging_mvcc_watcher_total gauge
+etcd_debugging_mvcc_watcher_total %d
+`, watchers)
+		err := testutil.CollectAndCompare(watcherGauge, strings.NewReader(expected), "etcd_debugging_mvcc_watcher_total")
+		if err != nil {
+			t.Error(err)
+		}
+	}
+
+	t.Run("regular watch", func(t *testing.T) {
+		b, _ := betesting.NewDefaultTmpBackend(t)
+		s := New(zaptest.NewLogger(t), b, &lease.FakeLessor{}, StoreConfig{})
+		defer cleanup(s, b)
+
+		// watcherGauge is a package variable and its value may change depending on
+		// the execution of other tests
+		initialGaugeState := int(testutil.ToFloat64(watcherGauge))
+
+		testKey := []byte("foo")
+		testValue := []byte("bar")
+		s.Put(testKey, testValue, lease.NoLease)
+
+		// we expect the gauge state to still be in its initial state
+		expectWatchGauge(initialGaugeState)
+
+		w := s.NewWatchStream()
+		defer w.Close()
+
+		wt, _ := w.Watch(0, testKey, nil, 0)
+
+		// after creating watch, the gauge state should have increased
+		expectWatchGauge(initialGaugeState + 1)
+
+		if err := w.Cancel(wt); err != nil {
+			t.Error(err)
+		}
+
+		// after cancelling watch, the gauge state should have decreased
+		expectWatchGauge(initialGaugeState)
+
+		w.Cancel(wt)
+
+		// cancelling the watch twice shouldn't decrement the counter twice
+		expectWatchGauge(initialGaugeState)
+	})
+
+	t.Run("compacted watch", func(t *testing.T) {
+		b, _ := betesting.NewDefaultTmpBackend(t)
+		s := New(zaptest.NewLogger(t), b, &lease.FakeLessor{}, StoreConfig{})
+		defer cleanup(s, b)
+
+		// watcherGauge is a package variable and its value may change depending on
+		// the execution of other tests
+		initialGaugeState := int(testutil.ToFloat64(watcherGauge))
+
+		testKey := []byte("foo")
+		testValue := []byte("bar")
+
+		s.Put(testKey, testValue, lease.NoLease)
+		rev := s.Put(testKey, testValue, lease.NoLease)
+
+		// compact up to the revision of the key we just put
+		_, err := s.Compact(traceutil.TODO(), rev)
+		if err != nil {
+			t.Error(err)
+		}
+
+		// we expect the gauge state to still be in its initial state
+		expectWatchGauge(initialGaugeState)
+
+		w := s.NewWatchStream()
+		defer w.Close()
+
+		wt, _ := w.Watch(0, testKey, nil, rev-1)
+
+		// wait for the watcher to be marked as compacted
+		select {
+		case resp := <-w.Chan():
+			if resp.CompactRevision == 0 {
+				t.Errorf("resp.Compacted = %v, want %v", resp.CompactRevision, rev)
+			}
+		case <-time.After(time.Second):
+			t.Fatalf("failed to receive response (timeout)")
+		}
+
+		// after creating watch, the gauge state should have increased
+		expectWatchGauge(initialGaugeState + 1)
+
+		if err := w.Cancel(wt); err != nil {
+			t.Error(err)
+		}
+
+		// after cancelling watch, the gauge state should have decreased
+		expectWatchGauge(initialGaugeState)
+
+		w.Cancel(wt)
+
+		// cancelling the watch twice shouldn't decrement the counter twice
+		expectWatchGauge(initialGaugeState)
+	})
+
+	t.Run("compacted watch, close/cancel race", func(t *testing.T) {
+		b, _ := betesting.NewDefaultTmpBackend(t)
+		s := New(zaptest.NewLogger(t), b, &lease.FakeLessor{}, StoreConfig{})
+		defer cleanup(s, b)
+
+		// watcherGauge is a package variable and its value may change depending on
+		// the execution of other tests
+		initialGaugeState := int(testutil.ToFloat64(watcherGauge))
+
+		testKey := []byte("foo")
+		testValue := []byte("bar")
+
+		s.Put(testKey, testValue, lease.NoLease)
+		rev := s.Put(testKey, testValue, lease.NoLease)
+
+		// compact up to the revision of the key we just put
+		_, err := s.Compact(traceutil.TODO(), rev)
+		if err != nil {
+			t.Error(err)
+		}
+
+		// we expect the gauge state to still be in its initial state
+		expectWatchGauge(initialGaugeState)
+
+		w := s.NewWatchStream()
+
+		wt, _ := w.Watch(0, testKey, nil, rev-1)
+
+		// wait for the watcher to be marked as compacted
+		select {
+		case resp := <-w.Chan():
+			if resp.CompactRevision == 0 {
+				t.Errorf("resp.Compacted = %v, want %v", resp.CompactRevision, rev)
+			}
+		case <-time.After(time.Second):
+			t.Fatalf("failed to receive response (timeout)")
+		}
+
+		// after creating watch, the gauge state should have increased
+		expectWatchGauge(initialGaugeState + 1)
+
+		// now race cancelling and closing the watcher and watch stream.
+		// in rare scenarios the watcher cancel function can be invoked
+		// multiple times, leading to a potentially negative gauge state,
+		// see: https://github.com/etcd-io/etcd/issues/19577
+		wg := sync.WaitGroup{}
+		wg.Add(2)
+
+		go func() {
+			w.Cancel(wt)
+			wg.Done()
+		}()
+
+		go func() {
+			w.Close()
+			wg.Done()
+		}()
+
+		wg.Wait()
+
+		// the gauge should be decremented to its original state
+		expectWatchGauge(initialGaugeState)
+	})
+}
+
 // TestCancelUnsynced tests if running CancelFunc removes watchers from unsynced.
 func TestCancelUnsynced(t *testing.T) {
 	b, _ := betesting.NewDefaultTmpBackend(t)