cloudhypervisor.go

// Package cloudhypervisor implements the hypervisor.Hypervisor interface
// for Cloud Hypervisor VMM.
package cloudhypervisor

import (
	"context"
	"fmt"
	"time"

	"github.com/kernel/hypeman/lib/hypervisor"
	"github.com/kernel/hypeman/lib/logger"
	"github.com/kernel/hypeman/lib/vmm"
)

// CloudHypervisor implements hypervisor.Hypervisor for Cloud Hypervisor VMM.
type CloudHypervisor struct {
	client     *vmm.VMM
	socketPath string
	// serial is set on the Starter path (StartVM/RestoreVM) so Shutdown
	// can stop the reader explicitly. When the client is constructed via
	// the reconnect factory (New), there is no reader to own — the
	// goroutine from the original process exited with that process.
	serial *serialReader
}

var balloonTargetCache hypervisor.BalloonTargetCache

func clearBalloonTargetCache(socketPath string) {
	balloonTargetCache.Delete(socketPath)
}

// New creates a new Cloud Hypervisor client for an existing VMM socket.
func New(socketPath string) (*CloudHypervisor, error) {
	client, err := vmm.NewVMM(socketPath)
	if err != nil {
		return nil, fmt.Errorf("create vmm client: %w", err)
	}
	return &CloudHypervisor{
		client:     client,
		socketPath: socketPath,
	}, nil
}

// Verify CloudHypervisor implements the interface
var _ hypervisor.Hypervisor = (*CloudHypervisor)(nil)

// Capabilities returns the features supported by Cloud Hypervisor.
func (c *CloudHypervisor) Capabilities() hypervisor.Capabilities {
	return capabilities()
}

func capabilities() hypervisor.Capabilities {
	return CapabilitiesForVersion(vmm.DefaultVersion)
}

// CapabilitiesForVersion returns capabilities for a specific CH version.
// Use version-specific capabilities when new features are introduced that
// aren't supported on previous versions.
func CapabilitiesForVersion(v vmm.CHVersion) hypervisor.Capabilities {
	caps := hypervisor.Capabilities{
		SupportsSnapshot:            true,
		SupportsHotplugMemory:       true,
		SupportsBalloonControl:      true,
		SupportsPause:               true,
		SupportsVsock:               true,
		SupportsGPUPassthrough:      true,
		SupportsDiskIOLimit:         true,
		SupportsGracefulVMMShutdown: true,
		SupportsSnapshotBaseReuse:   false,
	}
	switch v {
	case vmm.V51_1:
		caps.SupportsDiskResize = true
		caps.SupportsConcurrentForkPrepare = true
		caps.SupportsSnapshotBaseReuse = experimentalDiffSnapshotsEnabled()
	}
	return caps
}

// DeleteVM removes the VM configuration from Cloud Hypervisor.
func (c *CloudHypervisor) DeleteVM(ctx context.Context) error {
	resp, err := c.client.DeleteVMWithResponse(ctx)
	if err != nil {
		return fmt.Errorf("delete vm: %w", err)
	}
	if resp.StatusCode() != 204 {
		return fmt.Errorf("delete vm failed with status %d: %s", resp.StatusCode(), string(resp.Body))
	}
	clearBalloonTargetCache(c.socketPath)
	return nil
}

// Shutdown stops the VMM process gracefully.
func (c *CloudHypervisor) Shutdown(ctx context.Context) error {
	resp, err := c.client.ShutdownVMMWithResponse(ctx)
	// Stop the serial reader regardless of API outcome — once Shutdown
	// has been requested the VM is going away and the reader has no
	// further work.
	c.serial.Close()
	if err != nil {
		return fmt.Errorf("shutdown vmm: %w", err)
	}
	if resp.StatusCode() != 204 {
		return fmt.Errorf("shutdown vmm failed with status %d", resp.StatusCode())
	}
	clearBalloonTargetCache(c.socketPath)
	return nil
}

