diff --git a/cmd/serve.go b/cmd/serve.go index 3bbf77820..6f4a40eb0 100644 --- a/cmd/serve.go +++ b/cmd/serve.go @@ -3,15 +3,19 @@ package cmd import ( "context" "encoding/base64" + "errors" "fmt" "log/slog" "net/http" "os" "os/signal" + "slices" "strings" "syscall" "time" + "golang.org/x/sync/errgroup" + "github.com/raystack/frontier/core/aggregates/orgbilling" "github.com/raystack/frontier/core/aggregates/orginvoices" "github.com/raystack/frontier/core/aggregates/orgpats" @@ -29,8 +33,6 @@ import ( "github.com/raystack/frontier/core/prospect" "github.com/raystack/frontier/core/userpat" - "slices" - "github.com/doug-martin/goqu/v9" "github.com/jackc/pgx/v4" "github.com/stripe/stripe-go/v79" @@ -311,16 +313,37 @@ func StartServer(logger *slog.Logger, cfg *config.Frontier) error { }() } - go func() { - if err := deps.LogListener.Listen(ctx); err != nil { + // gctx is cancelled when ctx is cancelled or when any member returns an + // error, so a connect server failure also winds down the UI and listener. + g, gctx := errgroup.WithContext(ctx) + + g.Go(func() error { + if err := deps.LogListener.Listen(gctx); err != nil && !errors.Is(err, context.Canceled) { logger.Warn("log listener failed", "err", err) } - }() + return nil + }) + + g.Go(func() error { + server.ServeUI(gctx, logger, cfg.UI, cfg.App) + return nil + }) + + g.Go(func() error { + return server.ServeConnect(gctx, logger, cfg.App, deps, promRegistry) + }) - go server.ServeUI(ctx, logger, cfg.UI, cfg.App) + // Once shutdown starts, unregister the signal handler so a second + // SIGINT/SIGTERM gets default handling and can kill a stuck drain + // instead of being swallowed. + go func() { + <-gctx.Done() + cancelFunc() + }() - // start connect server (blocking) - return server.ServeConnect(ctx, logger, cfg.App, deps, promRegistry) + // Wait for every server to finish draining before the deferred cleanup + // above closes the database and other dependencies under them. + return g.Wait() } func buildAPIDependencies( diff --git a/pkg/server/serve_ui_test.go b/pkg/server/serve_ui_test.go new file mode 100644 index 000000000..1da6df62c --- /dev/null +++ b/pkg/server/serve_ui_test.go @@ -0,0 +1,80 @@ +package server + +import ( + "context" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "testing" + "time" + + "github.com/raystack/frontier/web/apps/admin" +) + +func TestServeUIReturnsAfterShutdownOnContextCancel(t *testing.T) { + if f, err := admin.Assets.Open("dist/admin/index.html"); err != nil { + t.Skip("admin SPA assets are not built; ServeUI exits before starting the server") + } else { + f.Close() + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("could not find a free port: %v", err) + } + port := l.Addr().(*net.TCPAddr).Port + l.Close() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + done := make(chan struct{}) + go func() { + defer close(done) + ServeUI(ctx, logger, UIConfig{Port: port}, Config{}) + }() + + url := fmt.Sprintf("http://127.0.0.1:%d/configs", port) + deadline := time.Now().Add(10 * time.Second) + for { + if time.Now().After(deadline) { + t.Fatalf("ui server did not respond at %s in time", url) + } + resp, err := http.Get(url) + if err == nil { + resp.Body.Close() + if resp.StatusCode == http.StatusOK { + break + } + } + time.Sleep(50 * time.Millisecond) + } + + cancel() + + select { + case <-done: + case <-time.After(15 * time.Second): + t.Fatal("ServeUI did not return after context cancel") + } +} + +func TestServeUIReturnsWhenPortNotConfigured(t *testing.T) { + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + + done := make(chan struct{}) + go func() { + defer close(done) + ServeUI(context.Background(), logger, UIConfig{}, Config{}) + }() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("ServeUI did not return when no port is configured") + } +} diff --git a/pkg/server/server.go b/pkg/server/server.go index 315da19a5..879071531 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "log/slog" "net/http" "net/http/httputil" "net/url" @@ -13,8 +14,6 @@ import ( "golang.org/x/net/http2" "golang.org/x/net/http2/h2c" - "log/slog" - "github.com/prometheus/client_golang/prometheus" "github.com/raystack/salt/server/spa" @@ -66,48 +65,72 @@ func ServeUI(ctx context.Context, logger *slog.Logger, uiConfig UIConfig, apiSer if err != nil { logger.Warn("failed to load ui", "err", err) return - } else { - connectRemoteHost := fmt.Sprintf("http://%s:%d", apiServerConfig.Host, apiServerConfig.Connect.Port) - connectRemote, err := url.Parse(connectRemoteHost) - if err != nil { - logger.Error("ui server failed: unable to parse connect server host") - return + } + + connectRemoteHost := fmt.Sprintf("http://%s:%d", apiServerConfig.Host, apiServerConfig.Connect.Port) + connectRemote, err := url.Parse(connectRemoteHost) + if err != nil { + logger.Error("ui server failed: unable to parse connect server host", "err", err) + return + } + + connectProxyHandler := func(p *httputil.ReverseProxy) func(http.ResponseWriter, *http.Request) { + return func(w http.ResponseWriter, r *http.Request) { + r.URL.Path = strings.Replace(r.URL.Path, "/frontier-connect", "", -1) + r.Host = connectRemoteHost + p.ServeHTTP(w, r) } + } - connectProxyHandler := func(p *httputil.ReverseProxy) func(http.ResponseWriter, *http.Request) { - return func(w http.ResponseWriter, r *http.Request) { - r.URL.Path = strings.Replace(r.URL.Path, "/frontier-connect", "", -1) - r.Host = connectRemoteHost - p.ServeHTTP(w, r) - } + connectProxy := httputil.NewSingleHostReverseProxy(connectRemote) + + mux := http.NewServeMux() + mux.HandleFunc("/configs", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + confResp := UIConfigApiResponse{ + Title: uiConfig.Title, + Logo: uiConfig.Logo, + AppUrl: uiConfig.AppURL, + TokenProductId: uiConfig.TokenProductId, + OrganizationTypes: uiConfig.OrganizationTypes, + Webhooks: WebhooksConfigApiResponse{ + EnableDelete: uiConfig.Webhooks.EnableDelete, + }, + Terminology: uiConfig.Terminology, } + json.NewEncoder(w).Encode(confResp) + }) - connectProxy := httputil.NewSingleHostReverseProxy(connectRemote) - - http.HandleFunc("/configs", func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - confResp := UIConfigApiResponse{ - Title: uiConfig.Title, - Logo: uiConfig.Logo, - AppUrl: uiConfig.AppURL, - TokenProductId: uiConfig.TokenProductId, - OrganizationTypes: uiConfig.OrganizationTypes, - Webhooks: WebhooksConfigApiResponse{ - EnableDelete: uiConfig.Webhooks.EnableDelete, - }, - Terminology: uiConfig.Terminology, - } - json.NewEncoder(w).Encode(confResp) - }) + mux.HandleFunc("/frontier-connect/", connectProxyHandler(connectProxy)) + mux.Handle("/", http.StripPrefix("/", spaHandler)) - http.HandleFunc("/frontier-connect/", connectProxyHandler(connectProxy)) - http.Handle("/", http.StripPrefix("/", spaHandler)) + server := &http.Server{ + Addr: fmt.Sprintf(":%d", uiConfig.Port), + Handler: mux, } logger.Info("ui server starting", "http-port", uiConfig.Port) - if err := http.ListenAndServe(fmt.Sprintf(":%d", uiConfig.Port), nil); err != nil { + serveErr := make(chan error, 1) + go func() { + serveErr <- server.ListenAndServe() + }() + + select { + case err := <-serveErr: logger.Error("ui server failed", "err", err) + return + case <-ctx.Done(): + } + + ctxShutdown, cancel := context.WithTimeout(context.Background(), connectServerGracePeriod) + defer cancel() + + if err := server.Shutdown(ctxShutdown); err != nil { + logger.ErrorContext(ctxShutdown, "ui server shutdown error", "err", err) + return } + + logger.Info("Graceful shutdown of ui server complete") } func ServeConnect(ctx context.Context, logger *slog.Logger, cfg Config, deps api.Deps, promRegistry *prometheus.Registry) error {