From 218f52a9f5c70540bc04c22364906654e20c6923 Mon Sep 17 00:00:00 2001
From: Ayush Ranjan <ayushranjan@google.com>
Date: Tue, 27 Aug 2024 14:00:04 -0700
Subject: [PATCH] Parallelize MemoryFile save and kernel save.

This compliments 39730b714ca4 ("Load pgalloc.MemoryFile and kernel parallely
with compression=none mode.")

This is a performance optimization. Kernel and MemoryFile are saved
independently. The save can be done in parallel when using --compression=none
because the kernel and MemoryFile are being saved in different files.

PiperOrigin-RevId: 668128205
---
 pkg/sentry/kernel/kernel.go | 34 +++++++++++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go
index 14ea29934..e487d8749 100644
--- a/pkg/sentry/kernel/kernel.go
+++ b/pkg/sentry/kernel/kernel.go
@@ -659,6 +659,23 @@ func (k *Kernel) SaveTo(ctx context.Context, w, pagesMetadata io.Writer, pagesFi
 		mf.MarkSavable()
 	}
 
+	var (
+		mfSaveWg  sync.WaitGroup
+		mfSaveErr error
+	)
+	parallelMfSave := pagesMetadata != nil && pagesFile != nil
+	if parallelMfSave {
+		// Parallelize MemoryFile save and kernel save. Both are independent.
+		mfSaveWg.Add(1)
+		go func() {
+			defer mfSaveWg.Done()
+			mfSaveErr = k.saveMemoryFiles(ctx, w, pagesMetadata, pagesFile, mfsToSave, mfOpts)
+		}()
+		// Defer a Wait() so we wait for k.saveMemoryFiles() to complete even if we
+		// error out without reaching the other Wait() below.
+		defer mfSaveWg.Wait()
+	}
+
 	// Save the CPUID FeatureSet before the rest of the kernel so we can
 	// verify its compatibility on restore before attempting to restore the
 	// entire kernel, which may fail on an incompatible machine.
@@ -690,6 +707,20 @@ func (k *Kernel) SaveTo(ctx context.Context, w, pagesMetadata io.Writer, pagesFi
 	log.Infof("Kernel save stats: %s", stats.String())
 	log.Infof("Kernel save took [%s].", time.Since(kernelStart))
 
+	if parallelMfSave {
+		mfSaveWg.Wait()
+	} else {
+		mfSaveErr = k.saveMemoryFiles(ctx, w, pagesMetadata, pagesFile, mfsToSave, mfOpts)
+	}
+	if mfSaveErr != nil {
+		return mfSaveErr
+	}
+
+	log.Infof("Overall save took [%s].", time.Since(saveStart))
+	return nil
+}
+
+func (k *Kernel) saveMemoryFiles(ctx context.Context, w, pagesMetadata io.Writer, pagesFile *fd.FD, mfsToSave map[string]*pgalloc.MemoryFile, mfOpts pgalloc.SaveOpts) error {
 	// Save the memory files' state.
 	memoryStart := time.Now()
 	pmw := w
@@ -707,9 +738,6 @@ func (k *Kernel) SaveTo(ctx context.Context, w, pagesMetadata io.Writer, pagesFi
 		return err
 	}
 	log.Infof("Memory files save took [%s].", time.Since(memoryStart))
-
-	log.Infof("Overall save took [%s].", time.Since(saveStart))
-
 	return nil
 }