Add command to get the current memory usage

Signed-off-by: Anders F Björklund <[email protected]>

Author: afbjorklund

cmd/limactl/memory.go

@@ -51,8 +51,10 @@ func memoryGetAction(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	_ = inst
-	mem := 0 // TODO: implement
+	mem, err := memory.GetCurrent(ctx, inst)
+	if err != nil {
+		return err
+	}
 	fmt.Fprintf(cmd.OutOrStdout(), "%d\n", mem>>20)
 	return nil
 }

pkg/driver/driver.go

@@ -89,6 +89,8 @@ type Driver interface {
 
 	SSHAddress(ctx context.Context) (string, error)
 
+	GetCurrentMemory() (int64, error)
+
 	SetTargetMemory(memory int64) error
 }
 

pkg/driver/external/client/methods.go

@@ -324,6 +324,10 @@ func (d *DriverClient) BootScripts() (map[string][]byte, error) {
 	return resp.Scripts, nil
 }
 
+func (d *DriverClient) GetCurrentMemory() (int64, error) {
+	return 0, errors.New("unavailable")
+}
+
 func (d *DriverClient) SetTargetMemory(_ int64) error {
 	return errors.New("unavailable")
 }

pkg/driver/qemu/qemu_driver.go

@@ -722,6 +722,25 @@ func (l *LimaQemuDriver) ForwardGuestAgent() bool {
 	return l.vSockPort == 0 && l.virtioPort == ""
 }
 
