github.com/SagerNet/gvisor@v0.0.0-20210707092255-7731c139d75c/pkg/refs/refcounter_test.go (about)

     1  // Copyright 2018 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package refs
    16  
    17  import (
    18  	"reflect"
    19  	"testing"
    20  
    21  	"github.com/SagerNet/gvisor/pkg/context"
    22  	"github.com/SagerNet/gvisor/pkg/sync"
    23  )
    24  
    25  type testCounter struct {
    26  	AtomicRefCount
    27  
    28  	// mu protects the boolean below.
    29  	mu sync.Mutex
    30  
    31  	// destroyed indicates whether this was destroyed.
    32  	destroyed bool
    33  }
    34  
    35  func (t *testCounter) DecRef(ctx context.Context) {
    36  	t.AtomicRefCount.DecRefWithDestructor(ctx, t.destroy)
    37  }
    38  
    39  func (t *testCounter) destroy(context.Context) {
    40  	t.mu.Lock()
    41  	defer t.mu.Unlock()
    42  	t.destroyed = true
    43  }
    44  
    45  func (t *testCounter) IsDestroyed() bool {
    46  	t.mu.Lock()
    47  	defer t.mu.Unlock()
    48  	return t.destroyed
    49  }
    50  
    51  func newTestCounter() *testCounter {
    52  	return &testCounter{destroyed: false}
    53  }
    54  
    55  func TestOneRef(t *testing.T) {
    56  	tc := newTestCounter()
    57  	tc.DecRef(context.Background())
    58  
    59  	if !tc.IsDestroyed() {
    60  		t.Errorf("object should have been destroyed")
    61  	}
    62  }
    63  
    64  func TestTwoRefs(t *testing.T) {
    65  	tc := newTestCounter()
    66  	tc.IncRef()
    67  	ctx := context.Background()
    68  	tc.DecRef(ctx)
    69  	tc.DecRef(ctx)
    70  
    71  	if !tc.IsDestroyed() {
    72  		t.Errorf("object should have been destroyed")
    73  	}
    74  }
    75  
    76  func TestMultiRefs(t *testing.T) {
    77  	tc := newTestCounter()
    78  	tc.IncRef()
    79  	ctx := context.Background()
    80  	tc.DecRef(ctx)
    81  
    82  	tc.IncRef()
    83  	tc.DecRef(ctx)
    84  
    85  	tc.DecRef(ctx)
    86  
    87  	if !tc.IsDestroyed() {
    88  		t.Errorf("object should have been destroyed")
    89  	}
    90  }
    91  
    92  func TestWeakRef(t *testing.T) {
    93  	tc := newTestCounter()
    94  	w := NewWeakRef(tc, nil)
    95  	ctx := context.Background()
    96  
    97  	// Try resolving.
    98  	if x := w.Get(); x == nil {
    99  		t.Errorf("weak reference didn't resolve: expected %v, got nil", tc)
   100  	} else {
   101  		x.DecRef(ctx)
   102  	}
   103  
   104  	// Try resolving again.
   105  	if x := w.Get(); x == nil {
   106  		t.Errorf("weak reference didn't resolve: expected %v, got nil", tc)
   107  	} else {
   108  		x.DecRef(ctx)
   109  	}
   110  
   111  	// Shouldn't be destroyed yet. (Can't continue if this fails.)
   112  	if tc.IsDestroyed() {
   113  		t.Fatalf("original object destroyed earlier than expected")
   114  	}
   115  
   116  	// Drop the original reference.
   117  	tc.DecRef(ctx)
   118  
   119  	// Assert destroyed.
   120  	if !tc.IsDestroyed() {
   121  		t.Errorf("original object not destroyed as expected")
   122  	}
   123  
   124  	// Shouldn't be anything.
   125  	if x := w.Get(); x != nil {
   126  		t.Errorf("weak reference resolved: expected nil, got %v", x)
   127  	}
   128  }
   129  
   130  func TestWeakRefDrop(t *testing.T) {
   131  	tc := newTestCounter()
   132  	w := NewWeakRef(tc, nil)
   133  	ctx := context.Background()
   134  	w.Drop(ctx)
   135  
   136  	// Just assert the list is empty.
   137  	if !tc.weakRefs.Empty() {
   138  		t.Errorf("weak reference not dropped")
   139  	}
   140  
   141  	// Drop the original reference.
   142  	tc.DecRef(ctx)
   143  }
   144  
   145  type testWeakRefUser struct {
   146  	weakRefGone func()
   147  }
   148  
   149  func (u *testWeakRefUser) WeakRefGone(ctx context.Context) {
   150  	u.weakRefGone()
   151  }
   152  
   153  func TestCallback(t *testing.T) {
   154  	called := false
   155  	tc := newTestCounter()
   156  	var w *WeakRef
   157  	w = NewWeakRef(tc, &testWeakRefUser{func() {
   158  		called = true
   159  
   160  		// Check that the weak ref has been zapped.
   161  		rc := w.obj.Load().(RefCounter)
   162  		if v := reflect.ValueOf(rc); v != reflect.Zero(v.Type()) {
   163  			t.Fatalf("Callback called with non-nil ptr")
   164  		}
   165  
   166  		// Check that we're not holding the mutex by acquiring and
   167  		// releasing it.
   168  		tc.mu.Lock()
   169  		tc.mu.Unlock()
   170  	}})
   171  
   172  	// Drop the original reference, this must trigger the callback.
   173  	ctx := context.Background()
   174  	tc.DecRef(ctx)
   175  
   176  	if !called {
   177  		t.Fatalf("Callback not called")
   178  	}
   179  }