Fixed data races in PipeConns

2 files changed, 56 insertions, 4 deletions

diff --git a/fasthttputil/pipeconns.go b/fasthttputil/pipeconns.go
index bfcdbde..24f09d8 100644
--- a/fasthttputil/pipeconns.go
+++ b/fasthttputil/pipeconns.go
@@ -120,10 +120,11 @@ func (c *pipeConn) Read(p []byte) (int, error) {
 
 func (c *pipeConn) read(p []byte, mayBlock bool) (int, error) {
 	if len(c.bb) == 0 {
+		c.rlock.Lock()
+
 		releaseByteBuffer(c.b)
 		c.b = nil
 
-		c.rlock.Lock()
 		if c.rclosed {
 			c.rlock.Unlock()
 			return 0, io.EOF
@@ -145,9 +146,8 @@ func (c *pipeConn) read(p []byte, mayBlock bool) (int, error) {
 			c.rlock.Unlock()
 			return 0, io.EOF
 		}
-		c.rlock.Unlock()
-
 		c.bb = c.b.b
+		c.rlock.Unlock()
 	}
 	n := copy(p, c.bb)
 	c.bb = c.bb[n:]
@@ -181,7 +181,6 @@ func (c *pipeConn) release() {
 
 	releaseByteBuffer(c.b)
 	c.b = nil
-	c.bb = nil
 
 	if !c.rclosed {
 		c.rclosed = true
diff --git a/fasthttputil/pipeconns_test.go b/fasthttputil/pipeconns_test.go
index b4bab0e..c39ecc1 100644
--- a/fasthttputil/pipeconns_test.go
+++ b/fasthttputil/pipeconns_test.go
@@ -3,11 +3,64 @@ package fasthttputil
 import (
 	"fmt"
 	"io"
+	"io/ioutil"
 	"net"
 	"testing"
 	"time"
 )
 
+func TestPipeConnsCloseWhileReadWrite(t *testing.T) {
+	for i := 0; i < 10; i++ {
+		testPipeConnsCloseWhileReadWrite(t)
+	}
+}
+
+func testPipeConnsCloseWhileReadWrite(t *testing.T) {
+	pc := NewPipeConns()
+	readCh := make(chan struct{})
+	writeCh := make(chan struct{})
+	go func() {
+		_, err := io.Copy(ioutil.Discard, pc.Conn1())
+		if err != nil {
+			if err != errConnectionClosed {
+				t.Fatalf("unexpected error: %s", err)
+			}
+		}
+		close(readCh)
+	}()
+	go func() {
+		for {
+			_, err := pc.Conn2().Write([]byte("foobar"))
+			if err != nil {
+				if err != errConnectionClosed {
+					t.Fatalf("unexpected error: %s", err)
+				}
+				break
+			}
+		}
+		close(writeCh)
+	}()
+
+	time.Sleep(10 * time.Millisecond)
+	if err := pc.Conn1().Close(); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if err := pc.Conn2().Close(); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	select {
+	case <-readCh:
+	case <-time.After(time.Second):
+		t.Fatalf("timeout")
+	}
+	select {
+	case <-writeCh:
+	case <-time.After(time.Second):
+		t.Fatalf("timeout")
+	}
+}
+
 func TestPipeConnsReadWriteSerial(t *testing.T) {
 	testPipeConnsReadWriteSerial(t)
 }