Make sure process context is set last

License: MIT
Signed-off-by: rht <rhtbot@gmail.com>

p2p/net/swarm/swarm_dial.go

@@ -17,6 +17,7 @@ import (
 	ma "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
 	manet "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr-net"
 	process "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess"
+	processctx "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess/context"
 	ratelimit "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess/ratelimit"
 	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
 )
@@ -416,8 +417,7 @@ func (s *Swarm) dialAddrs(ctx context.Context, d *conn.Dialer, p peer.ID, remote
 			}
 		})
 
-		<-ctx.Done()
-		limiter.Close()
+		processctx.CloseAfterContext(limiter, ctx)
 	}()
 
 	// wair fot the results.

routing/dht/dht.go

@@ -73,7 +73,7 @@ func NewDHT(ctx context.Context, h host.Host, dstore ds.ThreadSafeDatastore) *Ip
 	// register for network notifs.
 	dht.host.Network().Notify((*netNotifiee)(dht))
 
-	dht.proc = goprocessctx.WithContextAndTeardown(ctx, func() error {
+	dht.proc = goprocess.WithTeardown(func() error {
 		// remove ourselves from network notifs.
 		dht.host.Network().StopNotify((*netNotifiee)(dht))
 		return nil
@@ -84,6 +84,7 @@ func NewDHT(ctx context.Context, h host.Host, dstore ds.ThreadSafeDatastore) *Ip
 	h.SetStreamHandler(ProtocolDHT, dht.handleNewStream)
 	dht.providers = NewProviderManager(dht.ctx, dht.self)
 	dht.proc.AddChild(dht.providers.proc)
+	goprocessctx.CloseAfterContext(dht.proc, ctx)
 
 	dht.routingTable = kb.NewRoutingTable(20, kb.ConvertPeerID(dht.self), time.Minute, dht.peerstore)
 	dht.birth = time.Now()

routing/dht/query.go

@@ -127,10 +127,7 @@ func (r *dhtQueryRunner) Run(ctx context.Context, peers []peer.ID) (*dhtQueryRes
 	// now, if the context finishes, close the proc.
 	// we have to do it here because the logic before is setup, which
 	// should run without closing the proc.
-	go func() {
-		<-ctx.Done()
-		r.proc.Close()
-	}()
+	ctxproc.CloseAfterContext(r.proc, ctx)
 
 	select {
 	case <-r.peersRemaining.Done():