Merge pull request #36 from libp2p/open-stream-context

add a context to Open() and OpenStream()

bench_test.go

@@ -1,6 +1,7 @@
 package yamux
 
 import (
+	"context"
 	"io"
 	"testing"
 )
@@ -37,7 +38,7 @@ func BenchmarkAccept(b *testing.B) {
 	}()
 
 	for i := 0; i < b.N; i++ {
-		stream, err := client.Open()
+		stream, err := client.Open(context.Background())
 		if err != nil {
 			b.Fatalf("err: %v", err)
 		}
@@ -69,7 +70,7 @@ func BenchmarkSendRecv(b *testing.B) {
 		}
 	}()
 
-	stream, err := client.Open()
+	stream, err := client.Open(context.Background())
 	if err != nil {
 		b.Fatalf("err: %v", err)
 	}
@@ -113,7 +114,7 @@ func BenchmarkSendRecvLarge(b *testing.B) {
 		}
 	}()
 
-	stream, err := client.Open()
+	stream, err := client.Open(context.Background())
 	if err != nil {
 		b.Fatalf("err: %v", err)
 	}

session.go

@@ -2,6 +2,7 @@ package yamux
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -156,16 +157,16 @@ func (s *Session) NumStreams() int {
 }
 
 // Open is used to create a new stream as a net.Conn
-func (s *Session) Open() (net.Conn, error) {
-	conn, err := s.OpenStream()
+func (s *Session) Open(ctx context.Context) (net.Conn, error) {
+	conn, err := s.OpenStream(ctx)
 	if err != nil {
 		return nil, err
 	}
 	return conn, nil
 }
 
 // OpenStream is used to create a new stream
-func (s *Session) OpenStream() (*Stream, error) {
+func (s *Session) OpenStream(ctx context.Context) (*Stream, error) {
 	if s.IsClosed() {
 		return nil, s.shutdownErr
 	}
@@ -176,6 +177,8 @@ func (s *Session) OpenStream() (*Stream, error) {
 	// Block if we have too many inflight SYNs
 	select {
 	case s.synCh <- struct{}{}:
+	case <-ctx.Done():
+		return nil, ctx.Err()
 	case <-s.shutdownCh:
 		return nil, s.shutdownErr
 	}

session_norace_test.go

@@ -4,6 +4,7 @@ package yamux
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"io/ioutil"
 	"sync"
@@ -93,7 +94,7 @@ func TestSendData_VeryLarge(t *testing.T) {
 	for i := 0; i < workers; i++ {
 		go func() {
 			defer wg.Done()
-			stream, err := client.Open()
+			stream, err := client.Open(context.Background())
 			if err != nil {
 				t.Errorf("err: %v", err)
 				return
@@ -142,7 +143,7 @@ func TestLargeWindow(t *testing.T) {
 	defer client.Close()
 	defer server.Close()
 
-	stream, err := client.Open()
+	stream, err := client.Open(context.Background())
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}