Add support for conditional image tarball imports

Normally K3s will import all tarballs in the image dir on startup, and
re-import any tarballs that change while it is running.

This change allows users to opt into only importing tarballs that have
changed since they were last imported, even across restarts.

This behavior is opted into by touching a `.cache.json` file in the
images dir. This file is used to track the size and mtime of the image
files when they are imported.

Signed-off-by: Brad Davidson <brad.davidson@rancher.com>
(cherry picked from commit 67291090ca)
Signed-off-by: Brad Davidson <brad.davidson@rancher.com>

pkg/agent/containerd/containerd.go

@@ -10,7 +10,6 @@ import (
 	"os/exec"
 	"path/filepath"
 	"strings"
-	"time"
 
 	"github.com/containerd/containerd"
 	"github.com/containerd/containerd/errdefs"
@@ -131,7 +130,6 @@ func PreloadImages(ctx context.Context, cfg *config.Node) error {
 		return err
 	}
 	defer criConn.Close()
-	imageClient := runtimeapi.NewImageServiceClient(criConn)
 
 	// Ensure that our images are imported into the correct namespace
 	ctx = namespaces.WithNamespace(ctx, constants.K8sContainerdNamespace)
@@ -146,44 +144,7 @@ func PreloadImages(ctx context.Context, cfg *config.Node) error {
 		return pkgerrors.WithMessage(err, "failed to clear pinned labels")
 	}
 
-	go watchImages(ctx, cfg)
-
-	// After setting the watcher, connections and everything, k3s will see if the images folder is already created
-	// if the folder its already created, it will load the images
-	fileInfo, err := os.Stat(cfg.Images)
-	if os.IsNotExist(err) {
-		return nil
-	} else if err != nil {
-		logrus.Errorf("Unable to find images in %s: %v", cfg.Images, err)
-		return nil
-	}
-
-	if !fileInfo.IsDir() {
-		return nil
-	}
-
-	fileInfos, err := os.ReadDir(cfg.Images)
-	if err != nil {
-		logrus.Errorf("Unable to read images in %s: %v", cfg.Images, err)
-		return nil
-	}
-
-	for _, fileInfo := range fileInfos {
-		if fileInfo.IsDir() {
-			continue
-		}
-
-		start := time.Now()
-		filePath := filepath.Join(cfg.Images, fileInfo.Name())
-
-		if err := preloadFile(ctx, cfg, client, imageClient, filePath); err != nil {
-			logrus.Errorf("Error encountered while importing %s: %v", filePath, err)
-			continue
-		}
-		logrus.Infof("Imported images from %s in %s", filePath, time.Since(start))
-	}
-
-	return nil
+	return importAndWatchImages(ctx, cfg)
 }
 
 // preloadFile handles loading images from a single tarball or pre-pull image list.

pkg/agent/containerd/watcher.go

