-
Notifications
You must be signed in to change notification settings - Fork 0
/
connector.go
118 lines (99 loc) · 2.48 KB
/
connector.go
1
2
3
4
5
6
7
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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
package core
import (
"context"
"time"
"github.com/libp2p/go-libp2p/core/host"
"github.com/libp2p/go-libp2p/core/network"
"github.com/libp2p/go-libp2p/core/peer"
ma "github.com/multiformats/go-multiaddr"
)
type Connector interface {
Need(proc string, p peer.AddrInfo)
Done(proc string, p peer.ID)
}
func NewConnector(h host.Host) Connector {
return newConnector(h)
}
var _ Connector = (*connector)(nil)
type connector struct {
h host.Host
needed *PeerSet
bctx context.Context
cancel context.CancelFunc
}
func newConnector(h host.Host) *connector {
c := connector{}
c.h = h
c.needed = NewPeerSet()
c.h.Network().Notify((*connectorNotifiee)(&c))
c.bctx = nil
c.cancel = nil
return &c
}
func (c *connector) background(ctx context.Context) {
ticker := time.NewTicker(5 * time.Second)
for {
select {
case t := <-ticker.C:
peers := c.needed.Turn(t)
for _, peer := range peers {
c.connect(peer)
}
case <-ctx.Done():
ticker.Stop()
return
}
}
}
func (c *connector) mayStart() {
if c.bctx == nil {
c.bctx, c.cancel = context.WithCancel(context.Background())
go c.background(c.bctx)
}
}
func (c *connector) mayStop() {
if c.needed.Empty() && c.bctx != nil {
c.cancel()
c.bctx = nil
c.cancel = nil
}
}
func (c *connector) connect(p peer.AddrInfo) {
log.Debugf("try connecting to %s", p)
if c.h.Network().Connectedness(p.ID) != network.Connected {
go func(pi peer.AddrInfo) {
err := c.h.Connect(context.Background(), pi)
if err != nil {
c.needed.Failed(pi.ID)
log.Debugf("failed connecting to %s", p)
return
}
}(p)
}
}
func (c *connector) Need(proc string, p peer.AddrInfo) {
c.h.ConnManager().Protect(p.ID, proc)
c.needed.Add(proc, p)
c.needed.Force(p.ID)
c.mayStart()
c.connect(p)
}
func (c *connector) Done(proc string, p peer.ID) {
c.h.ConnManager().Unprotect(p, proc)
c.needed.Remove(proc, p)
c.mayStop()
}
type connectorNotifiee connector
func (cn *connectorNotifiee) connector() *connector {
return (*connector)(cn)
}
func (cn *connectorNotifiee) Listen(network.Network, ma.Multiaddr) {}
func (cn *connectorNotifiee) ListenClose(network.Network, ma.Multiaddr) {}
func (cn *connectorNotifiee) Connected(n network.Network, c network.Conn) {
log.Debug("peer connected")
cn.connector().needed.Done(c.RemotePeer())
}
func (cn *connectorNotifiee) Disconnected(n network.Network, c network.Conn) {
log.Debugf("node %v peer %v disconnected ", cn.h.ID(), c.RemotePeer())
cn.connector().needed.Failed(c.RemotePeer())
}