Fix simultaneous dials

dial_sync.go

@@ -5,95 +5,82 @@ import (
 	"errors"
 	"sync"
 
+	"github.com/libp2p/go-libp2p-core/network"
 	"github.com/libp2p/go-libp2p-core/peer"
 )
 
 // TODO: change this text when we fix the bug
 var errDialCanceled = errors.New("dial was aborted internally, likely due to https://git.io/Je2wW")
 
 // DialFunc is the type of function expected by DialSync.
-type DialFunc func(context.Context, peer.ID) (*Conn, error)
+type DialWorkerFunc func(context.Context, peer.ID, <-chan DialRequest)
 
 // NewDialSync constructs a new DialSync
-func NewDialSync(dfn DialFunc) *DialSync {
+func NewDialSync(worker DialWorkerFunc) *DialSync {
 	return &DialSync{
-		dials:    make(map[peer.ID]*activeDial),
-		dialFunc: dfn,
+		dials:      make(map[peer.ID]*activeDial),
+		dialWorker: worker,
 	}
 }
 
 // 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 {
-	dials    map[peer.ID]*activeDial
-	dialsLk  sync.Mutex
-	dialFunc DialFunc
+	dials      map[peer.ID]*activeDial
+	dialsLk    sync.Mutex
+	dialWorker DialWorkerFunc
 }
 
 type activeDial struct {
-	id       peer.ID
-	refCnt   int
-	refCntLk sync.Mutex
-	cancel   func()
+	id     peer.ID
+	refCnt int
 
-	err    error
-	conn   *Conn
-	waitch chan struct{}
+	ctx    context.Context
+	cancel func()
 
-	ds *DialSync
-}
+	reqch chan DialRequest
 
