analyzer: add pre-container monitor to capture early file access

Signed-off-by: abushwang <[email protected]>

Author: abushwang

analyzer/analyzer.go

@@ -41,6 +41,7 @@ import (
 	"github.com/containerd/log"
 	"github.com/containerd/platforms"
 	"github.com/containerd/stargz-snapshotter/analyzer/fanotify"
+	"github.com/containerd/stargz-snapshotter/analyzer/fanotify/service"
 	"github.com/containerd/stargz-snapshotter/analyzer/recorder"
 	"github.com/opencontainers/go-digest"
 	"github.com/opencontainers/image-spec/identity"
@@ -96,6 +97,26 @@ func Analyze(ctx context.Context, client *containerd.Client, ref string, opts ..
 	}
 	defer cleanup()
 
+	var preMonitorEnabled bool
+	var preMonitor *service.PreContainerMonitor
+	if aOpts.preMonitor {
+		preMonitor = service.NewPreContainerMonitor(target)
+		if err := preMonitor.Start(); err != nil {
+			return "", fmt.Errorf("failed to start pre-container monitor: %w", err)
+		}
+		preMonitorEnabled = true
+		defer func() {
+			if err := preMonitor.Close(); err != nil {
+				log.G(ctx).WithError(err).Warnf("failed to close pre-container monitor")
+			}
+		}()
+		go func() {
+			if err := preMonitor.Monitor(); err != nil {
+				log.G(ctx).WithError(err).Warnf("pre-container monitor failed")
+			}
+		}()
+	}
+
 	// Spawn a fanotifier process in a new mount namespace and setup recorder.
 	fanotifier, err := fanotify.SpawnFanotifier("/proc/self/exe")
 	if err != nil {
@@ -195,6 +216,21 @@ func Analyze(ctx context.Context, client *containerd.Client, ref string, opts ..
 		return "", err
 	}
 	defer rc.Close()
+
+	if preMonitorEnabled {
+		prePaths := preMonitor.GetPaths()
+		for _, path := range prePaths {
+			cleanPath := path
+			if strings.HasPrefix(path, target) {
+				cleanPath = strings.TrimPrefix(path, target)
+			}
+			if err := rc.Record(cleanPath); err != nil {
+				log.G(ctx).WithError(err).Debugf("failed to record pre-container path %q", cleanPath)
+			}
+		}
+		log.G(ctx).Infof("recorded %d paths from pre-container monitor", len(prePaths))
+	}
+
 	if err := fanotifier.Start(); err != nil {
 		return "", fmt.Errorf("failed to start fanotifier: %w", err)
 	}

analyzer/fanotify/service/service.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/containerd/stargz-snapshotter/analyzer/fanotify/conn"
@@ -95,3 +96,120 @@ func Serve(target string, r io.Reader, w io.Writer) error {
 
 	return nil
 }
+
+type PreContainerMonitor struct {
+	targetDir string
+	file      *os.File
+	paths     map[string]struct{}
+	mu        sync.RWMutex
+	done      chan struct{}
+	closed    bool
+	closeMu   sync.Mutex
+}
+
+func NewPreContainerMonitor(targetDir string) *PreContainerMonitor {
+	return &PreContainerMonitor{
+		targetDir: targetDir,
+		paths:     make(map[string]struct{}),
+		done:      make(chan struct{}),
+	}
+}
+
+func (m *PreContainerMonitor) Start() error {
+	fd, err := unix.FanotifyInit(unix.FAN_CLASS_NOTIF, unix.O_RDONLY)
+	if err != nil {
+		return fmt.Errorf("fanotify_init: %w", err)
+	}
+	m.file = os.NewFile(uintptr(fd), "fanotify")
+
+	if err := unix.FanotifyMark(fd,
+		unix.FAN_MARK_ADD|unix.FAN_MARK_FILESYSTEM,
+		unix.FAN_ACCESS|unix.FAN_OPEN,
+		unix.AT_FDCWD,
+		m.targetDir,
+	); err != nil {
+		m.file.Close()
+		return fmt.Errorf("fanotify_mark: %w", err)
+	}
+
+	return nil
+}
+
+func (m *PreContainerMonitor) Monitor() error {
+	if m.file == nil {
+		return fmt.Errorf("monitor not started")
+	}
+
+	nr := bufio.NewReader(m.file)
+
+	for {
+		event := &unix.FanotifyEventMetadata{}
+		if err := binary.Read(nr, binary.LittleEndian, event); err != nil {
+			if err == io.EOF {
+				return nil
+			}
+			select {
+			case <-m.done:
+				return nil
+			default:
+			}
+			return fmt.Errorf("read fanotify fd: %w", err)
+		}
+
+		if event.Vers != unix.FANOTIFY_METADATA_VERSION {
+			return fmt.Errorf("fanotify version mismatch %d(got) != %d(want)",
+				event.Vers, unix.FANOTIFY_METADATA_VERSION)
+		}
+
+		if event.Fd < 0 {
+			fmt.Fprintf(os.Stderr, "Warn: queue overflow")
+			continue
+		}
+
+		path, err := os.Readlink(fmt.Sprintf("/proc/self/fd/%d", event.Fd))
+		if err != nil {
+			unix.Close(int(event.Fd))
+			continue
+		}
+
+		m.mu.Lock()
+		m.paths[path] = struct{}{}
+		m.mu.Unlock()
+
+		if err := unix.Close(int(event.Fd)); err != nil {
+			fmt.Fprintf(os.Stderr, "Warn: failed to close fd %d: %v", event.Fd, err)
+		}
+	}
+}
+
+func (m *PreContainerMonitor) GetPaths() []string {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	paths := make([]string, 0, len(m.paths))
+	for path := range m.paths {
+		paths = append(paths, path)
+	}
+	return paths
+}
+
+func (m *PreContainerMonitor) Close() error {
+	m.closeMu.Lock()
+	defer m.closeMu.Unlock()
+
+	if m.closed {
+		return nil
+	}
+	m.closed = true
+
+	close(m.done)
+
+	if m.file != nil {
+		if err := m.file.Close(); err != nil {
+			return fmt.Errorf("failed to close fanotify fd: %w", err)
+		}
+		m.file = nil
+	}
+
+	return nil
+}

analyzer/option.go

@@ -31,6 +31,7 @@ type analyzerOpts struct {
 	terminal     bool
 	stdin        bool
 	waitLineOut  string
+	preMonitor   bool
 }
 
 // Option is runtime configuration of analyzer container
@@ -88,3 +89,10 @@ func WithWaitLineOut(s string) Option {
 		opts.waitLineOut = s
 	}
 }
+
+// WithPreMonitor enables pre-container monitoring phase.
+func WithPreMonitor() Option {
+	return func(opts *analyzerOpts) {
+		opts.preMonitor = true
+	}
+}

cmd/ctr-remote/commands/optimize.go

@@ -441,6 +441,9 @@ func analyze(ctx context.Context, clicontext *cli.Context, client *containerd.Cl
 
 	// Analyze layers and get prioritized files
 	aOpts := []analyzer.Option{analyzer.WithSpecOpts(getSpecOpts(clicontext))}
+	if clicontext.IsSet("gpus") {
+		aOpts = append(aOpts, analyzer.WithPreMonitor())
+	}
 	if clicontext.Bool("wait-on-signal") && clicontext.Bool("terminal") {
 		return "", nil, nil, fmt.Errorf("wait-on-signal can't be used with terminal flag")
 	}