Better shutdown code, release listeners

socketserver/cmd/ffzsocketserver/socketserver.go

@@ -1,6 +1,7 @@
 package main // import "github.com/FrankerFaceZ/FrankerFaceZ/socketserver/cmd/ffzsocketserver"
 
 import (
+	"context"
 	"encoding/json"
 	"flag"
 	"fmt"
@@ -8,6 +9,10 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"os/signal"
+	"sync"
+	"syscall"
+	"time"
 
 	"github.com/FrankerFaceZ/FrankerFaceZ/socketserver/server"
 )
@@ -56,17 +61,63 @@ func main() {
 
 	go commandLineConsole()
 
+	var server1, server2 *http.Server
+
+	stopSig := make(chan os.Signal, 3)
+	signal.Notify(stopSig, os.Interrupt)
+	signal.Notify(stopSig, syscall.SIGUSR1)
+	signal.Notify(stopSig, syscall.SIGTERM)
+
 	if conf.UseSSL {
+		server1 = &http.Server{
+			Addr:    conf.SSLListenAddr,
+			Handler: http.DefaultServeMux,
+		}
 		go func() {
-			if err := http.ListenAndServeTLS(conf.SSLListenAddr, conf.SSLCertificateFile, conf.SSLKeyFile, http.DefaultServeMux); err != nil {
-				log.Fatal("ListenAndServeTLS: ", err)
+			if err := server1.ListenAndServeTLS(conf.SSLCertificateFile, conf.SSLKeyFile); err != nil {
+				log.Println("ListenAndServeTLS:", err)
+				stopSig <- os.Interrupt
 			}
 		}()
 	}
 
-	if err = http.ListenAndServe(conf.ListenAddr, http.DefaultServeMux); err != nil {
-		log.Fatal("ListenAndServe: ", err)
+	if true {
+		server2 = &http.Server{
+			Addr:    conf.ListenAddr,
+			Handler: http.DefaultServeMux,
+		}
+		go func() {
+			if err := server2.ListenAndServe(); err != nil {
+				log.Println("ListenAndServe: ", err)
+				stopSig <- os.Interrupt
+			}
+		}()
+
 	}
+
+	<-stopSig
+	log.Println("Shutting down...")
+
+	var wg sync.WaitGroup
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		if conf.UseSSL {
+			server1.Shutdown(ctx)
+		}
+	}()
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		server2.Shutdown(ctx)
+	}()
+	server.Shutdown(&wg)
+
+	time.Sleep(1 * time.Second)
+	wg.Wait()
 }
 
 func generateKeys(outputFile string) {

socketserver/server/backend.go

@@ -141,9 +141,11 @@ func (backend *backendInfo) SendRemoteCommand(remoteCommand, data string, auth A
 		return "", ErrAuthorizationNeeded
 	} else if resp.StatusCode < 200 || resp.StatusCode > 299 { // any non-2xx
 		// If the Content-Type header includes a charset, ignore it.
+		// typeStr, _, _ = mime.ParseMediaType(resp.Header.Get("Content-Type"))
+		// inline the part of the function we care about
 		typeStr := resp.Header.Get("Content-Type")
 		splitIdx := strings.IndexRune(typeStr, ';')
-		if ( splitIdx != -1 ) {
+		if splitIdx != -1 {
 			typeStr = strings.TrimSpace(typeStr[0:splitIdx])
 		}
 

socketserver/server/handlecore.go

@@ -139,30 +139,21 @@ func startJanitors() {
 	go pubsubJanitor()
 
 	go ircConnection()
-	go shutdownHandler()
 }
 
-// is_init_func
-func shutdownHandler() {
-	ch := make(chan os.Signal)
-	signal.Notify(ch, syscall.SIGUSR1)
-	signal.Notify(ch, syscall.SIGTERM)
-	<-ch
-	log.Println("Shutting down...")
-
-	var wg sync.WaitGroup
+// Shutdown disconnects all clients.
+func Shutdown(wg *sync.WaitGroup) {
 	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		writeHLL()
-		wg.Done()
 	}()
-
-	StopAcceptingConnections = true
-	close(StopAcceptingConnectionsCh)
-
-	time.Sleep(1 * time.Second)
-	wg.Wait()
-	os.Exit(0)
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		close(StopAcceptingConnectionsCh)
+		time.Sleep(2 * time.Second)
+	}()
 }
 
 // is_init_func +test
@@ -200,7 +191,6 @@ var BannerHTML []byte
 
 // StopAcceptingConnectionsCh is closed while the server is shutting down.
 var StopAcceptingConnectionsCh = make(chan struct{})
-var StopAcceptingConnections = false
 
 // HTTPHandleRootURL is the http.HandleFunc for requests on `/`.
 // It either uses the SocketUpgrader or writes out the BannerHTML.
@@ -211,13 +201,6 @@ func HTTPHandleRootURL(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// racy, but should be ok?
-	if StopAcceptingConnections {
-		w.WriteHeader(503)
-		fmt.Fprint(w, "server is shutting down")
-		return
-	}
-
 	if strings.Contains(strings.ToLower(r.Header.Get("Connection")), "upgrade") {
 		updateSysMem()
 
@@ -376,8 +359,10 @@ func RunSocketConnection(conn *websocket.Conn) {
 
 	// And done.
 
-	if !StopAcceptingConnections {
+	select {
+	case <-StopAcceptingConnectionsCh:
 		// Don't perform high contention operations when server is closing
+	default:
 		atomic.AddUint64(&Statistics.CurrentClientCount, NegativeOne)
 		atomic.AddUint64(&Statistics.ClientDisconnectsTotal, 1)
 
@@ -410,7 +395,9 @@ func runSocketReader(conn *websocket.Conn, client *ClientInfo, errorChan chan<-
 
 		msg = ClientMessage{}
 		msgErr := UnmarshalClientMessage(packet, messageType, &msg)
-		if _, ok := msgErr.(interface{IsFatal() bool}); ok {
+		if _, ok := msgErr.(interface {
+			IsFatal() bool
+		}); ok {
 			errorChan <- msgErr
 			continue
 		} else if msgErr != nil {