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 (s *VMService) GetVMStats(ctx context.Context, idOrName string) (model.VMRecord, model.VMStats, error) {
	vm, err := s.withVMLockByRef(ctx, idOrName, func(vm model.VMRecord) (model.VMRecord, error) {
		return s.getVMStatsLocked(ctx, vm)
	})
	if err != nil {
		return model.VMRecord{}, model.VMStats{}, err
	}
	return vm, vm.Stats, nil
}

func (s *VMService) HealthVM(ctx context.Context, idOrName string) (result api.VMHealthResult, err error) {
	_, err = s.withVMLockByRef(ctx, idOrName, func(vm model.VMRecord) (model.VMRecord, error) {
		result.Name = vm.Name
		if !s.vmAlive(vm) {
			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 := s.net.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, s.logger, vm.Runtime.VSockPath); err != nil {
			return model.VMRecord{}, err
		}
		result.Healthy = true
		return vm, nil
	})
	return result, err
}

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

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

func (s *VMService) pollStats(ctx context.Context) error {
	vms, err := s.store.ListVMs(ctx)
	if err != nil {
		return err
	}
	for _, vm := range vms {
		if err := s.withVMLockByIDErr(ctx, vm.ID, func(vm model.VMRecord) error {
			if !s.vmAlive(vm) {
				return nil
			}
			stats, err := s.collectStats(ctx, vm)
			if err != nil {
				if s.logger != nil {
					s.logger.Debug("vm stats collection failed", append(vmLogAttrs(vm), "error", err.Error())...)
				}
				return nil
			}
			vm.Stats = stats
			vm.UpdatedAt = model.Now()
			return s.store.UpsertVM(ctx, vm)
		}); err != nil {
			return err
		}
	}
	return nil
}

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

func (s *VMService) 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 s.vmAlive(vm) {
		if ps, err := system.ReadProcessStats(ctx, s.vmHandles(vm.ID).PID); err == nil {
			stats.CPUPercent = ps.CPUPercent
			stats.RSSBytes = ps.RSSBytes
			stats.VSZBytes = ps.VSZBytes
		}
	}
	return stats, nil
}