feat(sketches): per-metric sample_p knob on the CMS + HLL warm processors

asap-precompute-go/sketches/cms.go

@@ -35,15 +35,73 @@ type CMSWrapper struct {
 	rows       int
 	cols       int
 	useMsgpack bool
+	// sampleP is the per-sketch sampling probability in (0,1]. 1.0 (the
+	// default) disables sampling so the sketch is byte-identical to an
+	// unsampled one. Set via WithSampleP; preserved across the
+	// re-construction paths (Reset / Merge / ApplyDelta) so a sampled
+	// wrapper stays sampled for its whole lifetime.
+	sampleP float64
 }
 
+// cmsSampleSeed is the fixed seed handed to sketchlib-go's geometric
+// sampler. A constant seed keeps the admitted-subset reproducible across
+// runs (sketchlib-go's NewGeometricSampler doc) and across the wrapper's
+// internal re-constructions; the value is arbitrary and only matters for
+// determinism, never for correctness (any seed yields an unbiased sample).
+const cmsSampleSeed int64 = 0x5A4D_5043 // "ZMPC"
+
 // NewCMSWrapper builds an empty CMS with the configured rows / cols.
 // useMsgpack=true selects the legacy msgpack emit path (no delta
 // transmission supported in that mode); useMsgpack=false uses the
 // proto SerializeProtoBytesFO format and supports delta transmission.
+//
+// Sampling is disabled (sampleP=1.0) by default — call WithSampleP to
+// enable it. The default keeps the emitted wire bytes byte-identical to
+// the pre-sampling format.
 func NewCMSWrapper(rows, cols int, useMsgpack bool) *CMSWrapper {
-	sk, _ := cms.NewCountMinSketch(rows, cols)
-	return &CMSWrapper{sk: sk, rows: rows, cols: cols, useMsgpack: useMsgpack}
+	w := &CMSWrapper{rows: rows, cols: cols, useMsgpack: useMsgpack, sampleP: 1.0}
+	w.sk = w.newSketch()
+	return w
+}
+
+// WithSampleP enables per-sketch geometric admission sampling at
+// probability p in (0,1]. p>=1 (or NaN) disables sampling (exact, the
+// default); p<=0 is clamped to a tiny positive probability by sketchlib-go
+// rather than dropping the whole stream. Returns the receiver for fluent
+// construction. The probability is stamped on the SketchEnvelope by
+// sketchlib-go so the backend rescales frequency estimates by 1/p at
+// query time.
+func (w *CMSWrapper) WithSampleP(p float64) *CMSWrapper {
+	if p >= 1.0 || p != p { // p != p ⇒ NaN
+		w.sampleP = 1.0
+	} else {
+		w.sampleP = p
+	}
+	// Re-apply to the live sketch so a WithSampleP after construction
+	// takes effect immediately.
+	if w.sk != nil {
+		w.sk.WithSampleP(w.sampleP, cmsSampleSeed)
+	}
+	return w
+}
+
+// SampleP returns the configured sampling probability (1.0 when disabled).
+func (w *CMSWrapper) SampleP() float64 {
+	if w.sampleP <= 0 {
+		return 1.0
+	}
+	return w.sampleP
+}
+
+// newSketch builds a fresh sketchlib-go CMS carrying the wrapper's
+// configured sampling probability. Centralises the construction so every
+// re-creation path (New / Reset / Merge / ApplyDelta) keeps sampleP.
+func (w *CMSWrapper) newSketch() *cms.CountMinSketch {
+	sk, _ := cms.NewCountMinSketch(w.rows, w.cols)
+	if sk != nil && w.sampleP > 0 && w.sampleP < 1.0 {
+		sk.WithSampleP(w.sampleP, cmsSampleSeed)
+	}
+	return sk
 }
 
 // InsertHash mirrors the legacy CMS processor's
@@ -131,7 +189,7 @@ func (w *CMSWrapper) ApplyDelta(payload []byte) error {
 		return nil
 	}
 	if w.sk == nil {
-		w.sk, _ = cms.NewCountMinSketch(w.rows, w.cols)
+		w.sk = w.newSketch()
 	}
 	if other, err := cms.DeserializeCountMinSketchFromProtoBytes(payload); err == nil && other != nil {
 		return w.sk.Merge(other)
@@ -158,16 +216,17 @@ func (w *CMSWrapper) Merge(other precompute.Sketch) error {
 		return nil
 	}
 	if w.sk == nil {
-		w.sk, _ = cms.NewCountMinSketch(w.rows, w.cols)
+		w.sk = w.newSketch()
 	}
 	return w.sk.Merge(o.sk)
 }
 
 // Reset zeros the sketch in place by replacing it with a fresh
