From 306eabcec74aef3d6870e57b302a71aea481d635 Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Fri, 3 Jul 2020 16:13:15 -0700
Subject: [PATCH] stmtdiagnostics: break into chunks

Implement breaking support bundles into chunks. The bundles contain
traces which can be large, resulting in "command is too large" errors.

Release note (bug fix): better support for large statement diagnostic
bundles.
---
 pkg/sql/explain_bundle_test.go                |  4 +++
 .../stmtdiagnostics/statement_diagnostics.go  | 29 +++++++++++++------
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/pkg/sql/explain_bundle_test.go b/pkg/sql/explain_bundle_test.go
index 3c22c8df505b..01adace06cd6 100644
--- a/pkg/sql/explain_bundle_test.go
+++ b/pkg/sql/explain_bundle_test.go
@@ -16,12 +16,14 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"math/rand"
 	"regexp"
 	"sort"
 	"strings"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/sql/stmtdiagnostics"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
@@ -44,6 +46,8 @@ func TestExplainAnalyzeDebug(t *testing.T) {
 	plans := "schema.sql opt.txt opt-v.txt opt-vv.txt plan.txt"
 
 	t.Run("basic", func(t *testing.T) {
+		// Set a small chunk size to test splitting into chunks.
+		defer stmtdiagnostics.TestingSetBundleChunkSize(1 + rand.Intn(100))()
 		rows := r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM abc WHERE c=1")
 		checkBundle(
 			t, fmt.Sprint(rows),
diff --git a/pkg/sql/stmtdiagnostics/statement_diagnostics.go b/pkg/sql/stmtdiagnostics/statement_diagnostics.go
index 58b0065d2146..87bb8b2e9f3f 100644
--- a/pkg/sql/stmtdiagnostics/statement_diagnostics.go
+++ b/pkg/sql/stmtdiagnostics/statement_diagnostics.go
@@ -38,6 +38,15 @@ var stmtDiagnosticsPollingInterval = settings.RegisterDurationSetting(
 	"rate at which the stmtdiagnostics.Registry polls for requests, set to zero to disable",
 	10*time.Second)
 
+var bundleChunkSize = 128 * 1024
+
+// TestingSetBundleChunkSize is used to change the chunk size for testing.
+func TestingSetBundleChunkSize(val int) (restore func()) {
+	old := bundleChunkSize
+	bundleChunkSize = val
+	return func() { bundleChunkSize = old }
+}
+
 // Registry maintains a view on the statement fingerprints
 // on which data is to be collected (i.e. system.statement_diagnostics_requests)
 // and provides utilities for checking a query against this list and satisfying
@@ -394,27 +403,29 @@ func (r *Registry) insertStatementDiagnostics(
 			errorVal = tree.NewDString(collectionErr.Error())
 		}
 
-		bundleChunksVal := tree.DNull
-		if len(bundle) != 0 {
-			// Insert the bundle into system.statement_bundle_chunks.
-			// TODO(radu): split in chunks.
+		bundleChunksVal := tree.NewDArray(types.Int)
+		for len(bundle) > 0 {
+			chunk := bundle
+			if len(chunk) > bundleChunkSize {
+				chunk = chunk[:bundleChunkSize]
+			}
+			bundle = bundle[len(chunk):]
+
+			// Insert the chunk into system.statement_bundle_chunks.
 			row, err := r.ie.QueryRowEx(
 				ctx, "stmt-bundle-chunks-insert", txn,
 				sqlbase.InternalExecutorSessionDataOverride{User: security.RootUser},
 				"INSERT INTO system.statement_bundle_chunks(description, data) VALUES ($1, $2) RETURNING id",
 				"statement diagnostics bundle",
-				tree.NewDBytes(tree.DBytes(bundle)),
+				tree.NewDBytes(tree.DBytes(chunk)),
 			)
 			if err != nil {
 				return err
 			}
 			chunkID := row[0].(*tree.DInt)
-
-			array := tree.NewDArray(types.Int)
-			if err := array.Append(chunkID); err != nil {
+			if err := bundleChunksVal.Append(chunkID); err != nil {
 				return err
 			}
-			bundleChunksVal = array
 		}
 
 		collectionTime := timeutil.Now()