Speed up post-restore guest network reconfigure

lib/guest/client.go

@@ -82,8 +82,12 @@ func GetOrCreateConn(ctx context.Context, dialer hypervisor.VsockDialer) (*grpc.
 	}
 
 	// Create new connection using the VsockDialer
+	traceCtx := ctx
 	conn, err := grpc.Dial("passthrough:///vsock",
 		grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) {
+			if span := trace.SpanFromContext(traceCtx); span.SpanContext().IsValid() {
+				ctx = trace.ContextWithSpan(ctx, span)
+			}
 			netConn, err := dialer.DialVsock(ctx, vsockGuestPort)
 			if err != nil {
 				return nil, &AgentVSockDialError{Err: err}
@@ -146,6 +150,109 @@ type ExecOptions struct {
 	ResizeChan   <-chan *WindowSize // Optional: channel to receive resize events (pointer to avoid copying mutex)
 }
 
+type ReconfigureNetworkOptions struct {
+	InterfaceName string
+	MAC           string
+	IPv4          string
+	Prefix        uint32
+	Gateway       string
+	WaitForAgent  time.Duration
+}
+
+func ReconfigureNetworkInInstance(ctx context.Context, dialer hypervisor.VsockDialer, opts ReconfigureNetworkOptions) error {
+	if opts.WaitForAgent == 0 {
+		return reconfigureNetworkOnce(ctx, dialer, opts)
+	}
+
+	ctx, span := otel.Tracer("hypeman/guest").Start(ctx, "guest.reconfigure_network", trace.WithAttributes(
+		attribute.Bool("wait_for_agent", true),
+		attribute.Int64("wait_for_agent_ms", opts.WaitForAgent.Milliseconds()),
+	))
+	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 := reconfigureNetworkOnce(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 reconfigureNetworkOnce(ctx context.Context, dialer hypervisor.VsockDialer, opts ReconfigureNetworkOptions) 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.reconfigure_network.rpc")
+	_, err = client.ReconfigureNetwork(ctx, &ReconfigureNetworkRequest{
+		InterfaceName: opts.InterfaceName,
+		Mac:           opts.MAC,
+		Ipv4:          opts.IPv4,
+		Prefix:        opts.Prefix,
+		Gateway:       opts.Gateway,
+	})
+	finishGuestNetworkStepSpan(span, err)
+	if err != nil {
+		return fmt.Errorf("reconfigure network rpc: %w", err)
+	}
+	return nil
+}
+
+func finishGuestNetworkStepSpan(span trace.Span, err error) {
+	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(otelcodes.Error, err.Error())
+	} else {
+		span.SetStatus(otelcodes.Ok, "")
+	}
+	span.End()
+}
+
 // ExecIntoInstance executes command in instance via vsock using gRPC.
 // The dialer is a hypervisor-specific VsockDialer that knows how to connect to the guest.
 // If WaitForAgent is set, it will retry on connection errors until the timeout.

lib/guest/guest.proto

@@ -20,6 +20,9 @@ service GuestService {
 
   // Shutdown requests graceful VM shutdown by signaling init (PID 1)
   rpc Shutdown(ShutdownRequest) returns (ShutdownResponse);
+
+  // ReconfigureNetwork updates the guest network identity without spawning shell commands
+  rpc ReconfigureNetwork(ReconfigureNetworkRequest) returns (ReconfigureNetworkResponse);
 }
 
 // ExecRequest represents messages from client to server
@@ -154,3 +157,15 @@ message ShutdownRequest {
 
 // ShutdownResponse acknowledges the shutdown request
 message ShutdownResponse {}
+
+// ReconfigureNetworkRequest updates a guest network interface after snapshot restore
+message ReconfigureNetworkRequest {
+  string interface_name = 1;  // Interface to reconfigure, defaults to eth0
+  string mac = 2;             // New MAC address
+  string ipv4 = 3;            // New IPv4 address without prefix
+  uint32 prefix = 4;          // IPv4 prefix length
+  string gateway = 5;         // Default gateway IPv4 address
+}
+
+// ReconfigureNetworkResponse acknowledges the network reconfiguration request
+message ReconfigureNetworkResponse {}