An attempt to gracefully shutdown websockets on go using gorilla/websockets

websocket_graceful.go

package main

import (
	"context"
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"os"
	"os/signal"
)

var upgrader = websocket.Upgrader{}
var joining chan *websocket.Conn

func receive(w http.ResponseWriter, r *http.Request){
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer c.Close()

	joining <- c

	for {
		_, message, err := c.ReadMessage()
		if err != nil {
			log.Println("Error: ", err)
			break
		}
		fmt.Println(string(message))

	}
}

func manageConns(incoming chan *websocket.Conn, interrupt chan os.Signal, closeIt chan struct{}, srv *http.Server) {

	clients := make([]*websocket.Conn,0)

	for {
		select {

		case newClient := <- incoming:
			clients = append(clients, newClient)
			fmt.Printf("Incoming: %v - %v \n", newClient.LocalAddr(), len(clients))

		case <- interrupt:
			fmt.Println("ManageConns: Shutting down server.")
			if err := srv.Shutdown(context.Background()); err != nil {
				// Error from closing listeners, or context timeout:
				log.Printf("HTTP server Shutdown Error: %v", err)
			}

			// Server may be shutdown and won't accept new connections but will accept data from open connections
			// this is a point you could test this.
			// <- time.After(20 * time.Second) // wait for slow clients to close up

			for _, v := range clients {
				err := v.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
				if err != nil {
					log.Println("error closing: ", err)
					v.Close() // close abruptly
				}
			}
			close(closeIt)
		}
	}
}

func main() {
	fmt.Println("Starting...")
	srv := &http.Server{Addr:"localhost:8080", Handler: nil}

	joining = make(chan *websocket.Conn)
	allDead := make(chan struct{}, 1)

	sigint := make(chan os.Signal, 1)
	signal.Notify(sigint, os.Interrupt)

	go manageConns(joining, sigint, allDead, srv)

	http.HandleFunc("/receive", receive)
	err := srv.ListenAndServe()
	if err != nil {
		fmt.Println("Error:", err)
	}

	<-allDead
	fmt.Println("Exiting main")
}