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

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)