fix(agent): bound PageBroker terminal replies

agent/internal/executor/restore.go

@@ -151,6 +151,7 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 	transactionID := ""
 	var broker pagebroker.Client
 	committed := false
+	var pageBrokerStageDuration, pageBrokerMountDuration, pageBrokerCommitDuration time.Duration
 	if brokered {
 		transactionID = uuid.NewString()
 		broker = pagebroker.Client{ControlSocketPath: req.PageBrokerControlSocketPath}
@@ -161,11 +162,15 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 				_ = broker.Abort(abortCtx, transactionID)
 			}
 		}()
+		stageStart := time.Now()
 		staged, err := broker.StagedRestore(ctx, transactionID, artifactPath)
+		pageBrokerStageDuration = time.Since(stageStart)
 		if err != nil {
 			return 0, fmt.Errorf("stage PageBroker restore: %w", err)
 		}
+		mountStart := time.Now()
 		stagingMount, err := mounts.MountPageBroker(ctx, bundleMount, staged)
+		pageBrokerMountDuration = time.Since(mountStart)
 		if err != nil {
 			return 0, fmt.Errorf("mount PageBroker staging: %w", err)
 		}
@@ -190,11 +195,13 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 		return 0, fmt.Errorf("nsrestore failed: %w", err)
 	}
 	if brokered {
+		commitStart := time.Now()
 		if err := broker.Commit(ctx, transactionID); err != nil {
 			log.Error(err, "failed to commit PageBroker restore")
 		} else {
 			committed = true
 		}
+		pageBrokerCommitDuration = time.Since(commitStart)
 	}
 	if result.CleanupError != nil {
 		cleanupErr = errors.Join(cleanupErr, result.CleanupError)
@@ -206,6 +213,9 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 	cleanup()
 	wall := time.Since(restoreStart)
 	unaccounted := remainingDuration(wall,
+		pageBrokerStageDuration,
+		pageBrokerMountDuration,
+		pageBrokerCommitDuration,
 		gpuDeviceMapDuration,
 		result.OverlayCaptureDuration,
 		result.CRIUPrepareDuration,
@@ -215,6 +225,9 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 	summary := map[string]any{
 		"duration": wall.String(),
 		"phases": map[string]string{
+			"pagebroker_stage": pageBrokerStageDuration.String(),
+			"pagebroker_mount": pageBrokerMountDuration.String(),
+			"pagebroker_commit": pageBrokerCommitDuration.String(),
 			"gpu_device_map":  gpuDeviceMapDuration.String(),
 			"overlay_capture": result.OverlayCaptureDuration.String(),
 			"criu_prepare":    result.CRIUPrepareDuration.String(),

agent/internal/pagebroker/client.go

@@ -20,6 +20,7 @@ const (
 	// PageBroker control requests and responses are limited to 64 KiB.
 	maxMessageSize   = 64 << 10
 	commitRetryDelay = 100 * time.Millisecond
+	commitRetryLimit = 30 * time.Second
 )
 
 var errMessageTooLarge = fmt.Errorf("message exceeds %d bytes", maxMessageSize)
@@ -57,6 +58,7 @@ func imageDirectory(directory string) (string, error) {
 }
 
 func (c Client) Commit(ctx context.Context, transactionID string) error {
+	var retryDeadline time.Time
 	for {
 		response, err := c.request(ctx, transactionID, &Request_Commit{Commit: &CommitRequest{}})
 		if err == nil {
@@ -69,10 +71,17 @@ func (c Client) Commit(ctx context.Context, transactionID string) error {
 		if !errors.As(err, &transport) {
 			return err
 		}
+		if retryDeadline.IsZero() {
+			retryDeadline = time.Now().Add(commitRetryLimit)
+		}
+		delay := min(commitRetryDelay, time.Until(retryDeadline))
+		if delay <= 0 {
+			return err
+		}
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
-		case <-time.After(commitRetryDelay):
+		case <-time.After(delay):
 		}
 	}
 }

agent/internal/types/config.go

@@ -45,6 +45,9 @@ func (c *AgentConfig) Validate() error {
 		return &ConfigError{Field: "storage.basePath", Message: fmt.Sprintf("storage.basePath must be %q", CheckpointBasePath)}
 	}
 	c.Storage.BasePath = basePath
+	if c.PageBroker.Enabled && strings.TrimSpace(c.PageBroker.ControlSocketPath) == "" {
+		return &ConfigError{Field: "pageBroker.controlSocketPath", Message: "pageBroker.controlSocketPath is required when PageBroker is enabled"}
+	}
 	if c.CRIU.TcpClose && c.CRIU.TcpEstablished {
 		return &ConfigError{
 			Field:   "criu",

agent/internal/types/config_test.go

@@ -26,3 +26,12 @@ func TestAgentConfigValidateRequiresFixedStorageBasePath(t *testing.T) {
 		}
 	}
 }
+
+func TestAgentConfigValidateRequiresPageBrokerControlSocket(t *testing.T) {
+	cfg := validAgentConfig()
+	cfg.PageBroker.Enabled = true
+
+	if err := cfg.Validate(); err == nil {
+		t.Fatal("expected error for missing PageBroker control socket")
+	}
+}

agent/pagebroker/Makefile

@@ -1,7 +1,7 @@
 PROTO := v1/pagebroker.proto
 GTEST_FLAGS := $(shell pkg-config --cflags --libs gtest_main)
-BROKER_SOURCES := broker.cpp checkpoint_transaction_descriptor.cpp posix_copy_engine.cpp restore_transaction_descriptor.cpp transfer_engine.cpp
-DAEMON_SOURCES := $(BROKER_SOURCES) daemon.cpp file_descriptor.cpp
+BROKER_SOURCES := broker.cpp checkpoint_transaction_descriptor.cpp posix_copy_engine.cpp restore_transaction_descriptor.cpp transaction.cpp transfer_engine.cpp
+DAEMON_SOURCES := $(BROKER_SOURCES) daemon.cpp main.cpp file_descriptor.cpp
 
 .PHONY: daemon generate test