Merge pull request soheilhy#30 from soheilhy/fix-patricia-tree

Fix race in patricia tree

.travis.yml

@@ -24,4 +24,5 @@ before_script:
   - if [[ $TRAVIS_GO_VERSION == 1.6* ]]; then go tool vet --shadow .; fi
 
 script:
-  - go test -v ./...
+  - go test -bench . -v ./...
+  - go test -race -bench . -v ./...

bench_test.go

@@ -62,12 +62,14 @@ func BenchmarkCMuxConnHTTP1(b *testing.B) {
 	wg.Add(b.N)
 
 	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		c := &mockConn{
-			r: bytes.NewReader(benchHTTP1Payload),
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			wg.Add(1)
+			m.serve(&mockConn{
+				r: bytes.NewReader(benchHTTP1Payload),
+			}, donec, &wg)
 		}
-		m.serve(c, donec, &wg)
-	}
+	})
 }
 
 func BenchmarkCMuxConnHTTP2(b *testing.B) {
@@ -80,12 +82,14 @@ func BenchmarkCMuxConnHTTP2(b *testing.B) {
 	wg.Add(b.N)
 
 	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		c := &mockConn{
-			r: bytes.NewReader(benchHTTP2Payload),
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			wg.Add(1)
+			m.serve(&mockConn{
+				r: bytes.NewReader(benchHTTP2Payload),
+			}, donec, &wg)
 		}
-		m.serve(c, donec, &wg)
-	}
+	})
 }
 
 func BenchmarkCMuxConnHTTP1n2(b *testing.B) {
@@ -98,15 +102,16 @@ func BenchmarkCMuxConnHTTP1n2(b *testing.B) {
 
 	donec := make(chan struct{})
 	var wg sync.WaitGroup
-	wg.Add(b.N)
 
 	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		c := &mockConn{
-			r: bytes.NewReader(benchHTTP2Payload),
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			wg.Add(1)
+			m.serve(&mockConn{
+				r: bytes.NewReader(benchHTTP2Payload),
+			}, donec, &wg)
 		}
-		m.serve(c, donec, &wg)
-	}
+	})
 }
 
 func BenchmarkCMuxConnHTTP2n1(b *testing.B) {
@@ -119,13 +124,14 @@ func BenchmarkCMuxConnHTTP2n1(b *testing.B) {
 
 	donec := make(chan struct{})
 	var wg sync.WaitGroup
-	wg.Add(b.N)
 
 	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		c := &mockConn{
-			r: bytes.NewReader(benchHTTP1Payload),
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			wg.Add(1)
+			m.serve(&mockConn{
+				r: bytes.NewReader(benchHTTP1Payload),
+			}, donec, &wg)
 		}
-		m.serve(c, donec, &wg)
-	}
+	})
 }

patricia.go

@@ -22,8 +22,8 @@ import (
 // patriciaTree is a simple patricia tree that handles []byte instead of string
 // and cannot be changed after instantiation.
 type patriciaTree struct {
-	root *ptNode
-	buf  []byte // preallocated buffer to read data while matching
+	root     *ptNode
+	maxDepth int // max depth of the tree.
 }
 
 func newPatriciaTree(bs ...[]byte) *patriciaTree {
@@ -34,8 +34,8 @@ func newPatriciaTree(bs ...[]byte) *patriciaTree {
 		}
 	}
 	return &patriciaTree{
-		root: newNode(bs),
-		buf:  make([]byte, max+1),
+		root:     newNode(bs),
+		maxDepth: max + 1,
 	}
 }
 
@@ -48,13 +48,15 @@ func newPatriciaTreeString(strs ...string) *patriciaTree {
 }
 
 func (t *patriciaTree) matchPrefix(r io.Reader) bool {
-	n, _ := io.ReadFull(r, t.buf)
-	return t.root.match(t.buf[:n], true)
+	buf := make([]byte, t.maxDepth)
+	n, _ := io.ReadFull(r, buf)
+	return t.root.match(buf[:n], true)
 }
 
 func (t *patriciaTree) match(r io.Reader) bool {
-	n, _ := io.ReadFull(r, t.buf)
-	return t.root.match(t.buf[:n], false)
+	buf := make([]byte, t.maxDepth)
+	n, _ := io.ReadFull(r, buf)
+	return t.root.match(buf[:n], false)
 }
 
 type ptNode struct {