From dbcc0d04e6510c0d2c9241403c5c81c33b75272b Mon Sep 17 00:00:00 2001
From: maru <maru@myyahoo.com>
Date: Wed, 17 Apr 2024 19:31:12 -0400
Subject: [PATCH] Add new exit handler

---
 api/endpoints.go    |  6 +++++
 pokerogue-server.go | 55 +++++++++++++++++++++++++++++++++++++--------
 2 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/api/endpoints.go b/api/endpoints.go
index 3dc6e10..85c082b 100644
--- a/api/endpoints.go
+++ b/api/endpoints.go
@@ -8,6 +8,7 @@ import (
 	"log"
 	"net/http"
 	"strconv"
+	"sync"
 
 	"github.com/pagefaultgames/pokerogue-server/api/account"
 	"github.com/pagefaultgames/pokerogue-server/api/daily"
@@ -17,6 +18,7 @@ import (
 
 type Server struct {
 	Debug bool
+	Exit *sync.RWMutex
 }
 
 /*
@@ -26,6 +28,10 @@ type Server struct {
 */
 
 func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	// kind of misusing the RWMutex but it doesn't matter
+	s.Exit.RLock()
+	defer s.Exit.RUnlock()
+
 	gob.Register([]interface{}{})
 	gob.Register(map[string]interface{}{})
 
diff --git a/pokerogue-server.go b/pokerogue-server.go
index d9a144a..6d8af7c 100644
--- a/pokerogue-server.go
+++ b/pokerogue-server.go
@@ -6,12 +6,16 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"os/signal"
+	"sync"
+	"syscall"
 
 	"github.com/pagefaultgames/pokerogue-server/api"
 	"github.com/pagefaultgames/pokerogue-server/db"
 )
 
 func main() {
+	// flag stuff
 	debug := flag.Bool("debug", false, "debug mode")
 
 	proto := flag.String("proto", "tcp", "protocol for api to use (tcp, unix)")
@@ -25,28 +29,61 @@ func main() {
 
 	flag.Parse()
 
+	// get database connection
 	err := db.Init(*dbuser, *dbpass, *dbproto, *dbaddr, *dbname)
 	if err != nil {
 		log.Fatalf("failed to initialize database: %s", err)
 	}
 
-	if *proto == "unix" {
-		os.Remove(*addr)
-	}
-
-	listener, err := net.Listen(*proto, *addr)
+	// create listener
+	listener, err := createListener(*proto, *addr)
 	if err != nil {
 		log.Fatalf("failed to create net listener: %s", err)
 	}
 
-	if *proto == "unix" {
-		os.Chmod(*addr, 0777)
-	}
+	// create exit handler
+	var exit sync.RWMutex
+	createExitHandler(&exit)
 
+	// init api
 	api.Init()
 
-	err = http.Serve(listener, &api.Server{Debug: *debug})
+	// start web server
+	err = http.Serve(listener, &api.Server{Debug: *debug, Exit: &exit})
 	if err != nil {
 		log.Fatalf("failed to create http server or server errored: %s", err)
 	}
 }
+
+func createListener(proto, addr string) (net.Listener, error) {
+	if proto == "unix" {
+		os.Remove(addr)
+	}
+
+	listener, err := net.Listen(proto, addr)
+	if err != nil {
+		return nil, err
+	}
+
+	if proto == "unix" {
+		os.Chmod(addr, 0777)
+	}
+
+	return listener, nil
+}
+
+func createExitHandler(mtx *sync.RWMutex) {
+	s := make(chan os.Signal, 1)
+	signal.Notify(s, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		// wait for exit signal of some kind
+		<-s
+
+		// block new requests and wait for existing ones to finish
+		mtx.Lock()
+
+		// bail
+		os.Exit(0)
+	}()
+}