add subscriber and pool health metrics

Signed-off-by: satyamg1620 <Satyam.Gupta.3@ibm.com>

Author: satyamg1620

pkg/kvcache/metrics/collector.go

@@ -33,6 +33,10 @@ import (
 // aliases so existing scrapers keep working during the migration.
 const routerSubsystem = "llm_d_router_epp"
 
+// podIdentifierLabel is the label key carried by the per-pod kvevents metrics.
+// Keeping it as a constant localizes label-schema changes to this package.
+const podIdentifierLabel = "pod_identifier"
+
 // dualCounter emits a value to both the deprecated kvcache_* counter and the
 // current llm_d_router_epp_* counter. It satisfies prometheus.Collector so both
 // register, and exposes the recording/read methods the call sites use.
@@ -130,6 +134,53 @@ var (
 	DedupRemovedHashesForwarded = newDualCounter("kvevents", "dedup_removed_hashes_forwarded_total",
 		"kv_cache_events_dedup_removed_hashes_forwarded_total",
 		"Block hashes forwarded for eviction after the KV-event dedup filter (block hashes, not BlockRemoved events)")
+
+	// SubscriberActive tracks the number of ZMQ subscribers currently managed by
+	// the SubscriberManager. A subscriber is counted from creation until it is
+	// removed, so the gauge includes subscribers that are retrying a failed
+	// connection; it is a measure of intent, not of live connectivity. Pair it
+	// with SubscriberReconnections and ZMQErrors to spot unhealthy subscribers.
+	SubscriberActive = prometheus.NewGauge(prometheus.GaugeOpts{
+		Subsystem: routerSubsystem, Name: "kv_cache_events_active_subscribers",
+		Help: metricsutil.HelpMsgWithStability(
+			"Number of ZMQ subscribers currently managed, including those retrying a failed connection",
+			compbasemetrics.ALPHA),
+	})
+	// SubscriberReconnections counts ZMQ subscriber reconnection attempts,
+	// labeled by the pod identifier the subscriber is bound to.
+	SubscriberReconnections = prometheus.NewCounterVec(prometheus.CounterOpts{
+		Subsystem: routerSubsystem, Name: "kv_cache_events_subscriber_reconnections_total",
+		Help: metricsutil.HelpMsgWithStability(
+			"Total number of ZMQ subscriber reconnection attempts", compbasemetrics.ALPHA),
+	}, []string{podIdentifierLabel})
+	// MessagesReceived counts raw messages received from ZMQ subscribers. The
+	// rate of message ingestion can be derived via rate() in Prometheus.
+	MessagesReceived = prometheus.NewCounterVec(prometheus.CounterOpts{
+		Subsystem: routerSubsystem, Name: "kv_cache_events_messages_received_total",
+		Help: metricsutil.HelpMsgWithStability(
+			"Total number of messages received from ZMQ subscribers", compbasemetrics.ALPHA),
+	}, []string{podIdentifierLabel})
+	// ZMQErrors counts errors encountered by ZMQ subscribers, labeled by the
+	// pod identifier and the operation that failed (e.g. "connect", "recv").
+	ZMQErrors = prometheus.NewCounterVec(prometheus.CounterOpts{
+		Subsystem: routerSubsystem, Name: "kv_cache_events_zmq_errors_total",
+		Help: metricsutil.HelpMsgWithStability(
+			"Total number of ZMQ subscriber errors", compbasemetrics.ALPHA),
+	}, []string{podIdentifierLabel, "operation"})
+	// PoolQueueDepth tracks the total number of messages queued across all
+	// worker shards of the event processing pool.
+	PoolQueueDepth = prometheus.NewGauge(prometheus.GaugeOpts{
+		Subsystem: routerSubsystem, Name: "kv_cache_events_pool_queue_depth",
+		Help: metricsutil.HelpMsgWithStability(
+			"Total number of messages queued across all worker shards", compbasemetrics.ALPHA),
+	})
+	// PoolCapacity tracks the number of worker shards (capacity) configured for
+	// the event processing pool.
+	PoolCapacity = prometheus.NewGauge(prometheus.GaugeOpts{
+		Subsystem: routerSubsystem, Name: "kv_cache_events_pool_capacity",
+		Help: metricsutil.HelpMsgWithStability(
+			"Number of worker shards (capacity) in the event processing pool", compbasemetrics.ALPHA),
+	})
 )
 
 // Collectors returns a slice of all registered Prometheus collectors.
@@ -138,9 +189,21 @@ func Collectors() []prometheus.Collector {
 		Admissions, Evictions,
 		LookupRequests, LookupHits, LookupLatency, MaxPodHitCount,
 		DedupRemovedHashesSuppressed, DedupRemovedHashesForwarded,
+		SubscriberActive, SubscriberReconnections, MessagesReceived, ZMQErrors,
+		PoolQueueDepth, PoolCapacity,
 	}
 }
 
