fix: graceful shutdown (#54)

Carlos Alexandro Becker created 3 years ago

Signed-off-by: Carlos A Becker <[e-mail]>

Change summary

cmd/soft/main.go | 24 +++++++++++++++++++++---
server/server.go |  6 ++++++
2 files changed, 27 insertions(+), 3 deletions(-)

Detailed changes

cmd/soft/main.go

@@ -1,10 +1,14 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"log"
 	"os"
+	"os/signal"
+	"syscall"
+	"time"
 
 	"github.com/charmbracelet/soft-serve/config"
 	"github.com/charmbracelet/soft-serve/server"
@@ -49,9 +53,23 @@ func main() {
 
 	cfg := config.DefaultConfig()
 	s := server.NewServer(cfg)
-	log.Printf("Starting SSH server on %s:%d\n", cfg.Host, cfg.Port)
-	err := s.Start()
-	if err != nil {
+
+	done := make(chan os.Signal, 1)
+	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+
+	log.Printf("Starting SSH server on %s:%d", cfg.Host, cfg.Port)
+	go func() {
+		if err := s.Start(); err != nil {
+			log.Fatalln(err)
+		}
+	}()
+
+	<-done
+
+	log.Printf("Stopping SSH server on %s:%d", cfg.Host, cfg.Port)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer func() { cancel() }()
+	if err := s.Shutdown(ctx); err != nil {
 		log.Fatalln(err)
 	}
 }

server/server.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"log"
 
@@ -63,3 +64,8 @@ func (srv *Server) Reload() error {
 func (srv *Server) Start() error {
 	return srv.SSHServer.ListenAndServe()
 }
+
+// Shutdown lets the server gracefully shutdown.
+func (srv *Server) Shutdown(ctx context.Context) error {
+	return srv.SSHServer.Shutdown(ctx)
+}