allow multistream to have zero rtt stream opening

License: MIT Signed-off-by: Jeromy <[email protected]>
ipfs · Jul 15, 2015 · b5b27f1 · b5b27f1
1 parent e6dd70a
commit b5b27f1
Show file tree

Hide file tree

Showing 6 changed files with 279 additions and 18 deletions.
diff --git a/Godeps/Godeps.json b/Godeps/Godeps.json
diff --git a/Godeps/_workspace/src/github.com/whyrusleeping/go-multistream/lazy.go b/Godeps/_workspace/src/github.com/whyrusleeping/go-multistream/lazy.go
diff --git a/Godeps/_workspace/src/github.com/whyrusleeping/go-multistream/multistream.go b/Godeps/_workspace/src/github.com/whyrusleeping/go-multistream/multistream.go
diff --git a/Godeps/_workspace/src/github.com/whyrusleeping/go-multistream/multistream_test.go b/Godeps/_workspace/src/github.com/whyrusleeping/go-multistream/multistream_test.go
diff --git a/p2p/host/basic/basic_host.go b/p2p/host/basic/basic_host.go
@@ -170,12 +170,11 @@ func (h *BasicHost) NewStream(pid protocol.ID, p peer.ID) (inet.Stream, error) {
 
 	logStream := mstream.WrapStream(s, pid, h.bwc)
 
-	if err := msmux.SelectProtoOrFail(string(pid), logStream); err != nil {
-		logStream.Close()
-		return nil, err
-	}
-
-	return logStream, nil
+	lzcon := msmux.NewLazyHandshakeConn(logStream, string(pid))
+	return &streamWrapper{
+		Stream: logStream,
+		rw:     lzcon,
+	}, nil
 }
 
 // Connect ensures there is a connection between this host and the peer with
@@ -254,3 +253,16 @@ func (h *BasicHost) Close() error {
 func (h *BasicHost) GetBandwidthReporter() metrics.Reporter {
 	return h.bwc
 }
+
+type streamWrapper struct {
+	inet.Stream
+	rw io.ReadWriter
+}
+
+func (s *streamWrapper) Read(b []byte) (int, error) {
+	return s.rw.Read(b)
+}
+
+func (s *streamWrapper) Write(b []byte) (int, error) {
+	return s.rw.Write(b)
+}
diff --git a/p2p/test/backpressure/backpressure_test.go b/p2p/test/backpressure/backpressure_test.go
@@ -299,6 +299,12 @@ func TestStBackpressureStreamWrite(t *testing.T) {
 		}
 	}
 
+	// trigger lazy connection handshaking
+	_, err = s.Read(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
 	// 500ms rounds of lockstep write + drain
 	roundsStart := time.Now()
 	roundsTotal := 0