From 0e0111c6f5fe5d3bf3c91295b39c0ea340d24c14 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 23 Aug 2021 17:38:57 +0100
Subject: [PATCH] simplify the DialSync code

It's easier to reason about this code if activeDial doesn't contain a pointer
back to DialSync (which already has a map of activeDials). It also allows us to
remove the memory footprint of the activeDial struct, so this should be
(slightly) more efficient.
---
 p2p/net/swarm/dial_sync.go | 45 +++++++++++++++++---------------------
 1 file changed, 20 insertions(+), 25 deletions(-)

diff --git a/p2p/net/swarm/dial_sync.go b/p2p/net/swarm/dial_sync.go
index 6b3f2afe..325abdf5 100644
--- a/p2p/net/swarm/dial_sync.go
+++ b/p2p/net/swarm/dial_sync.go
@@ -22,35 +22,26 @@ func newDialSync(worker dialWorkerFunc) *DialSync {
 // DialSync is a dial synchronization helper that ensures that at most one dial
 // to any given peer is active at any given time.
 type DialSync struct {
+	mutex      sync.Mutex
 	dials      map[peer.ID]*activeDial
-	dialsLk    sync.Mutex
 	dialWorker dialWorkerFunc
 }
 
 type activeDial struct {
-	id     peer.ID
 	refCnt int
 
 	ctx    context.Context
 	cancel func()
 
 	reqch chan dialRequest
-
-	ds *DialSync
 }
 
-func (ad *activeDial) decref() {
-	ad.ds.dialsLk.Lock()
-	ad.refCnt--
-	if ad.refCnt == 0 {
-		ad.cancel()
-		close(ad.reqch)
-		delete(ad.ds.dials, ad.id)
-	}
-	ad.ds.dialsLk.Unlock()
+func (ad *activeDial) close() {
+	ad.cancel()
+	close(ad.reqch)
 }
 
-func (ad *activeDial) dial(ctx context.Context, p peer.ID) (*Conn, error) {
+func (ad *activeDial) dial(ctx context.Context) (*Conn, error) {
 	dialCtx := ad.ctx
 
 	if forceDirect, reason := network.GetForceDirectDial(ctx); forceDirect {
@@ -76,8 +67,8 @@ func (ad *activeDial) dial(ctx context.Context, p peer.ID) (*Conn, error) {
 }
 
 func (ds *DialSync) getActiveDial(p peer.ID) (*activeDial, error) {
-	ds.dialsLk.Lock()
-	defer ds.dialsLk.Unlock()
+	ds.mutex.Lock()
+	defer ds.mutex.Unlock()
 
 	actd, ok := ds.dials[p]
 	if !ok {
@@ -85,21 +76,17 @@ func (ds *DialSync) getActiveDial(p peer.ID) (*activeDial, error) {
 		// to Dial is canceled, subsequent dial calls will also be canceled.
 		// XXX: this also breaks direct connection logic. We will need to pipe the
 		// information through some other way.
-		adctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(context.Background())
 		actd = &activeDial{
-			id:     p,
-			ctx:    adctx,
+			ctx:    ctx,
 			cancel: cancel,
 			reqch:  make(chan dialRequest),
-			ds:     ds,
 		}
 		go ds.dialWorker(p, actd.reqch)
 		ds.dials[p] = actd
 	}
-
-	// increase ref count before dropping dialsLk
+	// increase ref count before dropping mutex
 	actd.refCnt++
-
 	return actd, nil
 }
 
@@ -111,6 +98,14 @@ func (ds *DialSync) Dial(ctx context.Context, p peer.ID) (*Conn, error) {
 		return nil, err
 	}
 
-	defer ad.decref()
-	return ad.dial(ctx, p)
+	defer func() {
+		ds.mutex.Lock()
+		defer ds.mutex.Unlock()
+		ad.refCnt--
+		if ad.refCnt == 0 {
+			ad.close()
+			delete(ds.dials, p)
+		}
+	}()
+	return ad.dial(ctx)
 }