-func (ad *activeDial) wait(ctx context.Context) (*Conn, error) {
-	defer ad.decref()
-	select {
-	case <-ad.waitch:
-		return ad.conn, ad.err
-	case <-ctx.Done():
-		return nil, ctx.Err()
-	}
+	ds *DialSync
 }
 
 func (ad *activeDial) incref() {
-	ad.refCntLk.Lock()
-	defer ad.refCntLk.Unlock()
 	ad.refCnt++
 }
 
 func (ad *activeDial) decref() {
-	ad.refCntLk.Lock()
+	ad.ds.dialsLk.Lock()
 	ad.refCnt--
-	maybeZero := (ad.refCnt <= 0)
-	ad.refCntLk.Unlock()
-
-	// make sure to always take locks in correct order.
-	if maybeZero {
-		ad.ds.dialsLk.Lock()
-		ad.refCntLk.Lock()
-		// check again after lock swap drop to make sure nobody else called incref
-		// in between locks
-		if ad.refCnt <= 0 {
-			ad.cancel()
-			delete(ad.ds.dials, ad.id)
-		}
-		ad.refCntLk.Unlock()
-		ad.ds.dialsLk.Unlock()
+	if ad.refCnt == 0 {
+		ad.cancel()
+		close(ad.reqch)
+		delete(ad.ds.dials, ad.id)
 	}
+	ad.ds.dialsLk.Unlock()
 }
 
-func (ad *activeDial) start(ctx context.Context) {
-	ad.conn, ad.err = ad.ds.dialFunc(ctx, ad.id)
-
-	// This isn't the user's context so we should fix the error.
-	switch ad.err {
-	case context.Canceled:
-		// The dial was canceled with `CancelDial`.
-		ad.err = errDialCanceled
-	case context.DeadlineExceeded:
-		// We hit an internal timeout, not a context timeout.
-		ad.err = ErrDialTimeout
+func (ad *activeDial) dial(ctx context.Context, p peer.ID) (*Conn, error) {
+	dialCtx := ad.ctx
+
+	if forceDirect, reason := network.GetForceDirectDial(ctx); forceDirect {
+		dialCtx = network.WithForceDirectDial(dialCtx, reason)
+	}
+	if simConnect, reason := network.GetSimultaneousConnect(ctx); simConnect {
+		dialCtx = network.WithSimultaneousConnect(dialCtx, reason)
+	}
+
+	resch := make(chan DialResponse, 1)
+	select {
+	case ad.reqch <- DialRequest{Ctx: dialCtx, Resch: resch}:
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+
+	select {
+	case res := <-resch:
+		return res.Conn, res.Err
+	case <-ctx.Done():
+		return nil, ctx.Err()
 	}
-	close(ad.waitch)
-	ad.cancel()
 }
 
 func (ds *DialSync) getActiveDial(p peer.ID) *activeDial {
@@ -109,13 +96,14 @@ func (ds *DialSync) getActiveDial(p peer.ID) *activeDial {
 		adctx, cancel := context.WithCancel(context.Background())
 		actd = &activeDial{
 			id:     p,
+			ctx:    adctx,
 			cancel: cancel,
-			waitch: make(chan struct{}),
+			reqch:  make(chan DialRequest),
 			ds:     ds,
 		}
 		ds.dials[p] = actd
 
-		go actd.start(adctx)
+		go ds.dialWorker(adctx, p, actd.reqch)
 	}
 
 	// increase ref count before dropping dialsLk
@@ -127,14 +115,8 @@ func (ds *DialSync) getActiveDial(p peer.ID) *activeDial {
 // DialLock initiates a dial to the given peer if there are none in progress
 // then waits for the dial to that peer to complete.
 func (ds *DialSync) DialLock(ctx context.Context, p peer.ID) (*Conn, error) {
-	return ds.getActiveDial(p).wait(ctx)
-}
+	ad := ds.getActiveDial(p)
+	defer ad.decref()
 
-// CancelDial cancels all in-progress dials to the given peer.
-func (ds *DialSync) CancelDial(p peer.ID) {
-	ds.dialsLk.Lock()
-	defer ds.dialsLk.Unlock()
-	if ad, ok := ds.dials[p]; ok {
-		ad.cancel()
-	}
+	return ad.dial(ctx, p)
 }

dial_sync_test.go

@@ -12,19 +12,33 @@ import (
 	"github.com/libp2p/go-libp2p-core/peer"
 )
 
-func getMockDialFunc() (DialFunc, func(), context.Context, <-chan struct{}) {
+func getMockDialFunc() (DialWorkerFunc, func(), context.Context, <-chan struct{}) {
 	dfcalls := make(chan struct{}, 512) // buffer it large enough that we won't care
 	dialctx, cancel := context.WithCancel(context.Background())
 	ch := make(chan struct{})
-	f := func(ctx context.Context, p peer.ID) (*Conn, error) {
+	f := func(ctx context.Context, p peer.ID, reqch <-chan DialRequest) {
 		dfcalls <- struct{}{}
-		defer cancel()
-		select {
-		case <-ch:
-			return new(Conn), nil
-		case <-ctx.Done():
-			return nil, ctx.Err()
-		}
+		go func() {
+			defer cancel()
+			for {
+				select {
+				case req, ok := <-reqch:
+					if !ok {
+						return
+					}
+
+					select {
+					case <-ch:
+						req.Resch <- DialResponse{Conn: new(Conn)}
+					case <-ctx.Done():
+						req.Resch <- DialResponse{Err: ctx.Err()}
+						return
+					}
+				case <-ctx.Done():
+					return
+				}
+			}
+		}()
 	}
 
 	o := new(sync.Once)
@@ -174,12 +188,25 @@ func TestDialSyncAllCancel(t *testing.T) {
 
 func TestFailFirst(t *testing.T) {
 	var count int
-	f := func(ctx context.Context, p peer.ID) (*Conn, error) {
-		if count > 0 {
-			return new(Conn), nil
+	f := func(ctx context.Context, p peer.ID, reqch <-chan DialRequest) {
+		for {
+			select {
+			case req, ok := <-reqch:
+				if !ok {
+					return
+				}
+
+				if count > 0 {
+					req.Resch <- DialResponse{Conn: new(Conn)}
+				} else {
+					req.Resch <- DialResponse{Err: fmt.Errorf("gophers ate the modem")}
+				}
+				count++
+
+			case <-ctx.Done():
+				return
+			}
 		}
-		count++
-		return nil, fmt.Errorf("gophers ate the modem")
 	}
 
 	ds := NewDialSync(f)
@@ -205,8 +232,19 @@ func TestFailFirst(t *testing.T) {
 }
 
 func TestStressActiveDial(t *testing.T) {
-	ds := NewDialSync(func(ctx context.Context, p peer.ID) (*Conn, error) {
-		return nil, nil
+	ds := NewDialSync(func(ctx context.Context, p peer.ID, reqch <-chan DialRequest) {
+		for {
+			select {
+			case req, ok := <-reqch:
+				if !ok {
+					return
+				}
+
+				req.Resch <- DialResponse{}
+			case <-ctx.Done():
+				return
+			}
+		}
 	})
 
 	wg := sync.WaitGroup{}

go.mod

@@ -7,7 +7,7 @@ require (
 	github.com/jbenet/goprocess v0.1.4
 	github.com/libp2p/go-addr-util v0.0.2
 	github.com/libp2p/go-conn-security-multistream v0.2.1
-	github.com/libp2p/go-libp2p-core v0.8.3
+	github.com/libp2p/go-libp2p-core v0.8.5
 	github.com/libp2p/go-libp2p-loggables v0.1.0
 	github.com/libp2p/go-libp2p-peerstore v0.2.6
 	github.com/libp2p/go-libp2p-quic-transport v0.10.0