79 lines
1.8 KiB
Go
79 lines
1.8 KiB
Go
|
package mux
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"io"
|
||
|
|
"log"
|
||
|
|
"log/slog"
|
||
|
|
"net"
|
||
|
|
"net/http"
|
||
|
|
"os/signal"
|
||
|
|
"syscall"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type ServeCB func(srv *http.Server) error
|
||
|
|
|
||
|
|
const (
|
||
|
|
shutdownDelay = time.Second * 10
|
||
|
|
shutdownHardDelay = time.Second * 5
|
||
|
|
drainDelay = time.Second
|
||
|
|
)
|
||
|
|
|
||
|
|
// Serve with graceful shutdown
|
||
|
|
func (m *Mux) Serve(cb ServeCB) {
|
||
|
|
rootCtx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||
|
|
defer stop()
|
||
|
|
|
||
|
|
// catch all options
|
||
|
|
// lets get it thorugh all middlewares
|
||
|
|
m.OPTIONS("/", func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
w.Header().Set("Content-Length", "0")
|
||
|
|
if r.ContentLength != 0 {
|
||
|
|
// Read up to 4KB of OPTIONS body (as mentioned in the
|
||
|
|
// spec as being reserved for future use), but anything
|
||
|
|
// over that is considered a waste of server resources
|
||
|
|
// (or an attack) and we abort and close the connection,
|
||
|
|
// courtesy of MaxBytesReader's EOF behavior.
|
||
|
|
mb := http.MaxBytesReader(w, r.Body, 4<<10)
|
||
|
|
io.Copy(io.Discard, mb)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
|
||
|
|
srvCtx, cancelSrvCtx := context.WithCancel(context.Background())
|
||
|
|
srv := &http.Server{
|
||
|
|
Handler: m,
|
||
|
|
BaseContext: func(_ net.Listener) context.Context {
|
||
|
|
return srvCtx
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
go func() {
|
||
|
|
if err := cb(srv); !errors.Is(err, http.ErrServerClosed) {
|
||
|
|
panic(err)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
|
||
|
|
// Wait for interrupt signal
|
||
|
|
<-rootCtx.Done()
|
||
|
|
|
||
|
|
stop()
|
||
|
|
m.IsShuttingDown.Store(true)
|
||
|
|
slog.Info("received interrupt singal, shutting down")
|
||
|
|
time.Sleep(drainDelay)
|
||
|
|
slog.Info("readiness check propagated, now waiting for ongoing requests to finish.")
|
||
|
|
|
||
|
|
shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownDelay)
|
||
|
|
defer cancel()
|
||
|
|
|
||
|
|
err := srv.Shutdown(shutdownCtx)
|
||
|
|
cancelSrvCtx()
|
||
|
|
if err != nil {
|
||
|
|
log.Println("failed to wait for ongoing requests to finish, waiting for forced cancellation")
|
||
|
|
time.Sleep(shutdownHardDelay)
|
||
|
|
}
|
||
|
|
|
||
|
|
slog.Info("seerver shut down gracefully")
|
||
|
|
}
|