1 files changed, 22 insertions, 16 deletions
diff --git a/pkg/refs/refcounter_test.go b/pkg/refs/refcounter_test.go
index 1ab4a4440..6d0dd1018 100644
--- a/pkg/refs/refcounter_test.go
+++ b/pkg/refs/refcounter_test.go
@@ -18,6 +18,7 @@ import (
 	"reflect"
 	"testing"
 
+	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/sync"
 )
 
@@ -31,11 +32,11 @@ type testCounter struct {
 	destroyed bool
 }
 
-func (t *testCounter) DecRef() {
-	t.AtomicRefCount.DecRefWithDestructor(t.destroy)
+func (t *testCounter) DecRef(ctx context.Context) {
+	t.AtomicRefCount.DecRefWithDestructor(ctx, t.destroy)
 }
 
-func (t *testCounter) destroy() {
+func (t *testCounter) destroy(context.Context) {
 	t.mu.Lock()
 	defer t.mu.Unlock()
 	t.destroyed = true
@@ -53,7 +54,7 @@ func newTestCounter() *testCounter {
 
 func TestOneRef(t *testing.T) {
 	tc := newTestCounter()
-	tc.DecRef()
+	tc.DecRef(context.Background())
 
 	if !tc.IsDestroyed() {
 		t.Errorf("object should have been destroyed")
@@ -63,8 +64,9 @@ func TestOneRef(t *testing.T) {
 func TestTwoRefs(t *testing.T) {
 	tc := newTestCounter()
 	tc.IncRef()
-	tc.DecRef()
-	tc.DecRef()
+	ctx := context.Background()
+	tc.DecRef(ctx)
+	tc.DecRef(ctx)
 
 	if !tc.IsDestroyed() {
 		t.Errorf("object should have been destroyed")
@@ -74,12 +76,13 @@ func TestTwoRefs(t *testing.T) {
 func TestMultiRefs(t *testing.T) {
 	tc := newTestCounter()
 	tc.IncRef()
-	tc.DecRef()
+	ctx := context.Background()
+	tc.DecRef(ctx)
 
 	tc.IncRef()
-	tc.DecRef()
+	tc.DecRef(ctx)
 
-	tc.DecRef()
+	tc.DecRef(ctx)
 
 	if !tc.IsDestroyed() {
 		t.Errorf("object should have been destroyed")
@@ -89,19 +92,20 @@ func TestMultiRefs(t *testing.T) {
 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()
+		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()
+		x.DecRef(ctx)
 	}
 
 	// Shouldn't be destroyed yet. (Can't continue if this fails.)
@@ -110,7 +114,7 @@ func TestWeakRef(t *testing.T) {
 	}
 
 	// Drop the original reference.
-	tc.DecRef()
+	tc.DecRef(ctx)
 
 	// Assert destroyed.
 	if !tc.IsDestroyed() {
@@ -126,7 +130,8 @@ func TestWeakRef(t *testing.T) {
 func TestWeakRefDrop(t *testing.T) {
 	tc := newTestCounter()
 	w := NewWeakRef(tc, nil)
-	w.Drop()
+	ctx := context.Background()
+	w.Drop(ctx)
 
 	// Just assert the list is empty.
 	if !tc.weakRefs.Empty() {
@@ -134,14 +139,14 @@ func TestWeakRefDrop(t *testing.T) {
 	}
 
 	// Drop the original reference.
-	tc.DecRef()
+	tc.DecRef(ctx)
 }
 
 type testWeakRefUser struct {
 	weakRefGone func()
 }
 
-func (u *testWeakRefUser) WeakRefGone() {
+func (u *testWeakRefUser) WeakRefGone(ctx context.Context) {
 	u.weakRefGone()
 }
 
@@ -165,7 +170,8 @@ func TestCallback(t *testing.T) {
 	}})
 
 	// Drop the original reference, this must trigger the callback.
-	tc.DecRef()
+	ctx := context.Background()
+	tc.DecRef(ctx)
 
 	if !called {
 		t.Fatalf("Callback not called")