Merge #154298 #154427 #154444

154298: crosscluster/physical: add alter connection watcher r=jeffswenson a=msbutler

PCR now automatically updates the uri used if it's apart of an external
connection that is altered.

Informs #153471

Release note: none

154427: admission: correctly reset diskReadTokensAllocated r=tbg a=tbg

Epic: none


154444: kvserver: deflake TestProcessSplitAfterRightHandSideHasBeenRemoved r=tbg a=pav-kv

The test relies on the fact that the LHS replica will be caught up via log, and run the split trigger. In the flake, the replica has been caught up by a snapshot. Likely cause is that the leader's log has been truncated while the replica has been partitioned away. The PR disables log truncation for this test, so that the catch up via log is guaranteed.

Fixes #154313

Co-authored-by: Michael Butler <[email protected]>
Co-authored-by: Tobias Grieger <[email protected]>
Co-authored-by: Pavel Kalinnikov <[email protected]>

Author: 4 people

pkg/crosscluster/logical/logical_replication_job.go

@@ -326,35 +326,13 @@ func (r *logicalReplicationResumer) ingest(
 		return err
 	}
 
-	refreshConn := func(ctx context.Context) error {
-		ingestionJob := r.job
-		details := ingestionJob.Details().(jobspb.LogicalReplicationDetails)
-		resolvedDest, err := resolveDest(ctx, jobExecCtx.ExecCfg(), details.SourceClusterConnUri)
-		if err != nil {
-			return err
-		}
-		pollingInterval := 2 * time.Minute
-		if knobs := jobExecCtx.ExecCfg().StreamingTestingKnobs; knobs != nil && knobs.ExternalConnectionPollingInterval != nil {
-			pollingInterval = *knobs.ExternalConnectionPollingInterval
-		}
-		t := time.NewTicker(pollingInterval)
-		defer t.Stop()
-		for {
-			select {
-			case <-ctx.Done():
-				return ctx.Err()
-			case <-confPoller:
-				return nil
-			case <-t.C:
-				newDest, err := reloadDest(ctx, ingestionJob.ID(), jobExecCtx.ExecCfg())
-				if err != nil {
-					log.Dev.Warningf(ctx, "failed to check for updated configuration: %v", err)
-				} else if newDest != resolvedDest {
-					return errors.Mark(errors.Newf("replan due to detail change: old=%s, new=%s", resolvedDest, newDest), sql.ErrPlanChanged)
-				}
-			}
-		}
-	}
+	refreshConn := replicationutils.GetAlterConnectionChecker(
+		r.job.ID(),
+		uris[0].Serialize(),
+		geURIFromLoadedJobDetails,
+		execCfg,
+		confPoller,
+	)
 
 	defer func() {
 		if l := payload.MetricsLabel; l != "" {
@@ -1089,29 +1067,8 @@ func getRetryPolicy(knobs *sql.StreamingTestingKnobs) retry.Options {
 	}
 }
 
-func resolveDest(
-	ctx context.Context, execCfg *sql.ExecutorConfig, sourceURI string,
-) (string, error) {
-	configUri, err := streamclient.ParseConfigUri(sourceURI)
-	if err != nil {
-		return "", err
-	}
-
-	clusterUri, err := configUri.AsClusterUri(ctx, execCfg.InternalDB)
-	if err != nil {
-		return "", err
-	}
-
-	return clusterUri.Serialize(), nil
-}
-
-func reloadDest(ctx context.Context, id jobspb.JobID, execCfg *sql.ExecutorConfig) (string, error) {
-	reloadedJob, err := execCfg.JobRegistry.LoadJob(ctx, id)
-	if err != nil {
-		return "", err
-	}
-	newDetails := reloadedJob.Details().(jobspb.LogicalReplicationDetails)
-	return resolveDest(ctx, execCfg, newDetails.SourceClusterConnUri)
+func geURIFromLoadedJobDetails(details jobspb.Details) string {
+	return details.(jobspb.LogicalReplicationDetails).SourceClusterConnUri
 }
 
 func init() {

pkg/crosscluster/physical/replication_stream_e2e_test.go

@@ -1800,3 +1800,62 @@ func splitPrimaryKeyIndexSpan(
 	require.NoError(t, db.AdminSplit(ctx, pkStartKey, hlc.MaxTimestamp))
 	require.NoError(t, db.AdminSplit(ctx, pkEndKey, hlc.MaxTimestamp))
 }
+
+func TestAlterExternalConnection(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	skip.UnderDeadlock(t)
+	skip.UnderRace(t)
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	pollingInterval := 100 * time.Millisecond
+
+	var alreadyReplanned atomic.Int32
+
+	args := replicationtestutils.DefaultTenantStreamingClustersArgs
+	args.TestingKnobs = &sql.StreamingTestingKnobs{
+		ExternalConnectionPollingInterval: &pollingInterval,
+		AfterRetryIteration: func(err error) {
+			alreadyReplanned.Add(1)
+		},
+	}
+
+	c, cleanup := replicationtestutils.CreateTenantStreamingClusters(ctx, t, args)
+	defer cleanup()
+
+	externalConnection := "replication-source-addr"
+	ogConnection := c.SrcURL.String()
+	c.DestSysSQL.Exec(c.T, fmt.Sprintf(`CREATE EXTERNAL CONNECTION "%s" AS "%s"`,
+		externalConnection, ogConnection))
+	c.DestSysSQL.Exec(c.T, c.BuildCreateTenantQuery(externalConnection))
+	streamProducerJobID, ingestionJobID := replicationtestutils.GetStreamJobIds(c.T, ctx, c.DestSysSQL, c.Args.DestTenantName)
+
+	jobutils.WaitForJobToRun(c.T, c.SrcSysSQL, jobspb.JobID(streamProducerJobID))
+	jobutils.WaitForJobToRun(c.T, c.DestSysSQL, jobspb.JobID(ingestionJobID))
+	c.WaitUntilStartTimeReached(jobspb.JobID(ingestionJobID))
+
+	// Alter the external connection to break the stream, to test that pcr watches for new connection changes
+	c.SrcSysSQL.Exec(t, fmt.Sprintf("CREATE USER %s", username.TestUser))
+	srcAppURL, cleanupSinkCert := pgurlutils.PGUrl(t, c.SrcSysServer.AdvSQLAddr(), t.Name(), url.User(username.TestUser))
+	defer cleanupSinkCert()
+
+	beforeReplanCount := alreadyReplanned.Load()
+	c.DestSysSQL.Exec(c.T, fmt.Sprintf(`ALTER EXTERNAL CONNECTION "%s" AS "%s"`,
+		externalConnection, &srcAppURL))
+
+	testutils.SucceedsSoon(t, func() error {
+		if alreadyReplanned.Load() <= beforeReplanCount+2 {
+			return errors.New("not yet replanned twice")
+		}
+		return nil
+	})
+
+	// Alter the external connection to fix the stream, and ensure replication resumes
+	c.DestSysSQL.Exec(c.T, fmt.Sprintf(`ALTER EXTERNAL CONNECTION "%s" AS "%s"`,
+		externalConnection, ogConnection))
+	c.DestSysSQL.Exec(c.T, fmt.Sprintf("RESUME JOB %d", ingestionJobID))
+	jobutils.WaitForJobToRun(c.T, c.DestSysSQL, jobspb.JobID(ingestionJobID))
+	// ensure the stream advances
+	srcTime := c.SrcCluster.Server(0).Clock().Now()
+	c.WaitUntilReplicatedTime(srcTime, jobspb.JobID(ingestionJobID))
+}

pkg/crosscluster/physical/stream_ingestion_dist.go

@@ -96,12 +96,16 @@ func startDistIngestion(
 		updateStatus(ctx, ingestionJob, jobspb.InitializingReplication, msg)
 	}
 
-	client, err := connectToActiveClient(ctx, ingestionJob, execCtx.ExecCfg().InternalDB,
-		streamclient.WithStreamID(streamID))
+	clusterUris, err := getClusterUris(ctx, ingestionJob, execCtx.ExecCfg().InternalDB)
+	if err != nil {
+		return err
+	}
+	client, err := streamclient.GetFirstActiveClient(ctx, clusterUris, execCtx.ExecCfg().InternalDB, streamclient.WithStreamID(streamID))
 	if err != nil {
 		return err
 	}
 	defer closeAndLog(ctx, client)
+
 	if err := waitUntilProducerActive(ctx, client, streamID, heartbeatTimestamp, ingestionJob.ID()); err != nil {
 		return err
 	}
@@ -197,11 +201,23 @@ func startDistIngestion(
 		}
 		return ingestor.ingestSpanConfigs(ctx, details.SourceTenantName)
 	}
+
+	refreshConnStopper := make(chan struct{})
+
+	refreshConn := replicationutils.GetAlterConnectionChecker(
+		ingestionJob.ID(),
+		clusterUris[0].Serialize(),
+		geURIFromIngestionJobDetails,
+		execCtx.ExecCfg(),
+		refreshConnStopper,
+	)
+
 	execInitialPlan := func(ctx context.Context) error {
 		defer func() {
 			stopReplanner()
 			close(tracingAggCh)
 			close(spanConfigIngestStopper)
+			close(refreshConnStopper)
 		}()
 		ctx = logtags.AddTag(ctx, "stream-ingest-distsql", nil)
 
@@ -273,13 +289,17 @@ func startDistIngestion(
 		return err
 	}
 
-	err = ctxgroup.GoAndWait(ctx, execInitialPlan, replanner, tracingAggLoop, streamSpanConfigs)
+	err = ctxgroup.GoAndWait(ctx, execInitialPlan, replanner, tracingAggLoop, streamSpanConfigs, refreshConn)
 	if errors.Is(err, sql.ErrPlanChanged) {
 		execCtx.ExecCfg().JobRegistry.MetricsStruct().StreamIngest.(*Metrics).ReplanCount.Inc(1)
 	}
 	return err
 }
 
+func geURIFromIngestionJobDetails(details jobspb.Details) string {
+	return details.(jobspb.StreamIngestionDetails).SourceClusterConnUri
+}
+
 func sortSpans(partitions []streamclient.PartitionInfo) roachpb.Spans {
 	spansToSort := make(roachpb.Spans, 0)
 	for i := range partitions {

pkg/crosscluster/replicationutils/BUILD.bazel

@@ -3,12 +3,14 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 go_library(
     name = "replicationutils",
     srcs = [
+        "connection_checker.go",
         "stats.go",
         "utils.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/crosscluster/replicationutils",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/crosscluster/streamclient",
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
         "//pkg/kv/kvpb",
@@ -33,6 +35,7 @@ go_library(
         "//pkg/testutils/fingerprintutils",
         "//pkg/util/ctxgroup",
         "//pkg/util/hlc",
+        "//pkg/util/log",
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",

pkg/crosscluster/replicationutils/connection_checker.go

@@ -0,0 +1,74 @@
+// Copyright 2022 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package replicationutils
+
+import (
+	"context"
+	"time"
+
+	"github.com/cockroachdb/cockroach/pkg/crosscluster/streamclient"
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/sql"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/errors"
+)
+
+// GetAlterConnectionChecker returns a function that will poll for an altered
+// external connection. The initial URI passed to this function should be the
+// same one passed to subsequent client creation calls.
+func GetAlterConnectionChecker(
+	id jobspb.JobID,
+	initialURI string,
+	uriGetter URIGetter,
+	execCfg *sql.ExecutorConfig,
+	stopper chan struct{},
+) func(ctx context.Context) error {
+	return func(ctx context.Context) error {
+		pollingInterval := 2 * time.Minute
+		if knobs := execCfg.StreamingTestingKnobs; knobs != nil && knobs.ExternalConnectionPollingInterval != nil {
+			pollingInterval = *knobs.ExternalConnectionPollingInterval
+		}
+		t := time.NewTicker(pollingInterval)
+		defer t.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-stopper:
+				return nil
+			case <-t.C:
+				reloadedJob, err := execCfg.JobRegistry.LoadJob(ctx, id)
+				if err != nil {
+					return err
+				}
+				newURI, err := resolveURI(ctx, execCfg, uriGetter(reloadedJob.Details()))
+				if err != nil {
+					log.Dev.Warningf(ctx, "failed to load uri: %v", err)
+				} else if newURI != initialURI {
+					return errors.Mark(errors.Newf("uri has been updated: old=%s, new=%s", errors.Redact(initialURI), errors.Redact(newURI)), sql.ErrPlanChanged)
+				}
+			}
+		}
+	}
+}
+
+type URIGetter func(details jobspb.Details) string
+
+func resolveURI(
+	ctx context.Context, execCfg *sql.ExecutorConfig, sourceURI string,
+) (string, error) {
+	configUri, err := streamclient.ParseConfigUri(sourceURI)
+	if err != nil {
+		return "", err
+	}
+
+	clusterUri, err := configUri.AsClusterUri(ctx, execCfg.InternalDB)
+	if err != nil {
+		return "", err
+	}
+
+	return clusterUri.Serialize(), nil
+}

pkg/kv/kvserver/client_raft_test.go

@@ -33,6 +33,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/rditer"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/stateloader"
 	raft "github.com/cockroachdb/cockroach/pkg/raft"
@@ -5326,13 +5327,13 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 			t *testing.T, store *kvserver.Store, rangeID roachpb.RangeID,
 		) raftpb.HardState {
 			t.Helper()
-			hs, err := stateloader.Make(rangeID).LoadHardState(ctx, store.TODOEngine())
+			hs, err := logstore.NewStateLoader(rangeID).LoadHardState(ctx, store.LogEngine())
 			require.NoError(t, err)
 			return hs
 		}
 		partitionReplicaOnSplit := func(t *testing.T, tc *testcluster.TestCluster, key roachpb.Key, basePartition *testClusterPartitionedRange, partRange **testClusterPartitionedRange) {
-			// Set up a hook to partition the RHS range at its initial range ID
-			// before proposing the split trigger.
+			// Set up a hook to partition away the first store of the RHS range at the
+			// first opportunity (when the split trigger is proposed).
 			var setupOnce sync.Once
 			f := kvserverbase.ReplicaProposalFilter(func(args kvserverbase.ProposalFilterArgs) *kvpb.Error {
 				req, ok := args.Req.GetArg(kvpb.EndTxn)
@@ -5363,9 +5364,10 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 			proposalFilter.Store(f)
 		}
 
-		// The basic setup for all of these tests are that we have a LHS range on 3
-		// nodes and we've partitioned store 0 for the LHS range. The tests will now
-		// perform a split, remove the RHS, add it back and validate assumptions.
+		// The basic setup for all of these tests are that we have an LHS range on 3
+		// nodes (lease on the last one), and we've partitioned store 0 for the LHS
+		// range. The tests will now perform a split, remove the RHS, add it back
+		// and validate assumptions.
 		//
 		// Different outcomes will occur depending on whether and how the RHS is
 		// partitioned and whether the server is killed. In all cases we want the
@@ -5400,6 +5402,11 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 							// Newly-started stores (including the "rogue" one) should not GC
 							// their replicas. We'll turn this back on when needed.
 							DisableReplicaGCQueue: true,
+							// Some subtests, e.g. (4), expect that n1 catches up on the raft
+							// log after a partition, and runs the split trigger. Disable raft
+							// log truncation to make sure that it doesn't miss the split
+							// trigger by being caught up by a later snapshot. See #154313.
+							DisableRaftLogQueue:   true,
 							TestingProposalFilter: testingProposalFilter,
 						},
 					},
@@ -5486,9 +5493,9 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 			target := tc.Target(0)
 			log.KvExec.Infof(ctx, "removing voter: %v", target)
 			tc.RemoveVotersOrFatal(t, keyB, target)
-			// Unsuccessful because the RHS will not accept the learner snapshot
-			// and will be rolled back. Nevertheless it will have learned that it
-			// has been removed at the old replica ID.
+			// Unsuccessful because the RHS will not accept the learner snapshot and
+			// will be rolled back. Nevertheless, it will have learned that it has
+			// been removed at the old replica ID.
 			_, err = tc.Servers[0].DB().AdminChangeReplicas(
 				ctx, keyB, tc.LookupRangeOrFatal(t, keyB),
 				kvpb.MakeReplicationChanges(roachpb.ADD_VOTER, target),
@@ -5701,8 +5708,8 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 			target := tc.Target(0)
 			log.KvExec.Infof(ctx, "removing voter: %v", target)
 			tc.RemoveVotersOrFatal(t, keyB, target)
-			// Unsuccessfuly because the RHS will not accept the learner snapshot
-			// and will be rolled back. Nevertheless it will have learned that it
+			// Unsuccessful because the RHS will not accept the learner snapshot
+			// and will be rolled back. Nevertheless, it will have learned that it
 			// has been removed at the old replica ID.
 			//
 			// Not using tc.AddVoters because we expect an error, but that error
@@ -5722,7 +5729,7 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 			log.KvExec.Infof(ctx, "added %d to RHS partition: %s", rhsInfo.Desc.NextReplicaID, rhsPartition)
 
 			// We do all of this incrementing to ensure that nobody will ever
-			// succeed in sending a message the new RHS replica after we restart
+			// succeed in sending a message to the new RHS replica after we restart
 			// the store. Previously there were races which could happen if we
 			// stopped the store immediately. Sleeps worked but this feels somehow
 			// more principled.

pkg/sql/exec_util.go

@@ -2176,8 +2176,6 @@ type StreamingTestingKnobs struct {
 	// for whether the job record is updated on a progress update.
 	CutoverProgressShouldUpdate func() bool
 
-	//ExternalConnectionPollingInterval override the LDR alter
-	// connection polling frequency.
 	ExternalConnectionPollingInterval *time.Duration
 
 	DistSQLRetryPolicy *retry.Options

pkg/util/admission/io_load_listener.go

@@ -745,7 +745,7 @@ func (io *ioLoadListener) adjustTokens(ctx context.Context, metrics StoreMetrics
 		io.diskWriteTokens = tokens.writeByteTokens
 		io.diskWriteTokensAllocated = 0
 		io.diskReadTokens = tokens.readByteTokens
-		io.diskWriteTokensAllocated = 0
+		io.diskReadTokensAllocated = 0
 	}
 	io.diskBandwidthLimiter.unlimitedTokensOverride = false
 	if metrics.DiskStats.ProvisionedBandwidth == 0 ||