Add experimental memory command for balloons

cmd/limactl/main.go

@@ -203,6 +203,7 @@ func newApp() *cobra.Command {
 		newNetworkCommand(),
 		newCloneCommand(),
 		newRenameCommand(),
+		newMemoryCommand(),
 	)
 	addPluginCommands(rootCmd)
 

cmd/limactl/memory.go

@@ -0,0 +1,93 @@
+// SPDX-FileCopyrightText: Copyright The Lima Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package main
+
+import (
+	"fmt"
+	"strconv"
+
+	"github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+
+	"github.com/lima-vm/lima/v2/pkg/memory"
+	"github.com/lima-vm/lima/v2/pkg/store"
+)
+
+func newMemoryCommand() *cobra.Command {
+	memoryCmd := &cobra.Command{
+		Use:   "memory",
+		Short: "Manage instance memory",
+		PersistentPreRun: func(*cobra.Command, []string) {
+			logrus.Warn("`limactl memory` is experimental")
+		},
+		GroupID: advancedCommand,
+	}
+	memoryCmd.AddCommand(newMemoryGetCommand())
+	memoryCmd.AddCommand(newMemorySetCommand())
+
+	return memoryCmd
+}
+
+func newMemoryGetCommand() *cobra.Command {
+	getCmd := &cobra.Command{
+		Use:               "get INSTANCE",
+		Short:             "Get current memory",
+		Long:              "Get the currently used total memory of an instance, in MiB",
+		Args:              cobra.MinimumNArgs(1),
+		RunE:              memoryGetAction,
+		ValidArgsFunction: memoryBashComplete,
+	}
+
+	return getCmd
+}
+
+func memoryGetAction(cmd *cobra.Command, args []string) error {
+	ctx := cmd.Context()
+	instName := args[0]
+
+	inst, err := store.Inspect(ctx, instName)
+	if err != nil {
+		return err
+	}
+
+	mem, err := memory.GetCurrent(ctx, inst)
+	if err != nil {
+		return err
+	}
+	fmt.Fprintf(cmd.OutOrStdout(), "%d\n", mem>>20)
+	return nil
+}
+
+func newMemorySetCommand() *cobra.Command {
+	setCmd := &cobra.Command{
+		Use:               "set INSTANCE memory AMOUNT",
+		Short:             "Set target memory",
+		Long:              "Set the target total memory of an instance, in MiB",
+		Args:              cobra.MinimumNArgs(2),
+		RunE:              memorySetAction,
+		ValidArgsFunction: memoryBashComplete,
+	}
+
+	return setCmd
+}
+
+func memorySetAction(cmd *cobra.Command, args []string) error {
+	ctx := cmd.Context()
+	instName := args[0]
+	meg, err := strconv.Atoi(args[1])
+	if err != nil {
+		return err
+	}
+
+	inst, err := store.Inspect(ctx, instName)
+	if err != nil {
+		return err
+	}
+
+	return memory.SetTarget(ctx, inst, int64(meg)<<20)
+}
+
+func memoryBashComplete(cmd *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
+	return bashCompleteInstanceNames(cmd)
+}

pkg/driver/driver.go

