sentry: optimize dirent weakref map save / restore.

Weak references save / restore involves multiple interface indirection
and cause material latency overhead when there are lots of dirents, each
containing a weak reference map. The nil entries in the map should also
be purged.

PiperOrigin-RevId: 210593727
Change-Id: Ied6f4c3c0726fcc53a24b983d9b3a79121b6b758

2 files changed, 24 insertions, 1 deletions

diff --git a/pkg/sentry/fs/dirent.go b/pkg/sentry/fs/dirent.go
index 5587582b5..9417e808f 100644
--- a/pkg/sentry/fs/dirent.go
+++ b/pkg/sentry/fs/dirent.go
@@ -201,7 +201,7 @@ type Dirent struct {
 	mu sync.Mutex `state:"nosave"`
 
 	// children are cached via weak references.
-	children map[string]*refs.WeakRef
+	children map[string]*refs.WeakRef `state:".(map[string]*Dirent)"`
 }
 
 // NewDirent returns a new root Dirent, taking the caller's reference on inode. The caller
diff --git a/pkg/sentry/fs/dirent_state.go b/pkg/sentry/fs/dirent_state.go
index c6a1b5e38..fb81e7d54 100644
--- a/pkg/sentry/fs/dirent_state.go
+++ b/pkg/sentry/fs/dirent_state.go
@@ -17,6 +17,8 @@ package fs
 import (
 	"fmt"
 	"sync/atomic"
+
+	"gvisor.googlesource.com/gvisor/pkg/refs"
 )
 
 // beforeSave is invoked by stateify.
@@ -36,6 +38,27 @@ func (d *Dirent) beforeSave() {
 	}
 }
 
+// saveChildren is invoked by stateify.
+func (d *Dirent) saveChildren() map[string]*Dirent {
+	c := make(map[string]*Dirent)
+	for name, w := range d.children {
+		if rc := w.Get(); rc != nil {
+			// Drop the reference count obtain in w.Get()
+			rc.DecRef()
+			c[name] = rc.(*Dirent)
+		}
+	}
+	return c
+}
+
+// loadChildren is invoked by stateify.
+func (d *Dirent) loadChildren(children map[string]*Dirent) {
+	d.children = make(map[string]*refs.WeakRef)
+	for name, c := range children {
+		d.children[name] = refs.NewWeakRef(c, nil)
+	}
+}
+
 // afterLoad is invoked by stateify.
 func (d *Dirent) afterLoad() {
 	if d.userVisible {