*: add drpc client and server metrics

Add metrics collection for drpc client connections, server connections,
and stream pool.
- Introduce Counter, Observer, and Gauge metric interfaces in the drpc package,
allowing us to plug in the metrics implementation from cockroach
- Client metrics: bytes sent/received
- Server metrics: bytes sent/received, TLS handshake errors, inactivity
timeouts, and per-RPC stream lifecycle metrics (total, active, duration,
errors) with labels for service/method/type
- Pool metrics: pool size, connection hits/misses, connection age

Author: suj-krishnan

drpcconn/conn.go

@@ -19,14 +19,25 @@ import (
 	"storj.io/drpc/internal/drpcopts"
 )
 
+// ClientMetrics holds optional metrics that the client connection will
+// populate during operation.
+type ClientMetrics struct {
+	BytesSent drpc.Counter
+	BytesRecv drpc.Counter
+}
+
 // Options controls configuration settings for a conn.
 type Options struct {
 	// Manager controls the options we pass to the manager of this conn.
 	Manager drpcmanager.Options
 
-	// CollectStats controls whether the server should collect stats on the
+	// CollectStats controls whether the client should collect stats on the
 	// rpcs it creates.
 	CollectStats bool
+
+	// Metrics holds optional metrics the conn will populate. If nil, no
+	// metrics are recorded.
+	Metrics *ClientMetrics
 }
 
 // Conn is a drpc client connection.
@@ -51,6 +62,15 @@ func NewWithOptions(tr drpc.Transport, opts Options) *Conn {
 		tr: tr,
 	}
 