-// CountMinSketch of the same dimensions. Window rotation calls this
-// when the runtime decides to recycle entries.
+// CountMinSketch of the same dimensions (and the same sampling
+// probability). Window rotation calls this when the runtime decides to
+// recycle entries.
 func (w *CMSWrapper) Reset() {
-	w.sk, _ = cms.NewCountMinSketch(w.rows, w.cols)
+	w.sk = w.newSketch()
 }
 
 // EstimateCount returns the estimated frequency for a hashed key.

asap-precompute-go/sketches/hll.go

@@ -29,14 +29,63 @@ import (
 // legacy emit-side mode that only applies when DeltaTransmission=false.
 type HLLWrapper struct {
 	sk *hll.HyperLogLog
+	// sampleP is the per-sketch hash-threshold sampling probability in
+	// (0,1]. 1.0 (the default) disables sampling so the sketch is
+	// byte-identical to an unsampled one. Set via WithSampleP; preserved
+	// across the re-construction paths (Reset / Merge / ApplyDelta) so a
+	// sampled wrapper stays sampled for its whole lifetime.
+	sampleP float64
 }
 
 // NewHLLWrapper builds an empty HLL sketch. The sketchlib-go
 // constructor is parameterless (precision is hard-coded to
 // hll.HLLPrecision = 14); the adapter's encoding choice is honored
 // at the encode layer, not here.
+//
+// Sampling is disabled (sampleP=1.0) by default — call WithSampleP to
+// enable it. The default keeps the emitted wire bytes byte-identical to
+// the pre-sampling format.
 func NewHLLWrapper() *HLLWrapper {
-	return &HLLWrapper{sk: hll.NewHyperLogLog()}
+	w := &HLLWrapper{sampleP: 1.0}
+	w.sk = w.newSketch()
+	return w
+}
+
+// WithSampleP enables per-sketch hash-threshold element sampling at
+// probability p in (0,1]. p>=1 (or NaN) disables sampling (exact, the
+// default); p<=0 keeps nothing, which sketchlib-go clamps to disabled.
+// Returns the receiver for fluent construction. The probability is
+// stamped on the SketchEnvelope by sketchlib-go so the backend rescales
+// cardinality by 1/p at query time.
+func (w *HLLWrapper) WithSampleP(p float64) *HLLWrapper {
+	if p >= 1.0 || p != p { // p != p ⇒ NaN
+		w.sampleP = 1.0
+	} else {
+		w.sampleP = p
+	}
+	if w.sk != nil {
+		w.sk.WithSampleP(w.sampleP)
+	}
+	return w
+}
+
+// SampleP returns the configured sampling probability (1.0 when disabled).
+func (w *HLLWrapper) SampleP() float64 {
+	if w.sampleP <= 0 {
+		return 1.0
+	}
+	return w.sampleP
+}
+
+// newSketch builds a fresh sketchlib-go HLL carrying the wrapper's
+// configured sampling probability. Centralises the construction so every
+// re-creation path (New / Reset / Merge / ApplyDelta) keeps sampleP.
+func (w *HLLWrapper) newSketch() *hll.HyperLogLog {
+	sk := hll.NewHyperLogLog()
+	if sk != nil && w.sampleP > 0 && w.sampleP < 1.0 {
+		sk.WithSampleP(w.sampleP)
+	}
+	return sk
 }
 
 // UpdateValue feeds a single observation into the underlying HLL
@@ -97,7 +146,7 @@ func (w *HLLWrapper) ApplyDelta(payload []byte) error {
 		return nil
 	}
 	if w.sk == nil {
-		w.sk = hll.NewHyperLogLog()
+		w.sk = w.newSketch()
 	}
 	// Try delta first: RegisterDelta is the more constrained shape;
 	// proto-encoded HyperLogLogState envelopes won't decode as a
@@ -128,16 +177,16 @@ func (w *HLLWrapper) Merge(other precompute.Sketch) error {
 		return nil
 	}
 	if w.sk == nil {
-		w.sk = hll.NewHyperLogLog()
+		w.sk = w.newSketch()
 	}
 	return w.sk.Merge(o.sk)
 }
 
