Merge #160229

160229: crosscluster/physical: slim down TestProtectedTimestampManagement r=dt a=dt

This test was doing far too much for a unit test, running multiple test clusters and also waiting for multiple systems like PCR, GC, etc to all interact. In particular, verifying GC behaviors whenGC  TTLs have a minimum duration measured in seconds forced this test to have outsized delays in it -- it often took >30s, even when run on its own, let alone in a larger, more contended test run.

This slashes the test's scope to just verify PCR's direct behavior, cutting its runtime substantially. 

Release note: none.
Epic: none.

Co-authored-by: David Taylor <[email protected]>

Author: craig[bot]

pkg/crosscluster/physical/BUILD.bazel

@@ -151,7 +151,6 @@ go_test(
         "//pkg/kv/kvpb",
         "//pkg/kv/kvserver",
         "//pkg/kv/kvserver/protectedts",
-        "//pkg/kv/kvserver/protectedts/ptpb",
         "//pkg/repstream/streampb",
         "//pkg/roachpb",
         "//pkg/security/securityassets",

pkg/crosscluster/physical/replication_stream_e2e_test.go

@@ -26,7 +26,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts"
-	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts/ptpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/server"
@@ -952,172 +951,105 @@ func TestStreamingReplanOnLag(t *testing.T) {
 	c.WaitUntilReplicatedTime(cutoverTime, jobspb.JobID(ingestionJobID))
 }
 
-// TestProtectedTimestampManagement tests the active protected
-// timestamps management on the destination tenant's keyspan.
+// TestProtectedTimestampManagement tests that the destination tenant's
+// protected timestamp is created when the ingestion job starts and
+// released when the job completes or is canceled.
+//
+// Note: The PTS formula (protectedTS = replicatedTime - TTL) is unit tested
+// in TestResolveHeartbeatTime. Producer PTS management is tested in
+// TestTenantStreamingCancelIngestion.
 func TestProtectedTimestampManagement(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
-	skip.UnderRace(t, "slow test") // takes >1mn under race.
-
 	ctx := context.Background()
 	args := replicationtestutils.DefaultTenantStreamingClustersArgs
-	// Override the replication job details ReplicationTTLSeconds to a small value
-	// so that every progress update results in a protected timestamp update.
-	args.RetentionTTLSeconds = 1
-
-	testutils.RunTrueAndFalse(t, "pause-before-terminal", func(t *testing.T, pauseBeforeTerminal bool) {
-		testutils.RunTrueAndFalse(t, "complete-replication", func(t *testing.T, completeReplication bool) {
-
-			// waitForProducerProtection asserts that there is a PTS record protecting
-			// the source tenant. We ensure the PTS record is protecting a timestamp
-			// greater or equal to the frontier we know we have replicated up until.
-			waitForProducerProtection := func(c *replicationtestutils.TenantStreamingClusters, frontier hlc.Timestamp, producerJobID int) {
-				testutils.SucceedsSoon(t, func() error {
-					srv := c.SrcSysServer
-					job, err := srv.JobRegistry().(*jobs.Registry).LoadJob(ctx, jobspb.JobID(producerJobID))
-					if err != nil {
-						return err
-					}
-					ptsRecordID := job.Payload().Details.(*jobspb.Payload_StreamReplication).StreamReplication.ProtectedTimestampRecordID
-					ptsProvider := srv.ExecutorConfig().(sql.ExecutorConfig).ProtectedTimestampProvider
-
-					var ptsRecord *ptpb.Record
-					if err := srv.InternalDB().(descs.DB).Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
-						var err error
-						ptsRecord, err = ptsProvider.WithTxn(txn).GetRecord(ctx, ptsRecordID)
-						return err
-					}); err != nil {
-						return err
-					}
-					if ptsRecord.Timestamp.Less(frontier) {
-						return errors.Newf("protection is at %s, expected to be >= %s",
-							ptsRecord.Timestamp.String(), frontier.String())
-					}
-					return nil
-				})
-			}
-
-			// checkNoDestinationProtections asserts that there is no PTS record
-			// protecting the destination tenant.
-			checkNoDestinationProtection := func(c *replicationtestutils.TenantStreamingClusters, replicationJobID int) {
-				execCfg := c.DestSysServer.ExecutorConfig().(sql.ExecutorConfig)
-				require.NoError(t, c.DestSysServer.InternalDB().(isql.DB).Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
-					j, err := execCfg.JobRegistry.LoadJobWithTxn(ctx, jobspb.JobID(replicationJobID), txn)
-					require.NoError(t, err)
-					payload := j.Payload()
-					replicationDetails := payload.GetStreamIngestion()
-					ptp := execCfg.ProtectedTimestampProvider.WithTxn(txn)
-					_, err = ptp.GetRecord(ctx, *replicationDetails.ProtectedTimestampRecordID)
-					require.EqualError(t, err, protectedts.ErrNotExists.Error())
-					return nil
-				}))
-			}
-			checkDestinationProtection := func(c *replicationtestutils.TenantStreamingClusters, frontier hlc.Timestamp, replicationJobID int) {
-				execCfg := c.DestSysServer.ExecutorConfig().(sql.ExecutorConfig)
-				ptp := execCfg.ProtectedTimestampProvider
-				require.NoError(t, c.DestSysServer.InternalDB().(isql.DB).Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
-					j, err := execCfg.JobRegistry.LoadJobWithTxn(ctx, jobspb.JobID(replicationJobID), txn)
-					if err != nil {
-						return err
-					}
-					payload := j.Payload()
-					progress := j.Progress()
-					replicationDetails := payload.GetStreamIngestion()
-
-					require.NotNil(t, replicationDetails.ProtectedTimestampRecordID)
-					rec, err := ptp.WithTxn(txn).GetRecord(ctx, *replicationDetails.ProtectedTimestampRecordID)
-					if err != nil {
-						return err
-					}
-
-					replicatedTime := replicationutils.ReplicatedTimeFromProgress(&progress)
-					require.True(t, frontier.LessEq(replicatedTime))
-
-					roundedReplicatedTime := replicatedTime.GoTime().Round(time.Millisecond)
-					roundedProtectedTime := rec.Timestamp.GoTime().Round(time.Millisecond)
-					window := roundedReplicatedTime.Sub(roundedProtectedTime)
-					require.Equal(t, time.Second, window)
-					return nil
-				}))
-			}
+	// Disable metamorphic external connection since we reuse clusters across
+	// subtests and the external connection would already exist on subsequent runs.
+	args.NoMetamorphicExternalConnection = true
 
