From 191ac62c12fab7d996de0a06b185a9fc5572e6ee Mon Sep 17 00:00:00 2001
From: Jeromy <jeromyj@gmail.com>
Date: Tue, 14 Jul 2015 14:04:56 -0700
Subject: [PATCH] making the daemon shutdown quicker

License: MIT
Signed-off-by: Jeromy <jeromyj@gmail.com>
---
 exchange/bitswap/wantmanager.go |  2 ++
 p2p/host/basic/natmgr.go        | 16 ++++++++++------
 p2p/net/swarm/swarm.go          |  2 +-
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/exchange/bitswap/wantmanager.go b/exchange/bitswap/wantmanager.go
index a8eeb58e283..3b4626a4d0c 100644
--- a/exchange/bitswap/wantmanager.go
+++ b/exchange/bitswap/wantmanager.go
@@ -140,6 +140,8 @@ func (mq *msgQueue) runQueue(ctx context.Context) {
 			mq.doWork(ctx)
 		case <-mq.done:
 			return
+		case <-ctx.Done():
+			return
 		}
 	}
 }
diff --git a/p2p/host/basic/natmgr.go b/p2p/host/basic/natmgr.go
index 334d8b5c45b..04fa5a140df 100644
--- a/p2p/host/basic/natmgr.go
+++ b/p2p/host/basic/natmgr.go
@@ -75,18 +75,22 @@ func (nmgr *natManager) discoverNAT() {
 		// to avoid leaking resources in a non-obvious way. the only case
 		// this affects is when the daemon is being started up and _immediately_
 		// asked to close. other services are also starting up, so ok to wait.
-		nat := inat.DiscoverNAT()
-		if nat == nil { // no nat, or failed to get it.
-			return
-		}
+		discoverdone := make(chan struct{})
+		var nat *inat.NAT
+		go func() {
+			defer close(discoverdone)
+			nat = inat.DiscoverNAT()
+		}()
 
 		// by this point -- after finding the NAT -- we may have already
 		// be closing. if so, just exit.
 		select {
 		case <-worker.Closing():
-			nat.Close()
 			return
-		default:
+		case <-discoverdone:
+			if nat == nil { // no nat, or failed to get it.
+				return
+			}
 		}
 
 		// wire up the nat to close when nmgr closes.
diff --git a/p2p/net/swarm/swarm.go b/p2p/net/swarm/swarm.go
index bf3d0395f7f..9fe112c7c0a 100644
--- a/p2p/net/swarm/swarm.go
+++ b/p2p/net/swarm/swarm.go
@@ -187,7 +187,7 @@ func (s *Swarm) NewStreamWithPeer(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(context.Background(), p); err != nil {
+		if _, err := s.Dial(s.Context(), p); err != nil {
 			return nil, err
 		}
 	}