|
5 | 5 | "errors" |
6 | 6 | "io" |
7 | 7 | "log" |
| 8 | + "net" |
8 | 9 | "net/http" |
9 | 10 | "sync" |
10 | 11 | "time" |
@@ -69,14 +70,7 @@ func (cs *chatServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { |
69 | 70 | // subscribeHandler accepts the WebSocket connection and then subscribes |
70 | 71 | // it to all future messages. |
71 | 72 | func (cs *chatServer) subscribeHandler(w http.ResponseWriter, r *http.Request) { |
72 | | - c, err := websocket.Accept(w, r, nil) |
73 | | - if err != nil { |
74 | | - cs.logf("%v", err) |
75 | | - return |
76 | | - } |
77 | | - defer c.CloseNow() |
78 | | - |
79 | | - err = cs.subscribe(r.Context(), c) |
| 73 | + err := cs.subscribe(r.Context(), w, r) |
80 | 74 | if errors.Is(err, context.Canceled) { |
81 | 75 | return |
82 | 76 | } |
@@ -117,18 +111,39 @@ func (cs *chatServer) publishHandler(w http.ResponseWriter, r *http.Request) { |
117 | 111 | // |
118 | 112 | // It uses CloseRead to keep reading from the connection to process control |
119 | 113 | // messages and cancel the context if the connection drops. |
120 | | -func (cs *chatServer) subscribe(ctx context.Context, c *websocket.Conn) error { |
121 | | - ctx = c.CloseRead(ctx) |
122 | | - |
| 114 | +func (cs *chatServer) subscribe(ctx context.Context, w http.ResponseWriter, r *http.Request) error { |
| 115 | + var mu sync.Mutex |
| 116 | + var c *websocket.Conn |
| 117 | + var closed bool |
123 | 118 | s := &subscriber{ |
124 | 119 | msgs: make(chan []byte, cs.subscriberMessageBuffer), |
125 | 120 | closeSlow: func() { |
126 | | - c.Close(websocket.StatusPolicyViolation, "connection too slow to keep up with messages") |
| 121 | + mu.Lock() |
| 122 | + defer mu.Unlock() |
| 123 | + closed = true |
| 124 | + if c != nil { |
| 125 | + c.Close(websocket.StatusPolicyViolation, "connection too slow to keep up with messages") |
| 126 | + } |
127 | 127 | }, |
128 | 128 | } |
129 | 129 | cs.addSubscriber(s) |
130 | 130 | defer cs.deleteSubscriber(s) |
131 | 131 |
|
| 132 | + c2, err := websocket.Accept(w, r, nil) |
| 133 | + if err != nil { |
| 134 | + return err |
| 135 | + } |
| 136 | + mu.Lock() |
| 137 | + if closed { |
| 138 | + mu.Unlock() |
| 139 | + return net.ErrClosed |
| 140 | + } |
| 141 | + c = c2 |
| 142 | + mu.Unlock() |
| 143 | + defer c.CloseNow() |
| 144 | + |
| 145 | + ctx = c.CloseRead(ctx) |
| 146 | + |
132 | 147 | for { |
133 | 148 | select { |
134 | 149 | case msg := <-s.msgs: |
|
0 commit comments