-			c, cleanup := replicationtestutils.CreateTenantStreamingClusters(ctx, t, args)
-			defer cleanup()
-
-			c.DestSysSQL.Exec(t, "SET CLUSTER SETTING kv.closed_timestamp.target_duration = '100ms'")
-			c.DestSysSQL.Exec(t, "SET CLUSTER SETTING kv.protectedts.reconciliation.interval = '1ms';")
-
-			producerJobID, replicationJobID := c.StartStreamReplication(ctx)
-
-			jobutils.WaitForJobToRun(c.T, c.SrcSysSQL, jobspb.JobID(producerJobID))
-			jobutils.WaitForJobToRun(c.T, c.DestSysSQL, jobspb.JobID(replicationJobID))
-
-			// Ensure that we wait at least a second so that the gap between the first
-			// time we write the protected timestamp (t1) during replication job
-			// startup, and the first progress update (t2) is greater than 1s. This is
-			// important because if `frontier@t2 - ReplicationTTLSeconds < t1` then we
-			// will not update the PTS record.
-			now := c.SrcCluster.Server(0).SystemLayer().Clock().Now().Add(int64(time.Second)*2, 0)
-			c.WaitUntilReplicatedTime(now, jobspb.JobID(replicationJobID))
-
-			// Check that the producer and replication job have written a protected
-			// timestamp.
-			waitForProducerProtection(c, now, producerJobID)
-			checkDestinationProtection(c, now, replicationJobID)
-
-			now2 := now.Add(time.Second.Nanoseconds(), 0)
-			c.WaitUntilReplicatedTime(now2, jobspb.JobID(replicationJobID))
-			// Let the replication progress for a second before checking that the
-			// protected timestamp record has also been updated on the destination
-			// cluster. This update happens in the same txn in which we update the
-			// replication job's progress.
-			waitForProducerProtection(c, now2, producerJobID)
-			checkDestinationProtection(c, now2, replicationJobID)
-
-			if pauseBeforeTerminal {
-				c.DestSysSQL.Exec(t, fmt.Sprintf("PAUSE JOB %d", replicationJobID))
-				jobutils.WaitForJobToPause(c.T, c.DestSysSQL, jobspb.JobID(replicationJobID))
-			}
+	// Create clusters once and reuse them across all subtests.
+	c, cleanup := replicationtestutils.CreateTenantStreamingClusters(ctx, t, args)
+	defer cleanup()
 
