fixup! drpcpool: add support for ShouldRecord

Author: suj-krishnan

drpcconn/conn.go

@@ -61,10 +61,11 @@ func NewWithOptions(tr drpc.Transport, opts Options) *Conn {
 		tr: tr,
 	}
 
+	shouldWrap := opts.ShouldRecord != nil
 	if opts.ShouldRecord == nil {
 		opts.ShouldRecord = func() bool { return false }
 	}
-	if opts.ShouldRecord() {
+	if shouldWrap {
 		mt := drpcmetrics.ToMeteredTransport(tr, opts.Metrics.BytesSent,
 			opts.Metrics.BytesRecv, opts.ShouldRecord)
 		c.tr = mt

drpcpool/pool.go

@@ -76,9 +76,8 @@ func (p *Pool[K, V]) initPoolMetrics() {
 	if p.opts.ShouldRecord == nil {
 		p.opts.ShouldRecord = func() bool { return false }
 	}
-	if !p.opts.ShouldRecord() {
-		return
-	}
+
+	// initialize the metrics with no-op implementations if they are nil
 	metrics := &p.opts.Metrics
 	if metrics.PoolSize == nil {
 		metrics.PoolSize = drpcmetrics.NoOpGauge{}

drpcpool/pool_test.go

@@ -588,6 +588,64 @@ func TestPoolMetrics_ShouldRecordFalse(t *testing.T) {
 	})
 }
 
+func TestPoolMetrics_ShouldRecordDynamic(t *testing.T) {
+	recording := false
+	poolSize := &testPoolGauge{}
+	hits := &testPoolCounter{}
+	misses := &testPoolCounter{}
+
+	pool := New[string, Conn](Options{
+		Capacity: 10,
+		Metrics: PoolMetrics{
+			PoolSize:              poolSize,
+			ConnectionHitsTotal:   hits,
+			ConnectionMissesTotal: misses,
+		},
+		ShouldRecord: func() bool { return recording },
+	})
+	defer func() { _ = pool.Close() }()
+
+	newConn := func() *callbackConn {
+		return &callbackConn{
+			ClosedFn:    func() <-chan struct{} { return nil },
+			UnblockedFn: func() <-chan struct{} { return closedCh },
+		}
+	}
+
+	// Recording is off — no metrics should be collected.
+	pool.Put("key", newConn())
+	pool.Take("key")
+	pool.Take("miss")
+	assert.Equal(t, poolSize.total, 0.0)
+	assert.Equal(t, hits.total, 0.0)
+	assert.Equal(t, misses.total, 0.0)
+
+	// Turn recording on — metrics should now be collected.
+	recording = true
+	pool.Put("key", newConn())
+	assert.Equal(t, poolSize.total, 1.0)
+
+	_, ok := pool.Take("key")
+	assert.That(t, ok)
+	assert.Equal(t, hits.total, 1.0)
+	assert.Equal(t, poolSize.total, 0.0)
+
+	_, ok = pool.Take("miss")
+	assert.That(t, !ok)
+	assert.Equal(t, misses.total, 1.0)
+
+	// Turn recording back off — counters should stop incrementing.
+	recording = false
+	pool.Put("key", newConn())
+	assert.Equal(t, poolSize.total, 0.0) // unchanged
+
+	pool.Take("key")
+	assert.Equal(t, hits.total, 1.0) // unchanged
+
+	pool.Take("miss2")
+	assert.Equal(t, misses.total, 1.0) // unchanged
+}
+
 func BenchmarkPool(b *testing.B) {
 	ctx := drpctest.NewTracker(b)
 	defer ctx.Close()

internal/integration/metrics_test.go

@@ -165,6 +165,90 @@ func TestClientByteMetricsNotCollected(t *testing.T) {
 	_ = conn.Close()
 }
 
+func TestClientByteMetricsShouldRecordFalse(t *testing.T) {
+	ctx := drpctest.NewTracker(t)
+	defer ctx.Close()
+
+	sent := &testCounter{}
+	recv := &testCounter{}
+
+	c1, c2 := net.Pipe()
+	mux := drpcmux.New()
+	assert.NoError(t, DRPCRegisterService(mux, standardImpl))
+	srv := drpcserver.New(mux)
+	ctx.Run(func(ctx2 context.Context) { _ = srv.ServeOne(ctx2, c1) })
+	conn := drpcconn.NewWithOptions(c2, drpcconn.Options{
+		Metrics: drpcmetrics.ClientMetrics{
+			BytesSent: sent,
+			BytesRecv: recv,
+		},
+		ShouldRecord: func() bool { return false },
+	})
+	cli := NewDRPCServiceClient(conn)
+
+	out, err := cli.Method1(ctx, in(1))
+	assert.NoError(t, err)
+	assert.True(t, Equal(out, &Out{Out: 1}))
+
+	// ShouldRecord returns false, so no metrics should be collected
+	// even though the transport is wrapped.
+	assert.Equal(t, sent.total(), 0.0)
+	assert.Equal(t, recv.total(), 0.0)
+
+	_ = conn.Close()
+}
+
+func TestClientByteMetricsShouldRecordDynamic(t *testing.T) {
+	ctx := drpctest.NewTracker(t)
+	defer ctx.Close()
+
+	recording := false
+	sent := &testCounter{}
+	recv := &testCounter{}
+
+	c1, c2 := net.Pipe()
+	mux := drpcmux.New()
+	assert.NoError(t, DRPCRegisterService(mux, standardImpl))
+	srv := drpcserver.New(mux)
+	ctx.Run(func(ctx2 context.Context) { _ = srv.ServeOne(ctx2, c1) })
+	conn := drpcconn.NewWithOptions(c2, drpcconn.Options{
+		Manager: drpcmanager.Options{},
+		Metrics: drpcmetrics.ClientMetrics{
+			BytesSent: sent,
+			BytesRecv: recv,
+		},
+		ShouldRecord: func() bool { return recording },
+	})
+	cli := NewDRPCServiceClient(conn)
+
+	// Recording off — no metrics.
+	out, err := cli.Method1(ctx, in(1))
+	assert.NoError(t, err)
+	assert.True(t, Equal(out, &Out{Out: 1}))
+	assert.Equal(t, sent.total(), 0.0)
+	assert.Equal(t, recv.total(), 0.0)
+
+	// Turn recording on — metrics should now be collected.
+	recording = true
+	out, err = cli.Method1(ctx, in(1))
+	assert.NoError(t, err)
+	assert.True(t, Equal(out, &Out{Out: 1}))
+	assert.That(t, sent.total() > 0)
+	assert.That(t, recv.total() > 0)
+
+	// Turn recording back off — counters should stop incrementing.
+	sentBefore := sent.total()
+	recvBefore := recv.total()
+	recording = false
+	out, err = cli.Method1(ctx, in(1))
+	assert.NoError(t, err)
+	assert.True(t, Equal(out, &Out{Out: 1}))
+	assert.Equal(t, sent.total(), sentBefore)
+	assert.Equal(t, recv.total(), recvBefore)
+
+	_ = conn.Close()
+}
+
 func TestServerTLSHandshakeErrorMetric(t *testing.T) {
 	tlsErrors := &testCounter{}