From 381e8be8f305df03ed51ee7777cad909f297b9cc Mon Sep 17 00:00:00 2001 From: E99p1ant Date: Mon, 5 May 2025 02:30:42 +0800 Subject: [PATCH 1/3] sse: add `closed` SelectCase to handle request closing signal --- sse.go | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/sse.go b/sse.go index 480261f..8b031b1 100644 --- a/sse.go +++ b/sse.go @@ -45,7 +45,7 @@ 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()) + go sse.handle(log, c) } } @@ -59,7 +59,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) { + w := ctx.ResponseWriter() ticker := time.NewTicker(c.PingInterval) defer func() { ticker.Stop() }() @@ -78,11 +79,13 @@ func (c *connection) handle(log *log.Logger, w flamego.ResponseWriter) { senderSend = iota tickerTick timeout + closed ) - cases := make([]reflect.SelectCase, 3) + cases := make([]reflect.SelectCase, 4) 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)} loop: for { @@ -112,6 +115,9 @@ loop: write("events: stream timeout\n\n") w.Flush() break loop + + case closed: + return } } From f64ccc0fed5c4f06153d204e15f5d8359621cef8 Mon Sep 17 00:00:00 2001 From: E99p1ant Date: Mon, 5 May 2025 03:21:15 +0800 Subject: [PATCH 2/3] add test case --- sse_test.go | 43 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/sse_test.go b/sse_test.go index 7c82483..717aa8c 100644 --- a/sse_test.go +++ b/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) + }) } From d441660ade1f46e59eace632c6b0f06a492f1f45 Mon Sep 17 00:00:00 2001 From: E99p1ant Date: Sat, 7 Feb 2026 22:20:32 +0800 Subject: [PATCH 3/3] fix test Signed-off-by: E99p1ant --- sse.go | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/sse.go b/sse.go index 8b031b1..5866ed8 100644 --- a/sse.go +++ b/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) + // 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,7 @@ func newOptions(opts []Options) Options { return opts[0] } -func (c *connection) handle(log *log.Logger, ctx flamego.Context) { +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() }() @@ -80,12 +97,14 @@ func (c *connection) handle(log *log.Logger, ctx flamego.Context) { tickerTick timeout closed + stopped ) - cases := make([]reflect.SelectCase, 4) + 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 { @@ -118,6 +137,9 @@ loop: case closed: return + + case stopped: + return } }