From 9abbe345d8d9b2fc7af962a6dd09bff9967ae157 Mon Sep 17 00:00:00 2001 From: STeve Huang Date: Mon, 25 Mar 2024 16:27:44 -0400 Subject: [PATCH] Fix websocketALPNServerConn flaky test --- lib/web/conn_upgrade.go | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/lib/web/conn_upgrade.go b/lib/web/conn_upgrade.go index 3cb10f64b3ff5..1bcbf96bb323e 100644 --- a/lib/web/conn_upgrade.go +++ b/lib/web/conn_upgrade.go @@ -326,3 +326,15 @@ func (c *websocketALPNServerConn) SetDeadline(t time.Time) error { c.Conn.SetWriteDeadline(t), ) } + +func (c *websocketALPNServerConn) SetWriteDeadline(t time.Time) error { + c.writeMutex.Lock() + defer c.writeMutex.Unlock() + return trace.Wrap(c.Conn.SetWriteDeadline(t)) +} + +func (c *websocketALPNServerConn) SetReadDeadline(t time.Time) error { + c.writeMutex.Lock() + defer c.writeMutex.Unlock() + return trace.Wrap(c.Conn.SetReadDeadline(t)) +}