4 files changed, 542 insertions, 0 deletions
diff --git a/pkg/refs/BUILD b/pkg/refs/BUILD
new file mode 100644
index 000000000..4b7c9345d
--- /dev/null
+++ b/pkg/refs/BUILD
@@ -0,0 +1,37 @@
+package(licenses = ["notice"])  # Apache 2.0
+
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+load("//tools/go_stateify:defs.bzl", "go_stateify")
+
+go_stateify(
+    name = "refs_state",
+    srcs = [
+        "refcounter.go",
+        "refcounter_state.go",
+    ],
+    out = "refs_state.go",
+    package = "refs",
+)
+
+go_library(
+    name = "refs",
+    srcs = [
+        "refcounter.go",
+        "refcounter_state.go",
+        "refs_state.go",
+    ],
+    importpath = "gvisor.googlesource.com/gvisor/pkg/refs",
+    visibility = ["//:sandbox"],
+    deps = [
+        "//pkg/ilist",
+        "//pkg/log",
+        "//pkg/state",
+    ],
+)
+
+go_test(
+    name = "refs_test",
+    size = "small",
+    srcs = ["refcounter_test.go"],
+    embed = [":refs"],
+)
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)
+}
diff --git a/pkg/refs/refcounter_state.go b/pkg/refs/refcounter_state.go
new file mode 100644
index 000000000..1be67f951
--- /dev/null
+++ b/pkg/refs/refcounter_state.go
@@ -0,0 +1,34 @@
+// 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
+
+type savedReference struct {
+	obj interface{}
+}
+
+func (w *WeakRef) saveObj() savedReference {
+	// We load the object directly, because it is typed. This will be
+	// serialized and loaded as a typed value.
+	return savedReference{w.obj.Load()}
+}
+
+func (w *WeakRef) loadObj(v savedReference) {
+	// See note above. This will be serialized and loaded typed. So we're okay
+	// as long as refs aren't changing during save and load (which they should
+	// not be).
+	//
+	// w.user is loaded before loadObj is called.
+	w.init(v.obj.(RefCounter), w.user)
+}
diff --git a/pkg/refs/refcounter_test.go b/pkg/refs/refcounter_test.go
new file mode 100644
index 000000000..cc11bcd71
--- /dev/null
+++ b/pkg/refs/refcounter_test.go
@@ -0,0 +1,172 @@
+// 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
+
+import (
+	"reflect"
+	"sync"
+	"testing"
+)
+
+type testCounter struct {
+	AtomicRefCount
+
+	// mu protects the boolean below.
+	mu sync.Mutex
+
+	// destroyed indicates whether this was destroyed.
+	destroyed bool
+}
+
+func (t *testCounter) DecRef() {
+	t.AtomicRefCount.DecRefWithDestructor(t.destroy)
+}
+
+func (t *testCounter) destroy() {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	t.destroyed = true
+}
+
+func (t *testCounter) IsDestroyed() bool {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.destroyed
+}
+
+func newTestCounter() *testCounter {
+	return &testCounter{destroyed: false}
+}
+
+func TestOneRef(t *testing.T) {
+	tc := newTestCounter()
+	tc.DecRef()
+
+	if !tc.IsDestroyed() {
+		t.Errorf("object should have been destroyed")
+	}
+}
+
+func TestTwoRefs(t *testing.T) {
+	tc := newTestCounter()
+	tc.IncRef()
+	tc.DecRef()
+	tc.DecRef()
+
+	if !tc.IsDestroyed() {
+		t.Errorf("object should have been destroyed")
+	}
+}
+
+func TestMultiRefs(t *testing.T) {
+	tc := newTestCounter()
+	tc.IncRef()
+	tc.DecRef()
+
+	tc.IncRef()
+	tc.DecRef()
+
+	tc.DecRef()
+
+	if !tc.IsDestroyed() {
+		t.Errorf("object should have been destroyed")
+	}
+}
+
+func TestWeakRef(t *testing.T) {
+	tc := newTestCounter()
+	w := NewWeakRef(tc, nil)
+
+	// Try resolving.
+	if x := w.Get(); x == nil {
+		t.Errorf("weak reference didn't resolve: expected %v, got nil", tc)
+	} else {
+		x.DecRef()
+	}
+
+	// Try resolving again.
+	if x := w.Get(); x == nil {
+		t.Errorf("weak reference didn't resolve: expected %v, got nil", tc)
+	} else {
+		x.DecRef()
+	}
+
+	// Shouldn't be destroyed yet. (Can't continue if this fails.)
+	if tc.IsDestroyed() {
+		t.Fatalf("original object destroyed earlier than expected")
+	}
+
+	// Drop the original reference.
+	tc.DecRef()
+
+	// Assert destroyed.
+	if !tc.IsDestroyed() {
+		t.Errorf("original object not destroyed as expected")
+	}
+
+	// Shouldn't be anything.
+	if x := w.Get(); x != nil {
+		t.Errorf("weak reference resolved: expected nil, got %v", x)
+	}
+}
+
+func TestWeakRefDrop(t *testing.T) {
+	tc := newTestCounter()
+	w := NewWeakRef(tc, nil)
+	w.Drop()
+
+	// Just assert the list is empty.
+	if !tc.weakRefs.Empty() {
+		t.Errorf("weak reference not dropped")
+	}
+
+	// Drop the original reference.
+	tc.DecRef()
+}
+
+type testWeakRefUser struct {
+	weakRefGone func()
+}
+
+func (u *testWeakRefUser) WeakRefGone() {
+	u.weakRefGone()
+}
+
+func TestCallback(t *testing.T) {
+	called := false
+	tc := newTestCounter()
+	var w *WeakRef
+	w = NewWeakRef(tc, &testWeakRefUser{func() {
+		called = true
+
+		// Check that the weak ref has been zapped.
+		rc := w.obj.Load().(RefCounter)
+		if v := reflect.ValueOf(rc); v != reflect.Zero(v.Type()) {
+			t.Fatalf("Callback called with non-nil ptr")
+		}
+
+		// Check that we're not holding the mutex by acquiring and
+		// releasing it.
+		tc.mu.Lock()
+		tc.mu.Unlock()
+	}})
+
+	// Drop the original reference, this must trigger the callback.
+	tc.DecRef()
+
+	if !called {
+		t.Fatalf("Callback not called")
+	}
+}