🐞 fix: zhttp sse context cancel failure

sohaha · Aug 26, 2024 · 3c2bc86 · 3c2bc86
1 parent 906d21e
commit 3c2bc86
Showing 1 changed file with 17 additions and 1 deletion.
diff --git a/zhttp/sse.go b/zhttp/sse.go
@@ -128,7 +128,23 @@ func (e *Engine) SSE(url string, opt func(*SSEOption), v ...interface{}) (*SSEEn
 	if opt != nil {
 		opt(&o)
 	}
-	ctx, cancel := context.WithCancel(context.TODO())
+	var (
+		ctx    context.Context
+		cancel context.CancelFunc
+	)
+
+	for i := range v {
+		if c, ok := v[i].(context.Context); ok {
+			ctx = c
+		}
+	}
+
+	if ctx == nil {
+		ctx, cancel = context.WithCancel(context.TODO())
+	} else {
+		ctx, cancel = context.WithCancel(ctx)
+	}
+
 	sse := &SSEEngine{
 		readyState: 0,
 		ctx:        ctx,