From a293989e7c4c82d2228bff8c8470a71f6f8d8df8 Mon Sep 17 00:00:00 2001 From: Steven Allen Date: Tue, 5 Sep 2017 11:48:46 -0700 Subject: [PATCH 1/4] Make close thread safe 1. Ensure we only close the connection once. Especially, don't call the done function multiple times and/or concurrently. 2. Call WriteControl instead of WriteMessage in Close. WriteControl is thread-safe, WriteMessage isn't. Also, this sets a 100ms deadline on gracefully closing connections. --- p2p/transport/websocket/conn.go | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/p2p/transport/websocket/conn.go b/p2p/transport/websocket/conn.go index 1d8de2a01f..0ab71ce220 100644 --- a/p2p/transport/websocket/conn.go +++ b/p2p/transport/websocket/conn.go @@ -3,11 +3,16 @@ package websocket import ( "io" "net" + "sync" "time" ws "github.com/gorilla/websocket" ) +// GracefulCloseTimeout is the time to wait trying to gracefully close a +// connection before simply cutting it. +var GracefulCloseTimeout = 100 * time.Millisecond + var _ net.Conn = (*Conn)(nil) // Conn implements net.Conn interface for gorilla/websocket. @@ -16,6 +21,7 @@ type Conn struct { DefaultMessageType int done func() reader io.Reader + closeOnce sync.Once } func (c *Conn) Read(b []byte) (int, error) { @@ -73,13 +79,22 @@ func (c *Conn) Write(b []byte) (n int, err error) { return len(b), nil } +// Close closes the connection. Only the first call to Close will receive the +// close error, subsequent and concurrent calls will return nil. +// This method is thread-safe. func (c *Conn) Close() error { - if c.done != nil { - c.done() - } + var err error = nil + c.closeOnce.Do(func() { + if c.done != nil { + c.done() + // Be nice to GC + c.done = nil + } - c.Conn.WriteMessage(ws.CloseMessage, nil) - return c.Conn.Close() + c.Conn.WriteControl(ws.CloseMessage, nil, time.Now().Add(GracefulCloseTimeout)) + err = c.Conn.Close() + }) + return err } func (c *Conn) LocalAddr() net.Addr { From b863f44dabd3c8a230a6c4f0b57cfeee4576abc8 Mon Sep 17 00:00:00 2001 From: Steven Allen Date: Tue, 5 Sep 2017 12:12:21 -0700 Subject: [PATCH 2/4] test concurrent connection closing --- p2p/transport/websocket/websocket_test.go | 38 +++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/p2p/transport/websocket/websocket_test.go b/p2p/transport/websocket/websocket_test.go index 7aab918576..d2e166a685 100644 --- a/p2p/transport/websocket/websocket_test.go +++ b/p2p/transport/websocket/websocket_test.go @@ -53,3 +53,41 @@ func TestWebsocketListen(t *testing.T) { t.Fatal("got wrong message", out, msg) } } + +func TestConcurrentClose(t *testing.T) { + zero, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/0/ws") + if err != nil { + t.Fatal(err) + } + + tpt := &WebsocketTransport{} + l, err := tpt.Listen(zero) + if err != nil { + t.Fatal(err) + } + defer l.Close() + + msg := []byte("HELLO WORLD") + + go func() { + d, _ := tpt.Dialer(nil) + for i := 0; i < 100; i++ { + c, err := d.Dial(l.Multiaddr()) + if err != nil { + t.Error(err) + return + } + + go c.Write(msg) + go c.Close() + } + }() + + for i := 0; i < 100; i++ { + c, err := l.Accept() + if err != nil { + t.Fatal(err) + } + c.Close() + } +} From a88b185db2f234ec5743b9ec9810437fa856e669 Mon Sep 17 00:00:00 2001 From: Steven Allen Date: Tue, 5 Sep 2017 12:25:29 -0700 Subject: [PATCH 3/4] Add write-zero test. Just to be thorough. --- p2p/transport/websocket/websocket_test.go | 52 +++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/p2p/transport/websocket/websocket_test.go b/p2p/transport/websocket/websocket_test.go index d2e166a685..61bfa14230 100644 --- a/p2p/transport/websocket/websocket_test.go +++ b/p2p/transport/websocket/websocket_test.go @@ -2,6 +2,7 @@ package websocket import ( "bytes" + "io" "io/ioutil" "testing" "testing/iotest" @@ -91,3 +92,54 @@ func TestConcurrentClose(t *testing.T) { c.Close() } } + +func TestWriteZero(t *testing.T) { + zero, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/0/ws") + if err != nil { + t.Fatal(err) + } + + tpt := &WebsocketTransport{} + l, err := tpt.Listen(zero) + if err != nil { + t.Fatal(err) + } + defer l.Close() + + msg := []byte(nil) + + go func() { + d, _ := tpt.Dialer(nil) + c, err := d.Dial(l.Multiaddr()) + defer c.Close() + if err != nil { + t.Error(err) + return + } + + for i := 0; i < 100; i++ { + n, err := c.Write(msg) + if n != 0 { + t.Errorf("expected to write 0 bytes, wrote %d", n) + } + if err != nil { + t.Error(err) + return + } + } + }() + + c, err := l.Accept() + defer c.Close() + if err != nil { + t.Fatal(err) + } + buf := make([]byte, 100) + n, err := c.Read(buf) + if n != 0 { + t.Errorf("read %d bytes, expected 0", n) + } + if err != io.EOF { + t.Errorf("expected EOF, got err: %s", err) + } +} From da59505d57802250e79733282e7a296c8f81db13 Mon Sep 17 00:00:00 2001 From: Steven Allen Date: Tue, 5 Sep 2017 12:27:45 -0700 Subject: [PATCH 4/4] don't explicitly nil err (make @whyrusleeping happy) --- p2p/transport/websocket/conn.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/p2p/transport/websocket/conn.go b/p2p/transport/websocket/conn.go index 0ab71ce220..168497bf39 100644 --- a/p2p/transport/websocket/conn.go +++ b/p2p/transport/websocket/conn.go @@ -83,7 +83,7 @@ func (c *Conn) Write(b []byte) (n int, err error) { // close error, subsequent and concurrent calls will return nil. // This method is thread-safe. func (c *Conn) Close() error { - var err error = nil + var err error c.closeOnce.Do(func() { if c.done != nil { c.done()