Merge #143220

143220: roachtest-operations: sink selection in changefeed operations r=nameisbhaskar a=vidit-bhat

This PR implements probabilistic sink selection based on environment variables. `SINK_CONFIG` defines sink percentages (e.g., `kafka:20,webhook:10,gs:30,null:40`), ensuring they sum to 100%. Each sink's URI is retrieved from `SINK_CONFIG_<SINK_NAME>`. Changefeeds use a weighted random selection to determine the sink.

Epic: none
Fixes: #138488
Release note: None

Co-authored-by: Vidit Bhat <[email protected]>

Author: craig[bot]

pkg/cmd/roachtest/operations/BUILD.bazel

@@ -55,7 +55,13 @@ go_library(
 
 go_test(
     name = "operations_test",
-    srcs = ["cluster_settings_test.go"],
+    srcs = [
+        "cluster_settings_test.go",
+        "utils_test.go",
+    ],
     embed = [":operations"],
-    deps = ["@com_github_stretchr_testify//require"],
+    deps = [
+        "//pkg/cmd/roachtest/operations/changefeeds",
+        "@com_github_stretchr_testify//require",
+    ],
 )

pkg/cmd/roachtest/operations/changefeeds/utils.go

@@ -9,6 +9,9 @@ import (
 	"context"
 	gosql "database/sql"
 	"fmt"
+	"math/rand"
+	"os"
+	"strconv"
 	"strings"
 	"time"
 
@@ -39,6 +42,11 @@ type jobDetails struct {
 	highWaterTimestamp hlc.Timestamp             // high watermark timestamp
 }
 
+// configSetter defines a callback function used by parseConfigs to apply each
+// parsed sink and its associated percentage. This allows custom handling of
+// each config entry during parsing (e.g., storing in a map or validating).
+type configSetter func(key string, value int) error
+
 // getJobsUpdatedWithPayload fetches additional changefeed payload details for specific jobs from the database.
 // This returns a new slice of jobDetails with the updated payload details without mutating the one in input.
 func getJobsUpdatedWithPayload(
@@ -338,10 +346,93 @@ func createChangefeed(
 	return err
 }
 
+// ParseConfigs exported for testing
+func ParseConfigs(config string, cb configSetter) error {
+	if config == "" {
+		return fmt.Errorf("config string cannot be empty")
+	}
+
+	configsArr := strings.Split(config, ",")
+	totalCount := 0
+
+	for _, c := range configsArr {
+		parts := strings.Split(c, ":")
+		if len(parts) != 2 {
+			return fmt.Errorf("invalid config format: %s", c)
+		}
+
+		key := parts[0]
+		valueStr := parts[1]
+
+		value, err := strconv.Atoi(valueStr)
+		if err != nil {
+			return fmt.Errorf("invalid percentage value in config '%s'", c)
+		}
+
+		if value < 0 || value > 100 {
+			return fmt.Errorf("percentage value out of range in config '%s': must be between 0 and 100", c)
+		}
+
+		err = cb(key, value)
+		if err != nil {
+			return err
+		}
+
+		totalCount += value
+	}
+
+	if totalCount != 100 {
+		return fmt.Errorf("all sinks must sum to 100%%, but total is %d%%", totalCount)
+	}
+
+	return nil
+}
+
+func selectSink(sinks map[string]int) string {
+	choice := rand.Intn(100)
+	sum := 0
+	for sink, pct := range sinks {
+		sum += pct
+		if choice < sum {
+			return sink
+		}
+	}
+	return "null" // Default fallback if no valid selection
+}
+
 // getSinkConfigs returns the sink uri along with the options for creating the changefeed.
 // this will be extended later for more sinks
 func getSinkConfigs(_ context.Context, _ []*jobDetails) (string, []string, error) {
-	return "null://", make([]string, 0), nil
+	sinkConfigEnv := os.Getenv("SINK_CONFIG")
+	if sinkConfigEnv == "" {
+		return "null://", []string{}, nil // Default to null sink if env is not set
+	}
+
+	sinks := make(map[string]int)
+	uris := make(map[string]string)
+
+	err := ParseConfigs(sinkConfigEnv, func(sink string, value int) error {
+		sinks[sink] = value
+		uriEnv := fmt.Sprintf("SINK_CONFIG_%s", strings.ToUpper(sink))
+		uri, exists := os.LookupEnv(uriEnv)
+		if !exists {
+			return fmt.Errorf("environment variable %s not found for sink %s", uriEnv, sink)
+		}
+		uris[sink] = uri
+		return nil
+	})
+	if err != nil {
+		// Default to null sink on parsing error
+		return "null://", []string{}, nil //nolint:returnerrcheck
+	}
+
+	selectedSink := selectSink(sinks)
+	selectedURI, exists := uris[selectedSink]
+	if !exists {
+		return "null://", []string{}, nil // Default to null sink if selection fails
+	}
+
+	return selectedURI, []string{}, nil
 }
 
 // calculateScanOption determines whether the new changefeed should have an initial scan based on existing jobs.

pkg/cmd/roachtest/operations/utils_test.go

@@ -0,0 +1,83 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+//
+
+package operations
+
+import (
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/operations/changefeeds"
+)
+
+func TestParseConfigs(t *testing.T) {
+	tests := []struct {
+		name    string
+		config  string
+		wantErr bool
+		wantMap map[string]int
+	}{
+		{
+			name:    "valid config",
+			config:  "kafka:20,webhook:30,null:50",
+			wantErr: false,
+			wantMap: map[string]int{"kafka": 20, "webhook": 30, "null": 50},
+		},
+		{
+			name:    "percentage does not add up to 100",
+			config:  "kafka:20,webhook:30,null:40",
+			wantErr: true,
+		},
+		{
+			name:    "invalid percentage value",
+			config:  "kafka:abc,webhook:30,null:70",
+			wantErr: true,
+		},
+		{
+			name:    "missing colon",
+			config:  "kafka20,webhook:30,null:70",
+			wantErr: true,
+		},
+		{
+			name:    "negative percentage",
+			config:  "kafka:-10,webhook:60,null:50",
+			wantErr: true,
+		},
+		{
+			name:    "empty config",
+			config:  "",
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			gotMap := make(map[string]int)
+			err := changefeeds.ParseConfigs(tt.config, func(key string, value int) error {
+				gotMap[key] = value
+				return nil
+			})
+
+			if (err != nil) != tt.wantErr {
+				t.Errorf("parseConfigs() error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if !tt.wantErr && !equalMaps(gotMap, tt.wantMap) {
+				t.Errorf("parseConfigs() gotMap = %v, wantMap %v", gotMap, tt.wantMap)
+			}
+		})
+	}
+}
+
+func equalMaps(a, b map[string]int) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for k, v := range a {
+		if b[k] != v {
+			return false
+		}
+	}
+	return true
+}