gateway: cache local mounts from fs operations

Previously, each call to `ReadFile`/`ReadDir`/`StatFile` would trigger
an entire `Mount`+`Unmount` operation. This could result in terrible
performance for repeated filesystem accesses - potentially worse that
one `Mount`+`Unmount` per file, if the file is read in ranges.

To resolve this, we now directly store the `localMounter` in the
gateway, and rely on the fact that repeated calls to it return the same
value. These mounts will only be removed when the gateway itself is
released.

Signed-off-by: Justin Chadwell <me@jedevc.com>

cache/util/fsutil.go

@@ -7,7 +7,6 @@ import (
 	"path/filepath"
 
 	"github.com/containerd/continuity/fs"
-	"github.com/moby/buildkit/snapshot"
 	"github.com/pkg/errors"
 	"github.com/tonistiigi/fsutil"
 	fstypes "github.com/tonistiigi/fsutil/types"
@@ -23,64 +22,34 @@ type FileRange struct {
 	Length int
 }
 
-func withMount(mount snapshot.Mountable, cb func(string) error) error {
-	lm := snapshot.LocalMounter(mount)
-
-	root, err := lm.Mount()
+func ReadFile(ctx context.Context, root string, req ReadRequest) ([]byte, error) {
+	fp, err := fs.RootPath(root, req.Filename)
 	if err != nil {
-		return err
+		return nil, errors.WithStack(err)
 	}
 
-	defer func() {
-		if lm != nil {
-			lm.Unmount()
+	f, err := os.Open(fp)
+	if err != nil {
+		// The filename here is internal to the mount, so we can restore
+		// the request base path for error reporting.
+		// See os.DirFS.Open for details.
+		pe := &os.PathError{}
+		if errors.As(err, &pe) {
+			pe.Path = req.Filename
 		}
-	}()
-
-	if err := cb(root); err != nil {
-		return err
+		return nil, errors.WithStack(err)
 	}
+	defer f.Close()
 
-	if err := lm.Unmount(); err != nil {
-		return err
+	var rdr io.Reader = f
+	if req.Range != nil {
+		rdr = io.NewSectionReader(f, int64(req.Range.Offset), int64(req.Range.Length))
 	}
-	lm = nil
-	return nil
-}
-
-func ReadFile(ctx context.Context, mount snapshot.Mountable, req ReadRequest) ([]byte, error) {
-	var dt []byte
-
-	err := withMount(mount, func(root string) error {
-		fp, err := fs.RootPath(root, req.Filename)
-		if err != nil {
-			return errors.WithStack(err)
-		}
-
-		f, err := os.Open(fp)
-		if err != nil {
-			// The filename here is internal to the mount, so we can restore
-			// the request base path for error reporting.
-			// See os.DirFS.Open for details.
-			pe := &os.PathError{}
-			if errors.As(err, &pe) {
-				pe.Path = req.Filename
-			}
-			return errors.WithStack(err)
-		}
-		defer f.Close()
-
-		var rdr io.Reader = f
-		if req.Range != nil {
-			rdr = io.NewSectionReader(f, int64(req.Range.Offset), int64(req.Range.Length))
-		}
-		dt, err = io.ReadAll(rdr)
-		if err != nil {
-			return errors.WithStack(err)
-		}
-		return nil
-	})
-	return dt, err
+	dt, err := io.ReadAll(rdr)
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	return dt, nil
 }
 
 type ReadDirRequest struct {
@@ -88,7 +57,7 @@ type ReadDirRequest struct {
 	IncludePattern string
 }
 
-func ReadDir(ctx context.Context, mount snapshot.Mountable, req ReadDirRequest) ([]*fstypes.Stat, error) {
+func ReadDir(ctx context.Context, root string, req ReadDirRequest) ([]*fstypes.Stat, error) {
 	var (
 		rd []*fstypes.Stat
 		fo fsutil.FilterOpt
@@ -96,48 +65,46 @@ func ReadDir(ctx context.Context, mount snapshot.Mountable, req ReadDirRequest)
 	if req.IncludePattern != "" {
 		fo.IncludePatterns = append(fo.IncludePatterns, req.IncludePattern)
 	}
-	err := withMount(mount, func(root string) error {
-		fp, err := fs.RootPath(root, req.Path)
+	fp, err := fs.RootPath(root, req.Path)
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	err = fsutil.Walk(ctx, fp, &fo, func(path string, info os.FileInfo, err error) error {
 		if err != nil {
-			return errors.WithStack(err)
+			return errors.Wrapf(err, "walking %q", root)
 		}
-		return fsutil.Walk(ctx, fp, &fo, func(path string, info os.FileInfo, err error) error {
-			if err != nil {
-				return errors.Wrapf(err, "walking %q", root)
-			}
-			stat, ok := info.Sys().(*fstypes.Stat)
-			if !ok {
-				// This "can't happen(tm)".
-				return errors.Errorf("expected a *fsutil.Stat but got %T", info.Sys())
-			}
-			rd = append(rd, stat)
-
-			if info.IsDir() {
-				return filepath.SkipDir
-			}
-			return nil
-		})
-	})
-	return rd, err
-}
-
-func StatFile(ctx context.Context, mount snapshot.Mountable, path string) (*fstypes.Stat, error) {
-	var st *fstypes.Stat
-	err := withMount(mount, func(root string) error {
-		fp, err := fs.RootPath(root, path)
-		if err != nil {
-			return errors.WithStack(err)
+		stat, ok := info.Sys().(*fstypes.Stat)
+		if !ok {
+			// This "can't happen(tm)".
+			return errors.Errorf("expected a *fsutil.Stat but got %T", info.Sys())
 		}
-		if st, err = fsutil.Stat(fp); err != nil {
-			// The filename here is internal to the mount, so we can restore
-			// the request base path for error reporting.
-			// See os.DirFS.Open for details.
-			replaceErrorPath(err, path)
-			return errors.WithStack(err)
+		rd = append(rd, stat)
+
+		if info.IsDir() {
+			return filepath.SkipDir
 		}
 		return nil
 	})
-	return st, err
+	if err != nil {
+		return nil, err
+	}
+	return rd, nil
+}
+
+func StatFile(ctx context.Context, root string, path string) (*fstypes.Stat, error) {
+	fp, err := fs.RootPath(root, path)
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	st, err := fsutil.Stat(fp)
+	if err != nil {
+		// The filename here is internal to the mount, so we can restore
+		// the request base path for error reporting.
+		// See os.DirFS.Open for details.
+		replaceErrorPath(err, path)
+		return nil, errors.WithStack(err)
+	}
+	return st, nil
 }
 
 // replaceErrorPath will override the path in an os.PathError in the error chain.

frontend/gateway/forwarder/forward.go

@@ -358,7 +358,14 @@ func (r *ref) ReadFile(ctx context.Context, req client.ReadRequest) ([]byte, err
 			Length: r.Length,
 		}
 	}
-	return cacheutil.ReadFile(ctx, m, newReq)
+
+	var dt []byte
+	err = withMount(m, func(root string) error {
+		var err error
+		dt, err = cacheutil.ReadFile(ctx, root, newReq)
+		return err
+	})
+	return dt, err
 }
 
 func (r *ref) ReadDir(ctx context.Context, req client.ReadDirRequest) ([]*fstypes.Stat, error) {
@@ -370,7 +377,14 @@ func (r *ref) ReadDir(ctx context.Context, req client.ReadDirRequest) ([]*fstype
 		Path:           req.Path,
 		IncludePattern: req.IncludePattern,
 	}
-	return cacheutil.ReadDir(ctx, m, newReq)
+
+	var rd []*fstypes.Stat
+	err = withMount(m, func(root string) error {
+		var err error
+		rd, err = cacheutil.ReadDir(ctx, root, newReq)
+		return err
+	})
+	return rd, err
 }
 
 func (r *ref) StatFile(ctx context.Context, req client.StatRequest) (*fstypes.Stat, error) {
@@ -378,7 +392,14 @@ func (r *ref) StatFile(ctx context.Context, req client.StatRequest) (*fstypes.St
 	if err != nil {
 		return nil, err
 	}
-	return cacheutil.StatFile(ctx, m, req.Path)
+
+	var st *fstypes.Stat
+	err = withMount(m, func(root string) error {
+		var err error
+		st, err = cacheutil.StatFile(ctx, root, req.Path)
+		return err
+	})
+	return st, err
 }
 
 func (r *ref) getMountable(ctx context.Context) (snapshot.Mountable, error) {
@@ -392,3 +413,28 @@ func (r *ref) getMountable(ctx context.Context) (snapshot.Mountable, error) {
 	}
 	return ref.ImmutableRef.Mount(ctx, true, r.session)
 }
+
+func withMount(mount snapshot.Mountable, cb func(string) error) error {
+	lm := snapshot.LocalMounter(mount)
+
+	root, err := lm.Mount()
+	if err != nil {
+		return err
+	}
+
+	defer func() {
+		if lm != nil {
+			lm.Unmount()
+		}
+	}()
+
+	if err := cb(root); err != nil {
+		return err
+	}
+
+	if err := lm.Unmount(); err != nil {
+		return err
+	}
+	lm = nil
+	return nil
+}

frontend/gateway/gateway.go

@@ -370,6 +370,10 @@ func (lbf *llbBridgeForwarder) Discard() {
 		})
 	}
 
+	for _, mount := range lbf.mounts {
+		mount.Unmount()
+	}
+
 	for id, workerRef := range lbf.workerRefByID {
 		workerRef.Release(context.TODO())
 		delete(lbf.workerRefByID, id)
@@ -440,6 +444,7 @@ func newBridgeForwarder(ctx context.Context, llbBridge frontend.FrontendLLBBridg
 		sid:           sid,
 		sm:            sm,
 		ctrs:          map[string]gwclient.Container{},
+		mounts:        map[string]snapshot.Mounter{},
 		executor:      exec,
 	}
 	return lbf
@@ -553,8 +558,10 @@ type llbBridgeForwarder struct {
 	sm                *session.Manager
 	executor          executor.Executor
 	*pipe
-	ctrs   map[string]gwclient.Container
-	ctrsMu sync.Mutex
+	ctrs     map[string]gwclient.Container
+	ctrsMu   sync.Mutex
+	mounts   map[string]snapshot.Mounter
+	mountsMu sync.Mutex
 }
 
 func (lbf *llbBridgeForwarder) ResolveSourceMeta(ctx context.Context, req *pb.ResolveSourceMetaRequest) (*pb.ResolveSourceMetaResponse, error) {
@@ -869,6 +876,37 @@ func (lbf *llbBridgeForwarder) getImmutableRef(ctx context.Context, id string) (
 	return workerRef.ImmutableRef, nil
 }
 
+func (lbf *llbBridgeForwarder) getMounter(ctx context.Context, id string, ref cache.ImmutableRef) (snapshot.Mounter, error) {
+	lbf.mountsMu.Lock()
+	defer lbf.mountsMu.Unlock()
+
+	mounter, ok := lbf.mounts[id]
+	if ok {
+		return mounter, nil
+	}
+	var mountable snapshot.Mountable
+	if ref != nil {
+		var err error
+		mountable, err = ref.Mount(ctx, true, session.NewGroup(lbf.sid))
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	mounter = snapshot.LocalMounter(mountable)
+	lbf.mounts[id] = mounter
+	return mounter, nil
+}
+
+func (lbf *llbBridgeForwarder) getMount(ctx context.Context, id string, ref cache.ImmutableRef) (string, error) {
+	mounter, err := lbf.getMounter(ctx, id, ref)
+	if err != nil {
+		return "", err
+	}
+	// corresponding Unmount call is made in Discard()
+	return mounter.Mount()
+}
+
 func (lbf *llbBridgeForwarder) ReadFile(ctx context.Context, req *pb.ReadFileRequest) (*pb.ReadFileResponse, error) {
 	ctx = tracing.ContextWithSpanFromContext(ctx, lbf.callCtx)
 
@@ -876,6 +914,10 @@ func (lbf *llbBridgeForwarder) ReadFile(ctx context.Context, req *pb.ReadFileReq
 	if err != nil {
 		return nil, err
 	}
+	root, err := lbf.getMount(ctx, req.Ref, ref)
+	if err != nil {
+		return nil, err
+	}
 
 	newReq := cacheutil.ReadRequest{
 		Filename: req.FilePath,
@@ -887,15 +929,7 @@ func (lbf *llbBridgeForwarder) ReadFile(ctx context.Context, req *pb.ReadFileReq
 		}
 	}
 
-	var m snapshot.Mountable
-	if ref != nil {
-		m, err = ref.Mount(ctx, true, session.NewGroup(lbf.sid))
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	dt, err := cacheutil.ReadFile(ctx, m, newReq)
+	dt, err := cacheutil.ReadFile(ctx, root, newReq)
 	if err != nil {
 		return nil, lbf.wrapSolveError(err)
 	}
@@ -910,19 +944,17 @@ func (lbf *llbBridgeForwarder) ReadDir(ctx context.Context, req *pb.ReadDirReque
 	if err != nil {
 		return nil, err
 	}
+	root, err := lbf.getMount(ctx, req.Ref, ref)
+	if err != nil {
+		return nil, err
+	}
 
 	newReq := cacheutil.ReadDirRequest{
 		Path:           req.DirPath,
 		IncludePattern: req.IncludePattern,
 	}
-	var m snapshot.Mountable
-	if ref != nil {
-		m, err = ref.Mount(ctx, true, session.NewGroup(lbf.sid))
-		if err != nil {
-			return nil, err
-		}
-	}
-	entries, err := cacheutil.ReadDir(ctx, m, newReq)
+
+	entries, err := cacheutil.ReadDir(ctx, root, newReq)
 	if err != nil {
 		return nil, lbf.wrapSolveError(err)
 	}
@@ -937,14 +969,12 @@ func (lbf *llbBridgeForwarder) StatFile(ctx context.Context, req *pb.StatFileReq
 	if err != nil {
 		return nil, err
 	}
-	var m snapshot.Mountable
-	if ref != nil {
-		m, err = ref.Mount(ctx, true, session.NewGroup(lbf.sid))
-		if err != nil {
-			return nil, err
-		}
+	root, err := lbf.getMount(ctx, req.Ref, ref)
+	if err != nil {
+		return nil, err
 	}
-	st, err := cacheutil.StatFile(ctx, m, req.Path)
+
+	st, err := cacheutil.StatFile(ctx, root, req.Path)
 	if err != nil {
 		return nil, err
 	}