feat: support adaptive sampling rate (#221)

jy-tan · web-flow · commit c27d6869f017 · 2026-04-09T22:44:29.000-07:00
diff --git a/docs/drift/configuration.md b/docs/drift/configuration.md
@@ -317,11 +317,29 @@ For authentication in cloud mode, either use:
     </tr>
   </thead>
   <tbody>
+    <tr>
+      <td><code>recording.sampling.mode</code></td>
+      <td>string</td>
+      <td><code>fixed</code></td>
+      <td>Sampling strategy for Drift SDK recording. Supported values: <code>fixed</code> and <code>adaptive</code>.</td>
+    </tr>
+    <tr>
+      <td><code>recording.sampling.base_rate</code></td>
+      <td>number</td>
+      <td>0.1</td>
+      <td>Base sampling fraction when recording traces. In <code>fixed</code> mode this is the effective rate. In <code>adaptive</code> mode the SDK may temporarily reduce below this base rate under pressure.</td>
+    </tr>
+    <tr>
+      <td><code>recording.sampling.min_rate</code></td>
+      <td>number</td>
+      <td><code>0.001</code> in <code>adaptive</code> mode; unset in <code>fixed</code> mode</td>
+      <td>Lower bound for adaptive sampling after load shedding is applied. This is only defaulted in <code>adaptive</code> mode and remains unset in <code>fixed</code> mode.</td>
+    </tr>
     <tr>
       <td><code>recording.sampling_rate</code></td>
       <td>number</td>
       <td>0.1</td>
-      <td>Target sampling fraction when recording traces.</td>
+      <td><i>Deprecated.</i> Legacy alias for <code>recording.sampling.base_rate</code>. Still accepted for backwards compatibility.</td>
     </tr>
     <tr>
       <td><code>recording.export_spans</code></td>
diff --git a/internal/config/config.go b/internal/config/config.go
@@ -89,10 +89,17 @@ type ComparisonConfig struct {
 	IgnoreEpochTimestamps *bool    `koanf:"ignore_epoch_timestamps"`
 }
 
+type RecordingSamplingConfig struct {
+	Mode     string   `koanf:"mode"`
+	BaseRate *float64 `koanf:"base_rate"`
+	MinRate  *float64 `koanf:"min_rate"`
+}
+
 type RecordingConfig struct {
-	SamplingRate          float64 `koanf:"sampling_rate"`
-	ExportSpans           *bool   `koanf:"export_spans"`
-	EnableEnvVarRecording *bool   `koanf:"enable_env_var_recording"`
+	SamplingRate          float64                 `koanf:"sampling_rate"`
+	Sampling              RecordingSamplingConfig `koanf:"sampling"`
+	ExportSpans           *bool                   `koanf:"export_spans"`
+	EnableEnvVarRecording *bool                   `koanf:"enable_env_var_recording"`
 }
 
 type ReplayConfig struct {
@@ -165,6 +172,12 @@ func Load(configFile string) error {
 			if err := k.Set(configKey, val); err != nil {
 				return fmt.Errorf("error setting %s from env: %w", envKey, err)
 			}
+
+			if envKey == "TUSK_RECORDING_SAMPLING_RATE" {
+				if err := k.Set("recording.sampling.base_rate", val); err != nil {
+					return fmt.Errorf("error setting %s nested base rate from env: %w", envKey, err)
+				}
+			}
 		}
 	}
 
@@ -215,9 +228,23 @@ func parseAndValidate() (*Config, error) {
 	if cfg.TestExecution.Timeout == "" {
 		cfg.TestExecution.Timeout = "30s"
 	}
-	if cfg.Recording.SamplingRate == 0 {
+	if cfg.Recording.Sampling.BaseRate != nil {
+		cfg.Recording.SamplingRate = *cfg.Recording.Sampling.BaseRate
+	}
+	if cfg.Recording.SamplingRate == 0 && cfg.Recording.Sampling.BaseRate == nil {
 		cfg.Recording.SamplingRate = 0.1
 	}
+	if cfg.Recording.Sampling.Mode == "" {
+		cfg.Recording.Sampling.Mode = "fixed"
+	}
+	if cfg.Recording.Sampling.BaseRate == nil {
+		baseRate := cfg.Recording.SamplingRate
+		cfg.Recording.Sampling.BaseRate = &baseRate
+	}
+	if cfg.Recording.Sampling.MinRate == nil && cfg.Recording.Sampling.Mode == "adaptive" {
+		minRate := 0.001
+		cfg.Recording.Sampling.MinRate = &minRate
+	}
 	if cfg.Recording.ExportSpans == nil {
 		defaultExportSpans := false
 		cfg.Recording.ExportSpans = &defaultExportSpans
@@ -298,6 +325,29 @@ func (cfg *Config) Validate() error {
 		errs = append(errs, fmt.Errorf("replay.sandbox.mode must be 'auto', 'strict', or 'off', got %s", cfg.Replay.Sandbox.Mode))
 	}
 
+	validSamplingModes := map[string]bool{"fixed": true, "adaptive": true}
+	if cfg.Recording.Sampling.Mode != "" && !validSamplingModes[cfg.Recording.Sampling.Mode] {
+		errs = append(errs, fmt.Errorf("recording.sampling.mode must be 'fixed' or 'adaptive', got %s", cfg.Recording.Sampling.Mode))
+	}
+
+	if cfg.Recording.Sampling.BaseRate != nil {
+		if *cfg.Recording.Sampling.BaseRate < 0 || *cfg.Recording.Sampling.BaseRate > 1 {
+			errs = append(errs, fmt.Errorf("recording.sampling.base_rate must be between 0.0 and 1.0, got %v", *cfg.Recording.Sampling.BaseRate))
+		}
+	}
+
+	if cfg.Recording.Sampling.MinRate != nil {
+		if *cfg.Recording.Sampling.MinRate < 0 || *cfg.Recording.Sampling.MinRate > 1 {
+			errs = append(errs, fmt.Errorf("recording.sampling.min_rate must be between 0.0 and 1.0, got %v", *cfg.Recording.Sampling.MinRate))
+		}
+	}
+
+	if cfg.Recording.Sampling.BaseRate != nil && cfg.Recording.Sampling.MinRate != nil {
+		if *cfg.Recording.Sampling.MinRate > *cfg.Recording.Sampling.BaseRate {
+			errs = append(errs, fmt.Errorf("recording.sampling.min_rate must be less than or equal to recording.sampling.base_rate (got min_rate=%v, base_rate=%v)", *cfg.Recording.Sampling.MinRate, *cfg.Recording.Sampling.BaseRate))
+		}
+	}
+
 	if len(errs) > 0 {
 		return errors.Join(errs...)
 	}
diff --git a/internal/config/config_test.go b/internal/config/config_test.go
@@ -20,6 +20,150 @@ func evalSymlinks(path string) string {
 	return resolved
 }
 
+func TestNestedRecordingSamplingConfig(t *testing.T) {
+	defer Invalidate()
+
+	tmpDir := t.TempDir()
+	configPath := filepath.Join(tmpDir, "config.yaml")
+	require.NoError(t, os.WriteFile(configPath, []byte(`
+recording:
+  sampling:
+    mode: adaptive
+    base_rate: 0.25
+    min_rate: 0.05
+`), 0o600))
+
+	require.NoError(t, Load(configPath))
+
+	cfg, err := Get()
+	require.NoError(t, err)
+	assert.Equal(t, "adaptive", cfg.Recording.Sampling.Mode)
+	require.NotNil(t, cfg.Recording.Sampling.BaseRate)
+	assert.Equal(t, 0.25, *cfg.Recording.Sampling.BaseRate)
+	require.NotNil(t, cfg.Recording.Sampling.MinRate)
+	assert.Equal(t, 0.05, *cfg.Recording.Sampling.MinRate)
+	assert.Equal(t, 0.25, cfg.Recording.SamplingRate)
+}
+
+func TestLegacyRecordingSamplingRateBackfillsNestedSamplingConfig(t *testing.T) {
+	defer Invalidate()
+
+	tmpDir := t.TempDir()
+	configPath := filepath.Join(tmpDir, "config.yaml")
+	require.NoError(t, os.WriteFile(configPath, []byte(`
+recording:
+  sampling_rate: 0.25
+`), 0o600))
+
+	require.NoError(t, Load(configPath))
+
+	cfg, err := Get()
+	require.NoError(t, err)
+	assert.Equal(t, 0.25, cfg.Recording.SamplingRate)
+	assert.Equal(t, "fixed", cfg.Recording.Sampling.Mode)
+	require.NotNil(t, cfg.Recording.Sampling.BaseRate)
+	assert.Equal(t, 0.25, *cfg.Recording.Sampling.BaseRate)
+	assert.Nil(t, cfg.Recording.Sampling.MinRate)
+}
+
+func TestRecordingSamplingRateEnvOverrideBeatsNestedBaseRate(t *testing.T) {
+	defer Invalidate()
+	t.Setenv("TUSK_RECORDING_SAMPLING_RATE", "0.5")
+
+	tmpDir := t.TempDir()
+	configPath := filepath.Join(tmpDir, "config.yaml")
+	require.NoError(t, os.WriteFile(configPath, []byte(`
+recording:
+  sampling:
+    mode: adaptive
+    base_rate: 0.25
+    min_rate: 0.05
+`), 0o600))
+
+	require.NoError(t, Load(configPath))
+
+	cfg, err := Get()
+	require.NoError(t, err)
+	require.NotNil(t, cfg.Recording.Sampling.BaseRate)
+	assert.Equal(t, 0.5, *cfg.Recording.Sampling.BaseRate)
+	assert.Equal(t, 0.5, cfg.Recording.SamplingRate)
+	assert.Equal(t, "adaptive", cfg.Recording.Sampling.Mode)
+	require.NotNil(t, cfg.Recording.Sampling.MinRate)
+	assert.Equal(t, 0.05, *cfg.Recording.Sampling.MinRate)
+}
+
+func TestValidateRejectsInvalidRecordingSamplingMode(t *testing.T) {
+	cfg := &Config{
+		Service: ServiceConfig{
+			Port: 3000,
+			Communication: CommunicationConfig{
+				Type:    "auto",
+				TCPPort: 9001,
+			},
+		},
+		Recording: RecordingConfig{
+			Sampling: RecordingSamplingConfig{
+				Mode: "adapttive",
+			},
+		},
+	}
+
+	err := cfg.Validate()
+	require.Error(t, err)
+	assert.ErrorContains(t, err, "recording.sampling.mode must be 'fixed' or 'adaptive'")
+}
+
+func TestValidateRejectsOutOfRangeRecordingSamplingRates(t *testing.T) {
+	baseRate := 5.0
+	minRate := -0.1
+	cfg := &Config{
+		Service: ServiceConfig{
+			Port: 3000,
+			Communication: CommunicationConfig{
+				Type:    "auto",
+				TCPPort: 9001,
+			},
+		},
+		Recording: RecordingConfig{
+			Sampling: RecordingSamplingConfig{
+				Mode:     "adaptive",
+				BaseRate: &baseRate,
+				MinRate:  &minRate,
+			},
+		},
+	}
+
+	err := cfg.Validate()
+	require.Error(t, err)
+	assert.ErrorContains(t, err, "recording.sampling.base_rate must be between 0.0 and 1.0")
+	assert.ErrorContains(t, err, "recording.sampling.min_rate must be between 0.0 and 1.0")
+}
+
+func TestValidateRejectsRecordingMinRateGreaterThanBaseRate(t *testing.T) {
+	baseRate := 0.1
+	minRate := 0.9
+	cfg := &Config{
+		Service: ServiceConfig{
+			Port: 3000,
+			Communication: CommunicationConfig{
+				Type:    "auto",
+				TCPPort: 9001,
+			},
+		},
+		Recording: RecordingConfig{
+			Sampling: RecordingSamplingConfig{
+				Mode:     "adaptive",
+				BaseRate: &baseRate,
+				MinRate:  &minRate,
+			},
+		},
+	}
+
+	err := cfg.Validate()
+	require.Error(t, err)
+	assert.ErrorContains(t, err, "recording.sampling.min_rate must be less than or equal to recording.sampling.base_rate")
+}
+
 func TestFindConfigFile_ParentTraversal(t *testing.T) {
 	wd, _ := os.Getwd()
 	defer func() { _ = os.Chdir(wd) }()
