Fix goroutine leak caused by updating rate quotas

changelog/11371.txt

@@ -0,0 +1,3 @@
+```release-note:bug
+core: Fix goroutine leak when updating rate limit quota
+```

go.mod

@@ -156,6 +156,7 @@ require (
 	go.etcd.io/etcd v0.5.0-alpha.5.0.20200425165423-262c93980547
 	go.mongodb.org/mongo-driver v1.4.6
 	go.uber.org/atomic v1.6.0
+	go.uber.org/goleak v1.1.10
 	golang.org/x/crypto v0.0.0-20210220033148-5ea612d1eb83
 	golang.org/x/net v0.0.0-20201110031124-69a78807bb2b
 	golang.org/x/oauth2 v0.0.0-20200107190931-bf48bf16ab8d

helper/metricsutil/wrapped_metrics.go

@@ -88,8 +88,9 @@ func (m *ClusterMetricSink) MeasureSinceWithLabels(key []string, start time.Time
 
 // BlackholeSink is a default suitable for use in unit tests.
 func BlackholeSink() *ClusterMetricSink {
-	sink, _ := metrics.New(metrics.DefaultConfig(""),
-		&metrics.BlackholeSink{})
+	conf := metrics.DefaultConfig("")
+	conf.EnableRuntimeMetrics = false
+	sink, _ := metrics.New(conf, &metrics.BlackholeSink{})
 	cms := &ClusterMetricSink{
 		ClusterName: atomic.Value{},
 		Sink:        sink,

vault/quotas/quotas.go

@@ -261,14 +261,16 @@ func NewManager(logger log.Logger, walkFunc leaseWalkFunc, ms *metricsutil.Clust
 	return manager, nil
 }
 
-// SetQuota adds a new quota rule to the db.
+// SetQuota adds or updates a quota rule.
 func (m *Manager) SetQuota(ctx context.Context, qType string, quota Quota, loading bool) error {
 	m.lock.Lock()
 	defer m.lock.Unlock()
 	return m.setQuotaLocked(ctx, qType, quota, loading)
 }
 
-// setQuotaLocked should be called with the manager's lock held
+// setQuotaLocked adds or updates a quota rule, modifying the db as well as
+// any runtime elements such as goroutines.
+// It should be called with the write lock held.
 func (m *Manager) setQuotaLocked(ctx context.Context, qType string, quota Quota, loading bool) error {
 	if qType == TypeLeaseCount.String() {
 		m.setIsPerfStandby(quota)
@@ -277,13 +279,17 @@ func (m *Manager) setQuotaLocked(ctx context.Context, qType string, quota Quota,
 	txn := m.db.Txn(true)
 	defer txn.Abort()
 
-	raw, err := txn.First(qType, "id", quota.quotaID())
+	raw, err := txn.First(qType, indexID, quota.quotaID())
 	if err != nil {
 		return err
 	}
 
 	// If there already exists an entry in the db, remove that first.
 	if raw != nil {
+		quota := raw.(Quota)
+		if err := quota.close(); err != nil {
+			return err
+		}
 		err = txn.Delete(qType, raw)
 		if err != nil {
 			return err

vault/quotas/quotas_rate_limit.go

@@ -319,6 +319,7 @@ func (rlq *RateLimitQuota) allow(req *Request) (Response, error) {
 }
 
 // close stops the current running client purge loop.
+// It should be called with the write lock held.
 func (rlq *RateLimitQuota) close() error {
 	if rlq.purgeBlocked {
 		close(rlq.closePurgeBlockedCh)

vault/quotas/quotas_rate_limit_test.go

@@ -1,6 +1,7 @@
 package quotas
 
 import (
+	"context"
 	"fmt"
 	"math"
 	"sync"
@@ -12,6 +13,7 @@ import (
 	"github.com/hashicorp/vault/sdk/helper/logging"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
+	"go.uber.org/goleak"
 )
 
 type clientResult struct {
@@ -34,6 +36,9 @@ func TestNewRateLimitQuota(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			err := tc.rlq.initialize(logging.NewVaultLogger(log.Trace), metricsutil.BlackholeSink())
 			require.Equal(t, tc.expectErr, err != nil, err)
+			if err == nil {
+				require.Nil(t, tc.rlq.close())
+			}
 		})
 	}
 }
@@ -61,6 +66,7 @@ func TestRateLimitQuota_Allow(t *testing.T) {
 	}
 
 	require.NoError(t, rlq.initialize(logging.NewVaultLogger(log.Trace), metricsutil.BlackholeSink()))
+	defer rlq.close()
 
 	var wg sync.WaitGroup
 
@@ -135,6 +141,7 @@ func TestRateLimitQuota_Allow_WithBlock(t *testing.T) {
 	}
 
 	require.NoError(t, rlq.initialize(logging.NewVaultLogger(log.Trace), metricsutil.BlackholeSink()))
+	defer rlq.close()
 	require.True(t, rlq.getPurgeBlocked())
 
 	var wg sync.WaitGroup
@@ -204,3 +211,15 @@ func TestRateLimitQuota_Allow_WithBlock(t *testing.T) {
 		}
 	}()
 }
+
+func TestRateLimitQuota_Update(t *testing.T) {
+	defer goleak.VerifyNone(t)
+	qm, err := NewManager(logging.NewVaultLogger(log.Trace), nil, metricsutil.BlackholeSink())
+	require.NoError(t, err)
+
+	quota := NewRateLimitQuota("quota1", "", "", 10, time.Second, 0)
+	require.NoError(t, qm.SetQuota(context.Background(), TypeRateLimit.String(), quota, true))
+	require.NoError(t, qm.SetQuota(context.Background(), TypeRateLimit.String(), quota, true))
+
+	require.Nil(t, quota.close())
+}