MM-33233: Fix double close of webconn pump (#17026)

Automatic Merge
Этот коммит содержится в:
Agniva De Sarker
2021-03-01 19:22:27 +05:30
коммит произвёл GitHub
родитель 0319daf9bb
Коммит 0f98620b65
2 изменённых файлов: 28 добавлений и 8 удалений

Просмотреть файл

@@ -58,6 +58,7 @@ type WebConn struct {
isWindows bool
endWritePump chan struct{}
pumpFinished chan struct{}
closeOnce sync.Once
}
// NewWebConn returns a new WebConn instance.
@@ -94,15 +95,20 @@ func (a *App) NewWebConn(ws net.Conn, session model.Session, t i18n.TranslateFun
}
// Close closes the WebConn.
// It is made idempotent in nature by using a sync.Once
// to avoid a race condition that happens when an EventReadHup event
// and a connection close event happens at the same time.
func (wc *WebConn) Close() {
wc.WebSocket.Close()
if !wc.isWindows {
// This triggers the pump exit.
// If the pump has already exited, this just becomes a noop.
close(wc.endWritePump)
}
// We wait for the pump to fully exit.
<-wc.pumpFinished
wc.closeOnce.Do(func() {
wc.WebSocket.Close()
if !wc.isWindows {
// This triggers the pump exit.
// If the pump has already exited, this just becomes a noop.
close(wc.endWritePump)
}
// We wait for the pump to fully exit.
<-wc.pumpFinished
})
}
// GetSessionExpiresAt returns the time at which the session expires.

Просмотреть файл

@@ -90,6 +90,20 @@ func TestHubStopWithMultipleConnections(t *testing.T) {
defer wc3.Close()
}
func TestWebConnDoubleClose(t *testing.T) {
th := Setup(t)
defer th.TearDown()
s := httptest.NewServer(dummyWebsocketHandler(t))
defer s.Close()
wc1 := registerDummyWebConn(t, th.App, s.Listener.Addr(), "userID")
wc1.Close()
require.NotPanics(t, func() {
wc1.Close()
})
}
// TestHubStopRaceCondition verifies that attempts to use the hub after it has shutdown does not
// block the caller indefinitely.
func TestHubStopRaceCondition(t *testing.T) {