-			if completeReplication {
-				c.DestSysSQL.Exec(t, fmt.Sprintf("RESUME JOB %d", replicationJobID))
-				jobutils.WaitForJobToRun(c.T, c.DestSysSQL, jobspb.JobID(replicationJobID))
-				var emptyCutoverTime time.Time
-				c.Cutover(ctx, producerJobID, replicationJobID, emptyCutoverTime, false)
-				c.SrcSysSQL.Exec(t, fmt.Sprintf(`ALTER TENANT '%s' SET REPLICATION SOURCE EXPIRATION WINDOW ='100ms'`, c.Args.SrcTenantName))
-			}
+	// checkDestinationPTSExists verifies that a PTS record exists for the
+	// destination tenant's keyspan.
+	checkDestinationPTSExists := func(t *testing.T, replicationJobID int) {
+		t.Helper()
+		execCfg := c.DestSysServer.ExecutorConfig().(sql.ExecutorConfig)
+		testutils.SucceedsSoon(t, func() error {
+			return c.DestSysServer.InternalDB().(isql.DB).Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+				j, err := execCfg.JobRegistry.LoadJobWithTxn(ctx, jobspb.JobID(replicationJobID), txn)
+				if err != nil {
+					return err
+				}
+				payload := j.Payload()
+				replicationDetails := payload.GetStreamIngestion()
+				if replicationDetails.ProtectedTimestampRecordID == nil {
+					return errors.New("ProtectedTimestampRecordID not yet set")
+				}
+				ptp := execCfg.ProtectedTimestampProvider.WithTxn(txn)
+				_, err = ptp.GetRecord(ctx, *replicationDetails.ProtectedTimestampRecordID)
+				return err
+			})
+		})
+	}
 
-			// Set GC TTL low, so that the GC job completes quickly in the test.
-			c.DestSysSQL.Exec(t, "ALTER RANGE tenants CONFIGURE ZONE USING gc.ttlseconds = 1;")
-			c.DestSysSQL.Exec(t, fmt.Sprintf("DROP TENANT %s", c.Args.DestTenantName))
+	// checkNoDestinationProtection asserts that the PTS record has been released.
+	checkNoDestinationProtection := func(t *testing.T, replicationJobID int) {
+		t.Helper()
+		execCfg := c.DestSysServer.ExecutorConfig().(sql.ExecutorConfig)
+		testutils.SucceedsSoon(t, func() error {
+			return c.DestSysServer.InternalDB().(isql.DB).Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+				j, err := execCfg.JobRegistry.LoadJobWithTxn(ctx, jobspb.JobID(replicationJobID), txn)
+				if err != nil {
+					return err
+				}
+				payload := j.Payload()
+				replicationDetails := payload.GetStreamIngestion()
+				ptp := execCfg.ProtectedTimestampProvider.WithTxn(txn)
+				_, err = ptp.GetRecord(ctx, *replicationDetails.ProtectedTimestampRecordID)
+				if err == nil {
+					return errors.New("PTS record still exists")
+				}
+				if !errors.Is(err, protectedts.ErrNotExists) {
+					return err
+				}
+				return nil
+			})
+		})
+	}
 