@@ -88,6 +88,10 @@ type Driver interface {
 	FillConfig(ctx context.Context, cfg *limatype.LimaYAML, filePath string) error
 
 	SSHAddress(ctx context.Context) (string, error)
+
+	GetCurrentMemory() (int64, error)
+
+	SetTargetMemory(memory int64) error
 }
 
 type ConfiguredDriver struct {

pkg/driver/external/client/methods.go

@@ -323,3 +323,11 @@ func (d *DriverClient) BootScripts() (map[string][]byte, error) {
 	d.logger.Debugf("Boot scripts retrieved successfully: %d scripts", len(resp.Scripts))
 	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.go

@@ -835,6 +835,8 @@ func Cmdline(ctx context.Context, cfg Config) (exe string, args []string, err er
 	// virtio-rng-pci accelerates starting up the OS, according to https://wiki.gentoo.org/wiki/QEMU/Options
 	args = append(args, "-device", "virtio-rng-pci")
 
+	args = append(args, "-device", "virtio-balloon")
+
 	// Input
 	input := "mouse"
 

pkg/driver/qemu/qemu_driver.go

@@ -25,6 +25,7 @@ import (
 	"github.com/coreos/go-semver/semver"
 	"github.com/digitalocean/go-qemu/qmp"
 	"github.com/digitalocean/go-qemu/qmp/raw"
+	"github.com/docker/go-units"
 	"github.com/sirupsen/logrus"
 
 	"github.com/lima-vm/lima/v2/pkg/driver"
@@ -720,3 +721,41 @@ func (l *LimaQemuDriver) ForwardGuestAgent() bool {
 	// if driver is not providing, use host agent
 	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)
+	if err != nil {
+		return err
+	}
+	if err := qmpClient.Connect(); err != nil {
+		return err
+	}
+	defer func() { _ = qmpClient.Disconnect() }()
+	rawClient := raw.NewMonitor(qmpClient)
+	logrus.Infof("Balloon target size: %s", units.BytesSize(float64(memory)))
+	err = rawClient.Balloon(memory)
+	if err != nil {
+		return err
+	}
+	return nil
+}

pkg/driver/vz/vz_driver_darwin.go

@@ -437,3 +437,39 @@ func (l *LimaVzDriver) ForwardGuestAgent() bool {
 	// If driver is not providing, use host agent
 	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")
+	}
+	balloons := l.machine.MemoryBalloonDevices()
+	if len(balloons) != 1 {
+		return fmt.Errorf("unexpected number of devices: %d", len(balloons))
+	}
+	balloon := vz.AsVirtioTraditionalMemoryBalloonDevice(balloons[0])
+	if balloon == nil {
+		return errors.New("unexpected type of balloon")
+	}
+	balloon.SetTargetVirtualMachineMemorySize(uint64(memory))
+	// avoid segfault, when trying to Release
+	runtime.SetFinalizer(balloon, nil)
+	return nil
+}

pkg/driver/wsl2/wsl_driver_windows.go

@@ -357,3 +357,11 @@ func (l *LimaWslDriver) ForwardGuestAgent() bool {
 	// If driver is not providing, use host agent
 	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,7 +10,10 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
+	"strconv"
+	"strings"
 
 	"github.com/lima-vm/lima/v2/pkg/hostagent/api"
 	"github.com/lima-vm/lima/v2/pkg/httpclientutil"
@@ -19,6 +22,8 @@ import (
 type HostAgentClient interface {
 	HTTPClient() *http.Client
 	Info(context.Context) (*api.Info, error)
+	GetCurrentMemory(context.Context) (int64, error)
+	SetTargetMemory(context.Context, int64) error
 }
 
 // NewHostAgentClient creates a client.
@@ -65,3 +70,33 @@ 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)
+	b := strings.NewReader(body)
+	resp, err := httpclientutil.Put(ctx, c.HTTPClient(), u, b)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	return nil
+}

pkg/hostagent/api/server/server.go

@@ -6,7 +6,9 @@ package server
 import (
 	"context"
 	"encoding/json"
+	"io"
 	"net/http"
+	"strconv"
 
 	"github.com/lima-vm/lima/v2/pkg/hostagent"
 	"github.com/lima-vm/lima/v2/pkg/httputil"
@@ -53,6 +55,53 @@ 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 {
+		w.WriteHeader(http.StatusMethodNotAllowed)
+		return
+	}
+
+	body, err := io.ReadAll(r.Body)
+	if err != nil {
+		b.onError(w, err, http.StatusInternalServerError)
+		return
+	}
+	memory, err := strconv.ParseInt(string(body), 10, 64)
+	if err != nil {
+		b.onError(w, err, http.StatusInternalServerError)
+		return
+	}
+
+	err = b.Agent.SetTargetMemory(memory)
+	if err != nil {
+		b.onError(w, err, http.StatusInternalServerError)
+		return
+	}
+	w.WriteHeader(http.StatusOK)
+}
+
 func AddRoutes(r *http.ServeMux, b *Backend) {
 	r.Handle("/v1/info", http.HandlerFunc(b.GetInfo))
+	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,14 @@ 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)
+}
+
 func (a *HostAgent) startRoutinesAndWait(ctx context.Context, errCh <-chan error) error {
 	stBase := events.Status{
 		SSHLocalPort: a.sshLocalPort,

pkg/httpclientutil/httpclientutil.go

@@ -67,6 +67,22 @@ func Post(ctx context.Context, c *http.Client, url string, body io.Reader) (*htt
 	return resp, nil
 }
 
+func Put(ctx context.Context, c *http.Client, url string, body io.Reader) (*http.Response, error) {
+	req, err := http.NewRequestWithContext(ctx, "PUT", url, body)
+	if err != nil {
+		return nil, err
+	}
+	resp, err := c.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	if err := Successful(resp); err != nil {
+		resp.Body.Close()
+		return nil, err
+	}
+	return resp, nil
+}
+
 func readAtMost(r io.Reader, maxBytes int) ([]byte, error) {
 	lr := &io.LimitedReader{
 		R: r,