46 lines
805 B
Go
46 lines
805 B
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)
|
|
|
|
ctx := 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
|
|
}
|
|
}
|
|
}
|