jobsprofiler: store DistSQL diagram of jobs in job info

pkg/BUILD.bazel

@@ -176,6 +176,7 @@ ALL_TESTS = [
     "//pkg/jobs/joberror:joberror_test",
     "//pkg/jobs/jobsauth:jobsauth_test",
     "//pkg/jobs/jobspb:jobspb_test",
+    "//pkg/jobs/jobsprofiler:jobsprofiler_test",
     "//pkg/jobs/jobsprotectedts:jobsprotectedts_test",
     "//pkg/jobs:jobs_test",
     "//pkg/keys:keys_test",
@@ -1159,6 +1160,8 @@ GO_TARGETS = [
     "//pkg/jobs/jobsauth:jobsauth_test",
     "//pkg/jobs/jobspb:jobspb",
     "//pkg/jobs/jobspb:jobspb_test",
+    "//pkg/jobs/jobsprofiler:jobsprofiler",
+    "//pkg/jobs/jobsprofiler:jobsprofiler_test",
     "//pkg/jobs/jobsprotectedts:jobsprotectedts",
     "//pkg/jobs/jobsprotectedts:jobsprotectedts_test",
     "//pkg/jobs/jobstest:jobstest",
@@ -2632,6 +2635,7 @@ GET_X_DATA_TARGETS = [
     "//pkg/jobs/joberror:get_x_data",
     "//pkg/jobs/jobsauth:get_x_data",
     "//pkg/jobs/jobspb:get_x_data",
+    "//pkg/jobs/jobsprofiler:get_x_data",
     "//pkg/jobs/jobsprotectedts:get_x_data",
     "//pkg/jobs/jobstest:get_x_data",
     "//pkg/jobs/metricspoller:get_x_data",

pkg/ccl/backupccl/BUILD.bazel

@@ -59,6 +59,7 @@ go_library(
         "//pkg/jobs",
         "//pkg/jobs/joberror",
         "//pkg/jobs/jobspb",
+        "//pkg/jobs/jobsprofiler",
         "//pkg/jobs/jobsprotectedts",
         "//pkg/keys",
         "//pkg/kv",

pkg/ccl/backupccl/backup_processor_planning.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/ccl/backupccl/backuppb"
 	"github.com/cockroachdb/cockroach/pkg/cloud"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobsprofiler"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -170,7 +171,11 @@ func distBackup(
 	// Setup a one-stage plan with one proc per input spec.
 	corePlacement := make([]physicalplan.ProcessorCorePlacement, len(backupSpecs))
 	i := 0
+	var jobID jobspb.JobID
 	for sqlInstanceID, spec := range backupSpecs {
+		if i == 0 {
+			jobID = jobspb.JobID(spec.JobID)
+		}
 		corePlacement[i].SQLInstanceID = sqlInstanceID
 		corePlacement[i].Core.BackupData = spec
 		i++
@@ -206,6 +211,9 @@ func distBackup(
 	defer recv.Release()
 
 	defer close(progCh)
+	execCfg := execCtx.ExecCfg()
+	jobsprofiler.StorePlanDiagram(ctx, execCfg.DistSQLSrv.Stopper, p, execCfg.InternalDB, jobID)
+
 	// Copy the evalCtx, as dsp.Run() might change it.
 	evalCtxCopy := *evalCtx
 	dsp.Run(ctx, planCtx, noTxn, p, recv, &evalCtxCopy, nil /* finishedSetupFn */)

pkg/ccl/backupccl/restore_job.go

@@ -393,7 +393,7 @@ func restore(
 		return distRestore(
 			ctx,
 			execCtx,
-			int64(job.ID()),
+			job.ID(),
 			dataToRestore,
 			endTime,
 			encryption,

pkg/ccl/backupccl/restore_processor_planning.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/cloud"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobsprofiler"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
@@ -60,7 +61,7 @@ var replanRestoreFrequency = settings.RegisterDurationSetting(
 func distRestore(
 	ctx context.Context,
 	execCtx sql.JobExecContext,
-	jobID int64,
+	jobID jobspb.JobID,
 	dataToRestore restorationData,
 	restoreTime hlc.Timestamp,
 	encryption *jobspb.BackupEncryptionOptions,
@@ -111,7 +112,7 @@ func distRestore(
 		p := planCtx.NewPhysicalPlan()
 
 		restoreDataSpec := execinfrapb.RestoreDataSpec{
-			JobID:        jobID,
+			JobID:        int64(jobID),
 			RestoreTime:  restoreTime,
 			Encryption:   fileEncryption,
 			TableRekeys:  dataToRestore.getRekeys(),
@@ -184,7 +185,7 @@ func distRestore(
 			NumNodes:                 int64(numNodes),
 			UseSimpleImportSpans:     useSimpleImportSpans,
 			UseFrontierCheckpointing: spanFilter.useFrontierCheckpointing,
-			JobID:                    jobID,
+			JobID:                    int64(jobID),
 		}
 		if spanFilter.useFrontierCheckpointing {
 			spec.CheckpointedSpans = persistFrontier(spanFilter.checkpointFrontier, 0)
@@ -292,6 +293,9 @@ func distRestore(
 		)
 		defer recv.Release()
 
+		execCfg := execCtx.ExecCfg()
+		jobsprofiler.StorePlanDiagram(ctx, execCfg.DistSQLSrv.Stopper, p, execCfg.InternalDB, jobID)
+
 		// Copy the evalCtx, as dsp.Run() might change it.
 		evalCtxCopy := *evalCtx
 		dsp.Run(ctx, planCtx, noTxn, p, recv, &evalCtxCopy, nil /* finishedSetupFn */)

pkg/ccl/changefeedccl/BUILD.bazel

@@ -63,6 +63,7 @@ go_library(
         "//pkg/jobs",
         "//pkg/jobs/jobsauth",
         "//pkg/jobs/jobspb",
+        "//pkg/jobs/jobsprofiler",
         "//pkg/jobs/jobsprotectedts",
         "//pkg/keys",
         "//pkg/kv",

pkg/ccl/changefeedccl/changefeed_dist.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/changefeedbase"
 	"github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/kvfeed"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobsprofiler"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvclient/kvcoord"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -314,6 +315,8 @@ func startDistChangefeed(
 			finishedSetupFn = func(flowinfra.Flow) { resultsCh <- tree.Datums(nil) }
 		}
 
+		jobsprofiler.StorePlanDiagram(ctx, execCfg.DistSQLSrv.Stopper, p, execCfg.InternalDB, jobID)
+
 		// Copy the evalCtx, as dsp.Run() might change it.
 		evalCtxCopy := *evalCtx
 		// p is the physical plan, recv is the distsqlreceiver

pkg/ccl/streamingccl/streamingest/BUILD.bazel

@@ -28,6 +28,7 @@ go_library(
         "//pkg/cloud/externalconn/connectionpb",
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
+        "//pkg/jobs/jobsprofiler",
         "//pkg/jobs/jobsprotectedts",
         "//pkg/keys",
         "//pkg/kv",

pkg/ccl/streamingccl/streamingest/stream_ingestion_processor_planning.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/ccl/streamingccl"
 	"github.com/cockroachdb/cockroach/pkg/ccl/streamingccl/streamclient"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobsprofiler"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/repstream/streampb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -110,7 +111,11 @@ func distStreamIngest(
 
 	// Setup a one-stage plan with one proc per input spec.
 	corePlacement := make([]physicalplan.ProcessorCorePlacement, len(streamIngestionSpecs))
+	var jobID jobspb.JobID
 	for i := range streamIngestionSpecs {
+		if i == 0 {
+			jobID = jobspb.JobID(streamIngestionSpecs[i].JobID)
+		}
 		corePlacement[i].SQLInstanceID = sqlInstanceIDs[i]
 		corePlacement[i].Core.StreamIngestionData = streamIngestionSpecs[i]
 	}
@@ -151,6 +156,8 @@ func distStreamIngest(
 	)
 	defer recv.Release()
 
+	jobsprofiler.StorePlanDiagram(ctx, execCfg.DistSQLSrv.Stopper, p, execCfg.InternalDB, jobID)
+
 	// Copy the evalCtx, as dsp.Run() might change it.
 	evalCtxCopy := *evalCtx
 	dsp.Run(ctx, planCtx, noTxn, p, recv, &evalCtxCopy, nil /* finishedSetupFn */)

pkg/jobs/jobsprofiler/BUILD.bazel

@@ -0,0 +1,49 @@
+load("//build/bazelutil/unused_checker:unused.bzl", "get_x_data")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "jobsprofiler",
+    srcs = ["profiler.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/jobs/jobsprofiler",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/jobs",
+        "//pkg/jobs/jobspb",
+        "//pkg/sql",
+        "//pkg/sql/execinfrapb",
+        "//pkg/sql/isql",
+        "//pkg/util/log",
+        "//pkg/util/stop",
+        "//pkg/util/timeutil",
+    ],
+)
+
+go_test(
+    name = "jobsprofiler_test",
+    srcs = [
+        "main_test.go",
+        "profiler_test.go",
+    ],
+    args = ["-test.timeout=295s"],
+    deps = [
+        "//pkg/base",
+        "//pkg/ccl",
+        "//pkg/jobs",
+        "//pkg/jobs/jobspb",
+        "//pkg/security/securityassets",
+        "//pkg/security/securitytest",
+        "//pkg/server",
+        "//pkg/sql",
+        "//pkg/sql/isql",
+        "//pkg/testutils",
+        "//pkg/testutils/serverutils",
+        "//pkg/testutils/testcluster",
+        "//pkg/util/leaktest",
+        "//pkg/util/log",
+        "//pkg/util/randutil",
+        "@com_github_cockroachdb_errors//:errors",
+        "@com_github_stretchr_testify//require",
+    ],
+)
+
+get_x_data(name = "get_x_data")

pkg/jobs/jobsprofiler/main_test.go

@@ -0,0 +1,35 @@
+// Copyright 2017 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package jobsprofiler_test
+
+import (
+	"os"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/ccl"
+	"github.com/cockroachdb/cockroach/pkg/security/securityassets"
+	"github.com/cockroachdb/cockroach/pkg/security/securitytest"
+	"github.com/cockroachdb/cockroach/pkg/server"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+)
+
+func TestMain(m *testing.M) {
+	defer ccl.TestingEnableEnterprise()()
+	securityassets.SetLoader(securitytest.EmbeddedAssets)
+	randutil.SeedForTests()
+	serverutils.InitTestServerFactory(server.TestServerFactory)
+	serverutils.InitTestClusterFactory(testcluster.TestClusterFactory)
+	os.Exit(m.Run())
+}
+
+//go:generate ../../util/leaktest/add-leaktest.sh *_test.go

pkg/jobs/jobsprofiler/profiler.go

@@ -0,0 +1,58 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package jobsprofiler
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/cockroachdb/cockroach/pkg/jobs"
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/sql"
+	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
+	"github.com/cockroachdb/cockroach/pkg/sql/isql"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/stop"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+)
+
+// StorePlanDiagram stores the DistSQL diagram generated from p in the job info
+// table. The generation of the plan diagram and persistence to the info table
+// are done asynchronously and this method does not block on their completion.
+func StorePlanDiagram(
+	ctx context.Context, stopper *stop.Stopper, p *sql.PhysicalPlan, db isql.DB, jobID jobspb.JobID,
+) {
+	if err := stopper.RunAsyncTask(ctx, "jobs-store-plan-diagram", func(ctx context.Context) {
+		var cancel func()
+		ctx, cancel = stopper.WithCancelOnQuiesce(ctx)
+		defer cancel()
+
+		err := db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+			flowSpecs := p.GenerateFlowSpecs()
+			_, diagURL, err := execinfrapb.GeneratePlanDiagramURL(fmt.Sprintf("job:%d", jobID), flowSpecs, execinfrapb.DiagramFlags{})
+			if err != nil {
+				return err
+			}
+
+			const infoKey = "dsp-diag-url-%d"
+			infoStorage := jobs.InfoStorageForJob(txn, jobID)
+			return infoStorage.Write(ctx, []byte(fmt.Sprintf(infoKey, timeutil.Now().UnixNano())),
+				[]byte(diagURL.String()))
+		})
+		if err != nil {
+			log.Warningf(ctx, "failed to generate and write DistSQL diagram for job %d: %v",
+				jobID, err.Error())
+		}
+	}); err != nil {
+		log.Warningf(ctx, "failed to spawn task to generate DistSQL plan diagram for job %d: %v",
+			jobID, err.Error())
+	}
+}