Timeout peers when the server is shutdown. (#55)

When our backend needs to restart peers are disconnected can could get
into the timeouts table, therefor they could reconnect. This PR will
make sure all connections are closed when the server shuts down allowing
peers to be timed out and reconnect later.

cmd/signaling/main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"math/rand"
+	"net"
 	"net/http"
 	"os"
 	"os/signal"
@@ -49,7 +50,7 @@ func main() {
 	)
 	go credentialsClient.Run(ctx)
 
-	mux := internal.Signaling(ctx, store, credentialsClient)
+	mux, cleanup := internal.Signaling(ctx, store, credentialsClient)
 
 	cors := cors.Default()
 	handler := logging.Middleware(cors.Handler(mux), logger)
@@ -64,6 +65,10 @@ func main() {
 		Addr:    addr,
 		Handler: handler,
 
+		BaseContext: func(net.Listener) context.Context {
+			return ctx
+		},
+
 		ReadTimeout:  5 * time.Second,
 		WriteTimeout: 10 * time.Second,
 		IdleTimeout:  650 * time.Second,
@@ -77,13 +82,16 @@ func main() {
 	logger.Info("listening", zap.String("addr", addr))
 
 	<-ctx.Done()
-	if flushed != nil {
-		<-flushed
-	}
+	logger.Info("shutting down")
 
-	ctx, cancel = context.WithTimeout(context.Background(), 30*time.Second)
-	defer cancel()
-	if err := server.Shutdown(ctx); err != nil {
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer shutdownCancel()
+	if err := server.Shutdown(shutdownCtx); err != nil {
 		logger.Fatal("failed to shutdown server", zap.Error(err))
 	}
+
+	cleanup()
+	if flushed != nil {
+		<-flushed
+	}
 }

internal/signaling.go

@@ -11,10 +11,15 @@ import (
 	"github.com/poki/netlib/internal/util"
 )
 
-func Signaling(ctx context.Context, store stores.Store, credentialsClient *cloudflare.CredentialsClient) http.Handler {
+func Signaling(ctx context.Context, store stores.Store, credentialsClient *cloudflare.CredentialsClient) (http.Handler, func()) {
 	mux := http.NewServeMux()
 
-	mux.Handle("/v0/signaling", signaling.Handler(ctx, store, credentialsClient))
+	openConnections, signaling := signaling.Handler(ctx, store, credentialsClient)
+
+	cleanup := func() {
+		openConnections.Wait()
+	}
+	mux.Handle("/v0/signaling", signaling)
 
 	hasCredentials := uint32(0)
 	mux.HandleFunc("/ready", func(w http.ResponseWriter, r *http.Request) {
@@ -43,5 +48,5 @@ func Signaling(ctx context.Context, store stores.Store, credentialsClient *cloud
 	mux.HandleFunc("/health", healthCheck)
 	mux.HandleFunc("/", healthCheck)
 
-	return mux
+	return mux, cleanup
 }

internal/signaling/handler.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/koenbollen/logging"
@@ -18,13 +19,14 @@ import (
 
 const MaxConnectionTime = 1 * time.Hour
 
-func Handler(ctx context.Context, store stores.Store, cloudflare *cloudflare.CredentialsClient) http.HandlerFunc {
+func Handler(ctx context.Context, store stores.Store, cloudflare *cloudflare.CredentialsClient) (*sync.WaitGroup, http.HandlerFunc) {
 	manager := &TimeoutManager{
 		Store: store,
 	}
 	go manager.Run(ctx)
 
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	wg := &sync.WaitGroup{}
+	return wg, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
 		logger := logging.GetLogger(ctx)
 		logger.Debug("upgrading connection")
@@ -48,6 +50,9 @@ func Handler(ctx context.Context, store stores.Store, cloudflare *cloudflare.Cre
 			util.ErrorAndAbort(w, r, http.StatusBadRequest, "", err)
 		}
 
+		wg.Add(1)
+		defer wg.Done()
+
 		peer := &Peer{
 			store: store,
 			conn:  conn,