4 files changed, 339 insertions, 221 deletions
diff --git a/pkg/refs/BUILD b/pkg/refs/BUILD
deleted file mode 100644
index 9888cce9c..000000000
--- a/pkg/refs/BUILD
+++ /dev/null
@@ -1,42 +0,0 @@
-load("//tools:defs.bzl", "go_library", "go_test")
-load("//tools/go_generics:defs.bzl", "go_template_instance")
-
-package(licenses = ["notice"])
-
-go_template_instance(
-    name = "weak_ref_list",
-    out = "weak_ref_list.go",
-    package = "refs",
-    prefix = "weakRef",
-    template = "//pkg/ilist:generic_list",
-    types = {
-        "Element": "*WeakRef",
-        "Linker": "*WeakRef",
-    },
-)
-
-go_library(
-    name = "refs",
-    srcs = [
-        "refcounter.go",
-        "refcounter_state.go",
-        "weak_ref_list.go",
-    ],
-    visibility = ["//:sandbox"],
-    deps = [
-        "//pkg/context",
-        "//pkg/log",
-        "//pkg/sync",
-    ],
-)
-
-go_test(
-    name = "refs_test",
-    size = "small",
-    srcs = ["refcounter_test.go"],
-    library = ":refs",
-    deps = [
-        "//pkg/context",
-        "//pkg/sync",
-    ],
-)
diff --git a/pkg/refs/refcounter_test.go b/pkg/refs/refcounter_test.go
deleted file mode 100644
index 6d0dd1018..000000000
--- a/pkg/refs/refcounter_test.go
+++ /dev/null
@@ -1,179 +0,0 @@
-// Copyright 2018 The gVisor Authors.
-//
-// 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"
-	"testing"
-
-	"gvisor.dev/gvisor/pkg/context"
-	"gvisor.dev/gvisor/pkg/sync"
-)
-
-type testCounter struct {
-	AtomicRefCount
-
-	// mu protects the boolean below.
-	mu sync.Mutex
-
-	// destroyed indicates whether this was destroyed.
-	destroyed bool
-}
-
-func (t *testCounter) DecRef(ctx context.Context) {
-	t.AtomicRefCount.DecRefWithDestructor(ctx, t.destroy)
-}
-
-func (t *testCounter) destroy(context.Context) {
-	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(context.Background())
-
-	if !tc.IsDestroyed() {
-		t.Errorf("object should have been destroyed")
-	}
-}
-
-func TestTwoRefs(t *testing.T) {
-	tc := newTestCounter()
-	tc.IncRef()
-	ctx := context.Background()
-	tc.DecRef(ctx)
-	tc.DecRef(ctx)
-
-	if !tc.IsDestroyed() {
-		t.Errorf("object should have been destroyed")
-	}
-}
-
-func TestMultiRefs(t *testing.T) {
-	tc := newTestCounter()
-	tc.IncRef()
-	ctx := context.Background()
-	tc.DecRef(ctx)
-
-	tc.IncRef()
-	tc.DecRef(ctx)
-
-	tc.DecRef(ctx)
-
-	if !tc.IsDestroyed() {
-		t.Errorf("object should have been destroyed")
-	}
-}
-
-func TestWeakRef(t *testing.T) {
-	tc := newTestCounter()
-	w := NewWeakRef(tc, nil)
-	ctx := context.Background()
-
-	// Try resolving.
-	if x := w.Get(); x == nil {
-		t.Errorf("weak reference didn't resolve: expected %v, got nil", tc)
-	} else {
-		x.DecRef(ctx)
-	}
-
-	// Try resolving again.
-	if x := w.Get(); x == nil {
-		t.Errorf("weak reference didn't resolve: expected %v, got nil", tc)
-	} else {
-		x.DecRef(ctx)
-	}
-
-	// 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(ctx)
-
-	// 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)
-	ctx := context.Background()
-	w.Drop(ctx)
-
-	// Just assert the list is empty.
-	if !tc.weakRefs.Empty() {
-		t.Errorf("weak reference not dropped")
-	}
-
-	// Drop the original reference.
-	tc.DecRef(ctx)
-}
-
-type testWeakRefUser struct {
-	weakRefGone func()
-}
-
-func (u *testWeakRefUser) WeakRefGone(ctx context.Context) {
-	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.
-	ctx := context.Background()
-	tc.DecRef(ctx)
-
-	if !called {
-		t.Fatalf("Callback not called")
-	}
-}
diff --git a/pkg/refs/refs_state_autogen.go b/pkg/refs/refs_state_autogen.go
new file mode 100644
index 000000000..ab4c21b09
--- /dev/null
+++ b/pkg/refs/refs_state_autogen.go
@@ -0,0 +1,146 @@
+// automatically generated by stateify.
+
+package refs
+
+import (
+	"gvisor.dev/gvisor/pkg/state"
+)
+
+func (w *WeakRef) StateTypeName() string {
+	return "pkg/refs.WeakRef"
+}
+
+func (w *WeakRef) StateFields() []string {
+	return []string{
+		"obj",
+		"user",
+	}
+}
+
+func (w *WeakRef) beforeSave() {}
+
+func (w *WeakRef) StateSave(stateSinkObject state.Sink) {
+	w.beforeSave()
+	var objValue savedReference = w.saveObj()
+	stateSinkObject.SaveValue(0, objValue)
+	stateSinkObject.Save(1, &w.user)
+}
+
+func (w *WeakRef) afterLoad() {}
+
+func (w *WeakRef) StateLoad(stateSourceObject state.Source) {
+	stateSourceObject.Load(1, &w.user)
+	stateSourceObject.LoadValue(0, new(savedReference), func(y interface{}) { w.loadObj(y.(savedReference)) })
+}
+
+func (r *AtomicRefCount) StateTypeName() string {
+	return "pkg/refs.AtomicRefCount"
+}
+
+func (r *AtomicRefCount) StateFields() []string {
+	return []string{
+		"refCount",
+		"name",
+		"stack",
+	}
+}
+
+func (r *AtomicRefCount) beforeSave() {}
+
+func (r *AtomicRefCount) StateSave(stateSinkObject state.Sink) {
+	r.beforeSave()
+	stateSinkObject.Save(0, &r.refCount)
+	stateSinkObject.Save(1, &r.name)
+	stateSinkObject.Save(2, &r.stack)
+}
+
+func (r *AtomicRefCount) afterLoad() {}
+
+func (r *AtomicRefCount) StateLoad(stateSourceObject state.Source) {
+	stateSourceObject.Load(0, &r.refCount)
+	stateSourceObject.Load(1, &r.name)
+	stateSourceObject.Load(2, &r.stack)
+}
+
+func (s *savedReference) StateTypeName() string {
+	return "pkg/refs.savedReference"
+}
+
+func (s *savedReference) StateFields() []string {
+	return []string{
+		"obj",
+	}
+}
+
+func (s *savedReference) beforeSave() {}
+
+func (s *savedReference) StateSave(stateSinkObject state.Sink) {
+	s.beforeSave()
+	stateSinkObject.Save(0, &s.obj)
+}
+
+func (s *savedReference) afterLoad() {}
+
+func (s *savedReference) StateLoad(stateSourceObject state.Source) {
+	stateSourceObject.Load(0, &s.obj)
+}
+
+func (l *weakRefList) StateTypeName() string {
+	return "pkg/refs.weakRefList"
+}
+
+func (l *weakRefList) StateFields() []string {
+	return []string{
+		"head",
+		"tail",
+	}
+}
+
+func (l *weakRefList) beforeSave() {}
+
+func (l *weakRefList) StateSave(stateSinkObject state.Sink) {
+	l.beforeSave()
+	stateSinkObject.Save(0, &l.head)
+	stateSinkObject.Save(1, &l.tail)
+}
+
+func (l *weakRefList) afterLoad() {}
+
+func (l *weakRefList) StateLoad(stateSourceObject state.Source) {
+	stateSourceObject.Load(0, &l.head)
+	stateSourceObject.Load(1, &l.tail)
+}
+
+func (e *weakRefEntry) StateTypeName() string {
+	return "pkg/refs.weakRefEntry"
+}
+
+func (e *weakRefEntry) StateFields() []string {
+	return []string{
+		"next",
+		"prev",
+	}
+}
+
+func (e *weakRefEntry) beforeSave() {}
+
+func (e *weakRefEntry) StateSave(stateSinkObject state.Sink) {
+	e.beforeSave()
+	stateSinkObject.Save(0, &e.next)
+	stateSinkObject.Save(1, &e.prev)
+}
+
+func (e *weakRefEntry) afterLoad() {}
+
+func (e *weakRefEntry) StateLoad(stateSourceObject state.Source) {
+	stateSourceObject.Load(0, &e.next)
+	stateSourceObject.Load(1, &e.prev)
+}
+
+func init() {
+	state.Register((*WeakRef)(nil))
+	state.Register((*AtomicRefCount)(nil))
+	state.Register((*savedReference)(nil))
+	state.Register((*weakRefList)(nil))
+	state.Register((*weakRefEntry)(nil))
+}
diff --git a/pkg/refs/weak_ref_list.go b/pkg/refs/weak_ref_list.go
new file mode 100644
index 000000000..31460bfc3
--- /dev/null
+++ b/pkg/refs/weak_ref_list.go
@@ -0,0 +1,193 @@
+package refs
+
+// ElementMapper provides an identity mapping by default.
+//
+// This can be replaced to provide a struct that maps elements to linker
+// objects, if they are not the same. An ElementMapper is not typically
+// required if: Linker is left as is, Element is left as is, or Linker and
+// Element are the same type.
+type weakRefElementMapper struct{}
+
+// linkerFor maps an Element to a Linker.
+//
+// This default implementation should be inlined.
+//
+//go:nosplit
+func (weakRefElementMapper) linkerFor(elem *WeakRef) *WeakRef { return elem }
+
+// List is an intrusive list. Entries can be added to or removed from the list
+// in O(1) time and with no additional memory allocations.
+//
+// The zero value for List is an empty list ready to use.
+//
+// To iterate over a list (where l is a List):
+//      for e := l.Front(); e != nil; e = e.Next() {
+// 		// do something with e.
+//      }
+//
+// +stateify savable
+type weakRefList struct {
+	head *WeakRef
+	tail *WeakRef
+}
+
+// Reset resets list l to the empty state.
+func (l *weakRefList) Reset() {
+	l.head = nil
+	l.tail = nil
+}
+
+// Empty returns true iff the list is empty.
+func (l *weakRefList) Empty() bool {
+	return l.head == nil
+}
+
+// Front returns the first element of list l or nil.
+func (l *weakRefList) Front() *WeakRef {
+	return l.head
+}
+
+// Back returns the last element of list l or nil.
+func (l *weakRefList) Back() *WeakRef {
+	return l.tail
+}
+
+// Len returns the number of elements in the list.
+//
+// NOTE: This is an O(n) operation.
+func (l *weakRefList) Len() (count int) {
+	for e := l.Front(); e != nil; e = (weakRefElementMapper{}.linkerFor(e)).Next() {
+		count++
+	}
+	return count
+}
+
+// PushFront inserts the element e at the front of list l.
+func (l *weakRefList) PushFront(e *WeakRef) {
+	linker := weakRefElementMapper{}.linkerFor(e)
+	linker.SetNext(l.head)
+	linker.SetPrev(nil)
+	if l.head != nil {
+		weakRefElementMapper{}.linkerFor(l.head).SetPrev(e)
+	} else {
+		l.tail = e
+	}
+
+	l.head = e
+}
+
+// PushBack inserts the element e at the back of list l.
+func (l *weakRefList) PushBack(e *WeakRef) {
+	linker := weakRefElementMapper{}.linkerFor(e)
+	linker.SetNext(nil)
+	linker.SetPrev(l.tail)
+	if l.tail != nil {
+		weakRefElementMapper{}.linkerFor(l.tail).SetNext(e)
+	} else {
+		l.head = e
+	}
+
+	l.tail = e
+}
+
+// PushBackList inserts list m at the end of list l, emptying m.
+func (l *weakRefList) PushBackList(m *weakRefList) {
+	if l.head == nil {
+		l.head = m.head
+		l.tail = m.tail
+	} else if m.head != nil {
+		weakRefElementMapper{}.linkerFor(l.tail).SetNext(m.head)
+		weakRefElementMapper{}.linkerFor(m.head).SetPrev(l.tail)
+
+		l.tail = m.tail
+	}
+	m.head = nil
+	m.tail = nil
+}
+
+// InsertAfter inserts e after b.
+func (l *weakRefList) InsertAfter(b, e *WeakRef) {
+	bLinker := weakRefElementMapper{}.linkerFor(b)
+	eLinker := weakRefElementMapper{}.linkerFor(e)
+
+	a := bLinker.Next()
+
+	eLinker.SetNext(a)
+	eLinker.SetPrev(b)
+	bLinker.SetNext(e)
+
+	if a != nil {
+		weakRefElementMapper{}.linkerFor(a).SetPrev(e)
+	} else {
+		l.tail = e
+	}
+}
+
+// InsertBefore inserts e before a.
+func (l *weakRefList) InsertBefore(a, e *WeakRef) {
+	aLinker := weakRefElementMapper{}.linkerFor(a)
+	eLinker := weakRefElementMapper{}.linkerFor(e)
+
+	b := aLinker.Prev()
+	eLinker.SetNext(a)
+	eLinker.SetPrev(b)
+	aLinker.SetPrev(e)
+
+	if b != nil {
+		weakRefElementMapper{}.linkerFor(b).SetNext(e)
+	} else {
+		l.head = e
+	}
+}
+
+// Remove removes e from l.
+func (l *weakRefList) Remove(e *WeakRef) {
+	linker := weakRefElementMapper{}.linkerFor(e)
+	prev := linker.Prev()
+	next := linker.Next()
+
+	if prev != nil {
+		weakRefElementMapper{}.linkerFor(prev).SetNext(next)
+	} else if l.head == e {
+		l.head = next
+	}
+
+	if next != nil {
+		weakRefElementMapper{}.linkerFor(next).SetPrev(prev)
+	} else if l.tail == e {
+		l.tail = prev
+	}
+
+	linker.SetNext(nil)
+	linker.SetPrev(nil)
+}
+
+// Entry is a default implementation of Linker. Users can add anonymous fields
+// of this type to their structs to make them automatically implement the
+// methods needed by List.
+//
+// +stateify savable
+type weakRefEntry struct {
+	next *WeakRef
+	prev *WeakRef
+}
+
+// Next returns the entry that follows e in the list.
+func (e *weakRefEntry) Next() *WeakRef {
+	return e.next
+}
+
+// Prev returns the entry that precedes e in the list.
+func (e *weakRefEntry) Prev() *WeakRef {
+	return e.prev
+}
+
+// SetNext assigns 'entry' as the entry that follows e in the list.
+func (e *weakRefEntry) SetNext(elem *WeakRef) {
+	e.next = elem
+}
+
+// SetPrev assigns 'entry' as the entry that precedes e in the list.
+func (e *weakRefEntry) SetPrev(elem *WeakRef) {
+	e.prev = elem
+}