cli: add ZSTD compression support for tsdump generation and upload

This change adds native ZSTD compression support to the tsdump command,
eliminating the need for manual compression/decompression of tsdump files.

- Added `--format=raw-zstd` option to generate ZSTD-compressed tsdump files.
- Enabled `tsdump upload` to automatically detect and decompress ZSTD files.

Fixes #151727
Epic: CRDB-56325
Part-of: CRDB-53445
Release note (cli change): The `cockroach debug tsdump` command now
supports ZSTD compression via `--format=raw-zstd`. This generates
compressed tsdump files that are ~85% smaller than raw format. The
`tsdump upload` command automatically detects and decompresses ZSTD
files, allowing direct upload without manual decompression.

Author: visheshbardia

pkg/cli/BUILD.bazel

@@ -285,6 +285,7 @@ go_library(
         "@com_github_gogo_protobuf//jsonpb",
         "@com_github_jackc_pgx_v5//pgconn",
         "@com_github_jackc_pgx_v5//pgtype",
+        "@com_github_klauspost_compress//zstd",
         "@com_github_kr_pretty//:pretty",
         "@com_github_lestrrat_go_jwx_v2//jwk",
         "@com_github_marusama_semaphore//:semaphore",
@@ -498,6 +499,7 @@ go_test(
         "@com_github_datadog_datadog_api_client_go_v2//api/datadog",
         "@com_github_datadog_datadog_api_client_go_v2//api/datadogV2",
         "@com_github_google_pprof//profile",
+        "@com_github_klauspost_compress//zstd",
         "@com_github_pmezard_go_difflib//difflib",
         "@com_github_spf13_cobra//:cobra",
         "@com_github_spf13_pflag//:pflag",

pkg/cli/testdata/tsdump_upload_e2e

@@ -39,3 +39,17 @@ cr.store.rocksdb.block.cache.usage,2021-01-01T00:00:00Z,2,75.2
 [{"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:4","dry_run":"false","duration":"0","estimated_cost":"0.000186986301369863","hostname":"hostname","message":"tsdump upload completed: uploaded 4 series overall","series_uploaded":"4","service":"tsdump_upload","success":"true"}]
 ----
 ----
+
+
+upload-datadog compression=zstd
+cr.node.admission.admitted.elastic-cpu,2025-05-26T08:32:00Z,1,1
+cr.node.sql.query.count,2021-01-01T00:00:00Z,1,100.5
+cr.node.sql.query.count,2021-01-01T00:00:10Z,1,102.3
+cr.store.rocksdb.block.cache.usage,2021-01-01T00:00:00Z,2,75.2
+----
+----
+{"series":[{"interval":10,"metric":"cockroachdb.admission.admitted.elastic_cpu","points":[{"timestamp":1748248320,"value":1}],"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},{"interval":10,"metric":"cockroachdb.sql.query.count","points":[{"timestamp":1609459200,"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},{"interval":10,"metric":"cockroachdb.sql.query.count","points":[{"timestamp":1609459210,"value":1.8}],"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},{"interval":10,"metric":"cockroachdb.rocksdb.block.cache.usage","points":[{"timestamp":1609459200,"value":75.2}],"tags":["store:2","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:4","dry_run":"false","duration":"0","estimated_cost":"0.000186986301369863","hostname":"hostname","message":"tsdump upload completed: uploaded 4 series overall","series_uploaded":"4","service":"tsdump_upload","success":"true"}]
+----
+----

pkg/cli/tsdump.go

@@ -36,6 +36,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/netutil/addr"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
+	"github.com/klauspost/compress/zstd"
 	"github.com/spf13/cobra"
 )
 
@@ -127,7 +128,7 @@ will then convert it to the --format requested in the current invocation.
 
 		var w tsWriter
 		switch cmd := debugTimeSeriesDumpOpts.format; cmd {
-		case tsDumpRaw:
+		case tsDumpRaw, tsDumpRawZstd:
 			if convertFile != "" {
 				return errors.Errorf("input file is already in raw format")
 			}
@@ -263,7 +264,7 @@ will then convert it to the --format requested in the current invocation.
 			}
 
 			tsClient := conn.NewTimeSeriesClient()
-			if debugTimeSeriesDumpOpts.format == tsDumpRaw {
+			if debugTimeSeriesDumpOpts.format == tsDumpRaw || debugTimeSeriesDumpOpts.format == tsDumpRawZstd {
 				// get the node details so that we can get the SQL port
 				statusClient := conn.NewStatusClient()
 				resp, err := statusClient.Details(ctx, &serverpb.DetailsRequest{NodeId: "local"})
@@ -298,7 +299,18 @@ will then convert it to the --format requested in the current invocation.
 
 				// Buffer the writes since we're going to
 				// be writing potentially a lot of data.
-				w := bufio.NewWriterSize(output, 1024*1024)
+				bufWriter := bufio.NewWriterSize(output, 1024*1024)
+
+				// Wrap with ZSTD compressor if raw-zstd format
+				var w io.Writer = bufWriter
+				var zstdWriter *zstd.Encoder
+				if debugTimeSeriesDumpOpts.format == tsDumpRawZstd {
+					zstdWriter, err = zstd.NewWriter(bufWriter)
+					if err != nil {
+						return errors.Wrap(err, "creating zstd writer")
+					}
+					w = zstdWriter
+				}
 
 				// Write embedded metadata first
 				if err := tsdumpmeta.Write(w, metadata); err != nil {
@@ -309,7 +321,14 @@ will then convert it to the --format requested in the current invocation.
 					return err
 				}
 
-				if err := w.Flush(); err != nil {
+				// Close ZSTD writer before flushing buffer
+				if zstdWriter != nil {
+					if err := zstdWriter.Close(); err != nil {
+						return err
+					}
+				}
+
+				if err := bufWriter.Flush(); err != nil {
 					return err
 				}
 
@@ -839,6 +858,8 @@ const (
 	// to push older timestamps. There's no way to enable historical
 	// ingestion if DD doesn't already know your metric name.
 	tsDumpDatadogInit
+	// tsDumpRawZstd is like tsDumpRaw but with ZSTD compression.
+	tsDumpRawZstd
 )
 
 // Type implements the pflag.Value interface.
@@ -863,6 +884,8 @@ func (m *tsDumpFormat) String() string {
 		return "datadog"
 	case tsDumpDatadogInit:
 		return "datadoginit"
+	case tsDumpRawZstd:
+		return "raw-zstd"
 	}
 	return ""
 }
@@ -886,6 +909,8 @@ func (m *tsDumpFormat) Set(s string) error {
 		*m = tsDumpDatadog
 	case "datadoginit":
 		*m = tsDumpDatadogInit
+	case "raw-zstd":
+		*m = tsDumpRawZstd
 
 	default:
 		return fmt.Errorf("invalid value for --format: %s", s)

pkg/cli/tsdump_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/datadriven"
+	"github.com/klauspost/compress/zstd"
 	"github.com/stretchr/testify/require"
 )
 
@@ -485,6 +486,54 @@ func TestTSDumpRawGenerationWithEmbeddedMetadata(t *testing.T) {
 	require.Equal(t, "1", readMetadata.StoreToNodeMap["1"])
 }
 
+// TestTSDumpRawZstdGenerationWithEmbeddedMetadata tests that raw-zstd format tsdump generation
+// produces valid ZSTD-compressed output with embedded metadata.
+func TestTSDumpRawZstdGenerationWithEmbeddedMetadata(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	c := NewCLITest(TestCLIParams{})
+	defer c.Cleanup()
+
+	tmpFile, err := os.CreateTemp("", "tsdump_*.gob.zst")
+	require.NoError(t, err)
+	defer func() {
+		_ = os.Remove(tmpFile.Name())
+	}()
+	tmpFile.Close()
+
+	_, err = c.RunWithCapture(fmt.Sprintf(
+		"debug tsdump --format=raw-zstd --output=%s --cluster-name=test-cluster-1 --disable-cluster-name-verification",
+		tmpFile.Name(),
+	))
+	require.NoError(t, err)
+
+	file, err := os.Open(tmpFile.Name())
+	require.NoError(t, err)
+	defer file.Close()
+
+	// Read first 4 bytes to verify ZSTD magic number
+	magic := make([]byte, 4)
+	_, err = file.Read(magic)
+	require.NoError(t, err)
+	require.Equal(t, []byte{0x28, 0xB5, 0x2F, 0xFD}, magic, "file should have ZSTD magic number")
+
+	_, err = file.Seek(0, io.SeekStart)
+	require.NoError(t, err)
+
+	zstdReader, err := zstd.NewReader(file)
+	require.NoError(t, err)
+	defer zstdReader.Close()
+
+	dec := gob.NewDecoder(zstdReader)
+	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 {

pkg/cli/tsdump_upload.go

@@ -35,6 +35,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/yamlutil"
 	"github.com/cockroachdb/errors"
+	"github.com/klauspost/compress/zstd"
 	"go.yaml.in/yaml/v4"
 )
 
@@ -61,6 +62,7 @@ var (
 		"tpl_var_cluster=%s&tpl_var_upload_id=%s&tpl_var_upload_day=%d&tpl_var_upload_month=%d&tpl_var_upload_year=%d&from_ts=%d&to_ts=%d"
 	zipFileSignature            = []byte{0x50, 0x4B, 0x03, 0x04}
 	gzipFileSignature           = []byte{0x1f, 0x8b}
+	zstdFileSignature           = []byte{0x28, 0xB5, 0x2F, 0xFD}
 	logMessageFormat            = "tsdump upload to datadog is partially failed for metric: %s"
 	partialFailureMessageFormat = "The Tsdump upload to Datadog succeeded but %d metrics partially failed to upload." +
 		" These failures can be due to transient network errors.\nMetrics:\n%s\n" +
@@ -1026,6 +1028,13 @@ func getFileReader(fileName string) (io.Reader, error) {
 		}
 		return gzipReader, nil
 
+	case bytes.HasPrefix(buf, zstdFileSignature):
+		zstdReader, err := zstd.NewReader(file)
+		if err != nil {
+			return nil, err
+		}
+		return zstdReader, nil
+
 	default:
 		return file, nil
 	}

pkg/cli/tsdump_upload_test.go

@@ -36,6 +36,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/datadriven"
+	"github.com/klauspost/compress/zstd"
 	"github.com/stretchr/testify/require"
 )
 
@@ -416,6 +417,15 @@ func generateMockTSDumpFromCSV(t *testing.T, csvInput string, options ...mockTSD
 		require.NoError(t, err)
 		encoder = gob.NewEncoder(writer)
 
+	case "zstd":
+		filePattern = "mock_tsdump_*.gob.zst"
+		zstdWriter, err := zstd.NewWriter(tmpFile)
+		require.NoError(t, err)
+		defer func() {
+			require.NoError(t, zstdWriter.Close(), "failed to close zstd writer")
+		}()
+		encoder = gob.NewEncoder(zstdWriter)
+
 	default:
 		filePattern = "mock_tsdump_*.gob"
 		encoder = gob.NewEncoder(tmpFile)