From d2742ae00cdae2c22377aad820e01ef22797ee36 Mon Sep 17 00:00:00 2001 From: Dwi Siswanto Date: Wed, 14 Aug 2024 23:20:14 +0700 Subject: [PATCH] feat: add `*Server.Stop` method (#1863) * feat: add `*Server.Stop` method Signed-off-by: Dwi Siswanto * feat: implement `*Server.Stop` Signed-off-by: Dwi Siswanto --------- Signed-off-by: Dwi Siswanto --- runner/apiendpoint.go | 28 ++++++++++++++++++++++++---- runner/runner.go | 4 ++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/runner/apiendpoint.go b/runner/apiendpoint.go index 5faa4b9..b3528aa 100644 --- a/runner/apiendpoint.go +++ b/runner/apiendpoint.go @@ -2,6 +2,9 @@ package runner import ( + "context" + "time" + "encoding/json" "net/http" ) @@ -12,8 +15,9 @@ type Concurrency struct { // Server represents the HTTP server that handles the concurrency settings endpoints. type Server struct { - addr string - config *Options + addr string + config *Options + httpServer *http.Server } // New creates a new instance of Server. @@ -26,10 +30,18 @@ func NewServer(addr string, config *Options) *Server { // Start initializes the server and its routes, then starts listening on the specified address. func (s *Server) Start() error { - http.HandleFunc("/api/concurrency", s.handleConcurrency) - if err := http.ListenAndServe(s.addr, nil); err != nil { + mux := http.NewServeMux() + mux.HandleFunc("/api/concurrency", s.handleConcurrency) + + s.httpServer = &http.Server{ + Addr: s.addr, + Handler: mux, + } + + if err := s.httpServer.ListenAndServe(); err != nil { return err } + return nil } @@ -71,3 +83,11 @@ func (s *Server) updateSettings(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) } + +// Stop gracefully shuts down the server. +func (s *Server) Stop() error { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + return s.httpServer.Shutdown(ctx) +} diff --git a/runner/runner.go b/runner/runner.go index 5ba8b70..6a3511a 100644 --- a/runner/runner.go +++ b/runner/runner.go @@ -671,6 +671,7 @@ func (r *Runner) Close() { r.hm.Close() r.hp.Dialer.Close() r.ratelimiter.Stop() + if r.options.HostMaxErrors >= 0 { r.HostErrorsCache.Purge() } @@ -680,6 +681,9 @@ func (r *Runner) Close() { if r.options.ShowStatistics { _ = r.stats.Stop() } + if r.options.HttpApiEndpoint != "" { + _ = r.httpApiEndpoint.Stop() + } if r.options.OnClose != nil { r.options.OnClose() }