Merge #155063 #155162

155063: server: return draining progress if the server controller is draining r=iskettaneh a=iskettaneh

This commit fixes a bug where if the server controller is being drained, the drain command would return a response without populating the IsDraining. The CLI interprets this as if the draining is completed.

server: return draining progress if the server controller is draining

This commit fixes a bug where if the server controller is being drained, the drain command would return a response without populating the IsDraining field. The CLI interprets this as if the draining is completed.

Output example when running:

```
./cockroach node drain 3 --insecure --url {pgurl:3} --logtostderr=INFO

I251008 15:54:42.613200 15 2@rpc/peer.go:613  [rnode=?,raddr=10.142.1.228:26257,class=system,rpc] 1  connection is now healthy
node is draining... remaining: 3
I251008 15:54:42.622586 1 cli/rpc_node_shutdown.go:184  [-] 2  drain details: tenant servers: 3
node is draining... remaining: 3
I251008 15:54:42.824526 1 cli/rpc_node_shutdown.go:184  [-] 3  drain details: tenant servers: 3
node is draining... remaining: 1
I251008 15:54:43.026405 1 cli/rpc_node_shutdown.go:184  [-] 4  drain details: tenant servers: 1
node is draining... remaining: 1
I251008 15:54:43.228596 1 cli/rpc_node_shutdown.go:184  [-] 5  drain details: tenant servers: 1
node is draining... remaining: 243
I251008 15:54:44.580413 1 cli/rpc_node_shutdown.go:184  [-] 6  drain details: liveness record: 1, range lease iterations: 175, descriptor leases: 67
node is draining... remaining: 0 (complete)
drain ok
```

Release note (bug fix): fixed a bug in the drain command where draining a node using virtual clusters (such as clusters running Physical Cluster Replication) could return before the drain was complete, possibly resulting in shutting down the node while it still had active SQL clients and ranges leases.

Epic: None

155162: sql: bump up cluster setting read timeout for benchmark tests r=fqazi a=fqazi

Previously, we would see flakes on TestBenchmarkExpectation related to updating cluster settings and not being able to see the updated values after SET CLUSTER SETTING. We suspect this flake happens because under load, it may take longer for the cluster setting update to be visible. To address this, this patch modifies the cluster setting timeout for benchmark tests to prevent these flakes.

Fixes: #154219

Release note: None

Co-authored-by: iskettaneh <[email protected]>
Co-authored-by: Faizan Qazi <[email protected]>

Author: 3 people

pkg/server/drain.go

@@ -389,10 +389,13 @@ func (s *drainServer) drainInner(
 	return s.drainNode(ctx, reporter, verbose)
 }
 
-// isDraining returns true if either SQL client connections are being drained
-// or if one of the stores on the node is not accepting replicas.
+// isDraining returns true if the SQL client connections are being drained,
+// if one of the stores on the node is not accepting replicas, or if the
+// server controller is being drained.
 func (s *drainServer) isDraining() bool {
-	return s.sqlServer.pgServer.IsDraining() || (s.kvServer.node != nil && s.kvServer.node.IsDraining())
+	return s.sqlServer.pgServer.IsDraining() ||
+		(s.kvServer.node != nil && s.kvServer.node.IsDraining()) ||
+		(s.serverCtl != nil && s.serverCtl.IsDraining())
 }
 
 // drainClients starts draining the SQL layer.