+// CleanupSubscriber drops every per-pod kvevents series for podIdentifier so
+// stale time series do not linger after a subscriber is removed. Callers must
+// invoke it only once the subscriber's goroutine has exited, otherwise a late
+// increment resurrects the series.
+func CleanupSubscriber(podIdentifier string) {
+	SubscriberReconnections.DeleteLabelValues(podIdentifier)
+	MessagesReceived.DeleteLabelValues(podIdentifier)
+	ZMQErrors.DeletePartialMatch(prometheus.Labels{podIdentifierLabel: podIdentifier})
+}
+
 var registerMetricsOnce = sync.Once{}
 
 // Register registers all metrics with K8s registry.

pkg/kvcache/metrics/collector_test.go

@@ -52,6 +52,12 @@ func TestCollectorsIncludesAllMetrics(t *testing.T) {
 		{"MaxPodHitCount", MaxPodHitCount},
 		{"DedupRemovedHashesSuppressed", DedupRemovedHashesSuppressed},
 		{"DedupRemovedHashesForwarded", DedupRemovedHashesForwarded},
+		{"SubscriberActive", SubscriberActive},
+		{"SubscriberReconnections", SubscriberReconnections},
+		{"MessagesReceived", MessagesReceived},
+		{"ZMQErrors", ZMQErrors},
+		{"PoolQueueDepth", PoolQueueDepth},
+		{"PoolCapacity", PoolCapacity},
 	}
 
 	for _, e := range expected {
@@ -61,6 +67,92 @@ func TestCollectorsIncludesAllMetrics(t *testing.T) {
 	}
 }
 
