From 9416a681c09d81bd8a62c51dfef6373a5c92ef68 Mon Sep 17 00:00:00 2001
From: Raul Jordan <raul@prysmaticlabs.com>
Date: Thu, 17 Sep 2020 14:15:43 -0500
Subject: [PATCH] socket improvements
---
shared/logutil/stream.go | 11 ++++++++++-
1 file changed, 10 insertions(+), 1 deletion(-)
diff --git a/shared/logutil/stream.go b/shared/logutil/stream.go
index 2dd0ce40ba16..3054477533e1 100644
--- a/shared/logutil/stream.go
+++ b/shared/logutil/stream.go
@@ -42,8 +42,14 @@ func (ss *StreamServer) Handler(w http.ResponseWriter, r *http.Request) {
log.Errorf("Could not write websocket message: %v", err)
return
}
+ defer func() {
+ if err := conn.Close(); err != nil {
+ log.Errorf("Could not close websocket connection: %v", err)
+ }
+ }()
- ch := make(chan []byte)
+ ch := make(chan []byte, 1)
+ defer close(ch)
sub := ss.feed.Subscribe(ch)
defer sub.Unsubscribe()
@@ -52,14 +58,17 @@ func (ss *StreamServer) Handler(w http.ResponseWriter, r *http.Request) {
case evt := <-ch:
if err := conn.WriteMessage(websocket.TextMessage, evt); err != nil {
log.Errorf("Could not write websocket message: %v", err)
+ return
}
case <-r.Context().Done():
if err := conn.WriteMessage(websocket.CloseNormalClosure, []byte("context canceled")); err != nil {
log.Error(err)
+ return
}
case err := <-sub.Err():
if err := conn.WriteMessage(websocket.CloseInternalServerErr, []byte(err.Error())); err != nil {
log.Error(err)
+ return
}
}
}