From 459d5a78cbfe19b48391cf407dea916114db862f Mon Sep 17 00:00:00 2001 From: Miccah Date: Wed, 1 Jul 2026 14:28:50 -0700 Subject: [PATCH] Add prometheus metrics for engine channels and workers (#5095) * Add prometheus metrics for engine channels and workers * Address potential memory leak by unregistering existing metrics The engine isn't meant to be initialized multiple times in the same process, but this should prevent a bug in case it is. --- pkg/engine/engine.go | 7 ++ pkg/engine/runtime_collector.go | 201 ++++++++++++++++++++++++++++++++ 2 files changed, 208 insertions(+) create mode 100644 pkg/engine/runtime_collector.go diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go index 9b013ef37..b5a8435a3 100644 --- a/pkg/engine/engine.go +++ b/pkg/engine/engine.go @@ -232,6 +232,10 @@ type Engine struct { verificationOverlapWorkerMultiplier int maxDecodeDepth int + + // runtimeCollector exposes live channel/worker/scan counters to Prometheus + // while the engine is running. Set in Start, cleared in Finish. + runtimeCollector *runtimeCollector } // NewEngine creates a new Engine instance with the provided configuration. @@ -639,6 +643,7 @@ func (e *Engine) DetectorAvgTime() map[string][]time.Duration { func (e *Engine) Start(ctx context.Context) { e.metrics = runtimeMetrics{Metrics: Metrics{scanStartTime: time.Now()}} e.sanityChecks(ctx) + e.registerRuntimeMetrics(ctx) e.startWorkers(ctx) } @@ -756,6 +761,8 @@ func (e *Engine) Finish(ctx context.Context) error { e.metrics.ScanDuration = time.Since(e.metrics.scanStartTime) + e.unregisterRuntimeMetrics() + return err } diff --git a/pkg/engine/runtime_collector.go b/pkg/engine/runtime_collector.go new file mode 100644 index 000000000..8bc40f0c3 --- /dev/null +++ b/pkg/engine/runtime_collector.go @@ -0,0 +1,201 @@ +package engine + +import ( + "errors" + "sync/atomic" + + "github.com/prometheus/client_golang/prometheus" + + "github.com/trufflesecurity/trufflehog/v3/pkg/common" + "github.com/trufflesecurity/trufflehog/v3/pkg/context" +) + +// Channel and worker pool label values exposed by runtimeCollector. These +// names appear in metric labels, so external dashboards depend on them being +// stable. +const ( + channelSourceChunks = "source_chunks_chan" + channelDetectableChunks = "detectable_chunks_chan" + channelVerificationOverlapChunks = "verification_overlap_chunks_chan" + channelResults = "results_chan" + + workerEngine = "engine_workers" + workerDetector = "detector_workers" + workerVerificationOverlap = "verification_overlap_workers" + workerNotifier = "notifier_workers" + workerSources = "source_workers" +) + +// runtimeCollector exposes live engine internals (channel queue depths, +// channel capacities, configured worker pool sizes, and aggregate scan +// counters) to Prometheus. Values are sampled lazily on each scrape via +// len/cap and atomic loads, so there is no background goroutine overhead +// between scrapes. +type runtimeCollector struct { + engine *Engine + + channelSize *prometheus.Desc + channelCapacity *prometheus.Desc + workerCount *prometheus.Desc + activeSources *prometheus.Desc + + bytesScanned *prometheus.Desc + chunksScanned *prometheus.Desc + verifiedSecretsFound *prometheus.Desc + unverifiedSecretsFound *prometheus.Desc +} + +// newRuntimeCollector constructs a runtimeCollector bound to the given +// Engine. The metric descriptors are built once here; each subsequent Collect +// call reuses them and only samples fresh values. +func newRuntimeCollector(e *Engine) *runtimeCollector { + fq := func(name string) string { + return prometheus.BuildFQName(common.MetricsNamespace, common.MetricsSubsystem, name) + } + return &runtimeCollector{ + engine: e, + channelSize: prometheus.NewDesc( + fq("engine_channel_size"), + "Current number of items buffered in an engine channel.", + []string{"channel"}, nil, + ), + channelCapacity: prometheus.NewDesc( + fq("engine_channel_capacity"), + "Buffer capacity of an engine channel.", + []string{"channel"}, nil, + ), + workerCount: prometheus.NewDesc( + fq("engine_worker_count"), + "Configured number of workers in an engine worker pool.", + []string{"worker_type"}, nil, + ), + activeSources: prometheus.NewDesc( + fq("engine_active_sources"), + "Number of sources currently running.", + nil, nil, + ), + bytesScanned: prometheus.NewDesc( + fq("engine_bytes_scanned"), + "Total bytes scanned by the engine since Start.", + nil, nil, + ), + chunksScanned: prometheus.NewDesc( + fq("engine_chunks_scanned"), + "Total chunks scanned by the engine since Start.", + nil, nil, + ), + verifiedSecretsFound: prometheus.NewDesc( + fq("engine_verified_secrets_found"), + "Total verified secrets found by the engine since Start.", + nil, nil, + ), + unverifiedSecretsFound: prometheus.NewDesc( + fq("engine_unverified_secrets_found"), + "Total unverified secrets found by the engine since Start.", + nil, nil, + ), + } +} + +// Describe implements prometheus.Collector by emitting every metric +// descriptor this collector will ever produce. The registry uses these +// descriptors to detect conflicts at registration time. +func (c *runtimeCollector) Describe(ch chan<- *prometheus.Desc) { + ch <- c.channelSize + ch <- c.channelCapacity + ch <- c.workerCount + ch <- c.activeSources + ch <- c.bytesScanned + ch <- c.chunksScanned + ch <- c.verifiedSecretsFound + ch <- c.unverifiedSecretsFound +} + +// Collect implements prometheus.Collector and is invoked by the registry on +// every scrape. It samples the engine's channels via len/cap, reads scan +// counters via atomic.Load, and emits one metric per (descriptor, label set) +// combination. +func (c *runtimeCollector) Collect(ch chan<- prometheus.Metric) { + e := c.engine + + sourceChunks := e.sourceManager.Chunks() + channels := []struct { + name string + size int + capacity int + }{ + {channelSourceChunks, len(sourceChunks), cap(sourceChunks)}, + {channelDetectableChunks, len(e.detectableChunksChan), cap(e.detectableChunksChan)}, + {channelVerificationOverlapChunks, len(e.verificationOverlapChunksChan), cap(e.verificationOverlapChunksChan)}, + {channelResults, len(e.results), cap(e.results)}, + } + for _, q := range channels { + ch <- prometheus.MustNewConstMetric(c.channelSize, prometheus.GaugeValue, float64(q.size), q.name) + ch <- prometheus.MustNewConstMetric(c.channelCapacity, prometheus.GaugeValue, float64(q.capacity), q.name) + } + + workers := []struct { + name string + count int + }{ + {workerEngine, e.concurrency}, + {workerDetector, e.concurrency * e.detectorWorkerMultiplier}, + {workerVerificationOverlap, e.concurrency * e.verificationOverlapWorkerMultiplier}, + {workerNotifier, e.concurrency * e.notificationWorkerMultiplier}, + // SourceManager treats its concurrency limit as a semaphore rather than a + // fixed worker pool, but reporting it here lets dashboards ratio it against + // activeSources to spot source-side saturation. + {workerSources, e.sourceManager.MaxConcurrentSources()}, + } + for _, w := range workers { + ch <- prometheus.MustNewConstMetric(c.workerCount, prometheus.GaugeValue, float64(w.count), w.name) + } + + ch <- prometheus.MustNewConstMetric( + c.activeSources, prometheus.GaugeValue, + float64(e.sourceManager.ConcurrentSources()), + ) + + // Scan counters are written via atomic.AddUint64 on runtimeMetrics; read + // them with atomic.LoadUint64 to stay cheap on scrape and avoid the RW mutex. + m := &e.metrics.Metrics + ch <- prometheus.MustNewConstMetric(c.bytesScanned, prometheus.CounterValue, float64(atomic.LoadUint64(&m.BytesScanned))) + ch <- prometheus.MustNewConstMetric(c.chunksScanned, prometheus.CounterValue, float64(atomic.LoadUint64(&m.ChunksScanned))) + ch <- prometheus.MustNewConstMetric(c.verifiedSecretsFound, prometheus.CounterValue, float64(atomic.LoadUint64(&m.VerifiedSecretsFound))) + ch <- prometheus.MustNewConstMetric(c.unverifiedSecretsFound, prometheus.CounterValue, float64(atomic.LoadUint64(&m.UnverifiedSecretsFound))) +} + +// registerRuntimeMetrics installs the engine's runtime collector into the +// default Prometheus registry. If a collector with identical descriptors is +// already registered (e.g. a previous engine in the same process that didn't +// call Finish), the stale collector is evicted and replaced. Without eviction +// the stale collector would pin a dead engine in memory and permanently block +// future engines from exposing metrics. +func (e *Engine) registerRuntimeMetrics(ctx context.Context) { + collector := newRuntimeCollector(e) + err := prometheus.DefaultRegisterer.Register(collector) + if err != nil { + var already prometheus.AlreadyRegisteredError + if !errors.As(err, &already) { + ctx.Logger().Error(err, "failed to register engine runtime metrics") + return + } + ctx.Logger().V(2).Info("evicting stale engine runtime metrics collector") + prometheus.DefaultRegisterer.Unregister(already.ExistingCollector) + if err := prometheus.DefaultRegisterer.Register(collector); err != nil { + ctx.Logger().Error(err, "failed to register engine runtime metrics after eviction") + return + } + } + e.runtimeCollector = collector +} + +// unregisterRuntimeMetrics removes the engine's runtime collector from the +// default Prometheus registry. Safe to call when registration was skipped. +func (e *Engine) unregisterRuntimeMetrics() { + if e.runtimeCollector == nil { + return + } + prometheus.DefaultRegisterer.Unregister(e.runtimeCollector) + e.runtimeCollector = nil +}