sql/inspect: add protected timestamp support to INSPECT

This change adds protected timestamp record management to INSPECT jobs
when using AS OF SYSTEM TIME clauses. Protected timestamps prevent
garbage collection of historical data during the inspection process.

When an INSPECT job specifies an AS OF timestamp, the job now:
- Creates a protected timestamp record before beginning inspection
- Stores the record ID in the job details for tracking
- Automatically cleans up the record on job completion or failure

Note: when we add support for running INSPECT without a fixed timestamp
(in #148675), PTS won't be acquired then.

Closes #148573

Release note: none
Epic: CRDB-30356

Author: spilchen

pkg/jobs/jobspb/jobs.proto

@@ -1475,6 +1475,14 @@ message InspectDetails {
 
   // AsOf specifies the timestamp at which the inspect checks should be performed.
   util.hlc.Timestamp as_of = 2 [(gogoproto.nullable) = false];
+
+  // ProtectedTimestampRecord is the id of the protected timestamp record that
+  // prevents garbage collection of data at the AsOf timestamp. This is only
+  // set if a fixed timestamp was used for INSPECT (e.g. AsOf timestamp is non-empty).
+  bytes protected_timestamp_record = 3 [
+    (gogoproto.customname) = "ProtectedTimestampRecord",
+    (gogoproto.customtype) = "github.com/cockroachdb/cockroach/pkg/util/uuid.UUID"
+  ];
 }
 
 message UpdateTableMetadataCacheDetails {}

pkg/jobs/jobsprotectedts/jobs_protected_ts_manager.go

@@ -55,6 +55,9 @@ func setProtectedTSOnJob(details jobspb.Details, u *uuid.UUID) jobspb.Details {
 	case jobspb.SchemaChangeDetails:
 		v.ProtectedTimestampRecord = u
 		return v
+	case jobspb.InspectDetails:
+		v.ProtectedTimestampRecord = u
+		return v
 	default:
 		panic(errors.AssertionFailedf("not supported %T", details))
 	}
@@ -68,6 +71,8 @@ func getProtectedTSOnJob(details jobspb.Details) *uuid.UUID {
 		return v.ProtectedTimestampRecord
 	case jobspb.SchemaChangeDetails:
 		return v.ProtectedTimestampRecord
+	case jobspb.InspectDetails:
+		return v.ProtectedTimestampRecord
 	default:
 		panic("not supported")
 	}

pkg/sql/exec_util.go

@@ -2101,6 +2101,9 @@ type InspectTestingKnobs struct {
 	// OnInspectJobStart is called just before the inspect job begins execution.
 	// If it returns an error, the job fails immediately.
 	OnInspectJobStart func() error
+	// OnInspectAfterProtectedTimestamp is called after the protected timestamp
+	// has been created (if applicable). If it returns an error, the job fails.
+	OnInspectAfterProtectedTimestamp func() error
 	// InspectIssueLogger is an override to the default issue logger.
 	InspectIssueLogger interface{}
 }

pkg/sql/inspect/BUILD.bazel

@@ -18,6 +18,7 @@ go_library(
     deps = [
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
+        "//pkg/kv/kvserver/protectedts/ptpb",
         "//pkg/roachpb",
         "//pkg/security/username",
         "//pkg/settings",
@@ -72,6 +73,7 @@ go_test(
         "//pkg/jobs/jobspb",
         "//pkg/keys",
         "//pkg/kv",
+        "//pkg/kv/kvserver/protectedts",
         "//pkg/roachpb",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
@@ -83,6 +85,7 @@ go_test(
         "//pkg/sql/catalog/desctestutils",
         "//pkg/sql/execinfra",
         "//pkg/sql/execinfrapb",
+        "//pkg/sql/isql",
         "//pkg/sql/rowenc",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sessiondata",

pkg/sql/inspect/index_consistency_check.go

@@ -196,7 +196,6 @@ func (c *indexConsistencyCheck) Start(
 	)
 
 	// Wrap the query with AS OF SYSTEM TIME to ensure it uses the specified timestamp
-	// TODO(148573): use a protected timestamp record for this timestamp.
 	queryWithAsOf := fmt.Sprintf("SELECT * FROM (%s) AS OF SYSTEM TIME %s", checkQuery, c.asOf.AsOfSystemTime())
 
 	// Store the query for error reporting

pkg/sql/inspect/inspect_job.go

@@ -10,9 +10,11 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts/ptpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
@@ -21,6 +23,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/errors"
 )
 
 type inspectResumer struct {
@@ -40,11 +43,24 @@ func (c *inspectResumer) Resume(ctx context.Context, execCtx interface{}) error
 		return err
 	}
 
+	details := c.job.Details().(jobspb.InspectDetails)
+	if len(details.Checks) == 0 {
+		return nil
+	}
+
 	pkSpans, err := c.getPrimaryIndexSpans(ctx, execCfg)
 	if err != nil {
 		return err
 	}
 
+	if err := c.maybeProtectTimestamp(ctx, execCfg, details); err != nil {
+		return err
+	}
+
+	if err := c.maybeRunAfterProtectedTimestampHook(execCfg); err != nil {
+		return err
+	}
+
 	// TODO(149460): add a goroutine that will replan the job on topology changes
 	plan, planCtx, err := c.planInspectProcessors(ctx, jobExecCtx, pkSpans)
 	if err != nil {
@@ -55,13 +71,18 @@ func (c *inspectResumer) Resume(ctx context.Context, execCtx interface{}) error
 		return err
 	}
 
+	c.maybeCleanupProtectedTimestamp(ctx, execCfg)
+
 	return c.markJobComplete(ctx)
 }
 
 // OnFailOrCancel implements the Resumer interface
 func (c *inspectResumer) OnFailOrCancel(
 	ctx context.Context, execCtx interface{}, jobErr error,
 ) error {
+	jobExecCtx := execCtx.(sql.JobExecContext)
+	execCfg := jobExecCtx.ExecCfg()
+	c.maybeCleanupProtectedTimestamp(ctx, execCfg)
 	return nil
 }
 
@@ -77,6 +98,13 @@ func (c *inspectResumer) maybeRunOnJobStartHook(execCfg *sql.ExecutorConfig) err
 	return execCfg.InspectTestingKnobs.OnInspectJobStart()
 }
 
+func (c *inspectResumer) maybeRunAfterProtectedTimestampHook(execCfg *sql.ExecutorConfig) error {
+	if execCfg.InspectTestingKnobs == nil || execCfg.InspectTestingKnobs.OnInspectAfterProtectedTimestamp == nil {
+		return nil
+	}
+	return execCfg.InspectTestingKnobs.OnInspectAfterProtectedTimestamp()
+}
+
 // getPrimaryIndexSpans returns the primary index spans for all tables involved in
 // the INSPECT job's checks.
 func (c *inspectResumer) getPrimaryIndexSpans(
@@ -202,6 +230,50 @@ func (c *inspectResumer) markJobComplete(ctx context.Context) error {
 	)
 }
 
+// maybeProtectTimestamp creates a protected timestamp record for the AsOf
+// timestamp to prevent garbage collection during the inspect operation.
+// If no AsOf timestamp is specified, this function does nothing.
+// The protection target includes all tables involved in the inspect checks.
+// Uses the jobsprotectedts.Manager to store the PTS ID in job details.
+func (c *inspectResumer) maybeProtectTimestamp(
+	ctx context.Context, execCfg *sql.ExecutorConfig, details jobspb.InspectDetails,
+) error {
+	// If we are not running a historical query, nothing to do here.
+	if details.AsOf.IsEmpty() {
+		return nil
+	}
+
+	// Create a target for the specific tables involved in the inspect checks
+	var tableIDSet catalog.DescriptorIDSet
+	for _, check := range details.Checks {
+		tableIDSet.Add(check.TableID)
+	}
+	target := ptpb.MakeSchemaObjectsTarget(tableIDSet.Ordered())
+
+	// Protect will store the PTS ID in job details.
+	_, err := execCfg.ProtectedTimestampManager.Protect(ctx, c.job, target, details.AsOf)
+	if err != nil {
+		return errors.Wrapf(err, "failed to protect timestamp %s for INSPECT job %d", details.AsOf, c.job.ID())
+	}
+
+	log.Dev.Infof(ctx, "protected timestamp created for INSPECT job %d at %s", c.job.ID(), details.AsOf)
+	return nil
+}
+
+// maybeCleanupProtectedTimestamp cleans up any protected timestamp record
+// associated with this job. If no protected timestamp was created, this
+// function does nothing.
+func (c *inspectResumer) maybeCleanupProtectedTimestamp(
+	ctx context.Context, execCfg *sql.ExecutorConfig,
+) {
+	details := c.job.Details().(jobspb.InspectDetails)
+	if details.ProtectedTimestampRecord != nil {
+		if err := execCfg.ProtectedTimestampManager.Unprotect(ctx, c.job); err != nil {
+			log.Dev.Warningf(ctx, "failed to clean up protected timestamp: %v", err)
+		}
+	}
+}
+
 func init() {
 	createResumerFn := func(job *jobs.Job, settings *cluster.Settings) jobs.Resumer {
 		return &inspectResumer{job: job}

pkg/sql/inspect/inspect_job_test.go

@@ -13,7 +13,10 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts"
 	"github.com/cockroachdb/cockroach/pkg/sql"
+	"github.com/cockroachdb/cockroach/pkg/sql/isql"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
@@ -145,3 +148,159 @@ func TestInspectJobImplicitTxnSemantics(t *testing.T) {
 		})
 	}
 }
+
+// TestInspectJobProtectedTimestamp verifies that INSPECT jobs properly create
+// and clean up protected timestamp records when using AS OF SYSTEM TIME.
+func TestInspectJobProtectedTimestamp(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	for _, tc := range []struct {
+		desc              string
+		forceJobFailure   bool
+		expectedJobStatus string
+		expectError       bool
+	}{
+		{
+			desc:              "job success with cleanup",
+			forceJobFailure:   false,
+			expectedJobStatus: "succeeded",
+			expectError:       false,
+		},
+		{
+			desc:              "job failure with cleanup",
+			forceJobFailure:   true,
+			expectedJobStatus: "failed",
+			expectError:       true,
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			var blockInspectExecution atomic.Bool
+			var protectedTimestampCreated atomic.Bool
+
+			ctx := context.Background()
+			s, db, _ := serverutils.StartServer(t, base.TestServerArgs{
+				Knobs: base.TestingKnobs{
+					Inspect: &sql.InspectTestingKnobs{
+						OnInspectAfterProtectedTimestamp: func() error {
+							protectedTimestampCreated.Store(true)
+							// Block execution until we've verified the protected timestamp
+							for blockInspectExecution.Load() {
+								time.Sleep(10 * time.Millisecond)
+							}
+							if tc.forceJobFailure {
+								return errors.New("forced job failure for testing")
+							}
+							return nil
+						},
+					},
+				},
+			})
+			defer s.Stopper().Stop(ctx)
+
+			runner := sqlutils.MakeSQLRunner(db)
+			runner.Exec(t, `
+				CREATE DATABASE db;
+				SET enable_scrub_job = true;
+				CREATE TABLE db.t (
+					id INT PRIMARY KEY,
+					val INT
+				);
+				CREATE INDEX i1 on db.t (val);
+				INSERT INTO db.t VALUES (1, 2), (2, 3);`)
+
+			// Start blocking inspection execution
+			blockInspectExecution.Store(true)
+
+			// Start INSPECT job with AS OF timestamp in a goroutine
+			errCh := make(chan error, 1)
+			go func() {
+				_, err := db.Exec("EXPERIMENTAL SCRUB TABLE db.t AS OF SYSTEM TIME '-1us'")
+				errCh <- err
+			}()
+
+			// Wait for the protected timestamp hook to be called
+			testutils.SucceedsSoon(t, func() error {
+				if !protectedTimestampCreated.Load() {
+					return errors.New("protected timestamp hook not called yet")
+				}
+				return nil
+			})
+
+			// Get the job ID
+			var jobID int64
+			runner.QueryRow(t, `
+				SELECT id
+				FROM crdb_internal.system_jobs
+				WHERE job_type = 'INSPECT' AND status = 'running'
+				ORDER BY created DESC
+				LIMIT 1
+			`).Scan(&jobID)
+
+			// Load the job and get protected timestamp record
+			execCfg := s.ApplicationLayer().ExecutorConfig().(sql.ExecutorConfig)
+			job, err := execCfg.JobRegistry.LoadJob(ctx, jobspb.JobID(jobID))
+			require.NoError(t, err)
+
+			details := job.Details().(jobspb.InspectDetails)
+			require.NotNil(t, details.ProtectedTimestampRecord, "protected timestamp record should be set")
+			protectedTSID := *details.ProtectedTimestampRecord
+
+			// Check that the protected timestamp record actually exists in the system
+			var recordExists bool
+			require.NoError(t, execCfg.InternalDB.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+				pts := execCfg.ProtectedTimestampProvider.WithTxn(txn)
+				_, err := pts.GetRecord(ctx, protectedTSID)
+				if err != nil {
+					if errors.Is(err, protectedts.ErrNotExists) {
+						recordExists = false
+						return nil
+					}
+					return err
+				}
+				recordExists = true
+				return nil
+			}))
+			require.True(t, recordExists, "protected timestamp record should exist in the system")
+
+			// Allow the job to complete
+			blockInspectExecution.Store(false)
+
+			// Wait for job to complete
+			select {
+			case err := <-errCh:
+				if tc.expectError {
+					require.Error(t, err, "INSPECT job should fail due to forced error")
+				} else {
+					require.NoError(t, err, "INSPECT job should complete successfully")
+				}
+			case <-time.After(30 * time.Second):
+				t.Fatal("INSPECT job did not complete within timeout")
+			}
+
+			// Verify job status
+			var status string
+			runner.QueryRow(t, `
+				SELECT status
+				FROM crdb_internal.system_jobs
+				WHERE id = $1
+			`, jobID).Scan(&status)
+			require.Equal(t, tc.expectedJobStatus, status, "job should have expected status")
+
+			// Verify protected timestamp record is cleaned up
+			testutils.SucceedsSoon(t, func() error {
+				return execCfg.InternalDB.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+					pts := execCfg.ProtectedTimestampProvider.WithTxn(txn)
+					_, err := pts.GetRecord(ctx, protectedTSID)
+					if err != nil {
+						if errors.Is(err, protectedts.ErrNotExists) {
+							return nil // This is what we want
+						}
+						return err
+					}
+					return errors.New("protected timestamp record still exists")
+				})
+			})
+		})
+	}
+}