imap-copier/internal/httpapi/ws.go

package httpapi

import (
	"context"
	"net/http"
	"strconv"
	"time"

	"github.com/coder/websocket"
	"github.com/coder/websocket/wsjson"
)

func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) {
	taskID, err := strconv.ParseInt(r.URL.Query().Get("task_id"), 10, 64)
	if err != nil {
		http.Error(w, "task_id required", http.StatusBadRequest)
		return
	}
	c, err := websocket.Accept(w, r, nil)
	if err != nil {
		return
	}
	defer c.CloseNow()

	subID, ch := s.hub.Subscribe(taskID)
	defer s.hub.Unsubscribe(taskID, subID)

	// websocket.Accept хайджекает соединение, поэтому r.Context() не отменяется
	// при обрыве связи клиентом. CloseRead запускает фоновое чтение control-фреймов
	// и отменяет возвращаемый контекст, когда соединение действительно умирает.
	ctx := c.CloseRead(r.Context())
	for {
		select {
		case ev, ok := <-ch:
			if !ok {
				return
			}
			wctx, cancel := context.WithTimeout(ctx, 5*time.Second)
			err := wsjson.Write(wctx, c, ev)
			cancel()
			if err != nil {
				return
			}
		case <-ctx.Done():
			return
		}
	}
}