diff options
| -rw-r--r-- | cmd/pkgsite/main.go | 33 | ||||
| -rw-r--r-- | internal/middleware/timeout/timeout.go | 10 |
2 files changed, 42 insertions, 1 deletions
diff --git a/cmd/pkgsite/main.go b/cmd/pkgsite/main.go index eae690e0..06968618 100644 --- a/cmd/pkgsite/main.go +++ b/cmd/pkgsite/main.go @@ -89,6 +89,10 @@ func main() { flag.BoolVar(&serverCfg.DevMode, "dev", false, "enable developer mode (reload templates on each page load, serve non-minified JS/CSS, etc.)") flag.StringVar(&serverCfg.DevModeStaticDir, "static", "static", "path to folder containing static files served") + var shutdownIdleDuration string + flag.StringVar(&shutdownIdleDuration, `shutdown-idle`, ``, + `Shutdown the server after idle on specific duration.`) + flag.Usage = func() { out := flag.CommandLine.Output() fmt.Fprintf(out, "usage: %s [flags] [PATHS ...]\n", os.Args[0]) @@ -98,6 +102,15 @@ func main() { } flag.Parse() + if shutdownIdleDuration != `` { + var err error + timeout.ShutdownIdleDuration, err = time.ParseDuration(shutdownIdleDuration) + if err != nil { + dief(`invalid shutdown-idle duration %s`, shutdownIdleDuration) + } + timeout.ShutdownIdleTimer = time.NewTimer(timeout.ShutdownIdleDuration) + } + serverCfg.UseLocalStdlib = true serverCfg.GoRepoPath = *goRepoPath serverCfg.Paths = collectPaths(flag.Args()) @@ -171,7 +184,25 @@ func main() { server.Install(router.Handle, nil, nil) mw := timeout.Timeout(54 * time.Second) srv := &http.Server{Addr: addr, Handler: mw(router)} - dief("%v", srv.Serve(ln)) + + if timeout.ShutdownIdleTimer == nil { + dief("%v", srv.Serve(ln)) + } else { + waitq := make(chan struct{}, 1) + go func() { + err = srv.Serve(ln) + if err != nil { + fmt.Fprintf(os.Stderr, "%s\n", err.Error()) + } + waitq <- struct{}{} + }() + select { + case <-waitq: + // Server stopped. + case <-timeout.ShutdownIdleTimer.C: + _ = srv.Shutdown(context.Background()) + } + } } func dief(format string, args ...any) { diff --git a/internal/middleware/timeout/timeout.go b/internal/middleware/timeout/timeout.go index 93d992da..391919c4 100644 --- a/internal/middleware/timeout/timeout.go +++ b/internal/middleware/timeout/timeout.go @@ -10,11 +10,21 @@ import ( "time" ) +var ShutdownIdleDuration time.Duration +var ShutdownIdleTimer *time.Timer + // Timeout returns a new Middleware that times out each request after the given // duration. func Timeout(d time.Duration) func(http.Handler) http.Handler { return func(h http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if ShutdownIdleTimer != nil { + ok := ShutdownIdleTimer.Reset(ShutdownIdleDuration) + if !ok { + // Timer had reached or closed + return + } + } ctx, cancel := context.WithTimeout(r.Context(), d) defer cancel() h.ServeHTTP(w, r.WithContext(ctx)) |