diff --git a/util/wsproxy/server/proxy.go b/util/wsproxy/server/proxy.go index 1ee55ed6f..512d42ac9 100644 --- a/util/wsproxy/server/proxy.go +++ b/util/wsproxy/server/proxy.go @@ -2,6 +2,7 @@ package server import ( "net/http" + "sync/atomic" "time" "github.com/coder/websocket" @@ -59,7 +60,7 @@ func (p *Proxy) Handler() http.Handler { type proxyHandler struct { metrics MetricsRecorder handler http.Handler - conn *wsConnAdapter + conn atomic.Pointer[wsConnAdapter] headersReadTimeout time.Duration } @@ -91,7 +92,7 @@ func (ph *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { _ = serverConn.Close() }() - ph.conn = serverConn // used in tests only + ph.conn.Store(serverConn) // used in tests only log.Debugf("WebSocket proxy established: %s -> gRPC handler", r.RemoteAddr) diff --git a/util/wsproxy/server/ws_conn_adapter_test.go b/util/wsproxy/server/ws_conn_adapter_test.go index 459ac35ef..b701f59cd 100644 --- a/util/wsproxy/server/ws_conn_adapter_test.go +++ b/util/wsproxy/server/ws_conn_adapter_test.go @@ -90,7 +90,7 @@ func TestAdapterHandlingConnectionClosures(t *testing.T) { case 0: clientconn.Close(websocket.StatusNormalClosure, "") case 1: - handler.conn.Close() + handler.conn.Load().Close() case 2: cancel() case 3: @@ -99,7 +99,7 @@ func TestAdapterHandlingConnectionClosures(t *testing.T) { } assert.EventuallyWithT(t, func(c *assert.CollectT) { - assert.True(c, handler.conn.IsClosed()) + assert.True(c, handler.conn.Load().IsClosed()) }, 3*time.Second, 100*time.Millisecond) }) } @@ -154,7 +154,7 @@ func TestAdapterHandlingHttpConnection_NoHeadersSent(t *testing.T) { assert.Error(t, err) assert.EventuallyWithT(t, func(c *assert.CollectT) { - assert.True(c, handler.conn.IsClosed()) + assert.True(c, handler.conn.Load().IsClosed()) }, 3*time.Second, 100*time.Millisecond) }