pkg/cli: embed metadata in tsdump files

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

Author: Abhinav1299

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 {