// GetVMInfo returns current VM state.
func (c *CloudHypervisor) GetVMInfo(ctx context.Context) (*hypervisor.VMInfo, error) {
	resp, err := c.client.GetVmInfoWithResponse(ctx)
	if err != nil {
		return nil, fmt.Errorf("get vm info: %w", err)
	}
	if resp.StatusCode() != 200 || resp.JSON200 == nil {
		return nil, fmt.Errorf("get vm info failed with status %d", resp.StatusCode())
	}

	// Map Cloud Hypervisor state to hypervisor.VMState
	var state hypervisor.VMState
	switch resp.JSON200.State {
	case vmm.Created:
		state = hypervisor.StateCreated
	case vmm.Running:
		state = hypervisor.StateRunning
	case vmm.Paused:
		state = hypervisor.StatePaused
	case vmm.Shutdown:
		state = hypervisor.StateShutdown
	default:
		return nil, fmt.Errorf("unknown vm state: %s", resp.JSON200.State)
	}

	return &hypervisor.VMInfo{
		State:            state,
		MemoryActualSize: resp.JSON200.MemoryActualSize,
	}, nil
}

// Pause suspends VM execution.
func (c *CloudHypervisor) Pause(ctx context.Context) error {
	resp, err := c.client.PauseVMWithResponse(ctx)
	if err != nil {
		return fmt.Errorf("pause vm: %w", err)
	}
	if resp.StatusCode() != 204 {
		return fmt.Errorf("pause vm failed with status %d", resp.StatusCode())
	}
	return nil
}

// Resume continues VM execution.
func (c *CloudHypervisor) Resume(ctx context.Context) error {
	resp, err := c.client.ResumeVMWithResponse(ctx)
	if err != nil {
		return fmt.Errorf("resume vm: %w", err)
	}
	if resp.StatusCode() != 204 {
		return fmt.Errorf("resume vm failed with status %d", resp.StatusCode())
	}
	return nil
}

// Snapshot creates a VM snapshot.
func (c *CloudHypervisor) Snapshot(ctx context.Context, destPath string) error {
	diff, err := prepareDiffSnapshotDestination(destPath)
	if err != nil {
		return fmt.Errorf("prepare diff snapshot destination: %w", err)
	}

	snapshotURL := "file://" + destPath
	snapshotConfig := vmm.VmSnapshotConfig{DestinationUrl: &snapshotURL}
	if experimentalDiffSnapshotsEnabled() {
		snapshotType := vmm.Full
		if diff {
			snapshotType = vmm.Diff
		}
		snapshotConfig.SnapshotType = &snapshotType
	}
	resp, err := c.client.PutVmSnapshotWithResponse(ctx, snapshotConfig)
	if err != nil {
		return fmt.Errorf("snapshot: %w", err)
	}
	if resp.StatusCode() != 204 {
		return fmt.Errorf("snapshot failed with status %d", resp.StatusCode())
	}
	if diff {
		stats, err := mergeCloudHypervisorDiff(destPath)
		if err != nil {
			return fmt.Errorf("merge diff snapshot: %w", err)
		}
		logger.FromContext(ctx).InfoContext(ctx, "merged Cloud Hypervisor diff snapshot",
			"snapshot_dir", destPath,
			"delta_bytes", stats.DeltaBytes,
			"extents", stats.ExtentCount,
			"reflinked_bytes", stats.ReflinkedBytes,
			"copied_bytes", stats.CopiedBytes,
		)
	}
	optimized, err := prepareSnapshotForKernelPaging(destPath)
	if err != nil {
		return fmt.Errorf("prepare kernel-paged snapshot: %w", err)
	}
	if optimized {
		logger.FromContext(ctx).DebugContext(ctx, "prepared Cloud Hypervisor snapshot for kernel paging", "snapshot_dir", destPath)
	}
	return nil
}

func (c *CloudHypervisor) ResizeVCPUs(ctx context.Context, vcpus int) error {
	resp, err := c.client.PutVmResizeWithResponse(ctx, vmm.VmResize{DesiredVcpus: &vcpus})
	if err != nil {
		return fmt.Errorf("resize vCPUs: %w", err)
	}
	if resp.StatusCode() != 204 {
		return fmt.Errorf("resize vCPUs failed with status %d", resp.StatusCode())
	}
	return nil
}

