Fix panic on sse-server shutdown

sse.go

@@ -15,7 +15,7 @@ type Server struct {
 	channels     map[string]*Channel
 	addClient    chan *Client
 	removeClient chan *Client
-	shutdown     chan bool
+	shutdown     chan struct{}
 	closeChannel chan string
 }
 
@@ -37,7 +37,7 @@ func NewServer(options *Options) *Server {
 		make(map[string]*Channel),
 		make(chan *Client),
 		make(chan *Client),
-		make(chan bool),
+		make(chan struct{}),
 		make(chan string),
 	}
 
@@ -76,14 +76,24 @@ func (s *Server) ServeHTTP(response http.ResponseWriter, request *http.Request)
 			channelName = s.options.ChannelNameFunc(request)
 		}
 
+		select {
+		case <-s.shutdown:
+			http.Error(response, "SSE server is in termination state.", http.StatusServiceUnavailable)
+			return
+		default:
+		}
+
 		lastEventID := request.Header.Get("Last-Event-ID")
 		c := newClient(lastEventID, channelName)
 		s.addClient <- c
 		closeNotify := request.Context().Done()
 
 		go func() {
-			<-closeNotify
-			s.removeClient <- c
+			select {
+			case <-s.shutdown:
+			case <-closeNotify:
+				s.removeClient <- c
+			}
 		}()
 
 		response.WriteHeader(http.StatusOK)
@@ -128,7 +138,7 @@ func (s *Server) Restart() {
 
 // Shutdown performs a graceful server shutdown.
 func (s *Server) Shutdown() {
-	s.shutdown <- true
+	close(s.shutdown)
 }
 
 // ClientCount returns the number of clients connected to this server.
@@ -255,7 +265,6 @@ func (s *Server) dispatch() {
 			close(s.addClient)
 			close(s.removeClient)
 			close(s.closeChannel)
-			close(s.shutdown)
 
 			s.options.Logger.Print("server stopped.")
 			return

tests/integration_test.go

@@ -0,0 +1,50 @@
+package tests
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/alexandrevicenzi/go-sse"
+)
+
+func TestServerShutdown(t *testing.T) {
+	sseServer := sse.NewServer(nil)
+	s := httptest.NewServer(sseServer)
+
+	ch := make(chan struct{})
+	ch2 := make(chan struct{})
+	defer close(ch)
+
+	go func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+		defer cancel()
+
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.URL, nil)
+		if err != nil {
+			t.Fail()
+		}
+		req.Header.Set("Cache-Control", "no-cache")
+		req.Header.Set("Accept", "text/event-stream")
+		req.Header.Set("Connection", "keep-alive")
+
+		_, err = http.DefaultClient.Do(req)
+		if err != nil {
+			t.Fail()
+		}
+
+		close(ch2)
+		// Wait until the main test routine is done
+		<-ch
+	}()
+	// Wait until the HTTP request is sent
+	<-ch2
+
+	sseServer.Shutdown()
+	s.Close()
+
+	// Wait a bit after closing HTTP server
+	time.Sleep(250 * time.Millisecond)
+}