@@ -2,8 +2,8 @@ package containerd
 
 import (
 	"context"
-	"fmt"
-	"io/fs"
+	"encoding/json"
+	"io"
 	"os"
 	"path/filepath"
 	"strings"
@@ -15,29 +15,44 @@ import (
 	"github.com/k3s-io/k3s/pkg/daemons/config"
 	pkgerrors "github.com/pkg/errors"
 	"github.com/rancher/wharfie/pkg/tarfile"
-	"github.com/rancher/wrangler/v3/pkg/merr"
 	"github.com/sirupsen/logrus"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/util/workqueue"
 	runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1"
 )
 
-type Watcher struct {
+type fileInfo struct {
+	Size    int64       `json:"size"`
+	ModTime metav1.Time `json:"modTime"`
+	seen    bool        // field is not serialized, and can be used to track if a file has been seen since the last restart
+}
+
+type watchqueue struct {
+	cfg        *config.Node
 	watcher    *fsnotify.Watcher
-	filesCache map[string]fs.FileInfo
+	filesCache map[string]*fileInfo
 	workqueue  workqueue.DelayingInterface
 }
 
-func CreateWatcher() (*Watcher, error) {
+func createWatcher(path string) (*fsnotify.Watcher, error) {
 	watcher, err := fsnotify.NewWatcher()
 	if err != nil {
 		return nil, err
 	}
 
-	return &Watcher{
-		watcher:    watcher,
-		filesCache: make(map[string]fs.FileInfo),
-		workqueue:  workqueue.NewDelayingQueue(),
-	}, nil
+	if err := watcher.Add(path); err != nil {
+		return nil, err
+	}
+
+	return watcher, nil
+}
+
+func mustCreateWatcher(path string) *fsnotify.Watcher {
+	watcher, err := createWatcher(path)
+	if err != nil {
+		panic("Failed to create image import watcher:" + err.Error())
+	}
+	return watcher
 }
 
 func isFileSupported(path string) bool {
@@ -50,55 +65,11 @@ func isFileSupported(path string) bool {
 	return false
 }
 
-func (w *Watcher) HandleWatch(path string) error {
-	if err := w.watcher.Add(path); err != nil {
-		return pkgerrors.WithMessage(err, fmt.Sprintf("failed to watch from %s directory: %v", path, err))
-	}
-
-	return nil
-}
-
-// Populate the state of the files in the directory
-// for the watcher to have infos about the file changing
-// this function need to break
-func (w *Watcher) Populate(path string) error {
-	var errs []error
-
-	fileInfos, err := os.ReadDir(path)
-	if err != nil {
-		logrus.Errorf("Unable to read files in %s: %v", path, err)
-		return err
-	}
-
-	for _, dirEntry := range fileInfos {
-		if dirEntry.IsDir() {
-			continue
-		}
-
-		// get the file info to add to the state map
-		fileInfo, err := dirEntry.Info()
-		if err != nil {
-			logrus.Errorf("Failed while getting the info from file: %v", err)
-			errs = append(errs, err)
-			continue
-		}
-
-		if isFileSupported(dirEntry.Name()) {
-			// insert the file into the state map that will have the state from the file
-			w.filesCache[filepath.Join(path, dirEntry.Name())] = fileInfo
-		}
-	}
-
-	return merr.NewErrors(errs...)
-}
-
-func (w *Watcher) ClearMap() {
-	w.filesCache = make(map[string]fs.FileInfo)
-}
-
-func (w *Watcher) runWorkerForImages(ctx context.Context, cfg *config.Node) {
+// runWorkerForImages connects to containerd and calls processNextEventForImages to process items from the workqueue.
+// This blocks until the workqueue is shut down.
+func (w *watchqueue) runWorkerForImages(ctx context.Context) {
 	// create the connections to not create every time when processing a event
-	client, err := Client(cfg.Containerd.Address)
+	client, err := Client(w.cfg.Containerd.Address)
 	if err != nil {
 		logrus.Errorf("Failed to create containerd client: %v", err)
 		w.watcher.Close()
@@ -107,7 +78,7 @@ func (w *Watcher) runWorkerForImages(ctx context.Context, cfg *config.Node) {
 
 	defer client.Close()
 
-	criConn, err := cri.Connection(ctx, cfg.Containerd.Address)
+	criConn, err := cri.Connection(ctx, w.cfg.Containerd.Address)
 	if err != nil {
 		logrus.Errorf("Failed to create CRI connection: %v", err)
 		w.watcher.Close()
@@ -118,25 +89,28 @@ func (w *Watcher) runWorkerForImages(ctx context.Context, cfg *config.Node) {
 
 	imageClient := runtimeapi.NewImageServiceClient(criConn)
 
-	for w.processNextEventForImages(ctx, cfg, client, imageClient) {
+	for w.processNextEventForImages(ctx, client, imageClient) {
 	}
 }
 
-func (w *Watcher) processNextEventForImages(ctx context.Context, cfg *config.Node, client *containerd.Client, imageClient runtimeapi.ImageServiceClient) bool {
-	key, shutdown := w.workqueue.Get()
+// processNextEventForImages retrieves a single event from the workqueue and processes it.
+// It returns a boolean that is true if the workqueue is still open and this function should be called again.
+func (w *watchqueue) processNextEventForImages(ctx context.Context, client *containerd.Client, imageClient runtimeapi.ImageServiceClient) bool {
+	obj, shutdown := w.workqueue.Get()
 
 	if shutdown {
 		return false
 	}
 
-	if err := w.processImageEvent(ctx, key, cfg, client, imageClient); err != nil {
+	if err := w.processImageEvent(ctx, obj, client, imageClient); err != nil {
 		logrus.Errorf("Failed to process image event: %v", err)
 	}
 
 	return true
 }
 
-func (w *Watcher) processImageEvent(ctx context.Context, obj interface{}, cfg *config.Node, client *containerd.Client, imageClient runtimeapi.ImageServiceClient) error {
+// processImageEvent processes a single item from the workqueue.
+func (w *watchqueue) processImageEvent(ctx context.Context, obj interface{}, client *containerd.Client, imageClient runtimeapi.ImageServiceClient) error {
 	var (
 		key string
 		ok  bool
@@ -149,13 +123,19 @@ func (w *Watcher) processImageEvent(ctx context.Context, obj interface{}, cfg *c
 		return nil
 	}
 
-	defer w.workqueue.Done(key)
+	// Watch is rooted at the parent dir of the images dir, but we only need to handle things within the images dir
+	if !strings.HasPrefix(key, w.cfg.Images) {
+		return nil
+	}
 
 	file, err := os.Stat(key)
+
 	// if the file does not exists, we assume that the event was RENAMED or REMOVED
 	if os.IsNotExist(err) {
-		if key == cfg.Images {
-			w.ClearMap()
+		// if the whole images dir was removed, reset the fileinfo cache
+		if key == w.cfg.Images {
+			w.filesCache = make(map[string]*fileInfo)
+			defer w.syncCache()
 			return nil
 		}
 
@@ -164,49 +144,26 @@ func (w *Watcher) processImageEvent(ctx context.Context, obj interface{}, cfg *c
 		}
 
 		delete(w.filesCache, key)
-		logrus.Debugf("File removed from the image watcher controller: %s", key)
+		defer w.syncCache()
 		return nil
 	} else if err != nil {
-		logrus.Errorf("Failed to get file %s info for image event: %v", key, err)
-		return err
+		return pkgerrors.Wrapf(err, "failed to get fileinfo for image event %s", key)
 	}
 
 	if file.IsDir() {
-		// only add the image watcher, populate and search for images when it is the images folder
-		if key == cfg.Images {
-			if err := w.HandleWatch(cfg.Images); err != nil {
-				logrus.Errorf("Failed to watch %s: %v", cfg.Images, err)
-				return err
-			}
-
-			if err := w.Populate(cfg.Images); err != nil {
-				logrus.Errorf("Failed to populate %s files: %v", cfg.Images, err)
-				return err
-			}
-
-			// Read the directory to see if the created folder has files inside
-			fileInfos, err := os.ReadDir(cfg.Images)
-			if err != nil {
-				logrus.Errorf("Unable to read images in %s: %v", cfg.Images, err)
-				return err
-			}
-
-			for _, fileInfo := range fileInfos {
-				if fileInfo.IsDir() {
-					continue
-				}
-
-				start := time.Now()
-				filePath := filepath.Join(cfg.Images, fileInfo.Name())
-
-				if err := preloadFile(ctx, cfg, client, imageClient, filePath); err != nil {
-					logrus.Errorf("Error encountered while importing %s: %v", filePath, err)
-					continue
-				}
-				logrus.Infof("Imported images from %s in %s", filePath, time.Since(start))
-			}
+		// Add to watch and list+enqueue directory contents, as notify is not recursive
+		if err := w.watcher.Add(key); err != nil {
+			return pkgerrors.Wrapf(err, "failed to add watch of %s", key)
 		}
 
+		fileInfos, err := os.ReadDir(key)
+		if err != nil {
+			return pkgerrors.Wrapf(err, "unable to list contents of %s", key)
+		}
+
+		for _, fileInfo := range fileInfos {
+			w.workqueue.Add(filepath.Join(key, fileInfo.Name()))
+		}
 		return nil
 	}
 
@@ -214,89 +171,173 @@ func (w *Watcher) processImageEvent(ctx context.Context, obj interface{}, cfg *c
 		return nil
 	}
 
-	lastStateFile := w.filesCache[key]
-	w.filesCache[key] = file
-	if lastStateFile == nil || (file.Size() != lastStateFile.Size()) && file.ModTime().After(lastStateFile.ModTime()) {
-		logrus.Debugf("File met the requirements for import to containerd image store: %s", key)
+	if lastFileState := w.filesCache[key]; lastFileState == nil || (file.Size() != lastFileState.Size && file.ModTime().After(lastFileState.ModTime.Time)) {
 		start := time.Now()
-		if err := preloadFile(ctx, cfg, client, imageClient, key); err != nil {
-			logrus.Errorf("Failed to import %s: %v", key, err)
-			return err
+		if err := preloadFile(ctx, w.cfg, client, imageClient, key); err != nil {
+			return pkgerrors.Wrapf(err, "failed to import %s", key)
 		}
 		logrus.Infof("Imported images from %s in %s", key, time.Since(start))
+		w.filesCache[key] = &fileInfo{Size: file.Size(), ModTime: metav1.NewTime(file.ModTime()), seen: true}
+		defer w.syncCache()
+	} else if lastFileState != nil && !lastFileState.seen {
+		lastFileState.seen = true
+		// no need to sync as the field is not serialized
 	}
 
 	return nil
 }
 
-func watchImages(ctx context.Context, cfg *config.Node) {
-	w, err := CreateWatcher()
+// pruneCache removes entries for all files that have not been seen since the last restart,
+// and syncs the cache to disk. This is done to ensure that the cache file does not grow without
+// bounds by continuing to track files that do not exist.
+func (w *watchqueue) pruneCache() {
+	for path, fileState := range w.filesCache {
+		if !fileState.seen {
+			delete(w.filesCache, path)
+		}
+	}
+	w.syncCache()
+}
+
+// syncCache writes the fileinfo cache to disk.
+// if the cache file does not exist, this is a no-op. The file must be manually
+// created by the user in order for the cache to be persisted across restarts.
+func (w *watchqueue) syncCache() {
+	filePath := filepath.Join(w.cfg.Images, ".cache.json")
+	f, err := os.OpenFile(filePath, os.O_WRONLY|os.O_TRUNC, 0644)
 	if err != nil {
-		logrus.Errorf("Failed to create image watcher: %v", err)
+		if !os.IsNotExist(err) {
+			logrus.Errorf("Failed to truncate image import fileinfo cache: %v", err)
+		}
+		return
+	}
+	defer f.Close()
+
+	b, err := json.Marshal(&w.filesCache)
+	if err != nil {
+		logrus.Errorf("Failed to marshal image import fileinfo cache: %v", err)
 		return
 	}
 
-	logrus.Debugf("Image Watcher created")
-	defer w.watcher.Close()
-
-	if err := w.HandleWatch(filepath.Dir(cfg.Images)); err != nil {
-		logrus.Errorf("Failed to watch %s: %v", filepath.Dir(cfg.Images), err)
-		return
-	}
-
-	_, err = os.Stat(cfg.Images)
-	if err == nil {
-		if err := w.HandleWatch(cfg.Images); err != nil {
-			logrus.Errorf("Failed to watch %s: %v", cfg.Images, err)
-			return
-		}
-
-		if err := w.Populate(cfg.Images); err != nil {
-			logrus.Errorf("Failed to populate %s files: %v", cfg.Images, err)
-			return
-		}
-	} else if os.IsNotExist(err) {
-		logrus.Debugf("Image dir %s does not exist", cfg.Images)
-	} else {
-		logrus.Debugf("Failed to stat image dir %s: %v", cfg.Images, err)
-	}
-
-	go w.runWorkerForImages(ctx, cfg)
-
-	for {
-		select {
-		case event, ok := <-w.watcher.Events:
-			if !ok {
-				logrus.Info("Image watcher channel closed, shutting down workqueue and retrying in 5 seconds")
-				w.workqueue.ShutDown()
-				select {
-				case <-time.After(time.Second * 5):
-					go watchImages(ctx, cfg)
-					return
-				case <-ctx.Done():
-					return
-				}
-
-			}
-
-			// this part is to specify to only get events that were from /agent/images
-			if strings.Contains(event.Name, "/agent/images") {
-				w.workqueue.AddAfter(event.Name, 2*time.Second)
-			}
-
-		case err, ok := <-w.watcher.Errors:
-			if !ok {
-				logrus.Info("Image watcher channel closed, shutting down workqueue and retrying in 5 seconds")
-				w.workqueue.ShutDown()
-				select {
-				case <-time.After(time.Second * 5):
-					go watchImages(ctx, cfg)
-					return
-				case <-ctx.Done():
-					return
-				}
-			}
-			logrus.Errorf("Image watcher received an error: %v", err)
-		}
+	if _, err := f.Write(b); err != nil {
+		logrus.Errorf("Failed to write image import fileinfo cache: %v", err)
 	}
 }
+
+// loadCaache reads the fileinfo cache from disk.
+// It is not an error if this file exists or is empty.
+func (w *watchqueue) loadCache() {
+	filePath := filepath.Join(w.cfg.Images, ".cache.json")
+	f, err := os.OpenFile(filePath, os.O_RDONLY, 0664)
+	if err != nil {
+		if !os.IsNotExist(err) {
+			logrus.Errorf("Failed to open image import fileinfo cache: %v", err)
+		}
+		return
+	}
+	defer f.Close()
+
+	b, err := io.ReadAll(f)
+	if err != nil {
+		logrus.Errorf("Failed to read image import fileinfo cache: %v", err)
+		return
+	}
+
+	// 0 byte file is fine, but don't try to load it - allows users to simply
+	// touch the file to enable it for use.
+	if len(b) == 0 {
+		return
+	}
+
+	if err := json.Unmarshal(b, &w.filesCache); err != nil {
+		logrus.Errorf("Failed to unmarshal image import fileinfo cache: %v", err)
+	}
+}
+
+// importAndWatchImages starts the image watcher and workqueue.
+// This function block until the workqueue is empty, indicating that all images
+// that currently exist have been imported.
+func importAndWatchImages(ctx context.Context, cfg *config.Node) error {
+	w, err := watchImages(ctx, cfg)
+	if err != nil {
+		return err
+	}
+
+	// Add images dir to workqueue; if it exists and contains images
+	// they will be recursively listed and enqueued.
+	w.workqueue.Add(cfg.Images)
+
+	// wait for the workqueue to empty before returning
+	for w.workqueue.Len() > 0 {
+		logrus.Debugf("Waiting for initial import of images from %s", cfg.Images)
+		time.Sleep(time.Second * 2)
+	}
+
+	// prune unseen entries from last run once all existing files have been processed
+	w.pruneCache()
+
+	return nil
+}
+
+// watchImages starts a watcher on the parent of the images dir, and a workqueue to process events
+// from the watch stream.
+func watchImages(ctx context.Context, cfg *config.Node) (*watchqueue, error) {
+	// watch the directory above the images dir, as it may not exist yet when the watch is started.
+	watcher, err := createWatcher(filepath.Dir(cfg.Images))
+	if err != nil {
+		return nil, pkgerrors.Wrapf(err, "failed to create image import watcher for %s", filepath.Dir(cfg.Images))
+	}
+
+	w := &watchqueue{
+		cfg:        cfg,
+		watcher:    watcher,
+		filesCache: make(map[string]*fileInfo),
+		workqueue:  workqueue.NewDelayingQueue(),
+	}
+	logrus.Debugf("Image import watcher created")
+
+	w.loadCache()
+
+	go func() {
+		<-ctx.Done()
+		w.watcher.Close()
+	}()
+
+	go w.runWorkerForImages(ctx)
+
+	go func() {
+		for {
+			select {
+			case event, ok := <-w.watcher.Events:
+				if !ok {
+					logrus.Info("Image import watcher event channel closed; retrying in 5 seconds")
+					select {
+					case <-time.After(time.Second * 5):
+						w.watcher = mustCreateWatcher(filepath.Dir(cfg.Images))
+					case <-ctx.Done():
+						return
+					}
+				}
+
+				// only enqueue event if it is for a path within the images dir - not the parent dir that we are watching
+				if strings.HasPrefix(event.Name, cfg.Images) {
+					w.workqueue.AddAfter(event.Name, time.Second*2)
+				}
+
+			case err, ok := <-w.watcher.Errors:
+				if !ok {
+					logrus.Info("Image import watcher error channel closed; retrying in 5 seconds")
+					select {
+					case <-time.After(time.Second * 5):
+						w.watcher = mustCreateWatcher(filepath.Dir(cfg.Images))
+					case <-ctx.Done():
+						return
+					}
+				}
+				logrus.Errorf("Image import watcher received an error: %v", err)
+			}
+		}
+	}()
+
+	return w, nil
+}