ccl/sqlproxyccl: add metrics to track the number of rebalance requests

pkg/ccl/sqlproxyccl/balancer/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
         "balancer.go",
         "conn.go",
         "conn_tracker.go",
+        "metrics.go",
         "pod.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/balancer",
@@ -14,6 +15,7 @@ go_library(
         "//pkg/ccl/sqlproxyccl/tenant",
         "//pkg/roachpb",
         "//pkg/util/log",
+        "//pkg/util/metric",
         "//pkg/util/randutil",
         "//pkg/util/stop",
         "//pkg/util/syncutil",

pkg/ccl/sqlproxyccl/balancer/balancer.go

@@ -105,6 +105,9 @@ type Balancer struct {
 	// balancer.
 	stopper *stop.Stopper
 
+	// metrics contains various counters reflecting the balancer operations.
+	metrics *Metrics
+
 	// directoryCache corresponds to the tenant directory cache, which will be
 	// used to lookup IP addresses of SQL pods for tenants.
 	directoryCache tenant.DirectoryCache
@@ -131,6 +134,7 @@ type Balancer struct {
 func NewBalancer(
 	ctx context.Context,
 	stopper *stop.Stopper,
+	metrics *Metrics,
 	directoryCache tenant.DirectoryCache,
 	connTracker *ConnTracker,
 	opts ...Option,
@@ -150,13 +154,14 @@ func NewBalancer(
 	// Ensure that ctx gets cancelled on stopper's quiescing.
 	ctx, _ = stopper.WithCancelOnQuiesce(ctx)
 
-	q, err := newRebalancerQueue(ctx)
+	q, err := newRebalancerQueue(ctx, metrics)
 	if err != nil {
 		return nil, err
 	}
 
 	b := &Balancer{
 		stopper:        stopper,
+		metrics:        metrics,
 		connTracker:    connTracker,
 		directoryCache: directoryCache,
 		queue:          q,
@@ -219,8 +224,7 @@ func (b *Balancer) processQueue(ctx context.Context) {
 		}
 
 		// TODO(jaylim-crl): implement enhancements:
-		//   1. Add metrics to track the number of active transfers.
-		//   2. Rate limit the number of transfers per connection (e.g. once
+		//   1. Rate limit the number of transfers per connection (e.g. once
 		//      every 5 minutes). This ensures that the connection isn't
 		//      ping-ponged between pods within a short interval. However, for
 		//      draining ones, we may want to move right away (or after 60 secs),
@@ -229,6 +233,9 @@ func (b *Balancer) processQueue(ctx context.Context) {
 		if err := b.stopper.RunAsyncTask(ctx, "processQueue-item", func(ctx context.Context) {
 			defer b.processSem.Release(1)
 
+			b.metrics.processRebalanceStart()
+			defer b.metrics.processRebalanceFinish()
+
 			// Each request is retried up to maxTransferAttempts.
 			for i := 0; i < maxTransferAttempts && ctx.Err() == nil; i++ {
 				// TODO(jaylim-crl): Once the TransferConnection API accepts a
@@ -377,14 +384,16 @@ type rebalanceRequest struct {
 // rebalancing requests. All methods on the queue are thread-safe.
 type rebalancerQueue struct {
 	mu       syncutil.Mutex
+	metrics  *Metrics
 	sem      semaphore.Semaphore
 	queue    *list.List
 	elements map[ConnectionHandle]*list.Element
 }
 
 // newRebalancerQueue returns a new instance of rebalancerQueue.
-func newRebalancerQueue(ctx context.Context) (*rebalancerQueue, error) {
+func newRebalancerQueue(ctx context.Context, metrics *Metrics) (*rebalancerQueue, error) {
 	q := &rebalancerQueue{
+		metrics:  metrics,
 		sem:      semaphore.New(math.MaxInt32),
 		queue:    list.New(),
 		elements: make(map[ConnectionHandle]*list.Element),
@@ -398,8 +407,7 @@ func newRebalancerQueue(ctx context.Context) (*rebalancerQueue, error) {
 }
 
 // enqueue puts the rebalance request into the queue. If a request for the
-// connection already exists, the newer of the two will be used. This returns
-// nil if the operation succeeded.
+// connection already exists, the newer of the two will be used.
 //
 // NOTE: req should not be nil.
 func (q *rebalancerQueue) enqueue(req *rebalanceRequest) {
@@ -412,11 +420,13 @@ func (q *rebalancerQueue) enqueue(req *rebalanceRequest) {
 		if e.Value.(*rebalanceRequest).createdAt.Before(req.createdAt) {
 			e.Value = req
 		}
-	} else {
-		e = q.queue.PushBack(req)
-		q.elements[req.conn] = e
-		q.sem.Release(1)
+		return
 	}
+
+	e = q.queue.PushBack(req)
+	q.elements[req.conn] = e
+	q.metrics.rebalanceReqQueued.Inc(1)
+	q.sem.Release(1)
 }
 
 // dequeue removes a request at the front of the queue, and returns that. If the
@@ -445,5 +455,6 @@ func (q *rebalancerQueue) dequeue(ctx context.Context) (*rebalanceRequest, error
 
 	req := q.queue.Remove(e).(*rebalanceRequest)
 	delete(q.elements, req.conn)
+	q.metrics.rebalanceReqQueued.Dec(1)
 	return req, nil
 }

pkg/ccl/sqlproxyccl/balancer/balancer_test.go

@@ -38,6 +38,7 @@ func TestBalancer_SelectTenantPod(t *testing.T) {
 	b, err := NewBalancer(
 		ctx,
 		stopper,
+		nil, /* metrics */
 		nil, /* directoryCache */
 		nil, /* connTracker */
 		NoRebalanceLoop(),
@@ -67,6 +68,7 @@ func TestRebalancer_processQueue(t *testing.T) {
 	b, err := NewBalancer(
 		ctx,
 		stopper,
+		NewMetrics(),
 		nil, /* directoryCache */
 		nil, /* connTracker */
 		MaxConcurrentRebalances(1),
@@ -99,6 +101,7 @@ func TestRebalancer_processQueue(t *testing.T) {
 			conn: &testBalancerConnHandle{
 				onTransferConnection: func() error {
 					count++
+					require.Equal(t, int64(1), b.metrics.rebalanceReqRunning.Value())
 					return errors.New("cannot transfer")
 				},
 			},
@@ -112,6 +115,7 @@ func TestRebalancer_processQueue(t *testing.T) {
 
 		// Ensure that we only retried up to 3 times.
 		require.Equal(t, 3, count)
+		require.Equal(t, int64(0), b.metrics.rebalanceReqRunning.Value())
 	})
 
 	t.Run("conn_was_transferred_by_other", func(t *testing.T) {
@@ -121,6 +125,7 @@ func TestRebalancer_processQueue(t *testing.T) {
 			count++
 			// Simulate that connection was transferred by someone else.
 			conn.remoteAddr = "foo"
+			require.Equal(t, int64(1), b.metrics.rebalanceReqRunning.Value())
 			return errors.New("cannot transfer")
 		}
 		req := &rebalanceRequest{
@@ -136,6 +141,7 @@ func TestRebalancer_processQueue(t *testing.T) {
 
 		// We should only retry once.
 		require.Equal(t, 1, count)
+		require.Equal(t, int64(0), b.metrics.rebalanceReqRunning.Value())
 	})
 
 	t.Run("conn_was_transferred", func(t *testing.T) {
@@ -144,6 +150,7 @@ func TestRebalancer_processQueue(t *testing.T) {
 		conn.onTransferConnection = func() error {
 			count++
 			conn.remoteAddr = "foo"
+			require.Equal(t, int64(1), b.metrics.rebalanceReqRunning.Value())
 			return nil
 		}
 		req := &rebalanceRequest{
@@ -159,13 +166,15 @@ func TestRebalancer_processQueue(t *testing.T) {
 
 		// We should only retry once.
 		require.Equal(t, 1, count)
+		require.Equal(t, int64(0), b.metrics.rebalanceReqRunning.Value())
 	})
 
 	t.Run("conn_was_closed", func(t *testing.T) {
 		count := 0
 		conn := &testBalancerConnHandle{}
 		conn.onTransferConnection = func() error {
 			count++
+			require.Equal(t, int64(1), b.metrics.rebalanceReqRunning.Value())
 			return context.Canceled
 		}
 		req := &rebalanceRequest{
@@ -181,9 +190,15 @@ func TestRebalancer_processQueue(t *testing.T) {
 
 		// We should only retry once.
 		require.Equal(t, 1, count)
+		require.Equal(t, int64(0), b.metrics.rebalanceReqRunning.Value())
 	})
 
 	t.Run("limit_concurrent_rebalances", func(t *testing.T) {
+		// Temporarily override metrics so we could test total counts
+		// independent of other tests.
+		defer func(oldMetrics *Metrics) { b.metrics = oldMetrics }(b.metrics)
+		b.metrics = NewMetrics()
+
 		const reqCount = 100
 
 		// Allow up to 2 concurrent rebalances.
@@ -215,6 +230,7 @@ func TestRebalancer_processQueue(t *testing.T) {
 						// concurrent rebalances defined.
 						newCount := atomic.AddInt32(&count, 1)
 						require.True(t, newCount <= 2)
+						require.True(t, b.metrics.rebalanceReqRunning.Value() <= 2)
 						return nil
 					},
 				},
@@ -231,6 +247,7 @@ func TestRebalancer_processQueue(t *testing.T) {
 
 		// We should only transfer once for every connection.
 		require.Equal(t, int32(0), count)
+		require.Equal(t, int64(reqCount), b.metrics.rebalanceReqTotal.Count())
 	})
 }
 
@@ -245,6 +262,7 @@ func TestRebalancer_rebalanceLoop(t *testing.T) {
 	stopper := stop.NewStopper()
 	defer stopper.Stop(ctx)
 
+	metrics := NewMetrics()
 	directoryCache := newTestDirectoryCache()
 	connTracker := NewConnTracker()
 
@@ -255,6 +273,7 @@ func TestRebalancer_rebalanceLoop(t *testing.T) {
 	_, err := NewBalancer(
 		ctx,
 		stopper,
+		metrics,
 		directoryCache,
 		connTracker,
 		TimeSource(timeSource),
@@ -294,6 +313,7 @@ func TestRebalancer_rebalance(t *testing.T) {
 	stopper := stop.NewStopper()
 	defer stopper.Stop(ctx)
 
+	metrics := NewMetrics()
 	directoryCache := newTestDirectoryCache()
 	connTracker := NewConnTracker()
 
@@ -304,6 +324,7 @@ func TestRebalancer_rebalance(t *testing.T) {
 	b, err := NewBalancer(
 		ctx,
 		stopper,
+		metrics,
 		directoryCache,
 		connTracker,
 		NoRebalanceLoop(),
@@ -467,7 +488,7 @@ func TestRebalancerQueue(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	q, err := newRebalancerQueue(ctx)
+	q, err := newRebalancerQueue(ctx, NewMetrics())
 	require.NoError(t, err)
 
 	// Use a custom time source for testing.
@@ -496,10 +517,13 @@ func TestRebalancerQueue(t *testing.T) {
 
 	// Enqueue in a specific order. req3 overrides req1; req2 is a no-op.
 	q.enqueue(req1)
+	require.Equal(t, int64(1), q.metrics.rebalanceReqQueued.Value())
 	q.enqueue(req3)
+	require.Equal(t, int64(1), q.metrics.rebalanceReqQueued.Value())
 	q.enqueue(req2)
 	require.Len(t, q.elements, 1)
 	require.Equal(t, 1, q.queue.Len())
+	require.Equal(t, int64(1), q.metrics.rebalanceReqQueued.Value())
 
 	// Create another request.
 	conn2 := &testBalancerConnHandle{}
@@ -509,16 +533,19 @@ func TestRebalancerQueue(t *testing.T) {
 		dst:       "bar1",
 	}
 	q.enqueue(req4)
+	require.Equal(t, int64(2), q.metrics.rebalanceReqQueued.Value())
 	require.Len(t, q.elements, 2)
 	require.Equal(t, 2, q.queue.Len())
 
 	// Dequeue the items.
 	item, err := q.dequeue(ctx)
 	require.NoError(t, err)
 	require.Equal(t, req3, item)
+	require.Equal(t, int64(1), q.metrics.rebalanceReqQueued.Value())
 	item, err = q.dequeue(ctx)
 	require.NoError(t, err)
 	require.Equal(t, req4, item)
+	require.Equal(t, int64(0), q.metrics.rebalanceReqQueued.Value())
 	require.Empty(t, q.elements)
 	require.Equal(t, 0, q.queue.Len())
 
@@ -527,6 +554,7 @@ func TestRebalancerQueue(t *testing.T) {
 	req4, err = q.dequeue(ctx)
 	require.EqualError(t, err, context.Canceled.Error())
 	require.Nil(t, req4)
+	require.Equal(t, int64(0), q.metrics.rebalanceReqQueued.Value())
 }
 
 func TestRebalancerQueueBlocking(t *testing.T) {
@@ -535,7 +563,7 @@ func TestRebalancerQueueBlocking(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	q, err := newRebalancerQueue(ctx)
+	q, err := newRebalancerQueue(ctx, NewMetrics())
 	require.NoError(t, err)
 
 	reqCh := make(chan *rebalanceRequest, 10)

pkg/ccl/sqlproxyccl/balancer/metrics.go

@@ -0,0 +1,65 @@
+// Copyright 2022 The Cockroach Authors.
+//
+// Licensed as a CockroachDB Enterprise file under the Cockroach Community
+// License (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+//     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
+
+package balancer
+
+import "github.com/cockroachdb/cockroach/pkg/util/metric"
+
+// Metrics contains pointers to the metrics for monitoring balancer-related
+// operations.
+type Metrics struct {
+	rebalanceReqRunning *metric.Gauge
+	rebalanceReqQueued  *metric.Gauge
+	rebalanceReqTotal   *metric.Counter
+}
+
+// MetricStruct implements the metrics.Struct interface.
+func (Metrics) MetricStruct() {}
+
+var _ metric.Struct = Metrics{}
+
+var (
+	metaRebalanceReqRunning = metric.Metadata{
+		Name:        "proxy.balancer.rebalance.running",
+		Help:        "Number of rebalance requests currently running",
+		Measurement: "Rebalance Requests",
+		Unit:        metric.Unit_COUNT,
+	}
+	metaRebalanceReqQueued = metric.Metadata{
+		Name:        "proxy.balancer.rebalance.queued",
+		Help:        "Number of rebalance requests currently queued",
+		Measurement: "Rebalance Requests",
+		Unit:        metric.Unit_COUNT,
+	}
+	metaRebalanceReqTotal = metric.Metadata{
+		Name:        "proxy.balancer.rebalance.total",
+		Help:        "Number of rebalance requests that were processed",
+		Measurement: "Rebalance Requests",
+		Unit:        metric.Unit_COUNT,
+	}
+)
+
+// NewMetrics instantiates the metrics holder for balancer monitoring.
+func NewMetrics() *Metrics {
+	return &Metrics{
+		rebalanceReqRunning: metric.NewGauge(metaRebalanceReqRunning),
+		rebalanceReqQueued:  metric.NewGauge(metaRebalanceReqQueued),
+		rebalanceReqTotal:   metric.NewCounter(metaRebalanceReqTotal),
+	}
+}
+
+// processRebalanceStart indicates the start of processing a rebalance request.
+func (m *Metrics) processRebalanceStart() {
+	m.rebalanceReqRunning.Inc(1)
+	m.rebalanceReqTotal.Inc(1)
+}
+
+// processRebalanceFinish indicates the end of processing a rebalance request.
+func (m *Metrics) processRebalanceFinish() {
+	m.rebalanceReqRunning.Dec(1)
+}

pkg/ccl/sqlproxyccl/connector_test.go

@@ -487,6 +487,7 @@ func TestConnector_lookupAddr(t *testing.T) {
 	balancer, err := balancer.NewBalancer(
 		ctx,
 		stopper,
+		nil, /* metrics */
 		nil, /* directoryCache */
 		nil, /* connTracker */
 		balancer.NoRebalanceLoop(),