-// Reset zeros the sketch in place by replacing it with a fresh HLL.
-// Window rotation calls this when the runtime decides to recycle
-// entries.
+// Reset zeros the sketch in place by replacing it with a fresh HLL
+// (carrying the same sampling probability). Window rotation calls this
+// when the runtime decides to recycle entries.
 func (w *HLLWrapper) Reset() {
-	w.sk = hll.NewHyperLogLog()
+	w.sk = w.newSketch()
 }
 
 // EstimateCardinality satisfies precompute.CardinalitySketch — adapter

asap-precompute-go/sketches/sampling_test.go

@@ -0,0 +1,120 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package sketches
+
+import (
+	"testing"
+
+	"github.com/ProjectASAP/sketchlib-go/common"
+)
+
+// TestCMSWrapper_SamplePReachesSketchBuilder proves a configured
+// sample_p<1 is threaded all the way into the underlying sketchlib-go
+// CountMinSketch builder (read back off w.sk, the real sketch the wrapper
+// owns — not just the wrapper's own field), and that the default 1.0 is
+// unchanged (exact, sampling disabled). This is the per-metric sampling
+// knob's last hop: control plane → agent processor config → NewCMSWrapper
+// → WithSampleP → sketchlib-go cms.WithSampleP.
+func TestCMSWrapper_SamplePReachesSketchBuilder(t *testing.T) {
+	t.Parallel()
+
+	// Default: no WithSampleP call ⇒ exact (1.0) on both the wrapper and
+	// the underlying sketchlib-go sketch. Byte-identical to pre-sampling.
+	def := NewCMSWrapper(4, 256, false)
+	if got := def.SampleP(); got != 1.0 {
+		t.Fatalf("default wrapper SampleP = %v, want 1.0", got)
+	}
+	if got := def.sk.SampleP(); got != 1.0 {
+		t.Fatalf("default underlying sketch SampleP = %v, want 1.0", got)
+	}
+
+	// WithSampleP(1.0) is also an exact no-op (the safe rollout state the
+	// control plane emits when sample_p is unset / normalised to 1.0).
+	exact := NewCMSWrapper(4, 256, false).WithSampleP(1.0)
+	if got := exact.sk.SampleP(); got != 1.0 {
+		t.Fatalf("WithSampleP(1.0) underlying SampleP = %v, want 1.0", got)
+	}
+
+	// Configured p<1 reaches the builder.
+	const p = 0.1
+	w := NewCMSWrapper(4, 256, false).WithSampleP(p)
+	if got := w.SampleP(); got != p {
+		t.Fatalf("wrapper SampleP = %v, want %v", got, p)
+	}
+	if got := w.sk.SampleP(); got != p {
+		t.Fatalf("underlying sketch SampleP = %v, want %v — p did not reach the sketchlib-go builder", got, p)
+	}
+
+	// Sampling probability survives the re-construction paths so a sampled
+	// wrapper stays sampled for its whole lifetime.
+	w.Reset()
+	if got := w.sk.SampleP(); got != p {
+		t.Fatalf("after Reset underlying SampleP = %v, want %v", got, p)
+	}
+}
+
+// TestHLLWrapper_SamplePReachesSketchBuilder is the HLL twin of the CMS
+// test above: it pins that a configured sample_p<1 reaches the underlying
+// sketchlib-go HyperLogLog builder and that the default 1.0 is unchanged.
+func TestHLLWrapper_SamplePReachesSketchBuilder(t *testing.T) {
+	t.Parallel()
+
+	def := NewHLLWrapper()
+	if got := def.SampleP(); got != 1.0 {
+		t.Fatalf("default wrapper SampleP = %v, want 1.0", got)
+	}
+	if got := def.sk.SampleP(); got != 1.0 {
+		t.Fatalf("default underlying sketch SampleP = %v, want 1.0", got)
+	}
+
+	exact := NewHLLWrapper().WithSampleP(1.0)
+	if got := exact.sk.SampleP(); got != 1.0 {
+		t.Fatalf("WithSampleP(1.0) underlying SampleP = %v, want 1.0", got)
+	}
+
+	const p = 0.1
+	w := NewHLLWrapper().WithSampleP(p)
+	if got := w.SampleP(); got != p {
+		t.Fatalf("wrapper SampleP = %v, want %v", got, p)
+	}
+	if got := w.sk.SampleP(); got != p {
+		t.Fatalf("underlying sketch SampleP = %v, want %v — p did not reach the sketchlib-go builder", got, p)
+	}
+
+	w.Reset()
+	if got := w.sk.SampleP(); got != p {
+		t.Fatalf("after Reset underlying SampleP = %v, want %v", got, p)
+	}
+}
+
+// TestCMSWrapper_SamplePThinsUpdates is a behavioural check: with p<1 the
+// geometric sampler admits only a fraction of inserts, so the estimated
+// frequency of a key inserted N times is materially below N (the RAW
+// sampled count — the backend applies the ×1/p rescale at query time). At
+// p=1.0 every insert is admitted, so the estimate matches N.
+func TestCMSWrapper_SamplePThinsUpdates(t *testing.T) {
+	t.Parallel()
+
+	key := []byte("sampled-key")
+	h := common.FromBytes(key).Hash
+	const n = 10000
+
+	exact := NewCMSWrapper(5, 4096, false) // p=1.0 default
+	for i := 0; i < n; i++ {
+		exact.InsertHash(h)
+	}
+	if got := exact.EstimateCount(key); got < n*0.95 {
+		t.Fatalf("p=1.0 estimate = %v, want ~%d", got, n)
+	}
+
+	sampled := NewCMSWrapper(5, 4096, false).WithSampleP(0.1)
+	for i := 0; i < n; i++ {
+		sampled.InsertHash(h)
+	}
+	// Raw sampled count ≈ 0.1*n = 1000; allow a wide margin for RNG, but
+	// it must be well below the exact count, proving sampling is live.
+	if got := sampled.EstimateCount(key); got >= n*0.5 {
+		t.Fatalf("p=0.1 raw estimate = %v, expected well below %d (sampling not applied?)", got, n)
+	}
+}

