diff options
| -rw-r--r-- | tvix/nar-bridge/cmd/nar_bridge/main.go | 34 | ||||
| -rw-r--r-- | tvix/nar-bridge/pkg/server/server.go | 12 |
2 files changed, 30 insertions, 16 deletions
diff --git a/tvix/nar-bridge/cmd/nar_bridge/main.go b/tvix/nar-bridge/cmd/nar_bridge/main.go index 482012f39b..a1732433c0 100644 --- a/tvix/nar-bridge/cmd/nar_bridge/main.go +++ b/tvix/nar-bridge/cmd/nar_bridge/main.go @@ -1,8 +1,10 @@ package main import ( + "context" "os" "os/signal" + "time" "github.com/alecthomas/kong" @@ -33,25 +35,17 @@ func main() { } logrus.SetLevel(logLevel) - c := make(chan os.Signal, 1) - signal.Notify(c, os.Interrupt) - - go func() { - for range c { - log.Info("Received Signal, shutting down…") - os.Exit(1) - } - }() + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() // connect to tvix-store log.Debugf("Dialing to %v", cli.StoreAddr) - conn, err := grpc.Dial(cli.StoreAddr, grpc.WithTransportCredentials(insecure.NewCredentials())) + conn, err := grpc.DialContext(ctx, cli.StoreAddr, grpc.WithTransportCredentials(insecure.NewCredentials())) if err != nil { log.Fatalf("did not connect: %v", err) } defer conn.Close() - log.Printf("Starting nar-bridge at %v", cli.ListenAddr) s := server.New( storev1pb.NewDirectoryServiceClient(conn), storev1pb.NewBlobServiceClient(conn), @@ -60,9 +54,21 @@ func main() { 30, ) - err = s.ListenAndServe(cli.ListenAddr) - if err != nil { - log.Error("Server failed: %w", err) + log.Printf("Starting nar-bridge at %v", cli.ListenAddr) + go s.ListenAndServe(cli.ListenAddr) + + // listen for the interrupt signal. + <-ctx.Done() + + // Restore default behaviour on the interrupt signal + stop() + log.Info("Received Signal, shutting down, press Ctl+C again to force.") + + timeoutCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + if s.Shutdown(timeoutCtx); err != nil { + log.WithError(err).Warn("failed to shutdown") os.Exit(1) } } diff --git a/tvix/nar-bridge/pkg/server/server.go b/tvix/nar-bridge/pkg/server/server.go index 0de38c84df..f58842bfa7 100644 --- a/tvix/nar-bridge/pkg/server/server.go +++ b/tvix/nar-bridge/pkg/server/server.go @@ -1,6 +1,7 @@ package server import ( + "context" "fmt" "net/http" "sync" @@ -13,6 +14,7 @@ import ( ) type Server struct { + srv *http.Server handler chi.Router directoryServiceClient storev1pb.DirectoryServiceClient @@ -73,8 +75,14 @@ func New( return s } +func (s *Server) Shutdown(ctx context.Context) error { + return s.srv.Shutdown(ctx) +} + +// ListenAndServer starts the webserver, and waits for it being closed or +// shutdown, after which it'll return ErrServerClosed. func (s *Server) ListenAndServe(addr string) error { - srv := &http.Server{ + s.srv = &http.Server{ Addr: addr, Handler: s.handler, ReadTimeout: 500 * time.Second, @@ -82,5 +90,5 @@ func (s *Server) ListenAndServe(addr string) error { IdleTimeout: 500 * time.Second, } - return srv.ListenAndServe() + return s.srv.ListenAndServe() } |