+func (l *LimaQemuDriver) GetCurrentMemory() (int64, error) {
+	qmpSockPath := filepath.Join(l.Instance.Dir, filenames.QMPSock)
+	qmpClient, err := qmp.NewSocketMonitor("unix", qmpSockPath, 5*time.Second)
+	if err != nil {
+		return 0, err
+	}
+	if err := qmpClient.Connect(); err != nil {
+		return 0, err
+	}
+	defer func() { _ = qmpClient.Disconnect() }()
+	rawClient := raw.NewMonitor(qmpClient)
+	info, err := rawClient.QueryBalloon()
+	if err != nil {
+		return 0, err
+	}
+	logrus.Infof("Balloon actual size: %s", units.BytesSize(float64(info.Actual)))
+	return info.Actual, nil
+}
+
 func (l *LimaQemuDriver) SetTargetMemory(memory int64) error {
 	qmpSockPath := filepath.Join(l.Instance.Dir, filenames.QMPSock)
 	qmpClient, err := qmp.NewSocketMonitor("unix", qmpSockPath, 5*time.Second)

pkg/driver/vz/vz_driver_darwin.go

@@ -438,6 +438,24 @@ func (l *LimaVzDriver) ForwardGuestAgent() bool {
 	return l.vSockPort == 0 && l.virtioPort == ""
 }
 
+func (l *LimaVzDriver) GetCurrentMemory() (int64, error) {
+	if l.machine == nil {
+		return 0, errors.New("no machine")
+	}
+	balloons := l.machine.MemoryBalloonDevices()
+	if len(balloons) != 1 {
+		return 0, fmt.Errorf("unexpected number of devices: %d", len(balloons))
+	}
+	balloon := vz.AsVirtioTraditionalMemoryBalloonDevice(balloons[0])
+	if balloon == nil {
+		return 0, errors.New("unexpected type of balloon")
+	}
+	// avoid segfault, when trying to Release
+	runtime.SetFinalizer(balloon, nil)
+	memory := balloon.GetTargetVirtualMachineMemorySize()
+	return int64(memory), nil
+}
+
 func (l *LimaVzDriver) SetTargetMemory(memory int64) error {
 	if l.machine == nil {
 		return errors.New("no machine")

pkg/driver/wsl2/wsl_driver_windows.go

@@ -358,6 +358,10 @@ func (l *LimaWslDriver) ForwardGuestAgent() bool {
 	return l.vSockPort == 0 && l.virtioPort == ""
 }
 
+func (l *LimaWslDriver) GetCurrentMemory() (int64, error) {
+	return 0, errUnimplemented
+}
+
 func (l *LimaWslDriver) SetTargetMemory(_ int64) error {
 	return errUnimplemented
 }

pkg/hostagent/api/client/client.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"strconv"
 	"strings"
@@ -21,6 +22,7 @@ import (
 type HostAgentClient interface {
 	HTTPClient() *http.Client
 	Info(context.Context) (*api.Info, error)
+	GetCurrentMemory(context.Context) (int64, error)
 	SetTargetMemory(context.Context, int64) error
 }
 
@@ -69,6 +71,24 @@ func (c *client) Info(ctx context.Context) (*api.Info, error) {
 	return &info, nil
 }
 
+func (c *client) GetCurrentMemory(ctx context.Context) (int64, error) {
+	u := fmt.Sprintf("http://%s/%s/memory", c.dummyHost, c.version)
+	resp, err := httpclientutil.Get(ctx, c.HTTPClient(), u)
+	if err != nil {
+		return 0, err
+	}
+	defer resp.Body.Close()
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return 0, err
+	}
+	memory, err := strconv.ParseInt(string(body), 10, 64)
+	if err != nil {
+		return 0, err
+	}
+	return memory, nil
+}
+
 func (c *client) SetTargetMemory(ctx context.Context, memory int64) error {
 	u := fmt.Sprintf("http://%s/%s/memory", c.dummyHost, c.version)
 	body := strconv.FormatInt(memory, 10)

pkg/hostagent/api/server/server.go

@@ -55,6 +55,25 @@ func (b *Backend) GetInfo(w http.ResponseWriter, r *http.Request) {
 	_, _ = w.Write(m)
 }
 
+// GetMemory is the handler for GET /v1/memory.
+func (b *Backend) GetMemory(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodGet {
+		w.WriteHeader(http.StatusMethodNotAllowed)
+		return
+	}
+
+	memory, err := b.Agent.GetCurrentMemory()
+	if err != nil {
+		b.onError(w, err, http.StatusInternalServerError)
+		return
+	}
+	s := strconv.FormatInt(memory, 10)
+
+	w.Header().Set("Content-Type", "text/plain")
+	w.WriteHeader(http.StatusOK)
+	_, _ = w.Write([]byte(s))
+}
+
 // SetMemory is the handler for PUT /v1/memory.
 func (b *Backend) SetMemory(w http.ResponseWriter, r *http.Request) {
 	if r.Method != http.MethodPut {
@@ -83,5 +102,6 @@ func (b *Backend) SetMemory(w http.ResponseWriter, r *http.Request) {
 
 func AddRoutes(r *http.ServeMux, b *Backend) {
 	r.Handle("/v1/info", http.HandlerFunc(b.GetInfo))
-	r.Handle("/v1/memory", http.HandlerFunc(b.SetMemory))
+	r.Handle("GET /v1/memory", http.HandlerFunc(b.GetMemory))
+	r.Handle("PUT /v1/memory", http.HandlerFunc(b.SetMemory))
 }

pkg/hostagent/hostagent.go

@@ -445,6 +445,10 @@ func (a *HostAgent) Run(ctx context.Context) error {
 	return a.startRoutinesAndWait(ctx, errCh)
 }
 
+func (a *HostAgent) GetCurrentMemory() (int64, error) {
+	return a.driver.GetCurrentMemory()
+}
+
 func (a *HostAgent) SetTargetMemory(memory int64) error {
 	return a.driver.SetTargetMemory(memory)
 }

pkg/memory/memory.go

@@ -5,15 +5,82 @@ package memory
 
 import (
 	"context"
+	"fmt"
+	"os/exec"
 	"path/filepath"
+	"strconv"
+	"strings"
 	"time"
 
 	hostagentclient "github.com/lima-vm/lima/v2/pkg/hostagent/api/client"
 	"github.com/lima-vm/lima/v2/pkg/limatype"
 	"github.com/lima-vm/lima/v2/pkg/limatype/filenames"
+	"github.com/lima-vm/lima/v2/pkg/sshutil"
 	"github.com/lima-vm/lima/v2/pkg/store"
 )
 
+func GetCurrent(ctx context.Context, inst *limatype.Instance) (int64, error) {
+	var memory int64
+	hostAgentPID, err := store.ReadPIDFile(filepath.Join(inst.Dir, filenames.HostAgentPID))
+	if err != nil {
+		return 0, err
+	}
+	if hostAgentPID != 0 {
+		haSock := filepath.Join(inst.Dir, filenames.HostAgentSock)
+		haClient, err := hostagentclient.NewHostAgentClient(haSock)
+		if err != nil {
+			return 0, err
+		}
+		ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
+		defer cancel()
+		memory, err = haClient.GetCurrentMemory(ctx)
+		if err != nil {
+			return 0, err
+		}
+	}
+
+	sshExe, err := sshutil.NewSSHExe()
+	if err != nil {
+		return 0, err
+	}
+	sshOpts, err := sshutil.CommonOpts(ctx, sshExe, false)
+	if err != nil {
+		return 0, err
+	}
+	sshArgs := append(sshutil.SSHArgsFromOpts(sshOpts),
+		"-p", fmt.Sprintf("%d", inst.SSHLocalPort),
+		fmt.Sprintf("%s@%s", *inst.Config.User.Name, inst.SSHAddress),
+	)
+
+	args := []string{"cat", "/proc/meminfo"}
+	sshCmd := exec.CommandContext(ctx, sshExe.Exe, append(sshArgs, args...)...)
+	out, err := sshCmd.Output()
+	if err != nil {
+		return 0, err
+	}
+
+	var available int64
+	for _, line := range strings.Split(string(out), "\n") {
+		if !strings.HasPrefix(line, "MemAvailable: ") {
+			continue
+		}
+		fields := strings.Fields(line)
+		if len(fields) < 3 {
+			return 0, fmt.Errorf("unexpected line: %s", line)
+		}
+		num, err := strconv.ParseInt(fields[1], 10, 64)
+		if err != nil {
+			return 0, err
+		}
+		if fields[2] == "kB" {
+			num *= 1024
+		}
+		available = num
+	}
+
+	return memory - available, nil
+}
+
 func SetTarget(ctx context.Context, inst *limatype.Instance, memory int64) error {
 	hostAgentPID, err := store.ReadPIDFile(filepath.Join(inst.Dir, filenames.HostAgentPID))
 	if err != nil {