1 files changed, 299 insertions, 0 deletions
diff --git a/pkg/refs/refcounter.go b/pkg/refs/refcounter.go
new file mode 100644
index 000000000..1036553c7
--- /dev/null
+++ b/pkg/refs/refcounter.go
@@ -0,0 +1,299 @@
+// Copyright 2018 Google Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package refs defines an interface for reference counted objects. It
+// also provides a drop-in implementation called AtomicRefCount.
+package refs
+
+import (
+	"reflect"
+	"sync"
+	"sync/atomic"
+
+	"gvisor.googlesource.com/gvisor/pkg/ilist"
+)
+
+// RefCounter is the interface to be implemented by objects that are reference
+// counted.
+type RefCounter interface {
+	// IncRef increments the reference counter on the object.
+	IncRef()
+
+	// DecRef decrements the reference counter on the object.
+	//
+	// Note that AtomicRefCounter.DecRef() does not support destructors.
+	// If a type has a destructor, it must implement its own DecRef()
+	// method and call AtomicRefCounter.DecRefWithDestructor(destructor).
+	DecRef()
+
+	// TryIncRef attempts to increase the reference counter on the object,
+	// but may fail if all references have already been dropped. This
+	// should be used only in special circumstances, such as WeakRefs.
+	TryIncRef() bool
+
+	// addWeakRef adds the given weak reference. Note that you should have a
+	// reference to the object when calling this method.
+	addWeakRef(*WeakRef)
+
+	// dropWeakRef drops the given weak reference. Note that you should have
+	// a reference to the object when calling this method.
+	dropWeakRef(*WeakRef)
+}
+
+// A WeakRefUser is notified when the last non-weak reference is dropped.
+type WeakRefUser interface {
+	// WeakRefGone is called when the last non-weak reference is dropped.
+	WeakRefGone()
+}
+
+// WeakRef is a weak reference.
+type WeakRef struct {
+	ilist.Entry `state:"nosave"`
+
+	// obj is an atomic value that points to the refCounter.
+	obj atomic.Value `state:".(savedReference)"`
+
+	// user is notified when the weak ref is zapped by the object getting
+	// destroyed.
+	user WeakRefUser
+}
+
+// weakRefPool is a pool of weak references to avoid allocations on the hot path.
+var weakRefPool = sync.Pool{
+	New: func() interface{} {
+		return &WeakRef{}
+	},
+}
+
+// NewWeakRef acquires a weak reference for the given object.
+//
+// An optional user will be notified when the last non-weak reference is
+// dropped.
+//
+// Note that you must hold a reference to the object prior to getting a weak
+// reference. (But you may drop the non-weak reference after that.)
+func NewWeakRef(rc RefCounter, u WeakRefUser) *WeakRef {
+	w := weakRefPool.Get().(*WeakRef)
+	w.init(rc, u)
+	return w
+}
+
+// get attempts to get a normal reference to the underlying object, and returns
+// the object. If this weak reference has already been zapped (the object has
+// been destroyed) then false is returned. If the object still exists, then
+// true is returned.
+func (w *WeakRef) get() (RefCounter, bool) {
+	rc := w.obj.Load().(RefCounter)
+	if v := reflect.ValueOf(rc); v == reflect.Zero(v.Type()) {
+		// This pointer has already been zapped by zap() below. We do
+		// this to ensure that the GC can collect the underlying
+		// RefCounter objects and they don't hog resources.
+		return nil, false
+	}
+	if !rc.TryIncRef() {
+		return nil, true
+	}
+	return rc, true
+}
+
+// Get attempts to get a normal reference to the underlying object, and returns
+// the object. If this fails (the object no longer exists), then nil will be
+// returned instead.
+func (w *WeakRef) Get() RefCounter {
+	rc, _ := w.get()
+	return rc
+}
+
+// Drop drops this weak reference. You should always call drop when you are
+// finished with the weak reference. You may not use this object after calling
+// drop.
+func (w *WeakRef) Drop() {
+	rc, ok := w.get()
+	if !ok {
+		// We've been zapped already. When the refcounter has called
+		// zap, we're guaranteed it's not holding references.
+		weakRefPool.Put(w)
+		return
+	}
+	if rc == nil {
+		// The object is in the process of being destroyed. We can't
+		// remove this from the object's list, nor can we return this
+		// object to the pool. It'll just be garbage collected. This is
+		// a rare edge case, so it's not a big deal.
+		return
+	}
+
+	// At this point, we have a reference on the object. So destruction
+	// of the object (and zapping this weak reference) can't race here.
+	rc.dropWeakRef(w)
+
+	// And now aren't on the object's list of weak references. So it won't
+	// zap us if this causes the reference count to drop to zero.
+	rc.DecRef()
+
+	// Return to the pool.
+	weakRefPool.Put(w)
+}
+
+// init initializes this weak reference.
+func (w *WeakRef) init(rc RefCounter, u WeakRefUser) {
+	// Reset the contents of the weak reference.
+	// This is important because we are reseting the atomic value type.
+	// Otherwise, we could panic here if obj is different than what it was
+	// the last time this was used.
+	*w = WeakRef{}
+	w.user = u
+	w.obj.Store(rc)
+
+	// In the load path, we may already have a nil value. So we need to
+	// check whether or not that is the case before calling addWeakRef.
+	if v := reflect.ValueOf(rc); v != reflect.Zero(v.Type()) {
+		rc.addWeakRef(w)
+	}
+}
+
+// zap zaps this weak reference.
+func (w *WeakRef) zap() {
+	// We need to be careful about types here.
+	// So reflect is involved. But it's not that bad.
+	rc := w.obj.Load()
+	typ := reflect.TypeOf(rc)
+	w.obj.Store(reflect.Zero(typ).Interface())
+}
+
+// AtomicRefCount keeps a reference count using atomic operations and calls the
+// destructor when the count reaches zero.
+//
+// N.B. To allow the zero-object to be initialized, the count is offset by
+//      1, that is, when refCount is n, there are really n+1 references.
+type AtomicRefCount struct {
+	// refCount is composed of two fields:
+	//
+	//	[32-bit speculative references]:[32-bit real references]
+	//
+	// Speculative references are used for TryIncRef, to avoid a
+	// CompareAndSwap loop. See IncRef, DecRef and TryIncRef for details of
+	// how these fields are used.
+	refCount int64
+
+	// mu protects the list below.
+	mu sync.Mutex `state:"nosave"`
+
+	// weakRefs is our collection of weak references.
+	weakRefs ilist.List `state:"nosave"`
+}
+
+// TestReadRefs returns the current reference count of r. Use only for tests.
+func (r *AtomicRefCount) TestReadRefs() int64 {
+	return atomic.LoadInt64(&r.refCount)
+}
+
+// IncRef increments this object's reference count. While the count is kept
+// greater than zero, the destructor doesn't get called.
+//
+// The sanity check here is limited to real references, since if they have
+// dropped beneath zero then the object should have been destroyed.
+func (r *AtomicRefCount) IncRef() {
+	if v := atomic.AddInt64(&r.refCount, 1); v <= 0 {
+		panic("Incrementing non-positive ref count")
+	}
+}
+
+// TryIncRef attempts to increment the reference count, *unless the count has
+// already reached zero*. If false is returned, then the object has already
+// been destroyed, and the weak reference is no longer valid. If true if
+// returned then a valid reference is now held on the object.
+//
+// To do this safely without a loop, a speculative reference is first acquired
+// on the object. This allows multiple concurrent TryIncRef calls to
+// distinguish other TryIncRef calls from genuine references held.
+func (r *AtomicRefCount) TryIncRef() bool {
+	const speculativeRef = 1 << 32
+	v := atomic.AddInt64(&r.refCount, speculativeRef)
+	if int32(v) < 0 {
+		// This object has already been freed.
+		atomic.AddInt64(&r.refCount, -speculativeRef)
+		return false
+	}
+
+	// Turn into a real reference.
+	atomic.AddInt64(&r.refCount, -speculativeRef+1)
+	return true
+}
+
+// addWeakRef adds the given weak reference.
+func (r *AtomicRefCount) addWeakRef(w *WeakRef) {
+	r.mu.Lock()
+	r.weakRefs.PushBack(w)
+	r.mu.Unlock()
+}
+
+// dropWeakRef drops the given weak reference.
+func (r *AtomicRefCount) dropWeakRef(w *WeakRef) {
+	r.mu.Lock()
+	r.weakRefs.Remove(w)
+	r.mu.Unlock()
+}
+
+// DecRefWithDestructor decrements the object's reference count. If the
+// resulting count is negative and the destructor is not nil, then the
+// destructor will be called.
+//
+// Note that speculative references are counted here. Since they were added
+// prior to real references reaching zero, they will successfully convert to
+// real references. In other words, we see speculative references only in the
+// following case:
+//
+//	A: TryIncRef [speculative increase => sees non-negative references]
+//	B: DecRef [real decrease]
+//	A: TryIncRef [transform speculative to real]
+//
+func (r *AtomicRefCount) DecRefWithDestructor(destroy func()) {
+	switch v := atomic.AddInt64(&r.refCount, -1); {
+	case v < -1:
+		panic("Decrementing non-positive ref count")
+
+	case v == -1:
+		// Zap weak references. Note that at this point, all weak
+		// references are already invalid. That is, TryIncRef() will
+		// return false due to the reference count check.
+		r.mu.Lock()
+		for !r.weakRefs.Empty() {
+			w := r.weakRefs.Front().(*WeakRef)
+			// Capture the callback because w cannot be touched
+			// after it's zapped -- the owner is free it reuse it
+			// after that.
+			user := w.user
+			r.weakRefs.Remove(w)
+			w.zap()
+
+			if user != nil {
+				r.mu.Unlock()
+				user.WeakRefGone()
+				r.mu.Lock()
+			}
+		}
+		r.mu.Unlock()
+
+		// Call the destructor.
+		if destroy != nil {
+			destroy()
+		}
+	}
+}
+
+// DecRef decrements this object's reference count.
+func (r *AtomicRefCount) DecRef() {
+	r.DecRefWithDestructor(nil)
+}