@@ -410,7 +413,10 @@ func (s *drainServer) drainClientsInternal(
 	var cancel context.CancelFunc
 	ctx, cancel = context.WithCancel(ctx)
 	defer cancel()
-	shouldDelayDraining := !s.isDraining()
+
+	// If this is the first time we are draining the clients, we delay draining
+	// to allow health probes to notice that the node is not ready.
+	shouldDelayDraining := !s.sqlServer.pgServer.IsDraining()
 
 	// Set the gRPC mode of the node to "draining" and mark the node as "not ready".
 	// Probes to /health?ready=1 will now notice the change in the node's readiness.

pkg/server/drain_test.go

@@ -7,6 +7,7 @@ package server_test
 
 import (
 	"context"
+	gosql "database/sql"
 	"io"
 	"testing"
 	"time"
@@ -20,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats/persistedsqlstats/sqlstatstestutil"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
 	"github.com/cockroachdb/cockroach/pkg/util/grpcutil"
@@ -34,15 +36,39 @@ import (
 func TestDrain(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
-	doTestDrain(t)
+	testutils.RunTrueAndFalse(t, "secondaryTenants", func(t *testing.T, secondaryTenants bool) {
+		doTestDrain(t, secondaryTenants)
+	})
 }
 
 // doTestDrain runs the drain test.
-func doTestDrain(tt *testing.T) {
+func doTestDrain(tt *testing.T, secondaryTenants bool) {
+	if secondaryTenants {
+		// Draining the tenant server takes a long time under stress. See #155229.
+		skip.UnderRace(tt)
+	}
 	var drainSleepCallCount = 0
 	t := newTestDrainContext(tt, &drainSleepCallCount)
 	defer t.Close()
 
+	var tenantConn *gosql.DB
+	if secondaryTenants {
+		_, err := t.tc.ServerConn(0).Exec("CREATE TENANT hello")
+		require.NoError(tt, err)
+		_, err = t.tc.ServerConn(0).Exec("ALTER TENANT hello START SERVICE SHARED")
+		require.NoError(tt, err)
+
+		// Wait for the tenant to be ready by establishing a test connection.
+		testutils.SucceedsSoon(tt, func() error {
+			tenantConn, err = t.tc.Server(0).SystemLayer().SQLConnE(
+				serverutils.DBName("cluster:hello/defaultdb"))
+			if err != nil {
+				return err
+			}
+			return tenantConn.Ping()
+		})
+	}
+
 	// Issue a probe. We're not draining yet, so the probe should
 	// reflect that.
 	resp := t.sendProbe()
@@ -54,7 +80,12 @@ func doTestDrain(tt *testing.T) {
 	resp = t.sendDrainNoShutdown()
 	t.assertDraining(resp, true)
 	t.assertRemaining(resp, true)
-	t.assertEqual(1, drainSleepCallCount)
+	if !secondaryTenants {
+		// If we have secondary tenants, we might not have waited before draining
+		// the clients at this point because we might have returned early if we are
+		// still draining the serverController.
+		t.assertEqual(1, drainSleepCallCount)
+	}
 
 	// Issue another probe. This checks that the server is still running
 	// (i.e. Shutdown: false was effective), the draining status is
@@ -64,12 +95,14 @@ func doTestDrain(tt *testing.T) {
 	t.assertDraining(resp, true)
 	// probe-only has no remaining.
 	t.assertRemaining(resp, false)
-	t.assertEqual(1, drainSleepCallCount)
-
+	if !secondaryTenants {
+		t.assertEqual(1, drainSleepCallCount)
+	}
 	// Repeat drain commands until we verify that there are zero remaining leases
 	// (i.e. complete). Also validate that the server did not sleep again.
 	testutils.SucceedsSoon(t, func() error {
 		resp = t.sendDrainNoShutdown()
+		t.Logf("drain response: %+v", resp)
 		if !resp.IsDraining {
 			return errors.Newf("expected draining")
 		}
@@ -79,6 +112,7 @@ func doTestDrain(tt *testing.T) {
 		}
 		return nil
 	})
+
 	t.assertEqual(1, drainSleepCallCount)
 
 	// Now issue a drain request without drain but with shutdown.
@@ -190,6 +224,7 @@ func newTestDrainContext(t *testing.T, drainSleepCallCount *int) *testDrainConte
 			// care about TLS settings for the RPC client connection.
 			ServerArgs: base.TestServerArgs{
 				DefaultDRPCOption: base.TestDRPCDisabled,
+				DefaultTestTenant: base.TestControlsTenantsExplicitly,
 				Knobs: base.TestingKnobs{
 					Server: &server.TestingKnobs{
 						DrainSleepFn: func(time.Duration) {

pkg/server/server_controller.go

@@ -172,6 +172,10 @@ func (s *serverController) SetDraining() {
 	}
 }
 
+func (s *serverController) IsDraining() bool {
+	return s.draining.Load()
+}
+
 // start monitors changes to the service mode and updates
 // the running servers accordingly.
 func (c *serverController) start(ctx context.Context, ie isql.Executor) error {

pkg/sql/set_cluster_setting.go

@@ -40,6 +40,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessioninit"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqltelemetry"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/log/eventpb"
@@ -701,7 +702,12 @@ func waitForSettingUpdate(
 	}
 	errNotReady := errors.New("setting updated but timed out waiting to read new value")
 	var observed string
-	err := retry.ForDuration(10*time.Second, func() error {
+	defaultDuration := 10 * time.Second
+	// Bench tests maybe more prone to flaking, so use a longer time limit for them.
+	if buildutil.CrdbBenchBuild {
+		defaultDuration *= 3
+	}
+	err := retry.ForDuration(defaultDuration, func() error {
 		observed = setting.Encoded(&execCfg.Settings.SV)
 		if observed != expectedEncodedValue {
 			return errNotReady