refactor(shim): persist create annotations in one config.json write

Defer writing com.urunc.internal.rootfs.params and .view until after
guest rootfs choice and optional view prepare, then patch both via a
single PatchConfigJSON call instead of separate config.json updates.

Signed-off-by: sidneychang <2190206983@qq.com>

Author: sidneychang

pkg/containerd-shim/containerd/annotations.go

@@ -86,7 +86,7 @@ func InjectUruncAnnotations(ctx context.Context, session *Session, bundlePath st
 		return nil
 	}
 
-	return patchConfigJSON(bundlePath, annotations)
+	return PatchConfigJSON(bundlePath, annotations)
 }
 
 func (f *annotationFetcher) fetchUruncAnnotations(ctx context.Context) (map[string]string, error) {
@@ -152,12 +152,12 @@ func readBlob(ctx context.Context, namespace string, contentClient contentapi.Co
 	return raw, nil
 }
 
-// patchConfigJSON injects missing annotations into the OCI runtime spec
-// stored in the bundle's config.json.
+// PatchConfigJSON injects missing annotations into the OCI runtime spec stored in
+// the bundle's config.json.
 //
 // Existing annotations in config.json are preserved. Only annotation keys that
 // are not already present in the runtime spec are added.
-func patchConfigJSON(bundlePath string, annotations map[string]string) error {
+func PatchConfigJSON(bundlePath string, annotations map[string]string) error {
 	configPath := filepath.Join(bundlePath, "config.json")
 
 	fi, err := os.Stat(configPath)

pkg/containerd-shim/containerd/rootfs_view.go

@@ -25,7 +25,6 @@ import (
 	cntrtypes "github.com/containerd/containerd/api/types"
 	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/mount"
-	"github.com/sirupsen/logrus"
 	"github.com/urunc-dev/urunc/pkg/unikontainers"
 	"github.com/urunc-dev/urunc/pkg/unikontainers/types"
 	"google.golang.org/grpc/metadata"
@@ -37,10 +36,7 @@ const (
 	rootfsViewAnnotation  = "com.urunc.internal.rootfs.view"
 )
 
-var (
-	ErrRootfsViewNotPrepared = errors.New("rootfs view not prepared")
-	rootfsViewLog            = logrus.WithField("subsystem", "containerd-shim-rootfs-view")
-)
+var ErrRootfsViewNotPrepared = errors.New("rootfs view not prepared")
 
 type rootfsViewState struct {
 	Snapshotter string        `json:"snapshotter"`
@@ -79,33 +75,34 @@ func (a *RootfsViewAccessor) leaseID() string {
 	return rootfsViewLeasePrefix + a.containerID
 }
 
-func (a *RootfsViewAccessor) ShouldPrepare(rootfs types.RootfsParams) bool {
+func (a *RootfsViewAccessor) ShouldPrepare(rootfs types.RootfsParams) (bool, error) {
 	if a == nil ||
 		a.snapshotter == "" ||
 		a.snapshotKey == "" ||
 		(a.snapshotter != "devmapper" && a.snapshotter != "blockfile") ||
 		rootfs.Type != "block" ||
 		rootfs.MountedPath == "" {
-		return false
+		return false, nil
 	}
 
 	uruncCfg, cfgErr := unikontainers.LoadUruncConfig(unikontainers.UruncConfigPath)
 	if cfgErr != nil {
-		rootfsViewLog.WithError(cfgErr).Warn("failed to load urunc config; rootfs view disabled")
-		return false
+		return false, cfgErr
 	}
-	return uruncCfg.RootfsView.Enabled
+	return uruncCfg.RootfsView.Enabled, nil
 }
 
 // Prepare records a read-only view of the committed rootfs snapshot for runtime use.
-func (a *RootfsViewAccessor) Prepare(ctx context.Context, bundle string) error {
+// On success it returns JSON-encoded view state for the caller to persist in bundle
+// config.json together with other shim Create annotations.
+func (a *RootfsViewAccessor) Prepare(ctx context.Context) (string, error) {
 	if a == nil {
-		return fmt.Errorf("rootfs view accessor is nil")
+		return "", fmt.Errorf("rootfs view accessor is nil")
 	}
 
 	snapshotKey, err := a.resolveCommittedSnapshotBase(ctx, a.snapshotter, a.snapshotKey)
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	viewKey := a.viewKey()
@@ -115,17 +112,15 @@ func (a *RootfsViewAccessor) Prepare(ctx context.Context, bundle string) error {
 	if _, err := a.leases.Create(nsCtx, &leasesapi.CreateRequest{ID: leaseID}); err != nil {
 		err = containerdErr(err)
 		if err != nil && !errdefs.IsAlreadyExists(err) {
-			return fmt.Errorf("create rootfs view lease %s: %w", leaseID, err)
+			return "", fmt.Errorf("create rootfs view lease %s: %w", leaseID, err)
 		}
 	}
 
 	leaseCtx := metadata.AppendToOutgoingContext(nsCtx, "containerd-lease", leaseID)
 	mounts, err := a.createRootfsView(leaseCtx, viewKey, snapshotKey)
 	if err != nil {
-		if cleanupErr := cleanupRootfsViewLease(ctx, a.namespace, a.leaseID(), a.leases); cleanupErr != nil {
-			rootfsViewLog.WithError(cleanupErr).Warn("failed to clean up rootfs view lease after prepare failure")
-		}
-		return err
+		_ = cleanupRootfsViewLease(ctx, a.namespace, a.leaseID(), a.leases)
+		return "", err
 	}
 
 	state := &rootfsViewState{
@@ -135,26 +130,28 @@ func (a *RootfsViewAccessor) Prepare(ctx context.Context, bundle string) error {
 
 	encoded, err := json.Marshal(state)
 	if err != nil {
-		return fmt.Errorf("marshal rootfs view state: %w", err)
+		return "", fmt.Errorf("marshal rootfs view state: %w", err)
 	}
-	if err := unikontainers.PatchBundleRootfsView(bundle, string(encoded)); err != nil {
-		if cleanupErr := cleanupRootfsView(ctx, a.namespace, a.containerID, a.snapshotter, a.snapshots, a.leases); cleanupErr != nil {
-			rootfsViewLog.WithError(cleanupErr).Warn("failed to clean up rootfs view after state persistence failure")
-			return fmt.Errorf("persist rootfs view state: %w (cleanup also failed: %v)", err, cleanupErr)
-		}
-		return fmt.Errorf("persist rootfs view state: %w", err)
-	}
-
-	return nil
+	return string(encoded), nil
 }
 
+// Cleanup removes the per-container rootfs view snapshot and lease.
+// snapshotter names the devmapper/blockfile plugin that owns the view; callers
+// on Delete should pass the value persisted in bundle config.json at Prepare
+// time (see ShouldCleanupRootfsView). After inner task Delete, containerd
+// container metadata may no longer be loadable, so do not rely on a.snapshotter
+// alone on that path. When snapshotter is empty, a.snapshotter from the session
+// is used (Create-time rollback while the container record still exists).
 func (a *RootfsViewAccessor) Cleanup(ctx context.Context, snapshotter string) error {
 	if a == nil {
 		return fmt.Errorf("rootfs view accessor is nil")
 	}
 	if a.containerID == "" {
 		return fmt.Errorf("container id is empty")
 	}
+	if snapshotter == "" {
+		snapshotter = a.snapshotter
+	}
 	if snapshotter == "" {
 		return fmt.Errorf("snapshotter is empty")
 	}
@@ -300,7 +297,6 @@ func cleanupRootfsView(
 		Key:         rootfsViewKey(containerID),
 	})
 	if err = containerdErr(err); err != nil && !errdefs.IsNotFound(err) {
-		rootfsViewLog.WithError(err).Warn("failed to remove rootfs view from containerd")
 		return err
 	}
 	return cleanupRootfsViewLease(ctx, namespace, rootfsViewLeaseID(containerID), leases)
@@ -312,7 +308,6 @@ func cleanupRootfsViewLease(ctx context.Context, namespace, leaseID string, leas
 	}
 	_, err := leases.Delete(withNamespace(ctx, namespace), &leasesapi.DeleteRequest{ID: leaseID})
 	if err = containerdErr(err); err != nil && !errdefs.IsNotFound(err) {
-		rootfsViewLog.WithError(err).Warn("failed to remove rootfs view lease from containerd")
 		return err
 	}
 	return nil

pkg/containerd-shim/guest_rootfs.go

@@ -27,42 +27,35 @@ import (
 	"github.com/urunc-dev/urunc/pkg/unikontainers/types"
 )
 
-const annotRootfsParams = "com.urunc.internal.rootfs.params"
-
 var errGuestRootfsChoiceSkipped = errors.New("guest rootfs choice skipped")
 
 // chooseGuestRootfs runs the same ChooseRootfs logic as runtime Exec after inner
-// task Create (#684) and records the result in annotRootfsParams so Exec knows
-// selection already happened.
-func chooseGuestRootfs(r *taskAPI.CreateTaskRequest) (types.RootfsParams, error) {
+// task Create (#684). The caller persists the JSON-encoded result in bundle
+// config.json so Exec can reuse the selection.
+func chooseGuestRootfs(r *taskAPI.CreateTaskRequest) (types.RootfsParams, string, error) {
 	configPath := filepath.Join(r.Bundle, "config.json")
-	info, err := os.Stat(configPath)
-	if err != nil {
-		return types.RootfsParams{}, fmt.Errorf("stat config.json: %w", err)
-	}
-
 	data, err := os.ReadFile(configPath)
 	if err != nil {
-		return types.RootfsParams{}, fmt.Errorf("read config.json: %w", err)
+		return types.RootfsParams{}, "", fmt.Errorf("read config.json: %w", err)
 	}
 
 	var spec specs.Spec
 	if err := json.Unmarshal(data, &spec); err != nil {
-		return types.RootfsParams{}, fmt.Errorf("unmarshal config.json: %w", err)
+		return types.RootfsParams{}, "", fmt.Errorf("unmarshal config.json: %w", err)
 	}
 	if spec.Root == nil {
-		return types.RootfsParams{}, fmt.Errorf("invalid OCI spec: root section is required")
+		return types.RootfsParams{}, "", fmt.Errorf("invalid OCI spec: root section is required")
 	}
 
 	config, err := unikontainers.GetUnikernelConfig(filepath.Clean(r.Bundle), &spec)
 	if err != nil {
-		return types.RootfsParams{}, fmt.Errorf("%w: %w", errGuestRootfsChoiceSkipped, err)
+		return types.RootfsParams{}, "", fmt.Errorf("%w: %w", errGuestRootfsChoiceSkipped, err)
 	}
 
 	annotations := config.Map()
 	uruncCfg, err := unikontainers.LoadUruncConfig(unikontainers.UruncConfigPath)
 	if err != nil && uruncCfg == nil {
-		return types.RootfsParams{}, err
+		return types.RootfsParams{}, "", err
 	}
 
 	rootfsParams, err := unikontainers.ChooseRootfs(
@@ -72,24 +65,12 @@ func chooseGuestRootfs(r *taskAPI.CreateTaskRequest) (types.RootfsParams, error)
 		uruncCfg,
 	)
 	if err != nil {
-		return types.RootfsParams{}, err
+		return types.RootfsParams{}, "", err
 	}
 
 	encoded, err := json.Marshal(rootfsParams)
 	if err != nil {
-		return types.RootfsParams{}, err
-	}
-	if spec.Annotations == nil {
-		spec.Annotations = make(map[string]string)
-	}
-	spec.Annotations[annotRootfsParams] = string(encoded)
-
-	patched, err := json.MarshalIndent(spec, "", "  ")
-	if err != nil {
-		return types.RootfsParams{}, fmt.Errorf("marshal config.json: %w", err)
-	}
-	if err := os.WriteFile(configPath, patched, info.Mode()); err != nil {
-		return types.RootfsParams{}, err
+		return types.RootfsParams{}, "", err
 	}
-	return rootfsParams, nil
+	return rootfsParams, string(encoded), nil
 }

pkg/containerd-shim/task_service.go

@@ -27,6 +27,12 @@ import (
 	containerdShim "github.com/urunc-dev/urunc/pkg/containerd-shim/containerd"
 )
 
+// Internal bundle annotations (duplicated in unikontainers; keep in sync).
+const (
+	annotRootfsParams = "com.urunc.internal.rootfs.params"
+	annotRootfsView   = "com.urunc.internal.rootfs.view"
+)
+
 // taskService is urunc's shim-side wrapper around containerd's runc task
 // service. It wires urunc task setup before forwarding calls to the wrapped
 // service.
@@ -58,7 +64,7 @@ func (s *taskService) Create(ctx context.Context, r *taskAPI.CreateTaskRequest)
 		return resp, err
 	}
 
-	rootfsChoice, err := chooseGuestRootfs(r)
+	rootfsChoice, rootfsParamsJSON, err := chooseGuestRootfs(r)
 	if err != nil {
 		if errors.Is(err, errGuestRootfsChoiceSkipped) {
 			log.G(ctx).WithError(err).Debug("urunc(shim): guest rootfs choice skipped")
@@ -68,17 +74,18 @@ func (s *taskService) Create(ctx context.Context, r *taskAPI.CreateTaskRequest)
 		return nil, err
 	}
 
-	log.G(ctx).WithFields(map[string]any{
-		"rootfs_type": rootfsChoice.Type,
-		"rootfs_path": rootfsChoice.Path,
-		"mon_rootfs":  rootfsChoice.MonRootfs,
-	}).Debug("urunc(shim): guest rootfs chosen")
-
+	rootfsViewJSON := ""
+	var rootfsViewAccessor *containerdShim.RootfsViewAccessor
 	if session != nil {
-		rootfsViewAccessor := containerdShim.NewRootfsViewAccessor(session)
-		if rootfsViewAccessor.ShouldPrepare(rootfsChoice) {
-			if err := rootfsViewAccessor.Prepare(ctx, r.Bundle); err != nil {
+		rootfsViewAccessor = containerdShim.NewRootfsViewAccessor(session)
+		shouldPrepare, shouldPrepareErr := rootfsViewAccessor.ShouldPrepare(rootfsChoice)
+		if shouldPrepareErr != nil {
+			log.G(ctx).WithError(shouldPrepareErr).Warn("urunc(shim): failed to load urunc config; rootfs view disabled")
+		} else if shouldPrepare {
+			rootfsViewJSON, err = rootfsViewAccessor.Prepare(ctx)
+			if err != nil {
 				log.G(ctx).WithError(err).Warn("urunc(shim): failed to prepare rootfs view; falling back to legacy boot artifact extraction")
+				rootfsViewJSON = ""
 			} else {
 				log.G(ctx).Debug("urunc(shim): rootfs view prepared")
 			}
@@ -87,6 +94,27 @@ func (s *taskService) Create(ctx context.Context, r *taskAPI.CreateTaskRequest)
 		}
 	}
 
+	var shimAnnotations map[string]string
+	if rootfsViewJSON != "" {
+		shimAnnotations = map[string]string{
+			annotRootfsParams: rootfsParamsJSON,
+			annotRootfsView:   rootfsViewJSON,
+		}
+	} else {
+		shimAnnotations = map[string]string{
+			annotRootfsParams: rootfsParamsJSON,
+		}
+	}
+	if err := containerdShim.PatchConfigJSON(r.Bundle, shimAnnotations); err != nil {
+		if rootfsViewJSON != "" && rootfsViewAccessor != nil {
+			if cleanupErr := rootfsViewAccessor.Cleanup(ctx, ""); cleanupErr != nil {
+				log.G(ctx).WithError(cleanupErr).Warn("urunc(shim): failed to clean up rootfs view after annotation persistence failure")
+			}
+		}
+		log.G(ctx).WithError(err).Warn("urunc(shim): failed to persist shim create annotations")
+		return nil, err
+	}
+
 	return resp, nil
 }
 

pkg/unikontainers/rootfs_view_boot.go

@@ -15,7 +15,6 @@
 package unikontainers
 
 import (
-	"encoding/json"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -32,37 +31,6 @@ type rootfsViewState struct {
 	Mounts      []mount.Mount `json:"mounts,omitempty"`
 }
 
-// PatchBundleRootfsView writes shim-prepared view state into bundle config.json.
-func PatchBundleRootfsView(bundleDir, rootfsViewJSON string) error {
-	configPath := filepath.Join(bundleDir, configFilename)
-	fi, err := os.Stat(configPath)
-	if err != nil {
-		return fmt.Errorf("stat config.json: %w", err)
-	}
-
-	spec, err := loadSpec(bundleDir)
-	if err != nil {
-		return fmt.Errorf("load bundle spec: %w", err)
-	}
-	if spec.Annotations == nil {
-		spec.Annotations = make(map[string]string)
-	}
-	if rootfsViewJSON == "" {
-		delete(spec.Annotations, rootfsViewAnnotation)
-	} else {
-		spec.Annotations[rootfsViewAnnotation] = rootfsViewJSON
-	}
-
-	patched, err := json.MarshalIndent(spec, "", "  ")
-	if err != nil {
-		return fmt.Errorf("marshal config.json: %w", err)
-	}
-	if err := os.WriteFile(configPath, patched, fi.Mode()); err != nil {
-		return fmt.Errorf("write config.json: %w", err)
-	}
-	return nil
-}
-
 func ReadBundleRootfsView(bundleDir string) (string, error) {
 	configPath := filepath.Join(bundleDir, configFilename)
 	if _, err := os.Stat(configPath); err != nil {