+	if opts.Metrics != nil && (opts.Metrics.BytesSent != nil || opts.Metrics.BytesRecv != nil) {
+		tr = &drpc.MeteredTransport{
+			Transport: tr,
+			BytesSent: opts.Metrics.BytesSent,
+			BytesRecv: opts.Metrics.BytesRecv,
+		}
+		c.tr = tr
+	}
+
 	if opts.CollectStats {
 		drpcopts.SetManagerStatsCB(&opts.Manager.Internal, c.getStats)
 		c.stats = make(map[string]*drpcstats.Stats)

drpcmux/mux.go

@@ -67,6 +67,13 @@ func (m *Mux) Register(srv interface{}, desc drpc.Description) error {
 	return nil
 }
 
+// IsUnitary reports whether the named RPC has unary input and output.
+// It returns false if the RPC is not registered or is a streaming RPC.
+func (m *Mux) IsUnitary(rpc string) bool {
+	data, ok := m.rpcs[rpc]
+	return ok && data.unitary
+}
+
 // registerOne does the work to register a single rpc.
 func (m *Mux) registerOne(
 	srv interface{}, rpc string, enc drpc.Encoding, receiver drpc.Receiver, method interface{},

drpcpool/entry.go

@@ -9,11 +9,12 @@ import (
 )
 
 type entry[K comparable, V Conn] struct {
-	key    K
-	val    V
-	exp    *time.Timer
-	global node[K, V]
-	local  node[K, V]
+	key     K
+	val     V
+	exp     *time.Timer
+	created time.Time
+	global  node[K, V]
+	local   node[K, V]
 }
 
 func (e *entry[K, V]) String() string {

drpcpool/pool.go

@@ -11,9 +11,18 @@ import (
 
 	"github.com/zeebo/errs"
 
+	"storj.io/drpc"
 	"storj.io/drpc/drpcdebug"
 )
 
+// PoolMetrics holds optional metrics for connection pool monitoring.
+type PoolMetrics struct {
+	PoolSize              drpc.Gauge
+	ConnectionHitsTotal   drpc.Counter
+	ConnectionMissesTotal drpc.Counter
+	ConnectionAgeSeconds  drpc.Observer
+}
+
 // Options contains the options to configure a pool.
 type Options struct {
 	// Expiration will remove any values from the Pool after the
@@ -28,6 +37,10 @@ type Options struct {
 	// the Pool holds unlimited for any single key. Negative means
 	// no values for any single key.
 	KeyCapacity int
+
+	// Metrics holds optional metrics the pool will populate. If nil,
+	// no metrics are recorded.
+	Metrics *PoolMetrics
 }
 
 // Pool is a connection pool with key type K. It maintains a cache of connections
@@ -49,6 +62,30 @@ func New[K comparable, V Conn](opts Options) *Pool[K, V] {
 	}
 }
 
+func (p *Pool[K, V]) recordHit(created time.Time) {
+	if p.opts.Metrics == nil {
+		return
+	}
+	if p.opts.Metrics.ConnectionHitsTotal != nil {
+		p.opts.Metrics.ConnectionHitsTotal.Add(nil, 1)
+	}
+	if p.opts.Metrics.ConnectionAgeSeconds != nil {
+		p.opts.Metrics.ConnectionAgeSeconds.Observe(nil, time.Since(created).Seconds())
+	}
+}
+
+func (p *Pool[K, V]) recordMiss() {
+	if p.opts.Metrics != nil && p.opts.Metrics.ConnectionMissesTotal != nil {
+		p.opts.Metrics.ConnectionMissesTotal.Add(nil, 1)
+	}
+}
+
+func (p *Pool[K, V]) recordPoolSizeChange(delta float64) {
+	if p.opts.Metrics != nil && p.opts.Metrics.PoolSize != nil {
+		p.opts.Metrics.PoolSize.Add(nil, delta)
+	}
+}
+
 func (p *Pool[K, V]) log(what string, cb func() string) {
 	if drpcdebug.Enabled {
 		drpcdebug.Log(func() (_, _, _ string) { return fmt.Sprintf("<pül %p>", p), what, cb() })
@@ -62,12 +99,14 @@ func (p *Pool[K, V]) Close() (err error) {
 	defer p.mu.Unlock()
 
 	var eg errs.Group
+	count := p.order.count
 	for ent := p.order.head; ent != nil; ent = ent.global.next {
 		eg.Add(p.closeEntry(ent))
 	}
 
 	p.entries = make(map[K]*list[K, V])
 	p.order = list[K, V]{}
+	p.recordPoolSizeChange(-float64(count))
 
 	return eg.Err()
 }
@@ -99,6 +138,7 @@ func (p *Pool[K, V]) removeEntry(ent *entry[K, V]) {
 
 	local.removeEntry(ent, (*entry[K, V]).localList)
 	p.order.removeEntry(ent, (*entry[K, V]).globalList)
+	p.recordPoolSizeChange(-1)
 
 	if local.count == 0 {
 		delete(p.entries, ent.key)
@@ -123,6 +163,7 @@ func (p *Pool[K, V]) Take(key K) (V, bool) {
 
 	local := p.entries[key]
 	if local == nil {
+		p.recordMiss()
 		return *new(V), false
 	}
 
@@ -137,6 +178,7 @@ func (p *Pool[K, V]) Take(key K) (V, bool) {
 
 		local.removeEntry(ent, (*entry[K, V]).localList)
 		p.order.removeEntry(ent, (*entry[K, V]).globalList)
+		p.recordPoolSizeChange(-1)
 
 		if ent.exp != nil && !ent.exp.Stop() {
 			continue
@@ -145,9 +187,11 @@ func (p *Pool[K, V]) Take(key K) (V, bool) {
 		}
 
 		p.log("TAKEN", ent.String)
+		p.recordHit(ent.created)
 		return ent.val, true
 	}
 
+	p.recordMiss()
 	return *new(V), false
 }
 
@@ -177,6 +221,7 @@ func (p *Pool[K, V]) Put(key K, val V) {
 
 		local.removeEntry(ent, (*entry[K, V]).localList)
 		p.order.removeEntry(ent, (*entry[K, V]).globalList)
+		p.recordPoolSizeChange(-1)
 	}
 
 	for p.opts.Capacity != 0 && p.order.count >= p.opts.Capacity {
@@ -187,15 +232,17 @@ func (p *Pool[K, V]) Put(key K, val V) {
 
 		local.removeEntry(ent, (*entry[K, V]).localList)
 		p.order.removeEntry(ent, (*entry[K, V]).globalList)
+		p.recordPoolSizeChange(-1)
 
 		if local.count == 0 {
 			delete(p.entries, ent.key)
 		}
 	}
 
-	ent := &entry[K, V]{key: key, val: val}
+	ent := &entry[K, V]{key: key, val: val, created: time.Now()}
 	local.appendEntry(ent, (*entry[K, V]).localList)
 	p.order.appendEntry(ent, (*entry[K, V]).globalList)
+	p.recordPoolSizeChange(1)
 	p.log("PUT", ent.String)
 
 	if p.opts.Expiration > 0 {

drpcpool/pool_test.go

@@ -419,6 +419,144 @@ func TestPool_StreamContext(t *testing.T) {
 	}
 }
 
+//
+// pool metrics tests
+//
+
+type testPoolCounter struct {
+	total float64
+}
+
+func (c *testPoolCounter) Add(_ map[string]string, value float64) {
+	c.total += value
+}
+
+type testPoolGauge struct {
+	total float64
+}
+
+func (g *testPoolGauge) Add(_ map[string]string, value float64) {
+	g.total += value
+}
+
+type testPoolObserver struct {
+	calls int
+	last  float64
+}
+
+func (o *testPoolObserver) Observe(_ map[string]string, value float64) {
+	o.calls++
+	o.last = value
+}
+
+func TestPoolMetrics_PutTakeClose(t *testing.T) {
+	ctx := drpctest.NewTracker(t)
+	defer ctx.Close()
+
+	poolSize := &testPoolGauge{}
+	hits := &testPoolCounter{}
+	misses := &testPoolCounter{}
+	age := &testPoolObserver{}
+
+	pool := New[string, Conn](Options{
+		Capacity: 10,
+		Metrics: &PoolMetrics{
+			PoolSize:              poolSize,
+			ConnectionHitsTotal:   hits,
+			ConnectionMissesTotal: misses,
+			ConnectionAgeSeconds:  age,
+		},
+	})
+
+	// Miss on empty pool.
+	_, ok := pool.Take("key")
+	assert.That(t, !ok)
+	assert.Equal(t, misses.total, 1.0)
+	assert.Equal(t, hits.total, 0.0)
+
+	// Put a connection.
+	conn := &callbackConn{
+		ClosedFn:    func() <-chan struct{} { return nil },
+		UnblockedFn: func() <-chan struct{} { return closedCh },
+	}
+	pool.Put("key", conn)
+	assert.Equal(t, poolSize.total, 1.0)
+
+	// Take it back — should be a hit.
+	val, ok := pool.Take("key")
+	assert.That(t, ok)
+	assert.Equal(t, val, Conn(conn))
+	assert.Equal(t, hits.total, 1.0)
+	assert.Equal(t, poolSize.total, 0.0) // decremented on take
+	assert.Equal(t, age.calls, 1)
+	assert.That(t, age.last >= 0)
+
+	// Miss again now that pool is empty.
+	_, ok = pool.Take("key")
+	assert.That(t, !ok)
+	assert.Equal(t, misses.total, 2.0)
+
+	// Put two, then Close.
+	conn2 := &callbackConn{
+		ClosedFn:    func() <-chan struct{} { return nil },
+		UnblockedFn: func() <-chan struct{} { return closedCh },
+	}
+	pool.Put("key", conn)
+	pool.Put("key", conn2)
+	assert.Equal(t, poolSize.total, 2.0)
+
+	assert.NoError(t, pool.Close())
+	assert.Equal(t, poolSize.total, 0.0)
+}
+
+func TestPoolMetrics_Eviction(t *testing.T) {
+	poolSize := &testPoolGauge{}
+	misses := &testPoolCounter{}
+
+	pool := New[string, Conn](Options{
+		Capacity:    1,
+		KeyCapacity: 1,
+		Metrics: &PoolMetrics{
+			PoolSize:              poolSize,
+			ConnectionMissesTotal: misses,
+		},
+	})
+	defer func() { _ = pool.Close() }()
+
+	conn1 := &callbackConn{
+		ClosedFn:    func() <-chan struct{} { return nil },
+		UnblockedFn: func() <-chan struct{} { return closedCh },
+	}
+	conn2 := &callbackConn{
+		ClosedFn:    func() <-chan struct{} { return nil },
+		UnblockedFn: func() <-chan struct{} { return closedCh },
+	}
+
+	pool.Put("key", conn1)
+	assert.Equal(t, poolSize.total, 1.0)
+
+	// Putting conn2 with KeyCapacity=1 should evict conn1.
+	pool.Put("key", conn2)
+	// One eviction (-1) + one put (+1) = net 1.
+	assert.Equal(t, poolSize.total, 1.0)
+}
+
+func TestPoolMetrics_NilFields(t *testing.T) {
+	// All PoolMetrics fields are nil — should not panic.
+	pool := New[string, Conn](Options{
+		Metrics: &PoolMetrics{},
+	})
+
+	conn := &callbackConn{
+		ClosedFn:    func() <-chan struct{} { return nil },
+		UnblockedFn: func() <-chan struct{} { return closedCh },
+	}
+	pool.Put("key", conn)
+	pool.Take("key")
+	pool.Take("miss")
+	assert.NoError(t, pool.Close())
+}
+
 func BenchmarkPool(b *testing.B) {
 	ctx := drpctest.NewTracker(b)
 	defer ctx.Close()