dial_sync.go 1.54 KB
Newer Older
1 2 3 4 5 6
package swarm

import (
	"context"
	"sync"

Jeromy's avatar
Jeromy committed
7
	peer "github.com/libp2p/go-libp2p-peer"
8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91
)

type DialFunc func(context.Context, peer.ID) (*Conn, error)

func NewDialSync(dfn DialFunc) *DialSync {
	return &DialSync{
		dials:    make(map[peer.ID]*activeDial),
		dialFunc: dfn,
	}
}

type DialSync struct {
	dials    map[peer.ID]*activeDial
	dialsLk  sync.Mutex
	dialFunc DialFunc
}

type activeDial struct {
	id       peer.ID
	refCnt   int
	refCntLk sync.Mutex
	cancel   func()

	err    error
	conn   *Conn
	waitch chan struct{}

	ds *DialSync
}

func (dr *activeDial) wait(ctx context.Context) (*Conn, error) {
	defer dr.decref()
	select {
	case <-dr.waitch:
		return dr.conn, dr.err
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (ad *activeDial) incref() {
	ad.refCntLk.Lock()
	defer ad.refCntLk.Unlock()
	ad.refCnt++
}

func (ad *activeDial) decref() {
	ad.refCntLk.Lock()
	defer ad.refCntLk.Unlock()
	ad.refCnt--
	if ad.refCnt <= 0 {
		ad.cancel()
		ad.ds.dialsLk.Lock()
		delete(ad.ds.dials, ad.id)
		ad.ds.dialsLk.Unlock()
	}
}

func (ds *DialSync) DialLock(ctx context.Context, p peer.ID) (*Conn, error) {
	ds.dialsLk.Lock()

	actd, ok := ds.dials[p]
	if !ok {
		ctx, cancel := context.WithCancel(context.Background())
		actd = &activeDial{
			id:     p,
			cancel: cancel,
			waitch: make(chan struct{}),
			ds:     ds,
		}
		ds.dials[p] = actd

		go func(ctx context.Context, p peer.ID, ad *activeDial) {
			ad.conn, ad.err = ds.dialFunc(ctx, p)
			close(ad.waitch)
			ad.cancel()
		}(ctx, p, actd)
	}

	actd.incref()
	ds.dialsLk.Unlock()

	return actd.wait(ctx)
}