diff --git a/pkg/ccl/sqlproxyccl/proxy_handler_test.go b/pkg/ccl/sqlproxyccl/proxy_handler_test.go index 23abc08eb643..966f7c521590 100644 --- a/pkg/ccl/sqlproxyccl/proxy_handler_test.go +++ b/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) + }) }) } diff --git a/pkg/ccl/sqlproxyccl/server.go b/pkg/ccl/sqlproxyccl/server.go index 1d2f92643143..deec01a6cd95 100644 --- a/pkg/ccl/sqlproxyccl/server.go +++ b/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