From 9e58d34e5c272c5f077cf12a8cebbc2407c44e3c Mon Sep 17 00:00:00 2001
From: sidneychang <2190206983@qq.com>
Date: Fri, 22 May 2026 17:22:54 +0800
Subject: [PATCH] feat(shim): choose guest rootfs after inner task create

Run ChooseRootfs in the shim after inner task Create, once the bundle
rootfs is mounted (#684). Persist JSON RootfsParams in config.json as
com.urunc.internal.rootfs.params so Exec reuses the choice; read state
then spec annotations because reexec may not yet mirror config.json.

Export ChooseRootfs from unikontainers. When the annotation is absent,
Exec runs the same selection as the former u.chooseRootfs() (podman and
urunc CLI). Ensure MonRootfs exists before rootfs setup in Exec.

Fixes: #684

Signed-off-by: sidneychang <2190206983@qq.com>
---
 pkg/containerd-shim/guest_rootfs.go | 92 +++++++++++++++++++++++++++++
 pkg/containerd-shim/task_service.go | 19 +++++-
 pkg/unikontainers/rootfs.go         |  4 ++
 pkg/unikontainers/unikontainers.go  | 49 +++++++++++----
 pkg/unikontainers/urunc_config.go   |  2 +-
 5 files changed, 152 insertions(+), 14 deletions(-)
 create mode 100644 pkg/containerd-shim/guest_rootfs.go

diff --git a/pkg/containerd-shim/guest_rootfs.go b/pkg/containerd-shim/guest_rootfs.go
new file mode 100644
index 00000000..f8982ecf
--- /dev/null
+++ b/pkg/containerd-shim/guest_rootfs.go
@@ -0,0 +1,92 @@
+// Copyright (c) 2023-2026, Nubificus LTD
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package containerdshim
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+
+	taskAPI "github.com/containerd/containerd/api/runtime/task/v2"
+	specs "github.com/opencontainers/runtime-spec/specs-go"
+	"github.com/urunc-dev/urunc/pkg/unikontainers"
+)
+
+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) error {
+	configPath := filepath.Join(r.Bundle, "config.json")
+	info, err := os.Stat(configPath)
+	if err != nil {
+		return fmt.Errorf("stat config.json: %w", err)
+	}
+
+	data, err := os.ReadFile(configPath)
+	if err != nil {
+		return fmt.Errorf("read config.json: %w", err)
+	}
+
+	var spec specs.Spec
+	if err := json.Unmarshal(data, &spec); err != nil {
+		return fmt.Errorf("unmarshal config.json: %w", err)
+	}
+	if spec.Root == nil {
+		return fmt.Errorf("invalid OCI spec: root section is required")
+	}
+
+	config, err := unikontainers.GetUnikernelConfig(filepath.Clean(r.Bundle), &spec)
+	if err != nil {
+		return fmt.Errorf("%w: %w", errGuestRootfsChoiceSkipped, err)
+	}
+
+	annotations := config.Map()
+	uruncCfg, err := unikontainers.LoadUruncConfig(unikontainers.UruncConfigPath)
+	if err != nil && uruncCfg == nil {
+		return err
+	}
+
+	rootfsParams, err := unikontainers.ChooseRootfs(
+		filepath.Clean(r.Bundle),
+		spec.Root.Path,
+		annotations,
+		uruncCfg,
+	)
+	if err != nil {
+		return err
+	}
+
+	encoded, err := json.Marshal(rootfsParams)
+	if err != nil {
+		return 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 fmt.Errorf("marshal config.json: %w", err)
+	}
+
+	return os.WriteFile(configPath, patched, info.Mode())
+}
diff --git a/pkg/containerd-shim/task_service.go b/pkg/containerd-shim/task_service.go
index 2c3e53b5..fb126c3f 100644
--- a/pkg/containerd-shim/task_service.go
+++ b/pkg/containerd-shim/task_service.go
@@ -16,6 +16,7 @@ package containerdshim
 
 import (
 	"context"
+	"errors"
 
 	taskAPI "github.com/containerd/containerd/api/runtime/task/v2"
 	"github.com/containerd/log"
@@ -47,7 +48,23 @@ func (s *taskService) Create(ctx context.Context, r *taskAPI.CreateTaskRequest)
 		}
 	}
 
