sse: add closed SelectCase to handle request closing signal (#25)

Signed-off-by: E99p1ant <i@github.red>

Author: wuhan005

sse.go

@@ -45,7 +45,24 @@ func Bind(obj interface{}, opts ...Options) flamego.Handler {
 		}
 		c.Set(reflect.ChanOf(reflect.SendDir, sse.sender.Type().Elem()), sse.sender)
 
-		go sse.handle(log, c.ResponseWriter())
+		// stopCh is closed when the next handler returns, signaling handle()
+		// to stop writing to the ResponseWriter.
+		stopCh := make(chan struct{})
+		// doneCh is closed when handle() exits, allowing the handler to wait
+		// for the goroutine to fully stop before returning.
+		doneCh := make(chan struct{})
+
+		go func() {
+			defer close(doneCh)
+			sse.handle(log, c, stopCh)
+		}()
+
+		// Call the next handler(s) in the chain. When they return, signal
+		// the handle goroutine to stop, then wait for it to finish before
+		// returning control to the HTTP server.
+		c.Next()
+		close(stopCh)
+		<-doneCh
 	}
 }
 
@@ -59,7 +76,8 @@ func newOptions(opts []Options) Options {
 	return opts[0]
 }
 
-func (c *connection) handle(log *log.Logger, w flamego.ResponseWriter) {
+func (c *connection) handle(log *log.Logger, ctx flamego.Context, stopCh <-chan struct{}) {
+	w := ctx.ResponseWriter()
 	ticker := time.NewTicker(c.PingInterval)
 	defer func() { ticker.Stop() }()
 
@@ -78,11 +96,15 @@ func (c *connection) handle(log *log.Logger, w flamego.ResponseWriter) {
 		senderSend = iota
 		tickerTick
 		timeout
+		closed
+		stopped
 	)
-	cases := make([]reflect.SelectCase, 3)
+	cases := make([]reflect.SelectCase, 5)
 	cases[senderSend] = reflect.SelectCase{Dir: reflect.SelectRecv, Chan: c.sender, Send: reflect.ValueOf(nil)}
 	cases[tickerTick] = reflect.SelectCase{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(ticker.C), Send: reflect.ValueOf(nil)}
 	cases[timeout] = reflect.SelectCase{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(time.After(time.Hour)), Send: reflect.ValueOf(nil)}
+	cases[closed] = reflect.SelectCase{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(ctx.Request().Context().Done()), Send: reflect.ValueOf(nil)}
+	cases[stopped] = reflect.SelectCase{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(stopCh), Send: reflect.ValueOf(nil)}
 
 loop:
 	for {
@@ -112,6 +134,12 @@ loop:
 			write("events: stream timeout\n\n")
 			w.Flush()
 			break loop
+
+		case closed:
+			return
+
+		case stopped:
+			return
 		}
 	}
 

sse_test.go

@@ -6,6 +6,7 @@ package sse
 
 import (
 	"bytes"
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"sync"
@@ -63,6 +64,27 @@ func TestBind(t *testing.T) {
 			time.Sleep(1 * time.Second)
 		},
 	)
+	f.Get("/ticker",
+		Bind(
+			object{},
+			Options{
+				100 * time.Millisecond,
+			},
+		),
+		func(ctx flamego.Context, msg chan<- *object) {
+			ticker := time.NewTicker(1 * time.Second)
+			defer func() { ticker.Stop() }()
+
+			for {
+				select {
+				case <-ticker.C:
+					msg <- &object{Message: "Flamego"}
+				case <-ctx.Request().Context().Done():
+					return
+				}
+			}
+		},
+	)
 
 	t.Run("normal", func(t *testing.T) {
 		resp := &mockResponseWriter{
@@ -118,4 +140,25 @@ data: {"Message":"Flamego"}
 `
 		assert.Equal(t, wantBody, resp.Body())
 	})
+
+	t.Run("close connection", func(t *testing.T) {
+		server := httptest.NewServer(f)
+
+		reqContext, cancel := context.WithCancel(context.Background())
+		req, err := http.NewRequestWithContext(reqContext, http.MethodGet, server.URL+"/ticker", nil)
+		require.NoError(t, err)
+
+		// Close request connection after 1 second.
+		go func() {
+			time.Sleep(1 * time.Second)
+			cancel()
+		}()
+		resp, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		err = resp.Body.Close()
+		require.NoError(t, err)
+
+		// Sleep for 3 seconds to wait for new responses that may be in a closed request.
+		time.Sleep(3 * time.Second)
+	})
 }