feat: implement polling for WebSocket messages
Refactor WebSocket message handling from streaming to polling mechanism. Add messageQueue, queueMu, and isClosed fields to WSConnection struct. Introduce manageQueue goroutine to enqueue messages and handle connection close. Add GetMessages() and IsClosed() methods for client polling. Update endpoint from /sw-cgi/websocket/stream/ to /sw-cgi/websocket/poll/ and replace handleStream with handlePoll. This change improves client control over message retrieval and reduces server-side streaming overhead.
lemon-mint committed
Nov 1, 2025 at 02:20 UTC
611b297cc29c4b3304fa9711b69740f1f3b95a29
2 files changed
+123
-115
cmd/webclient/main_js.go
+92
-68
@@ -62,12 +62,15 @@ type WebSocketManager struct {
62
}
63
64
type WSConnection struct {
65
- id string
66
- conn *websocket.Conn
67
- messageChan chan wsMessage
68
- closeChan chan struct{}
69
- closeOnce sync.Once
70
- mu sync.Mutex
65
+ id string
66
+ conn *websocket.Conn
67
+ messageChan chan wsMessage
68
+ closeChan chan struct{}
69
+ closeOnce sync.Once
70
+ mu sync.Mutex
71
+ messageQueue []StreamMessage
72
+ queueMu sync.Mutex
73
+ isClosed bool
74
}
75
76
type wsMessage struct {
@@ -138,16 +141,18 @@ func (m *WebSocketManager) CreateConnection(uri string, protocols []string) (*WS
141
}
142
143
wsConn := &WSConnection{
141
- id: generateConnID(),
142
- conn: conn,
143
- messageChan: make(chan wsMessage, 100),
144
- closeChan: make(chan struct{}),
144
+ id: generateConnID(),
145
+ conn: conn,
146
+ messageChan: make(chan wsMessage, 100),
147
+ closeChan: make(chan struct{}),
148
+ messageQueue: make([]StreamMessage, 0),
149
}
150
151
m.connections.Store(wsConn.id, wsConn)
152
149
- // Start message receiver
153
+ // Start message receiver and queue manager
154
go wsConn.receiveMessages()
155
+ go wsConn.manageQueue()
156
157
return wsConn, negotiatedProtocol, nil
158
}
@@ -171,6 +176,9 @@ func (c *WSConnection) receiveMessages() {
176
messageType, msg, err := c.conn.ReadMessage()
177
if err != nil {
178
log.Error().Err(err).Str("connId", c.id).Msg("Error receiving message")
179
+ c.queueMu.Lock()
180
+ c.isClosed = true
181
+ c.queueMu.Unlock()
182
return
183
}
184
@@ -192,6 +200,57 @@ func (c *WSConnection) receiveMessages() {
200
}
201
}
202
203
+func (c *WSConnection) manageQueue() {
204
+ for {
205
+ select {
206
+ case msg := <-c.messageChan:
207
+ c.queueMu.Lock()
208
+
209
+ // Use message type from WebSocket frame
210
+ messageType := "binary"
211
+ if msg.isText {
212
+ messageType = "text"
213
+ }
214
+
215
+ streamMsg := StreamMessage{
216
+ Type: "message",
217
+ Data: base64.StdEncoding.EncodeToString(msg.data),
218
+ MessageType: messageType,
219
+ }
220
+ c.messageQueue = append(c.messageQueue, streamMsg)
221
+ c.queueMu.Unlock()
222
+
223
+ case <-c.closeChan:
224
+ c.queueMu.Lock()
225
+ c.isClosed = true
226
+ c.messageQueue = append(c.messageQueue, StreamMessage{
227
+ Type: "close",
228
+ Code: 1000,
229
+ Reason: "Connection closed",
230
+ })
231
+ c.queueMu.Unlock()
232
+ return
233
+ }
234
+ }
235
+}
236
+
237
+func (c *WSConnection) GetMessages() []StreamMessage {
238
+ c.queueMu.Lock()
239
+ defer c.queueMu.Unlock()
240
+
241
+ messages := make([]StreamMessage, len(c.messageQueue))
242
+ copy(messages, c.messageQueue)
243
+ c.messageQueue = c.messageQueue[:0]
244
+
245
+ return messages
246
+}
247
+
248
+func (c *WSConnection) IsClosed() bool {
249
+ c.queueMu.Lock()
250
+ defer c.queueMu.Unlock()
251
+ return c.isClosed
252
+}
253
+
254
func (c *WSConnection) Send(data []byte, isText bool) error {
255
c.mu.Lock()
256
defer c.mu.Unlock()
@@ -294,9 +353,9 @@ func (p *Proxy) handleWebSocketPolyfill(w http.ResponseWriter, r *http.Request)
353
return
354
}
355
297
- if strings.HasPrefix(path, "/sw-cgi/websocket/stream/") && r.Method == http.MethodGet {
298
- connID := strings.TrimPrefix(path, "/sw-cgi/websocket/stream/")
299
- p.handleStream(w, r, connID)
356
+ if strings.HasPrefix(path, "/sw-cgi/websocket/poll/") && r.Method == http.MethodGet {
357
+ connID := strings.TrimPrefix(path, "/sw-cgi/websocket/poll/")
358
+ p.handlePoll(w, r, connID)
359
return
360
}
361
@@ -334,70 +393,35 @@ func (p *Proxy) handleConnect(w http.ResponseWriter, r *http.Request) {
393
json.NewEncoder(w).Encode(resp)
394
}
395
337
-func (p *Proxy) handleStream(w http.ResponseWriter, r *http.Request, connID string) {
396
+func (p *Proxy) handlePoll(w http.ResponseWriter, r *http.Request, connID string) {
397
wsConn, ok := p.wsManager.GetConnection(connID)
398
if !ok {
399
http.Error(w, "Connection not found", http.StatusNotFound)
400
return
401
}
402
344
- log.Info().Str("connId", connID).Msg("Starting message stream")
345
-
346
- // Cleanup on exit
347
- defer func() {
348
- p.wsManager.RemoveConnection(connID)
349
- wsConn.Close()
350
- }()
351
-
352
- // Set headers for streaming
353
- w.Header().Set("Content-Type", "text/plain; charset=utf-8")
354
- w.Header().Set("Cache-Control", "no-cache")
355
- w.Header().Set("Connection", "keep-alive")
356
- w.Header().Set("X-Content-Type-Options", "nosniff")
357
-
358
- flusher, ok := w.(http.Flusher)
359
- if !ok {
360
- http.Error(w, "Streaming not supported", http.StatusInternalServerError)
361
- return
362
- }
363
-
364
- // Send messages as newline-delimited JSON
365
- encoder := json.NewEncoder(w)
366
- for {
367
- select {
368
- case msg := <-wsConn.messageChan:
369
- // Use message type from WebSocket frame
370
- messageType := "binary"
371
- if msg.isText {
372
- messageType = "text"
373
- }
374
-
375
- streamMsg := StreamMessage{
376
- Type: "message",
377
- Data: base64.StdEncoding.EncodeToString(msg.data),
378
- MessageType: messageType,
379
- }
380
- if err := encoder.Encode(streamMsg); err != nil {
381
- log.Error().Err(err).Msg("Failed to encode message")
382
- return
383
- }
384
- flusher.Flush()
403
+ // Get queued messages
404
+ messages := wsConn.GetMessages()
405
386
- case <-wsConn.closeChan:
387
- streamMsg := StreamMessage{
388
- Type: "close",
389
- Code: 1000,
390
- Reason: "Connection closed",
406
+ // Check if connection is closed and cleanup if needed
407
+ if wsConn.IsClosed() && len(messages) > 0 {
408
+ // Check if close message is in the queue
409
+ for _, msg := range messages {
410
+ if msg.Type == "close" {
411
+ defer func() {
412
+ p.wsManager.RemoveConnection(connID)
413
+ wsConn.Close()
414
+ }()
415
+ break
416
}
392
- encoder.Encode(streamMsg)
393
- flusher.Flush()
394
- return
395
-
396
- case <-r.Context().Done():
397
- log.Info().Str("connId", connID).Msg("Stream context cancelled")
398
- return
417
}
418
}
419
+
420
+ w.Header().Set("Content-Type", "application/json")
421
+ w.WriteHeader(http.StatusOK)
422
+ json.NewEncoder(w).Encode(map[string]interface{}{
423
+ "messages": messages,
424
+ })
425
}
426
427
func (p *Proxy) handleSend(w http.ResponseWriter, r *http.Request, connID string) {
cmd/webclient/polyfill.js
+31
-47
@@ -25,7 +25,7 @@
25
this._connId = null;
26
this._sendQueue = [];
27
this._isSending = false;
28
- this._streamAbortController = null;
28
+ this._pollInterval = null;
29
this._isClosed = false;
30
31
// Initialize connection
@@ -63,66 +63,52 @@
63
}
64
this.dispatchEvent(new Event('open'));
65
66
- // Start receiving messages
67
- this._startStream();
66
+ // Start polling for messages
67
+ this._startPolling();
68
69
} catch (error) {
70
this._handleError(error);
71
}
72
}
73
74
- async _startStream() {
74
+ _startPolling() {
75
if (!this._connId || this._isClosed) return;
76
77
- try {
78
- this._streamAbortController = new AbortController();
79
-
80
- const response = await fetch(`/sw-cgi/websocket/stream/${this._connId}`, {
81
- method: 'GET',
82
- signal: this._streamAbortController.signal
83
- });
84
-
85
- if (!response.ok) {
86
- throw new Error(`Stream failed: ${response.status}`);
77
+ // Poll every 100ms
78
+ this._pollInterval = setInterval(async () => {
79
+ if (this._isClosed) {
80
+ clearInterval(this._pollInterval);
81
+ return;
82
}
83
89
- const reader = response.body.getReader();
90
- const decoder = new TextDecoder();
91
- let buffer = '';
92
-
93
- while (!this._isClosed) {
94
- const { done, value } = await reader.read();
84
+ try {
85
+ const response = await fetch(`/sw-cgi/websocket/poll/${this._connId}`, {
86
+ method: 'GET'
87
+ });
88
96
- if (done) {
97
- // Connection closed by server
98
- this._handleClose(1000, 'Connection closed by server');
99
- break;
89
+ if (!response.ok) {
90
+ throw new Error(`Poll failed: ${response.status}`);
91
}
92
102
- // Decode chunk
103
- buffer += decoder.decode(value, { stream: true });
93
+ const result = await response.json();
94
105
- // Process complete messages (assuming newline-delimited JSON)
106
- let newlineIndex;
107
- while ((newlineIndex = buffer.indexOf('\n')) !== -1) {
108
- const line = buffer.slice(0, newlineIndex);
109
- buffer = buffer.slice(newlineIndex + 1);
110
-
111
- if (line.trim()) {
112
- try {
113
- const message = JSON.parse(line);
114
- this._handleMessage(message);
115
- } catch (e) {
116
- console.error('Failed to parse message:', e);
117
- }
95
+ if (result.messages && Array.isArray(result.messages)) {
96
+ for (const message of result.messages) {
97
+ this._handleMessage(message);
98
}
99
}
120
- }
121
-
122
- } catch (error) {
123
- if (error.name !== 'AbortError') {
100
+
101
+ } catch (error) {
102
+ console.error('Polling error:', error);
103
this._handleError(error);
104
}
105
+ }, 100);
106
+ }
107
+
108
+ _stopPolling() {
109
+ if (this._pollInterval) {
110
+ clearInterval(this._pollInterval);
111
+ this._pollInterval = null;
112
}
113
}
114
@@ -194,10 +180,8 @@
180
this._isClosed = true;
181
this.readyState = WebSocket.CLOSED;
182
197
- // Abort stream
198
- if (this._streamAbortController) {
199
- this._streamAbortController.abort();
200
- }
183
+ // Stop polling
184
+ this._stopPolling();
185
186
// Create CloseEvent
187
const event = new CloseEvent('close', {