From cd4dbe6157fdd6942d92242d253f3d616840e548 Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Wed, 19 Jan 2022 11:08:07 +0200
Subject: [PATCH] make protocols and peers sticky when setting their limit

this is typically a user action, so prevent them from being garbage collected
and losing their limit.
---
 p2p/host/resource-manager/extapi.go | 10 +++++++++
 p2p/host/resource-manager/rcmgr.go  | 33 +++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/p2p/host/resource-manager/extapi.go b/p2p/host/resource-manager/extapi.go
index f2155d9283..acc548d0dc 100644
--- a/p2p/host/resource-manager/extapi.go
+++ b/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()
diff --git a/p2p/host/resource-manager/rcmgr.go b/p2p/host/resource-manager/rcmgr.go
index 2d403bc2f3..61d9c2badd 100644
--- a/p2p/host/resource-manager/rcmgr.go
+++ b/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)