// ResizeMemory changes the VM's memory allocation.
func (c *CloudHypervisor) ResizeMemory(ctx context.Context, bytes int64) error {
	if ExperimentalHotplugOverlayEnabled() {
		resp, err := c.client.PutVmResizeZoneWithResponse(ctx, vmm.VmResizeZone{
			Id:         ptr(kernelPagingMemoryZoneID),
			DesiredRam: ptr(bytes),
		})
		if err != nil {
			return fmt.Errorf("resize memory zone: %w", err)
		}
		if resp.StatusCode() != 204 {
			return fmt.Errorf("resize memory zone failed with status %d", resp.StatusCode())
		}
		return nil
	}

	resizeConfig := vmm.VmResize{DesiredRam: &bytes}
	resp, err := c.client.PutVmResizeWithResponse(ctx, resizeConfig)
	if err != nil {
		return fmt.Errorf("resize memory: %w", err)
	}
	if resp.StatusCode() != 204 {
		return fmt.Errorf("resize memory failed with status %d", resp.StatusCode())
	}
	return nil
}

// ResizeMemoryAndWait changes the VM's memory allocation and waits for it to stabilize.
// It polls until the actual memory size stabilizes (stops changing) or timeout is reached.
func (c *CloudHypervisor) ResizeMemoryAndWait(ctx context.Context, bytes int64, timeout time.Duration) error {
	// First, request the resize
	if err := c.ResizeMemory(ctx, bytes); err != nil {
		return err
	}

	// Poll until memory stabilizes
	const pollInterval = 20 * time.Millisecond
	deadline := time.Now().Add(timeout)

	var lastSize int64 = -1
	stableCount := 0
	const requiredStableChecks = 3 // Require 3 consecutive stable readings

	for time.Now().Before(deadline) {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		info, err := c.GetVMInfo(ctx)
		if err != nil {
			return fmt.Errorf("poll memory size: %w", err)
		}

		if info.MemoryActualSize == nil {
			// No memory info available, just return after resize
			return nil
		}

		currentSize := *info.MemoryActualSize

		if currentSize == lastSize {
			stableCount++
			if stableCount >= requiredStableChecks {
				// Memory has stabilized
				return nil
			}
		} else {
			stableCount = 0
			lastSize = currentSize
		}

		time.Sleep(pollInterval)
	}

	// Timeout reached, but resize was requested successfully
	return nil
}

func (c *CloudHypervisor) SetTargetGuestMemoryBytes(ctx context.Context, bytes int64) error {
	info, err := c.client.GetVmInfoWithResponse(ctx)
	if err != nil {
		return fmt.Errorf("get vm info for balloon update: %w", err)
	}
	if info.StatusCode() != 200 || info.JSON200 == nil {
		return fmt.Errorf("get vm info for balloon update failed with status %d", info.StatusCode())
	}
	if info.JSON200.Config.Balloon == nil {
		return hypervisor.ErrNotSupported
	}

	assigned := assignedGuestMemoryBytes(info.JSON200)
	if bytes < 0 || bytes > assigned {
		return fmt.Errorf("target guest memory %d is outside valid range [0,%d]", bytes, assigned)
	}

	desiredBalloon := assigned - bytes
	resp, err := c.client.PutVmResizeWithResponse(ctx, vmm.VmResize{DesiredBalloon: &desiredBalloon})
	if err != nil {
		return fmt.Errorf("set balloon target: %w", err)
	}
	if resp.StatusCode() != 204 {
		return fmt.Errorf("set balloon target failed with status %d", resp.StatusCode())
	}
	balloonTargetCache.Store(c.socketPath, bytes)
	return nil
}

func (c *CloudHypervisor) GetTargetGuestMemoryBytes(ctx context.Context) (int64, error) {
	if target, ok := balloonTargetCache.Load(c.socketPath); ok {
		return target, nil
	}

	info, err := c.client.GetVmInfoWithResponse(ctx)
	if err != nil {
		return 0, fmt.Errorf("get vm info for balloon read: %w", err)
	}
	if info.StatusCode() != 200 || info.JSON200 == nil {
		return 0, fmt.Errorf("get vm info for balloon read failed with status %d", info.StatusCode())
	}
	if info.JSON200.Config.Balloon == nil {
		return 0, hypervisor.ErrNotSupported
	}
	assigned := assignedGuestMemoryBytes(info.JSON200)
	return assigned - info.JSON200.Config.Balloon.Size, nil
}

func assignedGuestMemoryBytes(info *vmm.VmInfo) int64 {
	assigned := info.Config.Memory.Size
	if info.MemoryActualSize != nil && info.Config.Balloon != nil {
		assigned = *info.MemoryActualSize + info.Config.Balloon.Size
	}
	return assigned
}