opentelemetry-collector-contrib-patch/processor/countminsketchprocessor/config.go

@@ -96,6 +96,15 @@ type Config struct {
 	// DeltaThreshold is the minimum absolute cell change required to include a
 	// cell in the delta payload. Defaults to 1.0 when DeltaTransmission=true.
 	DeltaThreshold float64 `mapstructure:"delta_threshold"`
+
+	// SampleP is the per-sketch geometric admission sampling probability in
+	// (0,1]. The control plane sets it per metric from the workload spec.
+	// 1.0 (the default — 0/unset is normalised to 1.0 in Validate) disables
+	// sampling so the emitted wire bytes are byte-identical to the
+	// pre-sampling format. A value <1 admits a ~p fraction of updates into
+	// the sketch; sketchlib-go stamps p on the SketchEnvelope so the backend
+	// rescales frequency estimates by 1/p at query time.
+	SampleP float64 `mapstructure:"sample_p"`
 }
 
 var _ component.Config = (*Config)(nil)
@@ -137,6 +146,17 @@ func (c *Config) Validate() error {
 		}
 	}
 
+	// SampleP: 0/unset normalises to 1.0 (sampling disabled — the safe
+	// default). Reject out-of-range values (negative or >1) rather than
+	// silently clamping, so a typo in the wire config surfaces at agent
+	// boot instead of producing a mis-scaled sketch.
+	if c.SampleP == 0 {
+		c.SampleP = 1.0
+	}
+	if c.SampleP < 0 || c.SampleP > 1.0 {
+		return fmt.Errorf("sample_p must be in (0, 1] (got %v)", c.SampleP)
+	}
+
 	// Default Encoding to proto when unset. Accept both supported
 	// values; anything else is a config error rather than a silent
 	// fallback.

opentelemetry-collector-contrib-patch/processor/countminsketchprocessor/processor.go

@@ -217,10 +217,13 @@ func (p *cmsProcessor) precomputeForLocked(name string) precompute.Precompute {
 		return pp
 	}
 	useMsgpack := p.cfg.Encoding == EncodingMsgpack && !p.cfg.DeltaTransmission
+	sampleP := p.cfg.SampleP
 	pp := precompute.New(
 		p.cfg.toPrecomputeConfig(name),
 		func() precompute.Sketch {
-			return sketches.NewCMSWrapper(p.cfg.Rows, p.cfg.Columns, useMsgpack)
+			// WithSampleP(1.0) is an exact no-op, so the default path is
+			// byte-identical to the pre-sampling build.
+			return sketches.NewCMSWrapper(p.cfg.Rows, p.cfg.Columns, useMsgpack).WithSampleP(sampleP)
 		},
 		sketches.CMSObserver{},
 	)