make new stream calls accept a context

swarm.go

@@ -192,16 +192,17 @@ func (s *Swarm) SetStreamHandler(handler inet.StreamHandler) {
 }
 
 // NewStreamWithPeer creates a new stream on any available connection to p
-func (s *Swarm) NewStreamWithPeer(p peer.ID) (*Stream, error) {
+func (s *Swarm) NewStreamWithPeer(ctx context.Context, p peer.ID) (*Stream, error) {
 	// if we have no connections, try connecting.
 	if len(s.ConnectionsToPeer(p)) == 0 {
 		log.Debug("Swarm: NewStreamWithPeer no connections. Attempting to connect...")
-		if _, err := s.Dial(s.Context(), p); err != nil {
+		if _, err := s.Dial(ctx, p); err != nil {
 			return nil, err
 		}
 	}
 	log.Debug("Swarm: NewStreamWithPeer...")
 
+	// TODO: think about passing a context down to NewStreamWithGroup
 	st, err := s.swarm.NewStreamWithGroup(p)
 	return wrapStream(st), err
 }

swarm_net.go

@@ -132,9 +132,9 @@ func (n *Network) Connectedness(p peer.ID) inet.Connectedness {
 
 // NewStream returns a new stream to given peer p.
 // If there is no connection to p, attempts to create one.
-func (n *Network) NewStream(p peer.ID) (inet.Stream, error) {
+func (n *Network) NewStream(ctx context.Context, p peer.ID) (inet.Stream, error) {
 	log.Debugf("[%s] network opening stream to peer [%s]", n.local, p)
-	s, err := n.Swarm().NewStreamWithPeer(p)
+	s, err := n.Swarm().NewStreamWithPeer(ctx, p)
 	if err != nil {
 		return nil, err
 	}

swarm_test.go

@@ -129,7 +129,7 @@ func SubtestSwarm(t *testing.T, SwarmNum int, MsgNum int) {
 				defer wg.Done()
 
 				// first, one stream per peer (nice)
-				stream, err := s1.NewStreamWithPeer(p)
+				stream, err := s1.NewStreamWithPeer(ctx, p)
 				if err != nil {
 					errChan <- err
 					return