Merge #151518

151518: pkg/cli: embed metadata in tsdump files r=Abhinav1299 a=Abhinav1299

This commit introduces functionality for writing and reading embedded metadata in tsdump files.
Additionally, the upload process now prioritizes embedded metadata over external YAML
configurations for store-to-node mappings, enhancing the tsdump upload workflow.

Summary of changes
- Current flow of tsdump generation now adds metadata information to it by default.
Metadata consist of store-to-node mappings, CRDB version and creation timestamp.
- The `tsdump upload` and `roachprod` are now capable of decoding this metadata so
no more need to provide store-to-node mapping manually.
- Changes are backward compatible for both `tsdump upload` and `roachprod upload` so existing
flow of providing store-to-node mapping is not hampered. Recompiling of roachprod is required though.
- A small bug fix of leaked go routine when raw tsdump is converted to any other format like csv.

Part of: CRDB-51457
Epic: CRDB-52093
Release note: None

Co-authored-by: Abhinav1299 <[email protected]>

Author: craig[bot]

pkg/BUILD.bazel

@@ -709,6 +709,7 @@ ALL_TESTS = [
     "//pkg/testutils/zerofields:zerofields_test",
     "//pkg/testutils:testutils_test",
     "//pkg/ts/testmodel:testmodel_test",
+    "//pkg/ts/tsdumpmeta:tsdumpmeta_test",
     "//pkg/ts:ts_test",
     "//pkg/ui:ui_test",
     "//pkg/upgrade/upgradecluster:upgradecluster_test",
@@ -2524,6 +2525,8 @@ GO_TARGETS = [
     "//pkg/ts/catalog:catalog",
     "//pkg/ts/testmodel:testmodel",
     "//pkg/ts/testmodel:testmodel_test",
+    "//pkg/ts/tsdumpmeta:tsdumpmeta",
+    "//pkg/ts/tsdumpmeta:tsdumpmeta_test",
     "//pkg/ts/tspb:tspb",
     "//pkg/ts/tsutil:tsutil",
     "//pkg/ts:ts",

pkg/cli/BUILD.bazel

@@ -193,6 +193,7 @@ go_library(
         "//pkg/testutils/serverutils",
         "//pkg/ts",
         "//pkg/ts/catalog",
+        "//pkg/ts/tsdumpmeta",
         "//pkg/ts/tspb",
         "//pkg/ts/tsutil",
         "//pkg/upgrade/upgrades",
@@ -453,6 +454,7 @@ go_test(
         "//pkg/testutils/sqlutils",
         "//pkg/testutils/testcluster",
         "//pkg/ts",
+        "//pkg/ts/tsdumpmeta",
         "//pkg/ts/tspb",
         "//pkg/util",
         "//pkg/util/envutil",

pkg/cli/testdata/tsdump_upload_embedded_metadata

@@ -0,0 +1,53 @@
+# Test embedded metadata functionality for tsdump upload
+
+# Test case 1: embedded metadata only (no YAML file)
+# This tests that store metrics get proper node_id tags based on embedded store-to-node mapping
+upload-datadog-embedded-only
+cr.store.rocksdb.block.cache.usage,2021-01-01T00:00:00Z,1,75.2
+cr.store.rocksdb.num.entries,2021-01-01T00:00:10Z,3,42.0
+cr.node.sql.query.count,2021-01-01T00:00:20Z,1,100.5
+----
+----
+{"series":[{"interval":10,"metric":"crdb.tsdump.rocksdb.block.cache.usage","points":[{"timestamp":1609459200,"value":75.2}],"tags":["node_id:1","store:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":3},{"interval":10,"metric":"crdb.tsdump.rocksdb.num.entries","points":[{"timestamp":1609459210,"value":42}],"tags":["node_id:2","store:3","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":0},{"interval":10,"metric":"crdb.tsdump.sql.query.count","points":[{"timestamp":1609459220,"value":100.5}],"tags":["node_id:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":1}]}
+
+[{"ddsource":"tsdump_upload","ddtags":"cluster_type:SELF_HOSTED,cluster_label:\"test-cluster\",cluster_id:test-cluster-id,zendesk_ticket:zd-test,org_name:test-org,user_name:test-user,upload_id:\"test-cluster\"-20241114000000,upload_timestamp:2024-11-14 00:00:00,upload_year:2024,upload_month:11,upload_day:14,series_uploaded:3","dry_run":"false","duration":"0","estimated_cost":"0.000186986301369863","hostname":"hostname","message":"tsdump upload completed: uploaded 3 series overall","series_uploaded":"3","service":"tsdump_upload","success":"true"}]
+----
+----
+
+# Test case 2: YAML file only (no embedded metadata)
+# This tests that store metrics get proper node_id tags from external YAML when no embedded metadata is present
+upload-datadog-yaml-only
+cr.store.rocksdb.block.cache.usage,2021-01-01T00:00:00Z,1,75.2
+----
+----
+{"series":[{"interval":10,"metric":"crdb.tsdump.rocksdb.block.cache.usage","points":[{"timestamp":1609459200,"value":75.2}],"tags":["node_id:1","store:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":3}]}
+
+[{"ddsource":"tsdump_upload","ddtags":"cluster_type:SELF_HOSTED,cluster_label:\"test-cluster\",cluster_id:test-cluster-id,zendesk_ticket:zd-test,org_name:test-org,user_name:test-user,upload_id:\"test-cluster\"-20241114000000,upload_timestamp:2024-11-14 00:00:00,upload_year:2024,upload_month:11,upload_day:14,series_uploaded:1","dry_run":"false","duration":"0","estimated_cost":"6.232876712328767e-05","hostname":"hostname","message":"tsdump upload completed: uploaded 1 series overall","series_uploaded":"1","service":"tsdump_upload","success":"true"}]
+----
+----
+
+# Test case 3: both embedded metadata and YAML file (embedded takes priority)
+# This tests that embedded metadata is prioritized over external YAML when both are present, proving precedence
+upload-datadog-embedded-priority
+cr.store.rocksdb.block.cache.usage,2021-01-01T00:00:00Z,1,75.2
+cr.store.rocksdb.num.entries,2021-01-01T00:00:10Z,3,42.0
+----
+----
+{"series":[{"interval":10,"metric":"crdb.tsdump.rocksdb.block.cache.usage","points":[{"timestamp":1609459200,"value":75.2}],"tags":["node_id:1","store:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":3},{"interval":10,"metric":"crdb.tsdump.rocksdb.num.entries","points":[{"timestamp":1609459210,"value":42}],"tags":["node_id:2","store:3","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":0}]}
+
+[{"ddsource":"tsdump_upload","ddtags":"cluster_type:SELF_HOSTED,cluster_label:\"test-cluster\",cluster_id:test-cluster-id,zendesk_ticket:zd-test,org_name:test-org,user_name:test-user,upload_id:\"test-cluster\"-20241114000000,upload_timestamp:2024-11-14 00:00:00,upload_year:2024,upload_month:11,upload_day:14,series_uploaded:2","dry_run":"false","duration":"0","estimated_cost":"0.00012465753424657533","hostname":"hostname","message":"tsdump upload completed: uploaded 2 series overall","series_uploaded":"2","service":"tsdump_upload","success":"true"}]
+----
+----
+
+# Test case 4: no metadata and no YAML file (normal upload)
+# This tests normal upload behavior when no store-to-node mapping is available, store metrics get only store tags
+upload-datadog-no-mapping
+cr.store.rocksdb.block.cache.usage,2021-01-01T00:00:00Z,1,75.2
+cr.node.sql.query.count,2021-01-01T00:00:20Z,1,100.5
+----
+----
+{"series":[{"interval":10,"metric":"crdb.tsdump.rocksdb.block.cache.usage","points":[{"timestamp":1609459200,"value":75.2}],"tags":["store:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":3},{"interval":10,"metric":"crdb.tsdump.sql.query.count","points":[{"timestamp":1609459220,"value":100.5}],"tags":["node_id:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":1}]}
+
+[{"ddsource":"tsdump_upload","ddtags":"cluster_type:SELF_HOSTED,cluster_label:\"test-cluster\",cluster_id:test-cluster-id,zendesk_ticket:zd-test,org_name:test-org,user_name:test-user,upload_id:\"test-cluster\"-20241114000000,upload_timestamp:2024-11-14 00:00:00,upload_year:2024,upload_month:11,upload_day:14,series_uploaded:2","dry_run":"false","duration":"0","estimated_cost":"0.00012465753424657533","hostname":"hostname","message":"tsdump upload completed: uploaded 2 series overall","series_uploaded":"2","service":"tsdump_upload","success":"true"}]
+----
+----

pkg/cli/tsdump.go

@@ -22,11 +22,13 @@ import (
 	"sync"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/build"
 	"github.com/cockroachdb/cockroach/pkg/cli/clierrorplus"
 	"github.com/cockroachdb/cockroach/pkg/cli/clisqlclient"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/ts"
+	"github.com/cockroachdb/cockroach/pkg/ts/tsdumpmeta"
 	"github.com/cockroachdb/cockroach/pkg/ts/tspb"
 	"github.com/cockroachdb/cockroach/pkg/ts/tsutil"
 	"github.com/cockroachdb/cockroach/pkg/util/netutil/addr"
@@ -192,13 +194,6 @@ will then convert it to the --format requested in the current invocation.
 					return errors.Wrapf(err, "connecting to %s", target)
 				}
 
-				// Buffer the writes to os.Stdout since we're going to
-				// be writing potentially a lot of data to it.
-				w := bufio.NewWriterSize(os.Stdout, 1024*1024)
-				if err := tsutil.DumpRawTo(stream, w); err != nil {
-					return err
-				}
-
 				// get the node details so that we can get the SQL port
 				statusClient := conn.NewStatusClient()
 				resp, err := statusClient.Details(ctx, &serverpb.DetailsRequest{NodeId: "local"})
@@ -213,6 +208,32 @@ will then convert it to the --format requested in the current invocation.
 					return err
 				}
 
+				// Get store-to-node mapping for metadata
+				storeToNodeMap, err := getStoreToNodeMapping(ctx)
+				if err != nil {
+					return err
+				}
+
+				// Create metadata header
+				metadata := tsdumpmeta.Metadata{
+					Version:        build.BinaryVersion(),
+					StoreToNodeMap: storeToNodeMap,
+					CreatedAt:      timeutil.Now(),
+				}
+
+				// Buffer the writes to os.Stdout since we're going to
+				// be writing potentially a lot of data to it.
+				w := bufio.NewWriterSize(os.Stdout, 1024*1024)
+
+				// Write embedded metadata first
+				if err := tsdumpmeta.Write(w, metadata); err != nil {
+					return err
+				}
+
+				if err := tsutil.DumpRawTo(stream, w); err != nil {
+					return err
+				}
+
 				if err = createYAML(ctx); err != nil {
 					return err
 				}
@@ -228,12 +249,26 @@ will then convert it to the --format requested in the current invocation.
 			if err != nil {
 				return err
 			}
+			defer f.Close()
 			type tup struct {
 				data *tspb.TimeSeriesData
 				err  error
 			}
 
 			dec := gob.NewDecoder(f)
+
+			// Try to read embedded metadata first
+			embeddedMetadata, metadataErr := tsdumpmeta.Read(dec)
+			if metadataErr != nil {
+				// No embedded metadata, restart from beginning
+				if _, err := f.Seek(0, io.SeekStart); err != nil {
+					return err
+				}
+				dec = gob.NewDecoder(f) // Reset decoder to read from beginning
+			} else {
+				fmt.Printf("Found embedded store-to-node mapping with %d entries\n", len(embeddedMetadata.StoreToNodeMap))
+			}
+
 			decodeOne := func() (*tspb.TimeSeriesData, error) {
 				var v roachpb.KeyValue
 				err := dec.Decode(&v)
@@ -260,6 +295,10 @@ will then convert it to the --format requested in the current invocation.
 				for {
 					data, err := decodeOne()
 					ch <- tup{data, err}
+					// Exit the goroutine if we encounter EOF or any error
+					if err != nil {
+						break
+					}
 				}
 			}()
 
@@ -465,38 +504,57 @@ type openMetricsWriter struct {
 // createYAML generates and writes tsdump.yaml to default /tmp or to a specified path.
 // This file is used for staging the tsdump data into a local database for debugging
 func createYAML(ctx context.Context) (resErr error) {
+	// Write the YAML file for backward compatibility
 	file, err := os.OpenFile(debugTimeSeriesDumpOpts.yaml, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0666)
 	if err != nil {
 		return err
 	}
 	defer file.Close()
 
-	sqlConn, err := makeSQLClient(ctx, "cockroach tsdump", useSystemDb)
+	mapping, err := getStoreToNodeMapping(ctx)
 	if err != nil {
 		return err
 	}
-	defer func() { resErr = errors.CombineErrors(resErr, sqlConn.Close()) }()
+
+	for storeID, nodeID := range mapping {
+		_, err := fmt.Fprintf(file, "%s: %s\n", storeID, nodeID)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// getStoreToNodeMapping retrieves the store-to-node mapping from the database
+func getStoreToNodeMapping(ctx context.Context) (map[string]string, error) {
+	sqlConn, err := makeSQLClient(ctx, "cockroach tsdump", useSystemDb)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		if closeErr := sqlConn.Close(); closeErr != nil {
+			fmt.Fprintf(os.Stderr, "Warning: failed to close SQL connection: %v\n", closeErr)
+		}
+	}()
 
 	_, rows, err := sqlExecCtx.RunQuery(
 		ctx,
 		sqlConn,
-		clisqlclient.MakeQuery(`SELECT store_id || ': ' || node_id FROM crdb_internal.kv_store_status`), false)
+		clisqlclient.MakeQuery(`SELECT store_id, node_id FROM crdb_internal.kv_store_status`), false)
 
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	var strStoreNodeID string
+	mapping := make(map[string]string)
 	for _, row := range rows {
-		storeNodeID := row
-		strStoreNodeID = strings.Join(storeNodeID, " ")
-		strStoreNodeID += "\n"
-		_, err := file.WriteString(strStoreNodeID)
-		if err != nil {
-			return err
+		if len(row) >= 2 {
+			storeID := strings.TrimSpace(row[0])
+			nodeID := strings.TrimSpace(row[1])
+			mapping[storeID] = nodeID
 		}
 	}
-	return nil
+	return mapping, nil
 }
 
 func makeOpenMetricsWriter(out io.Writer) *openMetricsWriter {

pkg/cli/tsdump_test.go

@@ -8,6 +8,7 @@ package cli
 import (
 	"bytes"
 	"compress/gzip"
+	"encoding/gob"
 	"fmt"
 	"io"
 	"math/rand"
@@ -22,6 +23,7 @@ import (
 
 	"github.com/DataDog/datadog-api-client-go/v2/api/datadogV2"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
+	"github.com/cockroachdb/cockroach/pkg/ts/tsdumpmeta"
 	"github.com/cockroachdb/cockroach/pkg/ts/tspb"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -134,6 +136,113 @@ func TestMakeOpenMetricsWriter(t *testing.T) {
 	require.Equal(t, dataPointsNum+1 /* datapoints + EOF final line */, len(res))
 }
 
+// TestTSDumpConversionWithEmbeddedMetadata tests the conversion of tsdump files
+// containing embedded metadata to CSV format, verifying metadata extraction.
+func TestTSDumpConversionWithEmbeddedMetadata(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	tmpFile, err := os.CreateTemp("", "tsdump_with_metadata_*.gob")
+	require.NoError(t, err)
+	defer func(name string) {
+		err := os.Remove(name)
+		if err != nil {
+			t.Fatalf("failed to remove temporary file %s: %v", name, err)
+		}
+	}(tmpFile.Name())
+
+	metadata := tsdumpmeta.Metadata{
+		Version: "v23.1.0",
+		StoreToNodeMap: map[string]string{
+			"1": "1",
+			"2": "2",
+		},
+		CreatedAt: timeutil.Unix(1609459200, 0),
+	}
+	err = tsdumpmeta.Write(tmpFile, metadata)
+	require.NoError(t, err)
+
+	enc := gob.NewEncoder(tmpFile)
+	kv, err := createMockTimeSeriesKV("cr.node.sql.query.count", "1", 1609459200000000000, 100.5)
+	require.NoError(t, err)
+	err = enc.Encode(kv)
+	require.NoError(t, err)
+
+	tmpFile.Close()
+
+	c := NewCLITest(TestCLIParams{})
+	defer c.Cleanup()
+
+	// Convert to CSV format
+	out, err := c.RunWithCapture(fmt.Sprintf(
+		"debug tsdump --format=csv %s",
+		tmpFile.Name(),
+	))
+	require.NoError(t, err)
+
+	lines := strings.Split(strings.TrimSpace(out), "\n")
+
+	require.Contains(t, out, "Found embedded store-to-node mapping with 2 entries")
+
+	csvFound := false
+	for _, line := range lines {
+		if strings.Contains(line, "cr.node.sql.query.count") &&
+			strings.Contains(line, "2021-01-01T00:00:00Z") {
+			csvFound = true
+			break
+		}
+	}
+	require.True(t, csvFound, "should contain converted CSV data")
+}
+
+// TestTSDumpRawGenerationWithEmbeddedMetadata tests that raw format tsdump generation
+// automatically includes embedded store-to-node mapping metadata in the output.
+func TestTSDumpRawGenerationWithEmbeddedMetadata(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	c := NewCLITest(TestCLIParams{})
+	defer c.Cleanup()
+
+	out, err := c.RunWithCapture("debug tsdump --format=raw --cluster-name=test-cluster-1 --disable-cluster-name-verification")
+	require.NoError(t, err)
+	require.NotEmpty(t, out)
+
+	// Remove the command prefix from the output to get the actual binary data
+	// The output starts with something like "debug tsdump --format=raw..."
+	// Find the first occurrence of gob magic bytes or skip the command line
+	actualData := out
+	if idx := strings.Index(out, "\n"); idx != -1 {
+		actualData = out[idx+1:] // Skip the command line
+	}
+
+	// Write the cleaned binary data to file
+	tmpFile, err := os.CreateTemp("", "tsdump_*.gob")
+	require.NoError(t, err)
+	defer func(name string) {
+		err := os.Remove(name)
+		if err != nil {
+			t.Fatalf("failed to remove temporary file %s: %v", name, err)
+		}
+	}(tmpFile.Name())
+
+	_, err = tmpFile.Write([]byte(actualData))
+	require.NoError(t, err)
+	tmpFile.Close()
+
+	file, err := os.Open(tmpFile.Name())
+	require.NoError(t, err)
+	defer file.Close()
+
+	dec := gob.NewDecoder(file)
+	readMetadata, err := tsdumpmeta.Read(dec)
+	require.NoError(t, err)
+
+	// Verify store-to-node mapping is embedded
+	require.NotNil(t, readMetadata.StoreToNodeMap)
+	require.Equal(t, "1", readMetadata.StoreToNodeMap["1"])
+}
+
 func makeTS(name, source string, dataPointsNum int) *tspb.TimeSeriesData {
 	dps := make([]tspb.TimeSeriesDatapoint, dataPointsNum)
 	for i := range dps {