Add mailbox resume network handoff

cmd/api/api/instances.go

@@ -657,9 +657,10 @@ func (s *ApiService) ForkInstance(ctx context.Context, request oapi.ForkInstance
 	}
 
 	result, err := s.InstanceManager.ForkInstance(ctx, inst.Id, instances.ForkInstanceRequest{
-		Name:        request.Body.Name,
-		FromRunning: request.Body.FromRunning != nil && *request.Body.FromRunning,
-		TargetState: targetState,
+		Name:           request.Body.Name,
+		FromRunning:    request.Body.FromRunning != nil && *request.Body.FromRunning,
+		TargetState:    targetState,
+		WaitForNetwork: request.Body.WaitForNetwork,
 	})
 	if err != nil {
 		switch {

cmd/api/api/instances_test.go

@@ -1165,13 +1165,15 @@ func TestForkInstance_Success(t *testing.T) {
 		result:  forked,
 	}
 	svc.InstanceManager = mockMgr
+	waitForNetwork := false
 
 	resp, err := svc.ForkInstance(
 		mw.WithResolvedInstance(ctx(), source.Id, source),
 		oapi.ForkInstanceRequestObject{
 			Id: source.Id,
 			Body: &oapi.ForkInstanceRequest{
-				Name: "forked-instance",
+				Name:           "forked-instance",
+				WaitForNetwork: &waitForNetwork,
 			},
 		},
 	)
@@ -1185,6 +1187,8 @@ func TestForkInstance_Success(t *testing.T) {
 	assert.Equal(t, "forked-instance", mockMgr.lastReq.Name)
 	assert.False(t, mockMgr.lastReq.FromRunning)
 	assert.Equal(t, instances.State(""), mockMgr.lastReq.TargetState)
+	require.NotNil(t, mockMgr.lastReq.WaitForNetwork)
+	assert.False(t, *mockMgr.lastReq.WaitForNetwork)
 }
 
 func TestForkInstance_NotSupported(t *testing.T) {

cmd/api/api/snapshots.go

@@ -175,7 +175,10 @@ func (s *ApiService) ForkSnapshot(ctx context.Context, request oapi.ForkSnapshot
 		return oapi.ForkSnapshot400JSONResponse{Code: "invalid_request", Message: "request body is required"}, nil
 	}
 
-	domainReq := instances.ForkSnapshotRequest{Name: request.Body.Name}
+	domainReq := instances.ForkSnapshotRequest{
+		Name:           request.Body.Name,
+		WaitForNetwork: request.Body.WaitForNetwork,
+	}
 	if request.Body.TargetState != nil {
 		domainReq.TargetState = instances.State(*request.Body.TargetState)
 	}

cmd/api/api/snapshots_test.go

@@ -1,14 +1,35 @@
 package api
 
 import (
+	"context"
 	"testing"
 	"time"
 
+	"github.com/kernel/hypeman/lib/hypervisor"
 	"github.com/kernel/hypeman/lib/instances"
+	"github.com/kernel/hypeman/lib/oapi"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+type captureForkSnapshotManager struct {
+	instances.Manager
+	lastID  string
+	lastReq *instances.ForkSnapshotRequest
+	result  *instances.Instance
+	err     error
+}
+
+func (m *captureForkSnapshotManager) ForkSnapshot(ctx context.Context, snapshotID string, req instances.ForkSnapshotRequest) (*instances.Instance, error) {
+	reqCopy := req
+	m.lastID = snapshotID
+	m.lastReq = &reqCopy
+	if m.err != nil {
+		return nil, m.err
+	}
+	return m.result, nil
+}
+
 func TestSnapshotScheduleToOAPIPreservesZeroMaxCount(t *testing.T) {
 	t.Parallel()
 
@@ -30,3 +51,43 @@ func TestSnapshotScheduleToOAPIPreservesZeroMaxCount(t *testing.T) {
 	require.NotNil(t, out.Retention.MaxAge)
 	assert.Equal(t, "24h0m0s", *out.Retention.MaxAge)
 }
+
+func TestForkSnapshotMapsWaitForNetwork(t *testing.T) {
+	t.Parallel()
+	svc := newTestService(t)
+
+	forked := &instances.Instance{
+		StoredMetadata: instances.StoredMetadata{
+			Id:             "forked-instance",
+			Name:           "forked-instance",
+			Image:          "docker.io/library/alpine:latest",
+			CreatedAt:      time.Now(),
+			HypervisorType: hypervisor.TypeFirecracker,
+		},
+		State: instances.StateRunning,
+	}
+	mockMgr := &captureForkSnapshotManager{
+		Manager: svc.InstanceManager,
+		result:  forked,
+	}
+	svc.InstanceManager = mockMgr
+	waitForNetwork := false
+
+	resp, err := svc.ForkSnapshot(ctx(), oapi.ForkSnapshotRequestObject{
+		SnapshotId: "snap-123",
+		Body: &oapi.ForkSnapshotRequest{
+			Name:           "forked-instance",
+			WaitForNetwork: &waitForNetwork,
+		},
+	})
+	require.NoError(t, err)
+
+	created, ok := resp.(oapi.ForkSnapshot201JSONResponse)
+	require.True(t, ok, "expected 201 response")
+	assert.Equal(t, "forked-instance", created.Name)
+	assert.Equal(t, "snap-123", mockMgr.lastID)
+	require.NotNil(t, mockMgr.lastReq)
+	assert.Equal(t, "forked-instance", mockMgr.lastReq.Name)
+	require.NotNil(t, mockMgr.lastReq.WaitForNetwork)
+	assert.False(t, *mockMgr.lastReq.WaitForNetwork)
+}

lib/guest/client.go

@@ -159,6 +159,92 @@ type ReconfigureNetworkOptions struct {
 	WaitForAgent  time.Duration
 }
 
+type ArmResumeNetworkOptions struct {
+	PollIntervalMS uint32
+	WaitForAgent   time.Duration
+}
+
+func ArmResumeNetworkInInstance(ctx context.Context, dialer hypervisor.VsockDialer, opts ArmResumeNetworkOptions) error {
+	if opts.WaitForAgent == 0 {
+		return armResumeNetworkOnce(ctx, dialer, opts)
+	}
+
+	ctx, span := otel.Tracer("hypeman/guest").Start(ctx, "guest.arm_resume_network", trace.WithAttributes(
+		attribute.Bool("wait_for_agent", true),
+		attribute.Int64("wait_for_agent_ms", opts.WaitForAgent.Milliseconds()),
+		attribute.Int("poll_interval_ms", int(opts.PollIntervalMS)),
+	))
+	defer span.End()
+
+	deadline := time.Now().Add(opts.WaitForAgent)
+	start := time.Now()
+	attempts := 0
+	retryableAttempts := 0
+	firstRetryableErrorType := ""
+	lastRetryableErrorType := ""
+	lastRetryInterval := time.Duration(0)
+
+	for {
+		attempts++
+		err := armResumeNetworkOnce(ctx, dialer, opts)
+		if err == nil {
+			recordGuestExecWait(span, start, attempts, retryableAttempts, firstRetryableErrorType, lastRetryableErrorType, lastRetryInterval)
+			span.SetStatus(otelcodes.Ok, "")
+			return nil
+		}
+		if !isRetryableConnectionError(err) {
+			recordGuestExecWait(span, start, attempts, retryableAttempts, firstRetryableErrorType, lastRetryableErrorType, lastRetryInterval)
+			span.RecordError(err)
+			span.SetStatus(otelcodes.Error, err.Error())
+			return err
+		}
+
+		retryableAttempts++
+		errType := retryableConnectionErrorType(err)
+		if firstRetryableErrorType == "" {
+			firstRetryableErrorType = errType
+		}
+		lastRetryableErrorType = errType
+		CloseConn(dialer.Key())
+
+		if time.Now().After(deadline) {
+			recordGuestExecWait(span, start, attempts, retryableAttempts, firstRetryableErrorType, lastRetryableErrorType, lastRetryInterval)
+			span.RecordError(err)
+			span.SetStatus(otelcodes.Error, err.Error())
+			return err
+		}
+
+		retryInterval := guestExecRetryInterval(time.Since(start))
+		lastRetryInterval = retryInterval
+		select {
+		case <-ctx.Done():
+			recordGuestExecWait(span, start, attempts, retryableAttempts, firstRetryableErrorType, lastRetryableErrorType, lastRetryInterval)
+			span.RecordError(ctx.Err())
+			span.SetStatus(otelcodes.Error, ctx.Err().Error())
+			return ctx.Err()
+		case <-time.After(retryInterval):
+		}
+	}
+}
+
+func armResumeNetworkOnce(ctx context.Context, dialer hypervisor.VsockDialer, opts ArmResumeNetworkOptions) error {
+	grpcConn, err := GetOrCreateConn(ctx, dialer)
+	if err != nil {
+		return fmt.Errorf("get grpc connection: %w", err)
+	}
+	client := NewGuestServiceClient(grpcConn)
+
+	_, span := otel.Tracer("hypeman/guest").Start(ctx, "guest.arm_resume_network.rpc")
+	_, err = client.ArmResumeNetwork(ctx, &ArmResumeNetworkRequest{
+		PollIntervalMs: opts.PollIntervalMS,
+	})
+	finishGuestNetworkStepSpan(span, err)
+	if err != nil {
+		return fmt.Errorf("arm resume network rpc: %w", err)
+	}
+	return nil
+}
+
 func ReconfigureNetworkInInstance(ctx context.Context, dialer hypervisor.VsockDialer, opts ReconfigureNetworkOptions) error {
 	if opts.WaitForAgent == 0 {
 		return reconfigureNetworkOnce(ctx, dialer, opts)