Add optional PinnableMetricsQuerier interface

A MetricsQuerier whose data can change between queries (for example one
serving in-memory snapshots that are replaced on an interval) can
implement source.PinnableMetricsQuerier to hand out a view bound to one
state of the data. PinMetrics and PinDataSource pin a querier or data
source when it supports this and return it unchanged otherwise, so
existing MetricsQuerier and OpenCostDataSource implementations are
unaffected.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: peatey <warwick@automatic.systems>

core/pkg/source/pin.go

@@ -0,0 +1,55 @@
+package source
+
+// PinnableMetricsQuerier is optionally implemented by a MetricsQuerier whose underlying data can change
+// between queries, for example one that serves from in-memory snapshots replaced on an interval.
+//
+// Computations that issue many queries for one window (allocation, assets, kube model) pin the querier
+// once and issue every query through the pinned view, so that all results come from one consistent
+// state of the data rather than a mix of states from before and after an update.
+type PinnableMetricsQuerier interface {
+	MetricsQuerier
+
+	// Pin returns a MetricsQuerier bound to the current state of the data, and a release function
+	// which must be called once the pinned querier is no longer used. The pinned querier must remain
+	// valid until released, even if the underlying data is updated in the meantime.
+	Pin() (MetricsQuerier, func())
+}
+
+// PinMetrics pins the querier if it implements PinnableMetricsQuerier. Otherwise it returns the querier
+// unchanged and a no-op release function.
+func PinMetrics(q MetricsQuerier) (MetricsQuerier, func()) {
+	if p, ok := q.(PinnableMetricsQuerier); ok {
+		pinned, release := p.Pin()
+		if release == nil {
+			release = func() {}
+		}
+		return pinned, release
+	}
+	return q, func() {}
+}
+
+// PinDataSource returns a data source whose Metrics() always returns the same pinned querier, for use
+// by a computation that reads metrics through several helpers taking an OpenCostDataSource. If the data
+// source's querier does not implement PinnableMetricsQuerier, the data source is returned unchanged.
+//
+// The returned data source only forwards the OpenCostDataSource methods; optional interfaces
+// implemented by the original are not visible through it.
+func PinDataSource(ds OpenCostDataSource) (OpenCostDataSource, func()) {
+	q := ds.Metrics()
+	if _, ok := q.(PinnableMetricsQuerier); !ok {
+		return ds, func() {}
+	}
+
+	pinned, release := PinMetrics(q)
+	return &pinnedDataSource{OpenCostDataSource: ds, metrics: pinned}, release
+}
+
+// pinnedDataSource is an OpenCostDataSource whose Metrics() returns a pinned querier.
+type pinnedDataSource struct {
+	OpenCostDataSource
+	metrics MetricsQuerier
+}
+
+func (p *pinnedDataSource) Metrics() MetricsQuerier {
+	return p.metrics
+}

core/pkg/source/pin_test.go

@@ -0,0 +1,59 @@
+package source
+
+import "testing"
+
+type pinnableQuerier struct {
+	*MockMetricsQuerier
+	pinned   *MockMetricsQuerier
+	released int
+}
+
+func (p *pinnableQuerier) Pin() (MetricsQuerier, func()) {
+	return p.pinned, func() { p.released++ }
+}
+
+type metricsOnlyDataSource struct {
+	OpenCostDataSource
+	metrics MetricsQuerier
+}
+
+func (m *metricsOnlyDataSource) Metrics() MetricsQuerier { return m.metrics }
+
+func TestPinMetrics(t *testing.T) {
+	plain := NewMockMetricsQuerier()
+	got, release := PinMetrics(plain)
+	if got != plain {
+		t.Errorf("expected a non-pinnable querier to be returned unchanged")
+	}
+	release()
+
+	p := &pinnableQuerier{MockMetricsQuerier: NewMockMetricsQuerier(), pinned: NewMockMetricsQuerier()}
+	got, release = PinMetrics(p)
+	if got != p.pinned {
+		t.Errorf("expected the pinned querier")
+	}
+	release()
+	if p.released != 1 {
+		t.Errorf("expected release to be forwarded, got %d", p.released)
+	}
+}
+
+func TestPinDataSource(t *testing.T) {
+	plain := &metricsOnlyDataSource{metrics: NewMockMetricsQuerier()}
+	got, release := PinDataSource(plain)
+	if got != plain {
+		t.Errorf("expected a data source with a non-pinnable querier to be returned unchanged")
+	}
+	release()
+
+	p := &pinnableQuerier{MockMetricsQuerier: NewMockMetricsQuerier(), pinned: NewMockMetricsQuerier()}
+	ds := &metricsOnlyDataSource{metrics: p}
+	got, release = PinDataSource(ds)
+	if got.Metrics() != p.pinned || got.Metrics() != p.pinned {
+		t.Errorf("expected every Metrics() call to return the pinned querier")
+	}
+	release()
+	if p.released != 1 {
+		t.Errorf("expected release to be forwarded, got %d", p.released)
+	}
+}