-			if !completeReplication {
-				c.SrcSysSQL.Exec(t, fmt.Sprintf(`ALTER TENANT '%s' SET REPLICATION SOURCE EXPIRATION WINDOW ='1ms'`, c.Args.SrcTenantName))
-				jobutils.WaitForJobToCancel(c.T, c.DestSysSQL, jobspb.JobID(replicationJobID))
-				jobutils.WaitForJobToFail(c.T, c.SrcSysSQL, jobspb.JobID(producerJobID))
-			}
+	// Use a counter to assign unique destination tenant IDs for each subtest.
+	destTenantCounter := 0
+	testutils.RunTrueAndFalse(t, "complete-replication", func(t *testing.T, completeReplication bool) {
+		// Assign a unique destination tenant for this subtest.
+		destTenantCounter++
+		destTenantID := roachpb.MustMakeTenantID(uint64(destTenantCounter + 1))
+		destTenantName := roachpb.TenantName(fmt.Sprintf("destination%d", destTenantCounter))
+		c.Args.DestTenantID = destTenantID
+		c.Args.DestTenantName = destTenantName
 
-			requireReleasedProducerPTSRecord(t, ctx, c.SrcSysServer, jobspb.JobID(producerJobID))
+		producerJobID, replicationJobID := c.StartStreamReplication(ctx)
 
-			// Check if the replication job has released protected timestamp.
-			checkNoDestinationProtection(c, replicationJobID)
+		jobutils.WaitForJobToRun(t, c.SrcSysSQL, jobspb.JobID(producerJobID))
+		jobutils.WaitForJobToRun(t, c.DestSysSQL, jobspb.JobID(replicationJobID))
 
-			// Wait for the GC job to finish, this should happen once the protected
-			// timestamp has been released.
-			c.DestSysSQL.Exec(t, "SHOW JOBS WHEN COMPLETE SELECT job_id FROM [SHOW JOBS] WHERE job_type = 'SCHEMA CHANGE GC'")
+		// Verify PTS exists on the destination.
+		checkDestinationPTSExists(t, replicationJobID)
 
-			// Check if dest tenant key range is cleaned up.
-			destTenantSpan := keys.MakeTenantSpan(args.DestTenantID)
-			rows, err := c.DestSysServer.DB().
-				Scan(ctx, destTenantSpan.Key, destTenantSpan.EndKey, 10)
-			require.NoError(t, err)
-			require.Empty(t, rows)
+		if completeReplication {
+			// Complete replication via cutover.
+			var emptyCutoverTime time.Time
+			c.Cutover(ctx, producerJobID, replicationJobID, emptyCutoverTime, false)
+		} else {
+			// Cancel replication.
+			c.DestSysSQL.Exec(t, fmt.Sprintf("CANCEL JOB %d", replicationJobID))
+			jobutils.WaitForJobToCancel(t, c.DestSysSQL, jobspb.JobID(replicationJobID))
+		}
 
-			c.DestSysSQL.CheckQueryResults(t,
-				fmt.Sprintf("SELECT count(*) FROM system.tenants WHERE id = %s", args.DestTenantID),
-				[][]string{{"0"}})
-		})
+		// Verify PTS is released after job completion/cancellation.
+		checkNoDestinationProtection(t, replicationJobID)
 	})
 }
 

pkg/crosscluster/replicationutils/utils_test.go

@@ -134,3 +134,70 @@ func TestScanSST(t *testing.T) {
 	checkScan(roachpb.Span{Key: roachpb.Key("da"), EndKey: roachpb.Key("e")},
 		[]storage.MVCCKeyValue{}, []storage.MVCCRangeKey{})
 }
+
+func TestResolveHeartbeatTime(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	// Helper to create timestamps more easily.
+	ts := func(seconds int64) hlc.Timestamp {
+		return hlc.Timestamp{WallTime: seconds * 1e9}
+	}
+
+	testCases := []struct {
+		name                 string
+		replicatedTime       hlc.Timestamp
+		replicationStartTime hlc.Timestamp
+		cutoverTime          hlc.Timestamp
+		ttlSeconds           int32
+		expected             hlc.Timestamp
+	}{
+		{
+			name:                 "basic formula: replicated - TTL",
+			replicatedTime:       ts(100),
+			replicationStartTime: ts(50),
+			cutoverTime:          hlc.Timestamp{},
+			ttlSeconds:           10,
+			expected:             ts(90),
+		},
+		{
+			name:                 "clamped to start time when result would be earlier",
+			replicatedTime:       ts(55),
+			replicationStartTime: ts(50),
+			cutoverTime:          hlc.Timestamp{},
+			ttlSeconds:           10,
+			expected:             ts(50), // 55-10=45, but clamped to 50
+		},
+		{
+			name:                 "clamped to cutover time when cutover is earlier",
+			replicatedTime:       ts(100),
+			replicationStartTime: ts(50),
+			cutoverTime:          ts(85),
+			ttlSeconds:           10,
+			expected:             ts(85), // 100-10=90, but clamped to cutover 85
+		},
+		{
+			name:                 "cutover later than calculated time has no effect",
+			replicatedTime:       ts(100),
+			replicationStartTime: ts(50),
+			cutoverTime:          ts(95),
+			ttlSeconds:           10,
+			expected:             ts(90), // 100-10=90, cutover 95 doesn't affect
+		},
+		{
+			name:                 "zero TTL returns replicated time",
+			replicatedTime:       ts(100),
+			replicationStartTime: ts(50),
+			cutoverTime:          hlc.Timestamp{},
+			ttlSeconds:           0,
+			expected:             ts(100),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result := ResolveHeartbeatTime(tc.replicatedTime, tc.replicationStartTime, tc.cutoverTime, tc.ttlSeconds)
+			require.Equal(t, tc.expected, result)
+		})
+	}
+}