fix: prevent context cancellation for critical write (#1151)

* fix: prevent context cancellation for critical write

* fix

* cs

Author: dunglas

publish.go

@@ -105,10 +105,12 @@ func (h *Hub) PublishHandler(w http.ResponseWriter, r *http.Request) {
 		Debug:   h.debug,
 		Event:   Event{r.PostForm.Get("data"), r.PostForm.Get("id"), r.PostForm.Get("type"), retry},
 	}
+
 	ctx = context.WithValue(ctx, UpdateContextKey, u)
+	dispatchCtx := context.WithoutCancel(ctx)
 
 	// Broadcast the update
-	if err := h.transport.Dispatch(ctx, u); err != nil {
+	if err := h.transport.Dispatch(dispatchCtx, u); err != nil {
 		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
 
 		if h.logger.Enabled(ctx, slog.LevelError) {

subscribe.go

@@ -213,11 +213,12 @@ func (h *Hub) registerSubscriber(ctx context.Context, w http.ResponseWriter, r *
 
 	s.SetTopics(topics, privateTopics)
 
-	h.dispatchSubscriptionUpdate(ctx, s, true)
+	addCtx := context.WithoutCancel(ctx)
+	h.dispatchSubscriptionUpdate(addCtx, s, true)
 
-	if err := h.transport.AddSubscriber(ctx, s); err != nil {
+	if err := h.transport.AddSubscriber(addCtx, s); err != nil {
 		http.Error(w, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable)
-		h.dispatchSubscriptionUpdate(ctx, s, false)
+		h.dispatchSubscriptionUpdate(addCtx, s, false)
 
 		if h.logger.Enabled(ctx, slog.LevelError) {
 			h.logger.LogAttrs(ctx, slog.LevelError, "Unable to add subscriber", slog.Any("error", err))
@@ -331,6 +332,8 @@ func (h *Hub) shutdown(ctx context.Context, s *LocalSubscriber) {
 	// Notify that the client is closing the connection
 	s.Disconnect()
 
+	ctx = context.WithoutCancel(ctx)
+
 	if err := h.transport.RemoveSubscriber(ctx, s); err != nil && h.logger.Enabled(ctx, slog.LevelError) {
 		h.logger.LogAttrs(ctx, slog.LevelError, "Failed to remove subscriber on shutdown", slog.Any("error", err))
 	}