update implementation for interface changes

rcmgr.go

@@ -86,6 +86,7 @@ type ConnectionScope struct {
 }
 
 var _ network.ConnectionScope = (*ConnectionScope)(nil)
+var _ network.UserConnectionScope = (*ConnectionScope)(nil)
 
 type StreamScope struct {
 	*ResourceScope
@@ -102,6 +103,7 @@ type StreamScope struct {
 }
 
 var _ network.StreamScope = (*StreamScope)(nil)
+var _ network.UserStreamScope = (*StreamScope)(nil)
 
 func NewResourceManager(limits Limiter) *ResourceManager {
 	r := &ResourceManager{
@@ -206,6 +208,18 @@ func (r *ResourceManager) OpenConnection(dir network.Direction, usefd bool) (net
 	return conn, nil
 }
 
+func (r *ResourceManager) OpenStream(p peer.ID, dir network.Direction) (network.StreamScope, error) {
+	peer := r.getPeerScope(p)
+	stream := NewStreamScope(dir, r.limits.GetStreamLimits(p), peer)
+
+	err := stream.AddStream(dir)
+	if err != nil {
+		return nil, err
+	}
+
+	return stream, nil
+}
+
 func (r *ResourceManager) Close() error {
 	r.cancel()
 	r.wg.Wait()
@@ -318,16 +332,6 @@ func (s *PeerScope) Peer() peer.ID {
 	return s.peer
 }
 
-func (s *PeerScope) OpenStream(dir network.Direction) (network.StreamScope, error) {
-	stream := NewStreamScope(dir, s.rcmgr.limits.GetStreamLimits(s.peer), s)
-	err := stream.AddStream(dir)
-	if err != nil {
-		return nil, err
-	}
-
-	return stream, nil
-}
-
 func (s *ConnectionScope) PeerScope() network.PeerScope {
 	s.Lock()
 	defer s.Unlock()

scope.go

@@ -28,6 +28,9 @@ type ResourceScope struct {
 	done   bool
 	refCnt int
 
+	// for transactional scope reference counting
+	parent *ResourceScope
+
 	rc          *Resources
 	constraints []*ResourceScope
 }
@@ -48,6 +51,15 @@ func NewResourceScope(limit Limit, constraints []*ResourceScope) *ResourceScope
 	}
 }
 
+func NewTxnResourceScope(parent *ResourceScope, limit Limit, constraints []*ResourceScope) *ResourceScope {
+	parent.IncRef()
+	return &ResourceScope{
+		parent:      parent,
+		rc:          NewResources(limit),
+		constraints: constraints,
+	}
+}
+
 // Resources implementation
 func (rc *Resources) checkMemory(rsvp int64) error {
 	// overflow check; this also has the side-effect that we cannot reserve negative memory.
@@ -578,6 +590,21 @@ func (s *ResourceScope) RemoveFDForChild(count int) {
 	s.rc.removeFD(count)
 }
 
+func (s *ResourceScope) BeginTxn() (network.TransactionalScope, error) {
+	s.Lock()
+	defer s.Unlock()
+
+	if s.done {
+		return nil, ErrResourceScopeClosed
+	}
+
+	constraints := make([]*ResourceScope, len(s.constraints)+1)
+	constraints[0] = s
+	copy(constraints[1:], s.constraints)
+
+	return NewTxnResourceScope(s, s.rc.limit, constraints), nil
+}
+
 func (s *ResourceScope) Done() {
 	s.Lock()
 	defer s.Unlock()
@@ -593,6 +620,10 @@ func (s *ResourceScope) Done() {
 		cst.RemoveFDForChild(s.rc.nfd)
 	}
 
+	if s.parent != nil {
+		s.parent.DecRef()
+	}
+
 	s.rc.releaseBuffers()
 
 	s.rc.nstreamsIn = 0