Merge release-20201019.0-34-g9ca66ec59 (automated)

13 files changed, 134 insertions, 167 deletions

diff --git a/pkg/sentry/kernel/abstract_socket_namespace.go b/pkg/sentry/kernel/abstract_socket_namespace.go
index 1b9721534..0ddbe5ff6 100644
--- a/pkg/sentry/kernel/abstract_socket_namespace.go
+++ b/pkg/sentry/kernel/abstract_socket_namespace.go
@@ -19,7 +19,7 @@ import (
 	"syscall"
 
 	"gvisor.dev/gvisor/pkg/context"
-	"gvisor.dev/gvisor/pkg/refs_vfs2"
+	"gvisor.dev/gvisor/pkg/refsvfs2"
 	"gvisor.dev/gvisor/pkg/sentry/socket/unix/transport"
 	"gvisor.dev/gvisor/pkg/sync"
 )
@@ -27,7 +27,7 @@ import (
 // +stateify savable
 type abstractEndpoint struct {
 	ep     transport.BoundEndpoint
-	socket refs_vfs2.RefCounter
+	socket refsvfs2.RefCounter
 	name   string
 	ns     *AbstractSocketNamespace
 }
@@ -57,7 +57,7 @@ func NewAbstractSocketNamespace() *AbstractSocketNamespace {
 // its backing socket.
 type boundEndpoint struct {
 	transport.BoundEndpoint
-	socket refs_vfs2.RefCounter
+	socket refsvfs2.RefCounter
 }
 
 // Release implements transport.BoundEndpoint.Release.
@@ -89,7 +89,7 @@ func (a *AbstractSocketNamespace) BoundEndpoint(name string) transport.BoundEndp
 //
 // When the last reference managed by socket is dropped, ep may be removed from the
 // namespace.
-func (a *AbstractSocketNamespace) Bind(ctx context.Context, name string, ep transport.BoundEndpoint, socket refs_vfs2.RefCounter) error {
+func (a *AbstractSocketNamespace) Bind(ctx context.Context, name string, ep transport.BoundEndpoint, socket refsvfs2.RefCounter) error {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
@@ -109,7 +109,7 @@ func (a *AbstractSocketNamespace) Bind(ctx context.Context, name string, ep tran
 
 // Remove removes the specified socket at name from the abstract socket
 // namespace, if it has not yet been replaced.
-func (a *AbstractSocketNamespace) Remove(name string, socket refs_vfs2.RefCounter) {
+func (a *AbstractSocketNamespace) Remove(name string, socket refsvfs2.RefCounter) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
diff --git a/pkg/sentry/kernel/fd_table.go b/pkg/sentry/kernel/fd_table.go
index 0ec7344cd..81a998966 100644
--- a/pkg/sentry/kernel/fd_table.go
+++ b/pkg/sentry/kernel/fd_table.go
@@ -110,7 +110,7 @@ func (f *FDTable) saveDescriptorTable() map[int32]descriptor {
 
 func (f *FDTable) loadDescriptorTable(m map[int32]descriptor) {
 	ctx := context.Background()
-	f.init() // Initialize table.
+	f.initNoLeakCheck() // Initialize table.
 	f.used = 0
 	for fd, d := range m {
 		if file, fileVFS2 := f.setAll(ctx, fd, d.file, d.fileVFS2, d.flags); file != nil || fileVFS2 != nil {
diff --git a/pkg/sentry/kernel/fd_table_refs.go b/pkg/sentry/kernel/fd_table_refs.go
index ecba138ac..cbf2e85ed 100644
--- a/pkg/sentry/kernel/fd_table_refs.go
+++ b/pkg/sentry/kernel/fd_table_refs.go
@@ -2,11 +2,9 @@ package kernel
 
 import (
 	"fmt"
-	"runtime"
 	"sync/atomic"
 
-	"gvisor.dev/gvisor/pkg/log"
-	refs_vfs1 "gvisor.dev/gvisor/pkg/refs"
+	"gvisor.dev/gvisor/pkg/refsvfs2"
 )
 
 // ownerType is used to customize logging. Note that we use a pointer to T so
@@ -19,11 +17,6 @@ var FDTableownerType *FDTable
 // Note that the number of references is actually refCount + 1 so that a default
 // zero-value Refs object contains one reference.
 //
-// TODO(gvisor.dev/issue/1486): Store stack traces when leak check is enabled in
-// a map with 16-bit hashes, and store the hash in the top 16 bits of refCount.
-// This will allow us to add stack trace information to the leak messages
-// without growing the size of Refs.
-//
 // +stateify savable
 type FDTableRefs struct {
 	// refCount is composed of two fields:
@@ -36,24 +29,16 @@ type FDTableRefs struct {
 	refCount int64
 }
 
-func (r *FDTableRefs) finalize() {
-	var note string
-	switch refs_vfs1.GetLeakMode() {
-	case refs_vfs1.NoLeakChecking:
-		return
-	case refs_vfs1.UninitializedLeakChecking:
-		note = "(Leak checker uninitialized): "
-	}
-	if n := r.ReadRefs(); n != 0 {
-		log.Warningf("%sRefs %p owned by %T garbage collected with ref count of %d (want 0)", note, r, FDTableownerType, n)
+// EnableLeakCheck enables reference leak checking on r.
+func (r *FDTableRefs) EnableLeakCheck() {
+	if refsvfs2.LeakCheckEnabled() {
+		refsvfs2.Register(r, fmt.Sprintf("%T", FDTableownerType))
 	}
 }
 
-// EnableLeakCheck checks for reference leaks when Refs gets garbage collected.
-func (r *FDTableRefs) EnableLeakCheck() {
-	if refs_vfs1.GetLeakMode() != refs_vfs1.NoLeakChecking {
-		runtime.SetFinalizer(r, (*FDTableRefs).finalize)
-	}
+// LeakMessage implements refsvfs2.CheckedObject.LeakMessage.
+func (r *FDTableRefs) LeakMessage() string {
+	return fmt.Sprintf("%T %p: reference count of %d instead of 0", FDTableownerType, r, r.ReadRefs())
 }
 
 // ReadRefs returns the current number of references. The returned count is
@@ -68,7 +53,7 @@ func (r *FDTableRefs) ReadRefs() int64 {
 //go:nosplit
 func (r *FDTableRefs) IncRef() {
 	if v := atomic.AddInt64(&r.refCount, 1); v <= 0 {
-		panic(fmt.Sprintf("Incrementing non-positive ref count %p owned by %T", r, FDTableownerType))
+		panic(fmt.Sprintf("Incrementing non-positive count %p on %T", r, FDTableownerType))
 	}
 }
 
@@ -110,9 +95,18 @@ func (r *FDTableRefs) DecRef(destroy func()) {
 		panic(fmt.Sprintf("Decrementing non-positive ref count %p, owned by %T", r, FDTableownerType))
 
 	case v == -1:
+		if refsvfs2.LeakCheckEnabled() {
+			refsvfs2.Unregister(r, fmt.Sprintf("%T", FDTableownerType))
+		}
 
 		if destroy != nil {
 			destroy()
 		}
 	}
 }
+
+func (r *FDTableRefs) afterLoad() {
+	if refsvfs2.LeakCheckEnabled() && r.ReadRefs() > 0 {
+		r.EnableLeakCheck()
+	}
+}
diff --git a/pkg/sentry/kernel/fd_table_unsafe.go b/pkg/sentry/kernel/fd_table_unsafe.go
index da79e6627..3476551f3 100644
--- a/pkg/sentry/kernel/fd_table_unsafe.go
+++ b/pkg/sentry/kernel/fd_table_unsafe.go
@@ -31,14 +31,21 @@ type descriptorTable struct {
 	slice unsafe.Pointer `state:".(map[int32]*descriptor)"`
 }
 
-// init initializes the table.
+// initNoLeakCheck initializes the table without enabling leak checking.
 //
-// TODO(gvisor.dev/1486): Enable leak check for FDTable.
-func (f *FDTable) init() {
+// This is used when loading an FDTable after S/R, during which the ref count
+// object itself will enable leak checking if necessary.
+func (f *FDTable) initNoLeakCheck() {
 	var slice []unsafe.Pointer // Empty slice.
 	atomic.StorePointer(&f.slice, unsafe.Pointer(&slice))
 }
 
+// init initializes the table with leak checking.
+func (f *FDTable) init() {
+	f.initNoLeakCheck()
+	f.EnableLeakCheck()
+}
+
 // get gets a file entry.
 //
 // The boolean indicates whether this was in range.
diff --git a/pkg/sentry/kernel/fs_context.go b/pkg/sentry/kernel/fs_context.go
index 08ea2e09c..41fb2a784 100644
--- a/pkg/sentry/kernel/fs_context.go
+++ b/pkg/sentry/kernel/fs_context.go
@@ -130,13 +130,15 @@ func (f *FSContext) Fork() *FSContext {
 		f.root.IncRef()
 	}
 
-	return &FSContext{
+	ctx := &FSContext{
 		cwd:      f.cwd,
 		root:     f.root,
 		cwdVFS2:  f.cwdVFS2,
 		rootVFS2: f.rootVFS2,
 		umask:    f.umask,
 	}
+	ctx.EnableLeakCheck()
+	return ctx
 }
 
 // WorkingDirectory returns the current working directory.
diff --git a/pkg/sentry/kernel/fs_context_refs.go b/pkg/sentry/kernel/fs_context_refs.go
index fb2fde971..025f11faa 100644
--- a/pkg/sentry/kernel/fs_context_refs.go
+++ b/pkg/sentry/kernel/fs_context_refs.go
@@ -2,11 +2,9 @@ package kernel
 
 import (
 	"fmt"
-	"runtime"
 	"sync/atomic"
 
-	"gvisor.dev/gvisor/pkg/log"
-	refs_vfs1 "gvisor.dev/gvisor/pkg/refs"
+	"gvisor.dev/gvisor/pkg/refsvfs2"
 )
 
 // ownerType is used to customize logging. Note that we use a pointer to T so
@@ -19,11 +17,6 @@ var FSContextownerType *FSContext
 // Note that the number of references is actually refCount + 1 so that a default
 // zero-value Refs object contains one reference.
 //
-// TODO(gvisor.dev/issue/1486): Store stack traces when leak check is enabled in
-// a map with 16-bit hashes, and store the hash in the top 16 bits of refCount.
-// This will allow us to add stack trace information to the leak messages
-// without growing the size of Refs.
-//
 // +stateify savable
 type FSContextRefs struct {
 	// refCount is composed of two fields:
@@ -36,24 +29,16 @@ type FSContextRefs struct {
 	refCount int64
 }
 
-func (r *FSContextRefs) finalize() {
-	var note string
-	switch refs_vfs1.GetLeakMode() {
-	case refs_vfs1.NoLeakChecking:
-		return
-	case refs_vfs1.UninitializedLeakChecking:
-		note = "(Leak checker uninitialized): "
-	}
-	if n := r.ReadRefs(); n != 0 {
-		log.Warningf("%sRefs %p owned by %T garbage collected with ref count of %d (want 0)", note, r, FSContextownerType, n)
+// EnableLeakCheck enables reference leak checking on r.
+func (r *FSContextRefs) EnableLeakCheck() {
+	if refsvfs2.LeakCheckEnabled() {
+		refsvfs2.Register(r, fmt.Sprintf("%T", FSContextownerType))
 	}
 }
 
-// EnableLeakCheck checks for reference leaks when Refs gets garbage collected.
-func (r *FSContextRefs) EnableLeakCheck() {
-	if refs_vfs1.GetLeakMode() != refs_vfs1.NoLeakChecking {
-		runtime.SetFinalizer(r, (*FSContextRefs).finalize)
-	}
+// LeakMessage implements refsvfs2.CheckedObject.LeakMessage.
+func (r *FSContextRefs) LeakMessage() string {
+	return fmt.Sprintf("%T %p: reference count of %d instead of 0", FSContextownerType, r, r.ReadRefs())
 }
 
 // ReadRefs returns the current number of references. The returned count is
@@ -68,7 +53,7 @@ func (r *FSContextRefs) ReadRefs() int64 {
 //go:nosplit
 func (r *FSContextRefs) IncRef() {
 	if v := atomic.AddInt64(&r.refCount, 1); v <= 0 {
-		panic(fmt.Sprintf("Incrementing non-positive ref count %p owned by %T", r, FSContextownerType))
+		panic(fmt.Sprintf("Incrementing non-positive count %p on %T", r, FSContextownerType))
 	}
 }
 
@@ -110,9 +95,18 @@ func (r *FSContextRefs) DecRef(destroy func()) {
 		panic(fmt.Sprintf("Decrementing non-positive ref count %p, owned by %T", r, FSContextownerType))
 
 	case v == -1:
+		if refsvfs2.LeakCheckEnabled() {
+			refsvfs2.Unregister(r, fmt.Sprintf("%T", FSContextownerType))
+		}
 
 		if destroy != nil {
 			destroy()
 		}
 	}
 }
+
+func (r *FSContextRefs) afterLoad() {
+	if refsvfs2.LeakCheckEnabled() && r.ReadRefs() > 0 {
+		r.EnableLeakCheck()
+	}
+}
diff --git a/pkg/sentry/kernel/ipc_namespace.go b/pkg/sentry/kernel/ipc_namespace.go
index 3f34ee0db..b87e40dd1 100644
--- a/pkg/sentry/kernel/ipc_namespace.go
+++ b/pkg/sentry/kernel/ipc_namespace.go
@@ -55,7 +55,7 @@ func (i *IPCNamespace) ShmRegistry() *shm.Registry {
 	return i.shms
 }
 
-// DecRef implements refs_vfs2.RefCounter.DecRef.
+// DecRef implements refsvfs2.RefCounter.DecRef.
 func (i *IPCNamespace) DecRef(ctx context.Context) {
 	i.IPCNamespaceRefs.DecRef(func() {
 		i.shms.Release(ctx)
diff --git a/pkg/sentry/kernel/ipc_namespace_refs.go b/pkg/sentry/kernel/ipc_namespace_refs.go
index 263919299..aec0f7a41 100644
--- a/pkg/sentry/kernel/ipc_namespace_refs.go
+++ b/pkg/sentry/kernel/ipc_namespace_refs.go
@@ -2,11 +2,9 @@ package kernel
 
 import (
 	"fmt"
-	"runtime"
 	"sync/atomic"
 
-	"gvisor.dev/gvisor/pkg/log"
-	refs_vfs1 "gvisor.dev/gvisor/pkg/refs"
+	"gvisor.dev/gvisor/pkg/refsvfs2"
 )
 
 // ownerType is used to customize logging. Note that we use a pointer to T so
@@ -19,11 +17,6 @@ var IPCNamespaceownerType *IPCNamespace
 // Note that the number of references is actually refCount + 1 so that a default
 // zero-value Refs object contains one reference.
 //
-// TODO(gvisor.dev/issue/1486): Store stack traces when leak check is enabled in
-// a map with 16-bit hashes, and store the hash in the top 16 bits of refCount.
-// This will allow us to add stack trace information to the leak messages
-// without growing the size of Refs.
-//
 // +stateify savable
 type IPCNamespaceRefs struct {
 	// refCount is composed of two fields:
@@ -36,24 +29,16 @@ type IPCNamespaceRefs struct {
 	refCount int64
 }
 
-func (r *IPCNamespaceRefs) finalize() {
-	var note string
-	switch refs_vfs1.GetLeakMode() {
-	case refs_vfs1.NoLeakChecking:
-		return
-	case refs_vfs1.UninitializedLeakChecking:
-		note = "(Leak checker uninitialized): "
-	}
-	if n := r.ReadRefs(); n != 0 {
-		log.Warningf("%sRefs %p owned by %T garbage collected with ref count of %d (want 0)", note, r, IPCNamespaceownerType, n)
+// EnableLeakCheck enables reference leak checking on r.
+func (r *IPCNamespaceRefs) EnableLeakCheck() {
+	if refsvfs2.LeakCheckEnabled() {
+		refsvfs2.Register(r, fmt.Sprintf("%T", IPCNamespaceownerType))
 	}
 }
 
-// EnableLeakCheck checks for reference leaks when Refs gets garbage collected.
-func (r *IPCNamespaceRefs) EnableLeakCheck() {
-	if refs_vfs1.GetLeakMode() != refs_vfs1.NoLeakChecking {
-		runtime.SetFinalizer(r, (*IPCNamespaceRefs).finalize)
-	}
+// LeakMessage implements refsvfs2.CheckedObject.LeakMessage.
+func (r *IPCNamespaceRefs) LeakMessage() string {
+	return fmt.Sprintf("%T %p: reference count of %d instead of 0", IPCNamespaceownerType, r, r.ReadRefs())
 }
 
 // ReadRefs returns the current number of references. The returned count is
@@ -68,7 +53,7 @@ func (r *IPCNamespaceRefs) ReadRefs() int64 {
 //go:nosplit
 func (r *IPCNamespaceRefs) IncRef() {
 	if v := atomic.AddInt64(&r.refCount, 1); v <= 0 {
-		panic(fmt.Sprintf("Incrementing non-positive ref count %p owned by %T", r, IPCNamespaceownerType))
+		panic(fmt.Sprintf("Incrementing non-positive count %p on %T", r, IPCNamespaceownerType))
 	}
 }
 
@@ -110,9 +95,18 @@ func (r *IPCNamespaceRefs) DecRef(destroy func()) {
 		panic(fmt.Sprintf("Decrementing non-positive ref count %p, owned by %T", r, IPCNamespaceownerType))
 
 	case v == -1:
+		if refsvfs2.LeakCheckEnabled() {
+			refsvfs2.Unregister(r, fmt.Sprintf("%T", IPCNamespaceownerType))
+		}
 
 		if destroy != nil {
 			destroy()
 		}
 	}
 }
+
+func (r *IPCNamespaceRefs) afterLoad() {
+	if refsvfs2.LeakCheckEnabled() && r.ReadRefs() > 0 {
+		r.EnableLeakCheck()
+	}
+}
diff --git a/pkg/sentry/kernel/kernel_state_autogen.go b/pkg/sentry/kernel/kernel_state_autogen.go
index b7d6d4f1c..be3c71199 100644
--- a/pkg/sentry/kernel/kernel_state_autogen.go
+++ b/pkg/sentry/kernel/kernel_state_autogen.go
@@ -169,10 +169,9 @@ func (r *FDTableRefs) StateSave(stateSinkObject state.Sink) {
 	stateSinkObject.Save(0, &r.refCount)
 }
 
-func (r *FDTableRefs) afterLoad() {}
-
 func (r *FDTableRefs) StateLoad(stateSourceObject state.Source) {
 	stateSourceObject.Load(0, &r.refCount)
+	stateSourceObject.AfterLoad(r.afterLoad)
 }
 
 func (f *FSContext) StateTypeName() string {
@@ -230,10 +229,9 @@ func (r *FSContextRefs) StateSave(stateSinkObject state.Sink) {
 	stateSinkObject.Save(0, &r.refCount)
 }
 
-func (r *FSContextRefs) afterLoad() {}
-
 func (r *FSContextRefs) StateLoad(stateSourceObject state.Source) {
 	stateSourceObject.Load(0, &r.refCount)
+	stateSourceObject.AfterLoad(r.afterLoad)
 }
 
 func (i *IPCNamespace) StateTypeName() string {
@@ -285,10 +283,9 @@ func (r *IPCNamespaceRefs) StateSave(stateSinkObject state.Sink) {
 	stateSinkObject.Save(0, &r.refCount)
 }
 
-func (r *IPCNamespaceRefs) afterLoad() {}
-
 func (r *IPCNamespaceRefs) StateLoad(stateSourceObject state.Source) {
 	stateSourceObject.Load(0, &r.refCount)
+	stateSourceObject.AfterLoad(r.afterLoad)
 }
 
 func (k *Kernel) StateTypeName() string {
@@ -758,10 +755,9 @@ func (r *ProcessGroupRefs) StateSave(stateSinkObject state.Sink) {
 	stateSinkObject.Save(0, &r.refCount)
 }
 
-func (r *ProcessGroupRefs) afterLoad() {}
-
 func (r *ProcessGroupRefs) StateLoad(stateSourceObject state.Source) {
 	stateSourceObject.Load(0, &r.refCount)
+	stateSourceObject.AfterLoad(r.afterLoad)
 }
 
 func (p *ptraceOptions) StateTypeName() string {
@@ -932,10 +928,9 @@ func (r *SessionRefs) StateSave(stateSinkObject state.Sink) {
 	stateSinkObject.Save(0, &r.refCount)
 }
 
-func (r *SessionRefs) afterLoad() {}
-
 func (r *SessionRefs) StateLoad(stateSourceObject state.Source) {
 	stateSourceObject.Load(0, &r.refCount)
+	stateSourceObject.AfterLoad(r.afterLoad)
 }
 
 func (s *Session) StateTypeName() string {
diff --git a/pkg/sentry/kernel/process_group_refs.go b/pkg/sentry/kernel/process_group_refs.go
index 4ed6e6458..1f4486817 100644
--- a/pkg/sentry/kernel/process_group_refs.go
+++ b/pkg/sentry/kernel/process_group_refs.go
@@ -2,11 +2,9 @@ package kernel
 
 import (
 	"fmt"
-	"runtime"
 	"sync/atomic"
 
-	"gvisor.dev/gvisor/pkg/log"
-	refs_vfs1 "gvisor.dev/gvisor/pkg/refs"
+	"gvisor.dev/gvisor/pkg/refsvfs2"
 )
 
 // ownerType is used to customize logging. Note that we use a pointer to T so
@@ -19,11 +17,6 @@ var ProcessGroupownerType *ProcessGroup
 // Note that the number of references is actually refCount + 1 so that a default
 // zero-value Refs object contains one reference.
 //
-// TODO(gvisor.dev/issue/1486): Store stack traces when leak check is enabled in
-// a map with 16-bit hashes, and store the hash in the top 16 bits of refCount.
-// This will allow us to add stack trace information to the leak messages
-// without growing the size of Refs.
-//
 // +stateify savable
 type ProcessGroupRefs struct {
 	// refCount is composed of two fields:
@@ -36,24 +29,16 @@ type ProcessGroupRefs struct {
 	refCount int64
 }
 
-func (r *ProcessGroupRefs) finalize() {
-	var note string
-	switch refs_vfs1.GetLeakMode() {
-	case refs_vfs1.NoLeakChecking:
-		return
-	case refs_vfs1.UninitializedLeakChecking:
-		note = "(Leak checker uninitialized): "
-	}
-	if n := r.ReadRefs(); n != 0 {
-		log.Warningf("%sRefs %p owned by %T garbage collected with ref count of %d (want 0)", note, r, ProcessGroupownerType, n)
+// EnableLeakCheck enables reference leak checking on r.
+func (r *ProcessGroupRefs) EnableLeakCheck() {
+	if refsvfs2.LeakCheckEnabled() {
+		refsvfs2.Register(r, fmt.Sprintf("%T", ProcessGroupownerType))
 	}
 }
 
-// EnableLeakCheck checks for reference leaks when Refs gets garbage collected.
-func (r *ProcessGroupRefs) EnableLeakCheck() {
-	if refs_vfs1.GetLeakMode() != refs_vfs1.NoLeakChecking {
-		runtime.SetFinalizer(r, (*ProcessGroupRefs).finalize)
-	}
+// LeakMessage implements refsvfs2.CheckedObject.LeakMessage.
+func (r *ProcessGroupRefs) LeakMessage() string {
+	return fmt.Sprintf("%T %p: reference count of %d instead of 0", ProcessGroupownerType, r, r.ReadRefs())
 }
 
 // ReadRefs returns the current number of references. The returned count is
@@ -68,7 +53,7 @@ func (r *ProcessGroupRefs) ReadRefs() int64 {
 //go:nosplit
 func (r *ProcessGroupRefs) IncRef() {
 	if v := atomic.AddInt64(&r.refCount, 1); v <= 0 {
-		panic(fmt.Sprintf("Incrementing non-positive ref count %p owned by %T", r, ProcessGroupownerType))
+		panic(fmt.Sprintf("Incrementing non-positive count %p on %T", r, ProcessGroupownerType))
 	}
 }
 
@@ -110,9 +95,18 @@ func (r *ProcessGroupRefs) DecRef(destroy func()) {
 		panic(fmt.Sprintf("Decrementing non-positive ref count %p, owned by %T", r, ProcessGroupownerType))
 
 	case v == -1:
+		if refsvfs2.LeakCheckEnabled() {
+			refsvfs2.Unregister(r, fmt.Sprintf("%T", ProcessGroupownerType))
+		}
 
 		if destroy != nil {
 			destroy()
 		}
 	}
 }
+
+func (r *ProcessGroupRefs) afterLoad() {
+	if refsvfs2.LeakCheckEnabled() && r.ReadRefs() > 0 {
+		r.EnableLeakCheck()
+	}
+}
diff --git a/pkg/sentry/kernel/session_refs.go b/pkg/sentry/kernel/session_refs.go
index f2e1bb797..86df93be1 100644
--- a/pkg/sentry/kernel/session_refs.go
+++ b/pkg/sentry/kernel/session_refs.go
@@ -2,11 +2,9 @@ package kernel
 
 import (
 	"fmt"
-	"runtime"
 	"sync/atomic"
 
-	"gvisor.dev/gvisor/pkg/log"
-	refs_vfs1 "gvisor.dev/gvisor/pkg/refs"
+	"gvisor.dev/gvisor/pkg/refsvfs2"
 )
 
 // ownerType is used to customize logging. Note that we use a pointer to T so
@@ -19,11 +17,6 @@ var SessionownerType *Session
 // Note that the number of references is actually refCount + 1 so that a default
 // zero-value Refs object contains one reference.
 //
-// TODO(gvisor.dev/issue/1486): Store stack traces when leak check is enabled in
-// a map with 16-bit hashes, and store the hash in the top 16 bits of refCount.
-// This will allow us to add stack trace information to the leak messages
-// without growing the size of Refs.
-//
 // +stateify savable
 type SessionRefs struct {
 	// refCount is composed of two fields:
@@ -36,24 +29,16 @@ type SessionRefs struct {
 	refCount int64
 }
 
-func (r *SessionRefs) finalize() {
-	var note string
-	switch refs_vfs1.GetLeakMode() {
-	case refs_vfs1.NoLeakChecking:
-		return
-	case refs_vfs1.UninitializedLeakChecking:
-		note = "(Leak checker uninitialized): "
-	}
-	if n := r.ReadRefs(); n != 0 {
-		log.Warningf("%sRefs %p owned by %T garbage collected with ref count of %d (want 0)", note, r, SessionownerType, n)
+// EnableLeakCheck enables reference leak checking on r.
+func (r *SessionRefs) EnableLeakCheck() {
+	if refsvfs2.LeakCheckEnabled() {
+		refsvfs2.Register(r, fmt.Sprintf("%T", SessionownerType))
 	}
 }
 
-// EnableLeakCheck checks for reference leaks when Refs gets garbage collected.
-func (r *SessionRefs) EnableLeakCheck() {
-	if refs_vfs1.GetLeakMode() != refs_vfs1.NoLeakChecking {
-		runtime.SetFinalizer(r, (*SessionRefs).finalize)
-	}
+// LeakMessage implements refsvfs2.CheckedObject.LeakMessage.
+func (r *SessionRefs) LeakMessage() string {
+	return fmt.Sprintf("%T %p: reference count of %d instead of 0", SessionownerType, r, r.ReadRefs())
 }
 
 // ReadRefs returns the current number of references. The returned count is
@@ -68,7 +53,7 @@ func (r *SessionRefs) ReadRefs() int64 {
 //go:nosplit
 func (r *SessionRefs) IncRef() {
 	if v := atomic.AddInt64(&r.refCount, 1); v <= 0 {
-		panic(fmt.Sprintf("Incrementing non-positive ref count %p owned by %T", r, SessionownerType))
+		panic(fmt.Sprintf("Incrementing non-positive count %p on %T", r, SessionownerType))
 	}
 }
 
@@ -110,9 +95,18 @@ func (r *SessionRefs) DecRef(destroy func()) {
 		panic(fmt.Sprintf("Decrementing non-positive ref count %p, owned by %T", r, SessionownerType))
 
 	case v == -1:
+		if refsvfs2.LeakCheckEnabled() {
+			refsvfs2.Unregister(r, fmt.Sprintf("%T", SessionownerType))
+		}
 
 		if destroy != nil {
 			destroy()
 		}
 	}
 }
+
+func (r *SessionRefs) afterLoad() {
+	if refsvfs2.LeakCheckEnabled() && r.ReadRefs() > 0 {
+		r.EnableLeakCheck()
+	}
+}
diff --git a/pkg/sentry/kernel/shm/shm_refs.go b/pkg/sentry/kernel/shm/shm_refs.go
index 51e07d0b3..58b0e80cf 100644
--- a/pkg/sentry/kernel/shm/shm_refs.go
+++ b/pkg/sentry/kernel/shm/shm_refs.go
@@ -2,11 +2,9 @@ package shm
 
 import (
 	"fmt"
-	"runtime"
 	"sync/atomic"
 
-	"gvisor.dev/gvisor/pkg/log"
-	refs_vfs1 "gvisor.dev/gvisor/pkg/refs"
+	"gvisor.dev/gvisor/pkg/refsvfs2"
 )
 
 // ownerType is used to customize logging. Note that we use a pointer to T so
@@ -19,11 +17,6 @@ var ShmownerType *Shm
 // Note that the number of references is actually refCount + 1 so that a default
 // zero-value Refs object contains one reference.
 //
-// TODO(gvisor.dev/issue/1486): Store stack traces when leak check is enabled in
-// a map with 16-bit hashes, and store the hash in the top 16 bits of refCount.
-// This will allow us to add stack trace information to the leak messages
-// without growing the size of Refs.
-//
 // +stateify savable
 type ShmRefs struct {
 	// refCount is composed of two fields:
@@ -36,24 +29,16 @@ type ShmRefs struct {
 	refCount int64
 }
 
-func (r *ShmRefs) finalize() {
-	var note string
-	switch refs_vfs1.GetLeakMode() {
-	case refs_vfs1.NoLeakChecking:
-		return
-	case refs_vfs1.UninitializedLeakChecking:
-		note = "(Leak checker uninitialized): "
-	}
-	if n := r.ReadRefs(); n != 0 {
-		log.Warningf("%sRefs %p owned by %T garbage collected with ref count of %d (want 0)", note, r, ShmownerType, n)
+// EnableLeakCheck enables reference leak checking on r.
+func (r *ShmRefs) EnableLeakCheck() {
+	if refsvfs2.LeakCheckEnabled() {
+		refsvfs2.Register(r, fmt.Sprintf("%T", ShmownerType))
 	}
 }
 
-// EnableLeakCheck checks for reference leaks when Refs gets garbage collected.
-func (r *ShmRefs) EnableLeakCheck() {
-	if refs_vfs1.GetLeakMode() != refs_vfs1.NoLeakChecking {
-		runtime.SetFinalizer(r, (*ShmRefs).finalize)
-	}
+// LeakMessage implements refsvfs2.CheckedObject.LeakMessage.
+func (r *ShmRefs) LeakMessage() string {
+	return fmt.Sprintf("%T %p: reference count of %d instead of 0", ShmownerType, r, r.ReadRefs())
 }
 
 // ReadRefs returns the current number of references. The returned count is
@@ -68,7 +53,7 @@ func (r *ShmRefs) ReadRefs() int64 {
 //go:nosplit
 func (r *ShmRefs) IncRef() {
 	if v := atomic.AddInt64(&r.refCount, 1); v <= 0 {
-		panic(fmt.Sprintf("Incrementing non-positive ref count %p owned by %T", r, ShmownerType))
+		panic(fmt.Sprintf("Incrementing non-positive count %p on %T", r, ShmownerType))
 	}
 }
 
@@ -110,9 +95,18 @@ func (r *ShmRefs) DecRef(destroy func()) {
 		panic(fmt.Sprintf("Decrementing non-positive ref count %p, owned by %T", r, ShmownerType))
 
 	case v == -1:
+		if refsvfs2.LeakCheckEnabled() {
+			refsvfs2.Unregister(r, fmt.Sprintf("%T", ShmownerType))
+		}
 
 		if destroy != nil {
 			destroy()
 		}
 	}
 }
+
+func (r *ShmRefs) afterLoad() {
+	if refsvfs2.LeakCheckEnabled() && r.ReadRefs() > 0 {
+		r.EnableLeakCheck()
+	}
+}
diff --git a/pkg/sentry/kernel/shm/shm_state_autogen.go b/pkg/sentry/kernel/shm/shm_state_autogen.go
index 8202c37d6..aca4c9b96 100644
--- a/pkg/sentry/kernel/shm/shm_state_autogen.go
+++ b/pkg/sentry/kernel/shm/shm_state_autogen.go
@@ -129,10 +129,9 @@ func (r *ShmRefs) StateSave(stateSinkObject state.Sink) {
 	stateSinkObject.Save(0, &r.refCount)
 }
 
-func (r *ShmRefs) afterLoad() {}
-
 func (r *ShmRefs) StateLoad(stateSourceObject state.Source) {
 	stateSourceObject.Load(0, &r.refCount)
+	stateSourceObject.AfterLoad(r.afterLoad)
 }
 
 func init() {