Merge pull request #7 from libp2p/feat/sticky-limits

make protocols and peers sticky when setting their limit

p2p/host/resource-manager/extapi.go

@@ -51,6 +51,16 @@ func (s *resourceScope) SetLimit(limit Limit) {
 	s.rc.limit = limit
 }
 
+func (s *protocolScope) SetLimit(limit Limit) {
+	s.rcmgr.setStickyProtocol(s.proto)
+	s.resourceScope.SetLimit(limit)
+}
+
+func (s *peerScope) SetLimit(limit Limit) {
+	s.rcmgr.setStickyPeer(s.peer)
+	s.resourceScope.SetLimit(limit)
+}
+
 func (r *resourceManager) ListServices() []string {
 	r.mx.Lock()
 	defer r.mx.Unlock()

p2p/host/resource-manager/rcmgr.go

@@ -32,6 +32,9 @@ type resourceManager struct {
 	proto map[protocol.ID]*protocolScope
 	peer  map[peer.ID]*peerScope
 
+	stickyProto map[protocol.ID]struct{}
+	stickyPeer  map[peer.ID]struct{}
+
 	connId, streamId int64
 }
 
@@ -200,6 +203,16 @@ func (r *resourceManager) getProtocolScope(proto protocol.ID) *protocolScope {
 	return s
 }
 
+func (r *resourceManager) setStickyProtocol(proto protocol.ID) {
+	r.mx.Lock()
+	defer r.mx.Unlock()
+
+	if r.stickyProto == nil {
+		r.stickyProto = make(map[protocol.ID]struct{})
+	}
+	r.stickyProto[proto] = struct{}{}
+}
+
 func (r *resourceManager) getPeerScope(p peer.ID) *peerScope {
 	r.mx.Lock()
 	defer r.mx.Unlock()
@@ -214,6 +227,17 @@ func (r *resourceManager) getPeerScope(p peer.ID) *peerScope {
 	return s
 }
 
+func (r *resourceManager) setStickyPeer(p peer.ID) {
+	r.mx.Lock()
+	defer r.mx.Unlock()
+
+	if r.stickyPeer == nil {
+		r.stickyPeer = make(map[peer.ID]struct{})
+	}
+
+	r.stickyPeer[p] = struct{}{}
+}
+
 func (r *resourceManager) nextConnId() int64 {
 	r.mx.Lock()
 	defer r.mx.Unlock()
@@ -285,6 +309,10 @@ func (r *resourceManager) gc() {
 	defer r.mx.Unlock()
 
 	for proto, s := range r.proto {
+		_, sticky := r.stickyProto[proto]
+		if sticky {
+			continue
+		}
 		if s.IsUnused() {
 			s.Done()
 			delete(r.proto, proto)
@@ -293,6 +321,11 @@ func (r *resourceManager) gc() {
 
 	var deadPeers []peer.ID
 	for p, s := range r.peer {
+		_, sticky := r.stickyPeer[p]
+		if sticky {
+			continue
+		}
+
 		if s.IsUnused() {
 			s.Done()
 			delete(r.peer, p)