-	return s.TaskService.Create(ctx, r)
+	resp, err := s.TaskService.Create(ctx, r)
+	if err != nil {
+		return resp, err
+	}
+
+	// ChooseRootfs after inner task Create so bundle rootfs is mounted;
+	// params are persisted in bundle config.json for runtime Exec.
+	if err := chooseGuestRootfs(r); err != nil {
+		if errors.Is(err, errGuestRootfsChoiceSkipped) {
+			log.G(ctx).WithError(err).Debug("urunc(shim): guest rootfs choice skipped")
+			return resp, nil
+		}
+		log.G(ctx).WithError(err).Warn("urunc(shim): failed to choose guest rootfs")
+		return nil, err
+	}
+
+	return resp, nil
 }
 
 func (s *taskService) Delete(ctx context.Context, r *taskAPI.DeleteRequest) (*taskAPI.DeleteResponse, error) {
diff --git a/pkg/unikontainers/rootfs.go b/pkg/unikontainers/rootfs.go
index 0fd7e38d..0d7c80c0 100644
--- a/pkg/unikontainers/rootfs.go
+++ b/pkg/unikontainers/rootfs.go
@@ -28,6 +28,10 @@ import (
 // TODO: Find and set the correct size for the tmpfs in the host
 const tmpfsSizeForNoRootfs = "65536k"
 
+// annotRootfsParams holds JSON RootfsParams after shim chooseGuestRootfs.
+// When present in bundle config.json, Exec reuses it; otherwise Exec runs ChooseRootfs.
+const annotRootfsParams = "com.urunc.internal.rootfs.params"
+
 type rootfsBuilder interface {
 	preSetup() error
 	postSetup() error
diff --git a/pkg/unikontainers/unikontainers.go b/pkg/unikontainers/unikontainers.go
index b7d61def..a3a00bb5 100644
--- a/pkg/unikontainers/unikontainers.go
+++ b/pkg/unikontainers/unikontainers.go
@@ -238,33 +238,37 @@ func (u *Unikontainer) SetupNet() (types.NetDevParams, error) {
 //  3. Container rootfs as block device (if MountRootfs=true and supported)
 //  4. Container rootfs as shared-fs: virtiofs > 9pfs (if MountRootfs=true and supported)
 //  5. No rootfs
-func (u *Unikontainer) chooseRootfs() (types.RootfsParams, error) {
-	bundleDir := filepath.Clean(u.State.Bundle)
-	rootfsDir := filepath.Clean(u.Spec.Root.Path)
+func ChooseRootfs(bundle, specRoot string, annot map[string]string, cfg *UruncConfig) (types.RootfsParams, error) {
+	bundleDir := filepath.Clean(bundle)
+	rootfsDir := filepath.Clean(specRoot)
 	rootfsDir, err := resolveAgainstBase(bundleDir, rootfsDir)
 	if err != nil {
 		uniklog.Errorf("could not resolve rootfs directory %s: %v", rootfsDir, err)
 		return types.RootfsParams{}, err
 	}
 
-	unikernelType := u.State.Annotations[annotType]
+	if cfg == nil {
+		return types.RootfsParams{}, fmt.Errorf("urunc config is required for guest rootfs selection")
+	}
+
+	unikernelType := annot[annotType]
 	unikernel, err := unikernels.New(unikernelType)
 	if err != nil {
 		return types.RootfsParams{}, err
 	}
 
-	vmmType := u.State.Annotations[annotHypervisor]
-	vmm, err := hypervisors.NewVMM(hypervisors.VmmType(vmmType), u.UruncCfg.Monitors)
+	vmmType := annot[annotHypervisor]
+	vmm, err := hypervisors.NewVMM(hypervisors.VmmType(vmmType), cfg.Monitors)
 	if err != nil {
 		return types.RootfsParams{}, err
 	}
 
-	virtiofsdConfig := u.UruncCfg.ExtraBins["virtiofsd"]
+	virtiofsdConfig := cfg.ExtraBins["virtiofsd"]
 
 	selector := &rootfsSelector{
 		bundle:     bundleDir,
 		cntrRootfs: rootfsDir,
-		annot:      u.State.Annotations,
+		annot:      annot,
 		unikernel:  unikernel,
 		vmm:        vmm,
 		vfsdPath:   virtiofsdConfig.Path,
@@ -426,11 +430,28 @@ func (u *Unikontainer) Exec(metrics m.Writer) error {
 	// if the respective annotation is set then, depending on the guest
 	// (supports block or 9pfs), it will use the supported option. In case
 	// both ae supported, then the block option will be used by default.
-	rootfsParams, err := u.chooseRootfs()
-	if err != nil {
-		uniklog.Errorf("could not choose guest rootfs: %v", err)
-		return err
+	var rootfsParams types.RootfsParams
+
+	// Read the rootfs choice written by the shim.
+	if rootfsParamsJSON := u.Spec.Annotations[annotRootfsParams]; rootfsParamsJSON != "" {
+		if err := json.Unmarshal([]byte(rootfsParamsJSON), &rootfsParams); err != nil {
+			return fmt.Errorf("could not decode guest rootfs params: %w", err)
+		}
+	}
+
+	// If there is no shim choice, the runtime chooses rootfs here.
+	if rootfsParams.MonRootfs == "" {
+		rootfsParams, err = ChooseRootfs(u.State.Bundle, u.Spec.Root.Path, u.State.Annotations, u.UruncCfg)
+		if err != nil {
+			uniklog.Errorf("could not choose guest rootfs: %v", err)
+			return err
+		}
 	}
+	uniklog.WithFields(logrus.Fields{
+		"rootfs_type": rootfsParams.Type,
+		"rootfs_path": rootfsParams.Path,
+		"mon_rootfs":  rootfsParams.MonRootfs,
+	}).Debug("guest rootfs params")
 
 	// TODO: Add support for using both an existing
 	// block based snapshot of the container's rootfs
@@ -479,6 +500,10 @@ func (u *Unikontainer) Exec(metrics m.Writer) error {
 		}
 	}
 
+	if err = os.MkdirAll(rootfsParams.MonRootfs, 0o755); err != nil {
+		return fmt.Errorf("failed to create monitor rootfs directory %s: %w", rootfsParams.MonRootfs, err)
+	}
+
 	err = rfsBuilder.preSetup()
 	if err != nil {
 		return fmt.Errorf("pre setup step for rootfs failed: %w", err)
diff --git a/pkg/unikontainers/urunc_config.go b/pkg/unikontainers/urunc_config.go
index 87164504..5f21d106 100644
--- a/pkg/unikontainers/urunc_config.go
+++ b/pkg/unikontainers/urunc_config.go
@@ -106,7 +106,7 @@ func defaultUruncConfig() *UruncConfig {
 // LoadUruncConfig loads the urunc configuration from the specified path.
 // If the file does not exist or is malformed, it returns the default configuration.
 func LoadUruncConfig(path string) (*UruncConfig, error) {
-	cfg := &UruncConfig{}
+	cfg := defaultUruncConfig()
 	_, err := toml.DecodeFile(path, cfg)
 	if err == nil {
 		return cfg, nil