Browse Source

fix data race in pipe

Darien Raymond 7 years ago
parent
commit
845606d86c
2 changed files with 33 additions and 3 deletions
  1. 10 3
      transport/pipe/impl.go
  2. 23 0
      transport/pipe/pipe_test.go

+ 10 - 3
transport/pipe/impl.go

@@ -1,6 +1,7 @@
 package pipe
 package pipe
 
 
 import (
 import (
+	"errors"
 	"io"
 	"io"
 	"sync"
 	"sync"
 	"time"
 	"time"
@@ -26,9 +27,14 @@ type pipe struct {
 	state       state
 	state       state
 }
 }
 
 
+var errBufferFull = errors.New("buffer full")
+
 func (p *pipe) getState(forRead bool) error {
 func (p *pipe) getState(forRead bool) error {
 	switch p.state {
 	switch p.state {
 	case open:
 	case open:
+		if !forRead && p.limit >= 0 && p.data.Len() > p.limit {
+			return errBufferFull
+		}
 		return nil
 		return nil
 	case closed:
 	case closed:
 		if forRead {
 		if forRead {
@@ -105,9 +111,10 @@ func (p *pipe) WriteMultiBuffer(mb buf.MultiBuffer) error {
 	}
 	}
 
 
 	for {
 	for {
-		if p.limit < 0 || p.data.Len()+mb.Len() <= p.limit {
-			defer p.readSignal.Signal()
-			return p.writeMultiBufferInternal(mb)
+		err := p.writeMultiBufferInternal(mb)
+		if err == nil || err != errBufferFull {
+			p.readSignal.Signal()
+			return err
 		}
 		}
 
 
 		<-p.writeSignal.Wait()
 		<-p.writeSignal.Wait()

+ 23 - 0
transport/pipe/pipe_test.go

@@ -0,0 +1,23 @@
+package pipe_test
+
+import (
+	"testing"
+
+	"v2ray.com/core/common/buf"
+	. "v2ray.com/core/transport/pipe"
+	. "v2ray.com/ext/assert"
+)
+
+func TestPipeReadWrite(t *testing.T) {
+	assert := With(t)
+
+	pReader, pWriter := New()
+	payload := []byte{'a', 'b', 'c', 'd'}
+	b := buf.New()
+	b.Append(payload)
+	assert(pWriter.WriteMultiBuffer(buf.NewMultiBufferValue(b)), IsNil)
+
+	rb, err := pReader.ReadMultiBuffer()
+	assert(err, IsNil)
+	assert(rb.String(), Equals, b.String())
+}