server: TestCheckRestartSafe omits descriptor leak check

These tests were flaking due to the descriptor leak detection during
draining that was implemented for tests. This check is unnecessary for
the endpoints that are being tested here.

A separate method was added to the test server to omit this assertion
called `DrainClientsWithoutLeakCheck` this calls a new internal helper
on the `drainServer` that omits the check.

Resolves: #151377
Resolves: #151455
Epic: None
Release note: None

Author: dhartunian

pkg/server/api_v2_test.go

@@ -547,7 +547,7 @@ func drain(ctx context.Context, ts1 serverutils.TestServerInterface, t *testing.
 	timeoutCtx, cancel := context.WithTimeout(ctx, testutils.SucceedsSoonDuration())
 	defer cancel()
 
-	err := ts1.DrainClients(ctx)
+	err := ts1.DrainClientsWithoutLeakCheck(ctx)
 	require.NoError(t, err)
 
 	for timeoutCtx.Err() == nil {

pkg/server/drain.go

@@ -398,6 +398,12 @@ func (s *drainServer) isDraining() bool {
 // drainClients starts draining the SQL layer.
 func (s *drainServer) drainClients(
 	ctx context.Context, reporter func(int, redact.SafeString),
+) error {
+	return s.drainClientsInternal(ctx, reporter, true /* assertOnLeakedDescriptor */)
+}
+
+func (s *drainServer) drainClientsInternal(
+	ctx context.Context, reporter func(int, redact.SafeString), assertOnLeakedDescriptor bool,
 ) error {
 	// Setup a cancelable context so that the logOpenConns goroutine exits when
 	// this function returns.
@@ -486,7 +492,7 @@ func (s *drainServer) drainClients(
 	// Drain all SQL table leases. This must be done after the pgServer has
 	// given sessions a chance to finish ongoing work and after the background
 	// tasks that may issue SQL statements have shut down.
-	s.sqlServer.leaseMgr.SetDraining(ctx, true /* drain */, reporter, true /*assertOnLeakedDescriptor*/)
+	s.sqlServer.leaseMgr.SetDraining(ctx, true /* drain */, reporter, assertOnLeakedDescriptor)
 
 	session, err := s.sqlServer.sqlLivenessProvider.Release(ctx)
 	if err != nil {

pkg/server/testserver.go

@@ -1242,6 +1242,11 @@ func (t *testTenant) DrainClients(ctx context.Context) error {
 	return t.drain.drainClients(ctx, nil /* reporter */)
 }
 
+// DrainClients exports the drainClientsInternal() method for use by tests.
+func (t *testTenant) DrainClientsWithoutLeakCheck(ctx context.Context) error {
+	return t.drain.drainClientsInternal(ctx, nil /* reporter */, false)
+}
+
 // Readiness is part of the serverutils.ApplicationLayerInterface.
 func (t *testTenant) Readiness(ctx context.Context) error {
 	return t.t.admin.checkReadinessForHealthCheck(ctx)
@@ -1943,6 +1948,13 @@ func (ts *testServer) DrainClients(ctx context.Context) error {
 	return ts.drain.drainClients(ctx, nil /* reporter */)
 }
 
+// DrainClientsWithoutLeakCheck exports the drainClients() method for use by tests
+// with descriptor leak checking disabled. This is useful for tests that focus on
+// restart safety logic rather than lease management correctness.
+func (ts *testServer) DrainClientsWithoutLeakCheck(ctx context.Context) error {
+	return ts.drain.drainClientsInternal(ctx, nil, false /* assertOnLeakedDescriptor */)
+}
+
 // Readiness is part of the serverutils.ApplicationLayerInterface.
 func (ts *testServer) Readiness(ctx context.Context) error {
 	return ts.admin.checkReadinessForHealthCheck(ctx)

pkg/testutils/serverutils/api.go

@@ -405,6 +405,11 @@ type ApplicationLayerInterface interface {
 	// DrainClients shuts down client connections.
 	DrainClients(ctx context.Context) error
 
+	// DrainClientsWithoutLeakCheck shuts down client connections but will
+	// skip the assertion that checks for descriptor leaks. This is meant
+	// to enable specific tests that may drain unsafely.
+	DrainClientsWithoutLeakCheck(ctx context.Context) error
+
 	// SystemConfigProvider provides access to the system config.
 	SystemConfigProvider() config.SystemConfigProvider