diff --git a/rcmgr.go b/rcmgr.go index 03a2eaa..fcae916 100644 --- a/rcmgr.go +++ b/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() diff --git a/scope.go b/scope.go index 6c4a9e8..d642f57 100644 --- a/scope.go +++ b/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