|
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 | | -iferr!=nil { |
74 | | -cs.logf("%v",err) |
75 | | -return |
76 | | -} |
77 | | -deferc.CloseNow() |
78 | | - |
79 | | -err=cs.subscribe(r.Context(),c) |
| 73 | +err:=cs.subscribe(r.Context(),w,r) |
80 | 74 | iferrors.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 | +varmu sync.Mutex |
| 116 | +varc*websocket.Conn |
| 117 | +varclosedbool |
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 | +defermu.Unlock() |
| 123 | +closed=true |
| 124 | +ifc!=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 | defercs.deleteSubscriber(s) |
131 | 131 |
|
| 132 | +c2,err:=websocket.Accept(w,r,nil) |
| 133 | +iferr!=nil { |
| 134 | +returnerr |
| 135 | +} |
| 136 | +mu.Lock() |
| 137 | +ifclosed { |
| 138 | +mu.Unlock() |
| 139 | +returnnet.ErrClosed |
| 140 | +} |
| 141 | +c=c2 |
| 142 | +mu.Unlock() |
| 143 | +deferc.CloseNow() |
| 144 | + |
| 145 | +ctx=c.CloseRead(ctx) |
| 146 | + |
132 | 147 | for { |
133 | 148 | select { |
134 | 149 | casemsg:=<-s.msgs: |
|