Merge #106549

106549: sqlproxyccl: deflake TestDirectoryConnect r=JeffSwenson a=JeffSwenson

This contains fixes to two sources of flakes in TestDirectoryConnect:
- sqlproxy http draining is now tied into the stopper. This avoids a
	source of goroutine leaks.
- The sql server is gracefully drained to work around #106537.

When combined with #106599, I was able to run the test for 25K
interations under stress with no flakes.

Fixes: #105402

Co-authored-by: Jeff <[email protected]>

pkg/ccl/sqlproxyccl/proxy_handler_test.go

@@ -1389,13 +1389,17 @@ func TestDirectoryConnect(t *testing.T) {
 		})
 	})
 
+	// Drain the tenant server gracefully. This is a workaround for #106537.
+	// Draining the server allows the server to delete the sql instance row.
+	require.NoError(t, tenants[0].DrainClients(ctx))
+
 	// Stop the directory server and the tenant SQL process started earlier.
 	// This tests whether the proxy can recover when the directory server and
 	// SQL pod restarts.
 	tds.Stop(ctx)
 	tenantStopper.Stop(ctx)
 
-	// Drain old pod and add a new one before starting the directory server.
+	// Drain the old pod and add a new one before starting the directory server.
 	tds.DrainPod(tenantID, tenants[0].SQLAddr())
 	tenants = startTestTenantPods(ctx, t, s, tenantID, 1, base.TestingKnobs{})
 	tds.AddPod(tenantID, &tenant.Pod{
@@ -1407,15 +1411,14 @@ func TestDirectoryConnect(t *testing.T) {
 	require.NoError(t, tds.Start(ctx))
 
 	t.Run("successful connection after restart", func(t *testing.T) {
-		require.Eventually(t, func() bool {
+		testutils.SucceedsSoon(t, func() error {
 			conn, err := pgx.Connect(ctx, connectionString)
 			if err != nil {
-				return false
+				return err
 			}
 			defer func() { _ = conn.Close(ctx) }()
-			require.NoError(t, runTestQuery(ctx, conn))
-			return true
-		}, 30*time.Second, 100*time.Millisecond)
+			return runTestQuery(ctx, conn)
+		})
 	})
 }
 

pkg/ccl/sqlproxyccl/server.go

@@ -181,7 +181,7 @@ func (s *Server) ServeHTTP(ctx context.Context, ln net.Listener) error {
 
 	srv := http.Server{Handler: s.mux}
 
-	go func() {
+	err := s.Stopper.RunAsyncTask(ctx, "sqlproxy-http-cleanup", func(ctx context.Context) {
 		<-ctx.Done()
 
 		// Wait up to 15 seconds for the HTTP server to shut itself
@@ -196,11 +196,13 @@ func (s *Server) ServeHTTP(ctx context.Context, ln net.Listener) error {
 				// Ignore any errors as this routine will only be called
 				// when the server is shutting down.
 				_ = srv.Shutdown(shutdownCtx)
-
 				return nil
 			},
 		)
-	}()
+	})
+	if err != nil {
+		return err
+	}
 
 	if err := srv.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
 		return err