Fix deadlock on testplanet shutdown (#798)

internal/testplanet/node.go

@@ -96,6 +96,9 @@ func (node *Node) Addr() string { return node.Info.Address.Address }
 // Shutdown shuts down all node dependencies
 func (node *Node) Shutdown() error {
 	var errs []error
+	if node.Kademlia != nil {
+		errs = append(errs, node.Kademlia.Disconnect())
+	}
 	if node.Provider != nil {
 		errs = append(errs, node.Provider.Close())
 	}
@@ -103,9 +106,6 @@ func (node *Node) Shutdown() error {
 	// if node.Listener != nil {
 	//    errs = append(errs, node.Listener.Close())
 	// }
-	if node.Kademlia != nil {
-		errs = append(errs, node.Kademlia.Disconnect())
-	}
 
 	for _, dep := range node.Dependencies {
 		err := dep.Close()

pkg/kademlia/kademlia.go

@@ -9,6 +9,8 @@ import (
 	"net"
 	"os"
 	"path/filepath"
+	"sync/atomic"
+	"unsafe"
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/zeebo/errs"
@@ -46,12 +48,13 @@ type discoveryOptions struct {
 
 // Kademlia is an implementation of kademlia adhering to the DHT interface.
 type Kademlia struct {
-	alpha          int // alpha is a system wide concurrency parameter
-	routingTable   *RoutingTable
-	bootstrapNodes []pb.Node
-	address        string
-	nodeClient     node.Client
-	identity       *provider.FullIdentity
+	alpha           int // alpha is a system wide concurrency parameter
+	routingTable    *RoutingTable
+	bootstrapNodes  []pb.Node
+	address         string
+	nodeClient      node.Client
+	identity        *provider.FullIdentity
+	bootstrapCancel unsafe.Pointer // context.CancelFunc
 }
 
 // NewKademlia returns a newly configured Kademlia instance
@@ -108,6 +111,12 @@ func NewKademliaWithRoutingTable(self pb.Node, bootstrapNodes []pb.Node, identit
 
 // Disconnect safely closes connections to the Kademlia network
 func (k *Kademlia) Disconnect() error {
+	// Cancel the bootstrap context
+	ptr := atomic.LoadPointer(&k.bootstrapCancel)
+	if ptr != nil {
+		(*(*context.CancelFunc)(ptr))()
+	}
+
 	return utils.CombineErrors(
 		k.nodeClient.Disconnect(),
 		k.routingTable.Close(),
@@ -172,7 +181,10 @@ func (k *Kademlia) Bootstrap(ctx context.Context) error {
 		return BootstrapErr.New("no bootstrap nodes provided")
 	}
 
-	return k.lookup(ctx, k.routingTable.self.Id, discoveryOptions{
+	bootstrapContext, bootstrapCancel := context.WithCancel(ctx)
+	atomic.StorePointer(&k.bootstrapCancel, unsafe.Pointer(&bootstrapCancel))
+
+	return k.lookup(bootstrapContext, k.routingTable.self.Id, discoveryOptions{
 		concurrency: k.alpha, retries: defaultRetries, bootstrap: true, bootstrapNodes: k.bootstrapNodes,
 	})
 }