4 files changed, 55 insertions, 51 deletions
diff --git a/pkg/sentry/fs/BUILD b/pkg/sentry/fs/BUILD
index dda6a0c9f..1742d3a65 100644
--- a/pkg/sentry/fs/BUILD
+++ b/pkg/sentry/fs/BUILD
@@ -32,7 +32,6 @@ go_library(
         "mock.go",
         "mount.go",
         "mount_overlay.go",
-        "mount_state.go",
         "mounts.go",
         "offset.go",
         "overlay.go",
diff --git a/pkg/sentry/fs/mount.go b/pkg/sentry/fs/mount.go
index dd6e64b4c..5cc777bef 100644
--- a/pkg/sentry/fs/mount.go
+++ b/pkg/sentry/fs/mount.go
@@ -119,7 +119,7 @@ type MountSource struct {
 
 	// fscache keeps Dirents pinned beyond application references to them.
 	// It must be flushed before kernel.SaveTo.
-	fscache *DirentCache `state:"nosave"`
+	fscache *DirentCache
 
 	// direntRefs is the sum of references on all Dirents in this MountSource.
 	//
diff --git a/pkg/sentry/fs/mount_state.go b/pkg/sentry/fs/mount_state.go
deleted file mode 100644
index 6344d5160..000000000
--- a/pkg/sentry/fs/mount_state.go
+++ /dev/null
@@ -1,25 +0,0 @@
-// Copyright 2018 Google LLC
-//
-// 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 fs
-
-// afterLoad is invoked by stateify.
-//
-// Beyond the cache, this method's existence is required to ensure that this
-// object is not marked "complete" until all dependent objects are also marked
-// "complete". Implementations (e.g. see gofer_state.go) reach into the
-// MountSourceOperations through this object, this is necessary on restore.
-func (msrc *MountSource) afterLoad() {
-	msrc.fscache = NewDirentCache(defaultDirentCacheSize)
-}
diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go
index 1cd2653ff..a9994f23b 100644
--- a/pkg/sentry/kernel/kernel.go
+++ b/pkg/sentry/kernel/kernel.go
@@ -311,7 +311,9 @@ func (k *Kernel) SaveTo(w io.Writer) error {
 	// Clear the dirent cache before saving because Dirents must be Loaded in a
 	// particular order (parents before children), and Loading dirents from a cache
 	// breaks that order.
-	k.mounts.FlushMountSourceRefs()
+	if err := k.flushMountSourceRefs(); err != nil {
+		return err
+	}
 
 	// Ensure that all pending asynchronous work is complete:
 	//   - inode and mount release
@@ -351,39 +353,67 @@ func (k *Kernel) SaveTo(w io.Writer) error {
 	return nil
 }
 
-func (ts *TaskSet) flushWritesToFiles(ctx context.Context) error {
+// flushMountSourceRefs flushes the MountSources for all mounted filesystems
+// and open FDs.
+func (k *Kernel) flushMountSourceRefs() error {
+	// Flush all mount sources for currently mounted filesystems.
+	k.mounts.FlushMountSourceRefs()
+
+	// There may be some open FDs whose filesystems have been unmounted. We
+	// must flush those as well.
+	return k.tasks.forEachFDPaused(func(desc descriptor) error {
+		desc.file.Dirent.Inode.MountSource.FlushDirentRefs()
+		return nil
+	})
+}
+
+// forEachFDPaused applies the given function to each open file descriptor in each
+// task.
+//
+// Precondition: Must be called with the kernel paused.
+func (ts *TaskSet) forEachFDPaused(f func(descriptor) error) error {
 	ts.mu.RLock()
 	defer ts.mu.RUnlock()
 	for t := range ts.Root.tids {
 		// We can skip locking Task.mu here since the kernel is paused.
-		if fdmap := t.fds; fdmap != nil {
-			for _, desc := range fdmap.files {
-				if flags := desc.file.Flags(); !flags.Write {
-					continue
-				}
-				if sattr := desc.file.Dirent.Inode.StableAttr; !fs.IsFile(sattr) && !fs.IsDir(sattr) {
-					continue
-				}
-				// Here we need all metadata synced.
-				syncErr := desc.file.Fsync(ctx, 0, fs.FileMaxOffset, fs.SyncAll)
-				if err := fs.SaveFileFsyncError(syncErr); err != nil {
-					name, _ := desc.file.Dirent.FullName(nil /* root */)
-					// Wrap this error in ErrSaveRejection
-					// so that it will trigger a save
-					// error, rather than a panic. This
-					// also allows us to distinguish Fsync
-					// errors from state file errors in
-					// state.Save.
-					return fs.ErrSaveRejection{
-						Err: fmt.Errorf("%q was not sufficiently synced: %v", name, err),
-					}
-				}
+		if t.fds == nil {
+			continue
+		}
+		for _, desc := range t.fds.files {
+			if err := f(desc); err != nil {
+				return err
 			}
 		}
 	}
 	return nil
 }
 
+func (ts *TaskSet) flushWritesToFiles(ctx context.Context) error {
+	return ts.forEachFDPaused(func(desc descriptor) error {
+		if flags := desc.file.Flags(); !flags.Write {
+			return nil
+		}
+		if sattr := desc.file.Dirent.Inode.StableAttr; !fs.IsFile(sattr) && !fs.IsDir(sattr) {
+			return nil
+		}
+		// Here we need all metadata synced.
+		syncErr := desc.file.Fsync(ctx, 0, fs.FileMaxOffset, fs.SyncAll)
+		if err := fs.SaveFileFsyncError(syncErr); err != nil {
+			name, _ := desc.file.Dirent.FullName(nil /* root */)
+			// Wrap this error in ErrSaveRejection
+			// so that it will trigger a save
+			// error, rather than a panic. This
+			// also allows us to distinguish Fsync
+			// errors from state file errors in
+			// state.Save.
+			return fs.ErrSaveRejection{
+				Err: fmt.Errorf("%q was not sufficiently synced: %v", name, err),
+			}
+		}
+		return nil
+	})
+}
+
 // Preconditions: The kernel must be paused.
 func (k *Kernel) invalidateUnsavableMappings(ctx context.Context) error {
 	invalidated := make(map[*mm.MemoryManager]struct{})