sync recent changes with upstream (#4)

regulating cache delay option which specifies the amount of time after which the processor starts regulating cache sizes. This can be used to avoid regulating cache sizes on an initial empty cache, as it can cause unnecessary cache size adjustments and high memory usage. More in docs.

Emitting span for not sampled trace. You can now specify on specific policies to emit a single span in place of the trace. This is useful if someone goes searching for a dropped trace in the system, it will tell them why it's dropped because a span appears with the policy.

Author: jamesmoessis

internal/ptraceutil/go.mod

@@ -1,6 +1,6 @@
 module github.com/atlassian-labs/atlassian-sampling-processor/internal/ptraceutil
 
-go 1.23.4
+go 1.23.5
 
 require (
 	github.com/open-telemetry/opentelemetry-collector-contrib/pkg/pdatatest v0.115.0

pkg/processor/atlassiansamplingprocessor/DESIGN.md

@@ -41,7 +41,7 @@ code significantly and potentially lead to bugs, as experienced in the upstream
 
 This is, of course, a trade-off. The processing throughput is limited by the capacity of a single goroutine, 
 creating a potential bottleneck. This can be alleviated by deploying more instances of the processor with reduced 
-memory allocation per instance (e.g., more pods, each with less memory). If the bottleneck becomes a significant issue, 
+memory allocation per instance (e.g., more nodes, each with less memory). If the bottleneck becomes a significant issue, 
 a future enhancement could involve sharding the processor. This would involve splitting the processing workload by trace 
 ID and maintaining separate caches and states for each shard.
 
@@ -107,7 +107,8 @@ they can only access cached metadata from the cache. This restriction is in plac
 1. **Cache Compression:** It allows for the compression of spans within the cache. We are guaranteed that a compressed 
 blob is decompressed at most once, which occurs in the case it is sampled and transmitted. 
 Restricting the reading of cached spans allows for this optimisation.
-1.**Performance Efficiency:** It ensures that policy evaluation remains efficient, adhering to O(n) complexity, 
+
+2. **Performance Efficiency:** It ensures that policy evaluation remains efficient, adhering to O(n) complexity, 
 where n represents the number of spans in the current arriving batch. This prevents the policies from becoming slow.
 
 ## Caches

pkg/processor/atlassiansamplingprocessor/README.md

@@ -44,7 +44,7 @@ The `primary_cache_size` value should be greater than 0, and should be set to a
 `processor_atlassian_sampling_trace_eviction_time` metric to tune how long you would like your traces to stay pending
 in memory before being considered not-sampled.
 
-The primary cache size is initially set to 80% of the `primary_cache_size` value.
+The primary cache size is initially set to 60% of the `primary_cache_size` value.
 It is automatically adjusted depending on heap memory usage at runtime, but will not exceed the `primary_cache_size` value.
 
 ### `secondary_cache_size`
@@ -103,6 +103,9 @@ used as the final decision.
 
 Policies include a `name`, `type`, and then further configuration depending on what the `type` was.
 
+`emit_single_span_for_not_sampled` is an optional field that can be set to `true` for a policy. If set, the processor will emit a single span for a trace that is not sampled, instead of dropping the trace entirely. 
+This span will have the same trace ID as the original trace. The span will have the name `TRACE NOT SAMPLED`, with policy name in its attribute.
+
 Current supported policy types are: 
 
 - `span_count` - samples the trace if it meets a minimum amount of spans. 

pkg/processor/atlassiansamplingprocessor/config.go

@@ -2,6 +2,7 @@ package atlassiansamplingprocessor // import "github.com/atlassian-labs/atlassia
 
 import (
 	"errors"
+	"time"
 
 	"go.opentelemetry.io/collector/component"
 	"go.uber.org/multierr"
@@ -17,6 +18,11 @@ type Config struct {
 	// A good starting point to set this is about 75% of overall memory resource allocation.
 	TargetHeapBytes uint64 `mapstructure:"target_heap_bytes"`
 
+	// RegulateCacheDelay is the amount of time after which the processor starts regulating cache sizes based on the set TargetHeapBytes (if specified).
+	// It is optional and defaults to 0s.
+	// This can be used to avoid regulating cache sizes on an initial empty cache, as it can cause unnecessary cache size adjustments and high memory usage.
+	RegulateCacheDelay time.Duration `mapstructure:"regulate_cache_delay"`
+
 	// PrimaryCacheSize sets the initial and maximum size of the primary cache that holds non-low priority traces.
 	PrimaryCacheSize int `mapstructure:"primary_cache_size"`
 

pkg/processor/atlassiansamplingprocessor/config_policy.go

@@ -31,6 +31,9 @@ type SharedPolicyConfig struct {
 	Name string `mapstructure:"name"`
 	// Type of the policy this will be used to match the proper configuration of the policy.
 	Type PolicyType `mapstructure:"type"`
+	// EmitSingleSpanForNotSampled enabled to emit a single span with same trace ID as the one that was dropped.
+	// So when people search for it, it comes up with “TRACE NOT SAMPLED” and then the sampling policy name attached as an attribute.
+	EmitSingleSpanForNotSampled bool `mapstructure:"emit_single_span_for_not_sampled"`
 	// Configs for probabilistic sampling policy evaluator.
 	ProbabilisticConfig `mapstructure:"probabilistic"`
 	// Configs for span count filter sampling policy evaluator.

pkg/processor/atlassiansamplingprocessor/config_policy_test.go

@@ -47,6 +47,7 @@ func TestPolicyCreationFromConfig(t *testing.T) {
 	assert.Equal(t, RemoteProbabilistic, policies[8].policyType)
 	assert.Equal(t, "test-policy-10", policies[9].name)
 	assert.Equal(t, Downgrader, policies[9].policyType)
+	assert.True(t, policies[9].emitSingleSpanForNotSampled)
 
 	assert.Equal(t, 10, len(policies), "wrong number of assertions")
 }

pkg/processor/atlassiansamplingprocessor/config_test.go

@@ -9,6 +9,7 @@ package atlassiansamplingprocessor
 import (
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -33,12 +34,16 @@ func TestLoadConfig(t *testing.T) {
 	require.NoError(t, err)
 	require.NoError(t, sub.Unmarshal(cfg))
 
+	delay, err := time.ParseDuration("5m")
+	require.NoError(t, err)
+
 	assert.Equal(t,
 		cfg,
 		&Config{
 			PrimaryCacheSize:   1000,
 			SecondaryCacheSize: 100,
 			TargetHeapBytes:    100_000_000,
+			RegulateCacheDelay: delay,
 			DecisionCacheCfg:   DecisionCacheCfg{SampledCacheSize: 1000, NonSampledCacheSize: 10000},
 			CompressionEnabled: true,
 			PolicyConfig: []PolicyConfig{
@@ -151,8 +156,9 @@ func TestLoadConfig(t *testing.T) {
 				},
 				{
 					SharedPolicyConfig: SharedPolicyConfig{
-						Name: "test-policy-10",
-						Type: "downgrader",
+						Name:                        "test-policy-10",
+						Type:                        "downgrader",
+						EmitSingleSpanForNotSampled: true,
 					},
 					DowngraderConfig: DowngraderConfig{
 						DowngradeTo: "NotSampled",

pkg/processor/atlassiansamplingprocessor/go.mod

@@ -1,6 +1,6 @@
 module github.com/atlassian-labs/atlassian-sampling-processor/pkg/processor/atlassiansamplingprocessor
 
-go 1.23.4
+go 1.23.5
 
 require (
 	github.com/atlassian-labs/atlassian-sampling-processor/internal/ptraceutil v0.7.6

pkg/processor/atlassiansamplingprocessor/policy.go

@@ -16,6 +16,8 @@ type policy struct {
 	name string
 	// type is used to identify this policy instance's policyType
 	policyType PolicyType
+	// emitSingleSpanForNotSampled is used to determine if a single span should be emitted for a trace that is not sampled.
+	emitSingleSpanForNotSampled bool
 	// evaluator that decides if a trace is sampled or not by this policy instance.
 	evaluator evaluators.PolicyEvaluator
 	// attribute to use in the telemetry to denote the policy.
@@ -38,10 +40,11 @@ func newPolicies(cfg []PolicyConfig, set component.TelemetrySettings) ([]*policy
 			return nil, err
 		}
 		p := &policy{
-			name:       policyCfg.Name,
-			policyType: policyCfg.Type,
-			evaluator:  eval,
-			attribute:  metric.WithAttributes(attribute.String("policy", policyCfg.Name)),
+			name:                        policyCfg.Name,
+			policyType:                  policyCfg.Type,
+			emitSingleSpanForNotSampled: policyCfg.EmitSingleSpanForNotSampled,
+			evaluator:                   eval,
+			attribute:                   metric.WithAttributes(attribute.String("policy", policyCfg.Name)),
 		}
 		pols[i] = p
 	}

pkg/processor/atlassiansamplingprocessor/processor.go

@@ -63,6 +63,8 @@ type atlassianSamplingProcessor struct {
 	memRegulator memory.RegulatorI
 	// memTicker controls how often the memRegulator is called
 	memTicker *time.Ticker
+	// regulatorStartTime is the time at which the regulator starts being used to regulate cache sizes
+	regulatorStartTime time.Time
 	// shutdownStart is a chan used to signal to the async goroutine to start shutdown.
 	// The deadline time is passed to the chan, so the async goroutine knows when to time out.
 	shutdownStart   chan time.Time
@@ -96,8 +98,8 @@ func newAtlassianSamplingProcessor(cCfg component.Config, set component.Telemetr
 		compress:        cfg.CompressionEnabled,
 	}
 
-	// Start with 80% of max cache size, the memory regulator will adjust the cache size as needed
-	initialPrimaryCacheSize := int(0.8 * float64(cfg.PrimaryCacheSize))
+	// Start with 60% of max cache size, the memory regulator will adjust the cache size as needed
+	initialPrimaryCacheSize := int(0.6 * float64(cfg.PrimaryCacheSize))
 
 	primaryCache, err := cache.NewLRUCache[*tracedata.TraceData](
 		initialPrimaryCacheSize,
@@ -143,7 +145,9 @@ func newAtlassianSamplingProcessor(cCfg component.Config, set component.Telemetr
 		if rErr != nil {
 			return nil, rErr
 		}
+
 		asp.memRegulator = memRegulator
+		asp.regulatorStartTime = time.Now().Add(cfg.RegulateCacheDelay)
 	}
 
 	pols, err := newPolicies(cfg.PolicyConfig, set)
@@ -238,9 +242,9 @@ func (asp *atlassianSamplingProcessor) consumeChan() {
 			for i := 0; i < resourceSpans.Len(); i++ {
 				asp.processTraces(ctx, resourceSpans.At(i))
 			}
-		case <-asp.memTicker.C:
+		case t := <-asp.memTicker.C:
 			// If ticker signals, call the memory regulator
-			if asp.memRegulator != nil {
+			if t.After(asp.regulatorStartTime) && asp.memRegulator != nil {
 				size := asp.memRegulator.RegulateCacheSize()
 				asp.telemetry.ProcessorAtlassianSamplingPrimaryCacheSize.Record(ctx, int64(size))
 			}
@@ -316,7 +320,7 @@ func (asp *atlassianSamplingProcessor) processTraces(ctx context.Context, resour
 			asp.sendSampledTraceData(ctx, currentTrace)
 			asp.telemetry.ProcessorAtlassianSamplingTracesSampled.Add(ctx, 1)
 		case evaluators.NotSampled:
-			asp.releaseNotSampledTrace(ctx, id)
+			asp.releaseNotSampledTrace(ctx, id, pol)
 		default:
 			if finalDecision == evaluators.LowPriority {
 				td.Metadata.Priority = priority.Low
@@ -382,7 +386,7 @@ func (asp *atlassianSamplingProcessor) cachedDecision(
 		}
 		asp.sampledDecisionCache.Put(id, time.Now())
 	} else {
-		asp.releaseNotSampledTrace(ctx, id)
+		asp.releaseNotSampledTrace(ctx, id, nil)
 	}
 	return true
 }
@@ -406,10 +410,25 @@ func (asp *atlassianSamplingProcessor) sendSampledTraceData(ctx context.Context,
 
 // releaseNotSampledTrace removes references to traces that have been decided to be not sampled.
 // It also caches the non sampled decision, and increments count of traces not sampled.
-func (asp *atlassianSamplingProcessor) releaseNotSampledTrace(ctx context.Context, id pcommon.TraceID) {
+func (asp *atlassianSamplingProcessor) releaseNotSampledTrace(ctx context.Context, id pcommon.TraceID, policy *policy) {
+	// Check if EmitSingleSpanForNotSampled is true in the policy config
+	if policy != nil && policy.emitSingleSpanForNotSampled {
+		// Create a placeholder trace with a single span containing the decision policy name
+		notSampledTrace := ptrace.NewTraces()
+		rs := notSampledTrace.ResourceSpans().AppendEmpty()
+		ss := rs.ScopeSpans().AppendEmpty()
+		span := ss.Spans().AppendEmpty()
+		span.SetTraceID(id)
+		span.SetName("TRACE NOT SAMPLED")
+		span.SetStartTimestamp(pcommon.NewTimestampFromTime(time.Now().Add(-time.Second)))
+		span.SetEndTimestamp(pcommon.NewTimestampFromTime(time.Now()))
+		span.Attributes().PutStr("sampling.policy", policy.name)
+
+		// Send the placeholder trace
+		asp.sendSampledTraceData(ctx, notSampledTrace)
+	}
 	asp.nonSampledDecisionCache.Put(id, time.Now())
 	asp.traceData.Delete(id)
-	asp.telemetry.ProcessorAtlassianSamplingTracesNotSampled.Add(ctx, 1)
 }
 
 func (asp *atlassianSamplingProcessor) flushAll(ctx context.Context) error {