refactor(shim): precompute guest rootfs after task create

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())
+}

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) {

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

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)

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