crosscluster/physical: slim down TestProtectedTimestampManagement

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. This slashes the test's scope to just verify PCR's direct
behavior.

Release note: none.
Epic: none.

Author: dt

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)
+		})
+	}
+}