banger/internal/daemon/vm_stats.go

package daemon

import (
	"context"
	"errors"
	"strings"
	"time"

	"banger/internal/api"
	"banger/internal/model"
	"banger/internal/system"
	"banger/internal/vsockagent"
)

func (d *Daemon) GetVMStats(ctx context.Context, idOrName string) (model.VMRecord, model.VMStats, error) {
	vm, err := d.withVMLockByRef(ctx, idOrName, func(vm model.VMRecord) (model.VMRecord, error) {
		return d.getVMStatsLocked(ctx, vm)
	})
	if err != nil {
		return model.VMRecord{}, model.VMStats{}, err
	}
	return vm, vm.Stats, nil
}

func (d *Daemon) HealthVM(ctx context.Context, idOrName string) (result api.VMHealthResult, err error) {
	_, err = d.withVMLockByRef(ctx, idOrName, func(vm model.VMRecord) (model.VMRecord, error) {
		result.Name = vm.Name
		if vm.State != model.VMStateRunning || !system.ProcessRunning(vm.Runtime.PID, vm.Runtime.APISockPath) {
			result.Healthy = false
			return vm, nil
		}
		if strings.TrimSpace(vm.Runtime.VSockPath) == "" {
			return model.VMRecord{}, errors.New("vm has no vsock path")
		}
		if vm.Runtime.VSockCID == 0 {
			return model.VMRecord{}, errors.New("vm has no vsock cid")
		}
		if err := d.ensureSocketAccess(ctx, vm.Runtime.VSockPath, "firecracker vsock socket"); err != nil {
			return model.VMRecord{}, err
		}
		pingCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
		defer cancel()
		if err := vsockagent.Health(pingCtx, d.logger, vm.Runtime.VSockPath); err != nil {
			return model.VMRecord{}, err
		}
		result.Healthy = true
		return vm, nil
	})
	return result, err
}

func (d *Daemon) PingVM(ctx context.Context, idOrName string) (result api.VMPingResult, err error) {
	health, err := d.HealthVM(ctx, idOrName)
	if err != nil {
		return api.VMPingResult{}, err
	}
	return api.VMPingResult{Name: health.Name, Alive: health.Healthy}, nil
}

func (d *Daemon) getVMStatsLocked(ctx context.Context, vm model.VMRecord) (model.VMRecord, error) {
	stats, err := d.collectStats(ctx, vm)
	if err == nil {
		vm.Stats = stats
		vm.UpdatedAt = model.Now()
		_ = d.store.UpsertVM(ctx, vm)
		if d.logger != nil {
			d.logger.Debug("vm stats collected", append(vmLogAttrs(vm), "rss_bytes", stats.RSSBytes, "vsz_bytes", stats.VSZBytes, "cpu_percent", stats.CPUPercent)...)
		}
	}
	return vm, nil
}

func (d *Daemon) pollStats(ctx context.Context) error {
	vms, err := d.store.ListVMs(ctx)
	if err != nil {
		return err
	}
	for _, vm := range vms {
		if err := d.withVMLockByIDErr(ctx, vm.ID, func(vm model.VMRecord) error {
			if vm.State != model.VMStateRunning || !system.ProcessRunning(vm.Runtime.PID, vm.Runtime.APISockPath) {
				return nil
			}
			stats, err := d.collectStats(ctx, vm)
			if err != nil {
				if d.logger != nil {
					d.logger.Debug("vm stats collection failed", append(vmLogAttrs(vm), "error", err.Error())...)
				}
				return nil
			}
			vm.Stats = stats
			vm.UpdatedAt = model.Now()
			return d.store.UpsertVM(ctx, vm)
		}); err != nil {
			return err
		}
	}
	return nil
}

func (d *Daemon) stopStaleVMs(ctx context.Context) (err error) {
	if d.config.AutoStopStaleAfter <= 0 {
		return nil
	}
	op := d.beginOperation("vm.stop_stale")
	defer func() {
		if err != nil {
			op.fail(err)
			return
		}
		op.done()
	}()
	vms, err := d.store.ListVMs(ctx)
	if err != nil {
		return err
	}
	now := model.Now()
	for _, vm := range vms {
		if err := d.withVMLockByIDErr(ctx, vm.ID, func(vm model.VMRecord) error {
			if vm.State != model.VMStateRunning || !system.ProcessRunning(vm.Runtime.PID, vm.Runtime.APISockPath) {
				return nil
			}
			if now.Sub(vm.LastTouchedAt) < d.config.AutoStopStaleAfter {
				return nil
			}
			op.stage("stopping_vm", vmLogAttrs(vm)...)
			_ = d.sendCtrlAltDel(ctx, vm)
			_ = d.waitForExit(ctx, vm.Runtime.PID, vm.Runtime.APISockPath, 10*time.Second)
			_ = d.cleanupRuntime(ctx, vm, true)
			vm.State = model.VMStateStopped
			vm.Runtime.State = model.VMStateStopped
			clearRuntimeHandles(&vm)
			vm.UpdatedAt = model.Now()
			return d.store.UpsertVM(ctx, vm)
		}); err != nil {
			return err
		}
	}
	return nil
}

func (d *Daemon) collectStats(ctx context.Context, vm model.VMRecord) (model.VMStats, error) {
	stats := model.VMStats{
		CollectedAt:        model.Now(),
		SystemOverlayBytes: system.AllocatedBytes(vm.Runtime.SystemOverlay),
		WorkDiskBytes:      system.AllocatedBytes(vm.Runtime.WorkDiskPath),
		MetricsRaw:         system.ParseMetricsFile(vm.Runtime.MetricsPath),
	}
	if vm.Runtime.PID > 0 && system.ProcessRunning(vm.Runtime.PID, vm.Runtime.APISockPath) {
		ps, err := system.ReadProcessStats(ctx, vm.Runtime.PID)
		if err == nil {
			stats.CPUPercent = ps.CPUPercent
			stats.RSSBytes = ps.RSSBytes
			stats.VSZBytes = ps.VSZBytes
		}
	}
	return stats, nil
}