p2p: fix connmgr use
License: MIT Signed-off-by: Łukasz Magiera <magik6k@gmail.com>
Łukasz Magiera committed
Sep 6, 2018 at 00:35 UTC
c5090508a86c150bb585a0a38f30bfdaadc902c4
5 files changed
+53
-34
core/commands/p2p.go
+1
-1
@@ -511,7 +511,7 @@ var p2pStreamCloseCmd = &cmds.Command{
511
n.P2P.Streams.Unlock()
512
513
for _, s := range toClose {
514
- s.Reset()
514
+ n.P2P.Streams.Reset(s)
515
}
516
},
517
}
p2p/local.go
+1
-7
@@ -77,23 +77,17 @@ func (l *localListener) setupStream(local manet.Conn) {
77
return
78
}
79
80
- cmgr := l.p2p.peerHost.ConnManager()
81
- cmgr.TagPeer(l.peer, CMGR_TAG, 20)
82
-
80
stream := &Stream{
81
Protocol: l.proto,
82
83
OriginAddr: local.RemoteMultiaddr(),
84
TargetAddr: l.TargetAddress(),
85
+ peer: l.peer,
86
87
Local: local,
88
Remote: remote,
89
90
Registry: l.p2p.Streams,
93
-
94
- cleanup: func() {
95
- cmgr.UntagPeer(l.peer, CMGR_TAG)
96
- },
91
}
92
93
l.p2p.Streams.Register(stream)
p2p/p2p.go
+3
-1
@@ -31,7 +31,9 @@ func NewP2P(identity peer.ID, peerHost p2phost.Host, peerstore pstore.Peerstore)
31
ListenersP2P: newListenerP2PRegistry(identity, peerHost),
32
33
Streams: &StreamRegistry{
34
- Streams: map[uint64]*Stream{},
34
+ Streams: map[uint64]*Stream{},
35
+ ConnManager: peerHost.ConnManager(),
36
+ conns: map[peer.ID]int{},
37
},
38
}
39
}
p2p/remote.go
+1
-7
@@ -57,23 +57,17 @@ func (l *remoteListener) handleStream(remote net.Stream) {
57
return
58
}
59
60
- cmgr := l.p2p.peerHost.ConnManager()
61
- cmgr.TagPeer(peer, CMGR_TAG, 20)
62
-
60
stream := &Stream{
61
Protocol: l.proto,
62
63
OriginAddr: peerMa,
64
TargetAddr: l.addr,
65
+ peer: peer,
66
67
Local: local,
68
Remote: remote,
69
70
Registry: l.p2p.Streams,
73
-
74
- cleanup: func() {
75
- cmgr.UntagPeer(peer, CMGR_TAG)
76
- },
71
}
72
73
l.p2p.Streams.Register(stream)
p2p/stream.go
+47
-18
@@ -5,9 +5,11 @@ import (
5
"sync"
6
7
net "gx/ipfs/QmQSbtGXCyNrj34LWL8EgXyNNYDZ8r3SwQcpW5pPxVhLnM/go-libp2p-net"
8
+ peer "gx/ipfs/QmQsErDt8Qgw1XrsXf2BpEzDgGWtB1YLsTAARBup5b6B9W/go-libp2p-peer"
9
manet "gx/ipfs/QmV6FjemM1K8oXjrvuq3wuVWWoU2TLDPmNnKrxHzY3v6Ai/go-multiaddr-net"
10
+ ifconnmgr "gx/ipfs/QmVz2p8ZVZ5GcWPNWGs2HZHiZyHumZcJpQdMRpxkMDhc2C/go-libp2p-interface-connmgr"
11
ma "gx/ipfs/QmYmsdtJ3HsodkePE3eU3TsCaP2YvPZJ4LoXnNkDE5Tpt7/go-multiaddr"
10
- "gx/ipfs/QmZNkThpqfVXs9GNbexPrfBbXSLNYeKrE7jwFM2oqHbyqN/go-libp2p-protocol"
12
+ protocol "gx/ipfs/QmZNkThpqfVXs9GNbexPrfBbXSLNYeKrE7jwFM2oqHbyqN/go-libp2p-protocol"
13
)
14
15
const CMGR_TAG = "stream-fwd"
@@ -20,29 +22,23 @@ type Stream struct {
22
23
OriginAddr ma.Multiaddr
24
TargetAddr ma.Multiaddr
25
+ peer peer.ID
26
27
Local manet.Conn
28
Remote net.Stream
29
30
Registry *StreamRegistry
28
-
29
- cleanup func()
31
}
32
32
-// Close closes stream endpoints and deregisters it
33
-func (s *Stream) Close() error {
34
- s.Local.Close()
35
- s.Remote.Close()
36
- s.cleanup()
37
- s.Registry.Deregister(s.id)
33
+// close closes stream endpoints and deregisters it
34
+func (s *Stream) close() error {
35
+ s.Registry.Close(s)
36
return nil
37
}
38
41
-// Reset closes stream endpoints and deregisters it
42
-func (s *Stream) Reset() error {
43
- s.Local.Close()
44
- s.Remote.Reset()
45
- s.Registry.Deregister(s.id)
39
+// reset closes stream endpoints and deregisters it
40
+func (s *Stream) reset() error {
41
+ s.Registry.Reset(s)
42
return nil
43
}
44
@@ -50,18 +46,18 @@ func (s *Stream) startStreaming() {
46
go func() {
47
_, err := io.Copy(s.Local, s.Remote)
48
if err != nil {
53
- s.Reset()
49
+ s.reset()
50
} else {
55
- s.Close()
51
+ s.close()
52
}
53
}()
54
55
go func() {
56
_, err := io.Copy(s.Remote, s.Local)
57
if err != nil {
62
- s.Reset()
58
+ s.reset()
59
} else {
64
- s.Close()
60
+ s.close()
61
}
62
}()
63
}
@@ -71,7 +67,10 @@ type StreamRegistry struct {
67
sync.Mutex
68
69
Streams map[uint64]*Stream
70
+ conns map[peer.ID]int
71
nextID uint64
72
+
73
+ ifconnmgr.ConnManager
74
}
75
76
// Register registers a stream to the registry
@@ -79,6 +78,9 @@ func (r *StreamRegistry) Register(streamInfo *Stream) {
78
r.Lock()
79
defer r.Unlock()
80
81
+ r.ConnManager.TagPeer(streamInfo.peer, CMGR_TAG, 20)
82
+ r.conns[streamInfo.peer]++
83
+
84
streamInfo.id = r.nextID
85
r.Streams[r.nextID] = streamInfo
86
r.nextID++
@@ -89,5 +91,32 @@ func (r *StreamRegistry) Deregister(streamID uint64) {
91
r.Lock()
92
defer r.Unlock()
93
94
+ s, ok := r.Streams[streamID]
95
+ if !ok {
96
+ return
97
+ }
98
+ p := s.peer
99
+ r.conns[p]--
100
+ if r.conns[p] < 1 {
101
+ delete(r.conns, p)
102
+ r.ConnManager.UntagPeer(p, CMGR_TAG)
103
+ }
104
+
105
delete(r.Streams, streamID)
106
}
107
+
108
+// close closes stream endpoints and deregisters it
109
+func (r *StreamRegistry) Close(s *Stream) error {
110
+ s.Local.Close()
111
+ s.Remote.Close()
112
+ s.Registry.Deregister(s.id)
113
+ return nil
114
+}
115
+
116
+// reset closes stream endpoints and deregisters it
117
+func (r *StreamRegistry) Reset(s *Stream) error {
118
+ s.Local.Close()
119
+ s.Remote.Reset()
120
+ s.Registry.Deregister(s.id)
121
+ return nil
122
+}