+func TestKVEventsMetricNames(t *testing.T) {
+	// The kvevents observability metrics must be emitted under the router EPP
+	// subsystem with the kv_cache_events prefix.
+	reg := prometheus.NewRegistry()
+	kvevents := []prometheus.Collector{
+		SubscriberActive, SubscriberReconnections, MessagesReceived,
+		ZMQErrors, PoolQueueDepth, PoolCapacity,
+	}
+	for _, c := range kvevents {
+		reg.MustRegister(c)
+	}
+
+	// Emit a sample for each labeled metric so it appears in the gather output.
+	SubscriberActive.Set(1)
+	SubscriberReconnections.WithLabelValues("pod-a").Inc()
+	MessagesReceived.WithLabelValues("pod-a").Inc()
+	ZMQErrors.WithLabelValues("pod-a", "recv").Inc()
+	PoolQueueDepth.Set(3)
+	PoolCapacity.Set(4)
+
+	mfs, err := reg.Gather()
+	if err != nil {
+		t.Fatalf("Gather() failed: %v", err)
+	}
+
+	got := make(map[string]bool, len(mfs))
+	for _, mf := range mfs {
+		got[mf.GetName()] = true
+	}
+
+	for _, name := range []string{
+		"llm_d_router_epp_kv_cache_events_active_subscribers",
+		"llm_d_router_epp_kv_cache_events_subscriber_reconnections_total",
+		"llm_d_router_epp_kv_cache_events_messages_received_total",
+		"llm_d_router_epp_kv_cache_events_zmq_errors_total",
+		"llm_d_router_epp_kv_cache_events_pool_queue_depth",
+		"llm_d_router_epp_kv_cache_events_pool_capacity",
+	} {
+		if !got[name] {
+			t.Errorf("expected metric %q to be registered, got names: %v", name, got)
+		}
+	}
+}
+
+func TestCleanupSubscriberDropsPerPodSeries(t *testing.T) {
+	// Removing a subscriber must drop only its own series, leaving other pods'
+	// series intact.
+	SubscriberReconnections.WithLabelValues("pod-a").Inc()
+	MessagesReceived.WithLabelValues("pod-a").Inc()
+	ZMQErrors.WithLabelValues("pod-a", "recv").Inc()
+	ZMQErrors.WithLabelValues("pod-a", "connect").Inc()
+	SubscriberReconnections.WithLabelValues("pod-b").Inc()
+	MessagesReceived.WithLabelValues("pod-b").Inc()
+	ZMQErrors.WithLabelValues("pod-b", "recv").Inc()
+
+	CleanupSubscriber("pod-a")
+
+	reg := prometheus.NewRegistry()
+	reg.MustRegister(SubscriberReconnections, MessagesReceived, ZMQErrors)
+	mfs, err := reg.Gather()
+	if err != nil {
+		t.Fatalf("Gather() failed: %v", err)
+	}
+
+	for _, mf := range mfs {
+		for _, m := range mf.GetMetric() {
+			for _, l := range m.GetLabel() {
+				if l.GetName() == podIdentifierLabel && l.GetValue() == "pod-a" {
+					t.Errorf("metric %q still has a series for pod-a", mf.GetName())
+				}
+			}
+		}
+	}
+
+	// pod-b must be untouched: three series across the three metrics.
+	remaining := 0
+	for _, mf := range mfs {
+		remaining += len(mf.GetMetric())
+	}
+	if remaining != 3 {
+		t.Errorf("expected 3 remaining pod-b series, got %d", remaining)
+	}
+
+	CleanupSubscriber("pod-b")
+}
+
 func TestLogMetrics(t *testing.T) {
 	// Set up a buffer to capture log output
 	var buf bytes.Buffer

pkg/kvevents/pool.go

@@ -20,6 +20,7 @@ import (
 	"hash/fnv"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"k8s.io/client-go/util/workqueue"
 	"sigs.k8s.io/controller-runtime/pkg/log"
@@ -114,11 +115,19 @@ type Pool struct {
 	// Index.Add succeeds — both of which only the Pool observes.
 	dedup *eventDedupFilter
 	wg    sync.WaitGroup
+	// queueDepth mirrors the number of tasks queued across all shards. It is
+	// tracked incrementally rather than by summing queue.Len() so that the
+	// depth gauge stays O(1) on the enqueue/dequeue hot path.
+	queueDepth atomic.Int64
 }
 
 // NewPool creates a Pool with a sharded worker setup.
 // Subscribers are managed by SubscriberManager which is controlled by the pod
 // reconciler.
+//
+// Side effect: it registers the kvcache metrics with the controller-runtime
+// registry so that the kvevents metrics are scraped wherever a pool runs.
+// Registration is idempotent (guarded by a sync.Once).
 func NewPool(cfg *Config, index kvblock.Index, tokenProcessor kvblock.TokenProcessor,
 	adapter EngineAdapter,
 ) *Pool {
@@ -140,9 +149,17 @@ func NewPool(cfg *Config, index kvblock.Index, tokenProcessor kvblock.TokenProce
 		p.queues[i] = workqueue.NewTypedRateLimitingQueue(workqueue.DefaultTypedControllerRateLimiter[*RawMessage]())
 	}
 
+	metrics.Register()
+
 	return p
 }
 
+// addQueueDepth adjusts the tracked queue depth by delta and publishes the new
+// total to the depth gauge.
+func (p *Pool) addQueueDepth(delta int64) {
+	metrics.PoolQueueDepth.Set(float64(p.queueDepth.Add(delta)))
+}
+
 // GroupCatalog returns the KV cache group metadata learned from events.
 func (p *Pool) GroupCatalog() *kvblock.GroupCatalog {
 	return p.groupCatalog
@@ -154,6 +171,8 @@ func (p *Pool) Start(ctx context.Context) {
 	logger := log.FromContext(ctx)
 	logger.Info("Starting sharded event processing pool", "workers", p.concurrency)
 
+	metrics.PoolCapacity.Set(float64(p.concurrency))
+
 	p.wg.Add(p.concurrency)
 	for i := 0; i < p.concurrency; i++ {
 		// Each worker is given its own dedicated queue shard.
@@ -171,6 +190,12 @@ func (p *Pool) Shutdown(ctx context.Context) {
 	}
 
 	p.wg.Wait()
+
+	// Tasks still queued at shutdown are dropped with the queues, so reset the
+	// depth rather than leaving the gauge pinned at the undrained count.
+	p.queueDepth.Store(0)
+	metrics.PoolQueueDepth.Set(0)
+
 	logger.Info("event processing pool shut down.")
 }
 
@@ -189,6 +214,7 @@ func (p *Pool) AddTask(task *RawMessage) {
 	//nolint:gosec // if concurrency overflows then the world is in trouble anyway
 	queueIndex := h.Sum32() % uint32(p.concurrency)
 	p.queues[queueIndex].Add(task)
+	p.addQueueDepth(1)
 }
 
 // worker is the main processing loop for a single worker goroutine.
@@ -209,6 +235,7 @@ func (p *Pool) worker(ctx context.Context, workerIndex int) {
 			// Task succeeded, remove it from the queue.
 			queue.Forget(task)
 		}(task)
+		p.addQueueDepth(-1)
 
 		// Check if context was cancelled after processing a task.
 		select {

pkg/kvevents/pool_test.go

@@ -3,6 +3,7 @@ package kvevents //nolint:testpackage // tests use unexported processEventBatch
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	dto "github.com/prometheus/client_model/go"
@@ -1096,6 +1097,64 @@ func TestPool_DedupMetricsCountBlockHashes(t *testing.T) {
 		"second remove must forward all 4 constituent block hashes")
 }
 
+// stubAdapter is a minimal EngineAdapter that shards every message onto the
+// same key and decodes to an empty batch, so tasks flow through the pool
+// without exercising any engine-specific parsing.
+type stubAdapter struct{}
+
+//nolint:gocritic // unnamed results match the EngineAdapter implementations
+func (stubAdapter) ParseMessage(_ *RawMessage) (string, string, EventBatch, error) {
+	return "pod-1", "model-1", EventBatch{}, nil
+}
+
+func (stubAdapter) ShardingKey(_ *RawMessage) string { return "pod-1" }
+
+// TestPool_QueueDepthAccounting verifies that the queue depth gauge tracks
+// enqueues and dequeues, and is reset once the pool shuts down.
+func TestPool_QueueDepthAccounting(t *testing.T) {
+	idx, err := kvblock.NewInMemoryIndex(kvblock.DefaultInMemoryIndexConfig())
+	require.NoError(t, err)
+
+	tp, err := kvblock.NewChunkedTokenDatabase(&kvblock.TokenProcessorConfig{
+		BlockSizeTokens: 4, HashSeed: "test",
+	})
+	require.NoError(t, err)
+
+	cfg := DefaultConfig()
+	cfg.Concurrency = 2
+	pool := NewPool(cfg, idx, tp, stubAdapter{})
+
+	const tasks = 3
+	for i := range uint64(tasks) {
+		pool.AddTask(&RawMessage{Topic: "kv@pod-1@model-1", Sequence: i})
+	}
+
+	assert.Equal(t, int64(tasks), pool.queueDepth.Load(), "every enqueued task must be counted")
+	assert.InDelta(t, float64(tasks), gaugeValue(t, metrics.PoolQueueDepth), 0.001)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	pool.Start(ctx)
+
+	assert.Eventually(t, func() bool {
+		return pool.queueDepth.Load() == 0
+	}, 5*time.Second, 10*time.Millisecond, "workers must decrement the depth as tasks drain")
+
+	pool.Shutdown(ctx)
+	assert.Equal(t, int64(0), pool.queueDepth.Load())
+	assert.InDelta(t, 0.0, gaugeValue(t, metrics.PoolQueueDepth), 0.001)
+	assert.InDelta(t, float64(cfg.Concurrency), gaugeValue(t, metrics.PoolCapacity), 0.001)
+}
+
+// gaugeValue reads the current value of a prometheus.Gauge without touching the
+// global registry.
+func gaugeValue(t *testing.T, g prometheus.Gauge) float64 {
+	t.Helper()
+	var m dto.Metric
+	require.NoError(t, g.Write(&m))
+	return m.GetGauge().GetValue()
+}
+
 // counterValue reads the current value of a plain prometheus.Counter without
 // touching the global registry, using the same dto.Metric.Write pattern as
 // pkg/kvcache/metrics.logMetrics.