Add experimental memory command for balloons

Make it possible to set target memory size, for an instance.

This way it should be possible to return memory to the host.

Signed-off-by: Anders F Björklund <anders.f.bjorklund@gmail.com>

Author: afbjorklund

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,91 @@
+// 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
+	}
+
+	_ = inst
+	mem := 0 // TODO: implement
+	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,8 @@ type Driver interface {
 	FillConfig(ctx context.Context, cfg *limatype.LimaYAML, filePath string) error
 
 	SSHAddress(ctx context.Context) (string, error)
+
+	SetTargetMemory(memory int64) error
 }
 
 type ConfiguredDriver struct {

pkg/driver/external/client/methods.go

@@ -323,3 +323,7 @@ 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) SetTargetMemory(_ int64) error {
+	return errors.New("unavailable")
+}

pkg/driver/qemu/errors.go

@@ -0,0 +1,8 @@
+// SPDX-FileCopyrightText: Copyright The Lima Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package qemu
+
+import "errors"
+
+var errUnimplemented = errors.New("unimplemented")

pkg/driver/qemu/qemu_driver.go

@@ -720,3 +720,7 @@ func (l *LimaQemuDriver) ForwardGuestAgent() bool {
 	// if driver is not providing, use host agent
 	return l.vSockPort == 0 && l.virtioPort == ""
 }
+
+func (l *LimaQemuDriver) SetTargetMemory(_ int64) error {
+	return errUnimplemented
+}

pkg/driver/vz/vz_driver_darwin.go

@@ -437,3 +437,19 @@ func (l *LimaVzDriver) ForwardGuestAgent() bool {
 	// If driver is not providing, use host agent
 	return l.vSockPort == 0 && l.virtioPort == ""
 }
+
+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))
+	return nil
+}

pkg/driver/wsl2/wsl_driver_windows.go

@@ -357,3 +357,7 @@ func (l *LimaWslDriver) ForwardGuestAgent() bool {
 	// If driver is not providing, use host agent
 	return l.vSockPort == 0 && l.virtioPort == ""
 }
+
+func (l *LimaWslDriver) SetTargetMemory(memory int64) error {
+	return errUnimplemented
+}

pkg/hostagent/api/client/client.go

@@ -11,6 +11,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"strconv"
+	"strings"
 
 	"github.com/lima-vm/lima/v2/pkg/hostagent/api"
 	"github.com/lima-vm/lima/v2/pkg/httpclientutil"
@@ -19,6 +21,7 @@ import (
 type HostAgentClient interface {
 	HTTPClient() *http.Client
 	Info(context.Context) (*api.Info, error)
+	SetTargetMemory(context.Context, int64) error
 }
 
 // NewHostAgentClient creates a client.
@@ -65,3 +68,15 @@ func (c *client) Info(ctx context.Context) (*api.Info, error) {
 	}
 	return &info, 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,33 @@ func (b *Backend) GetInfo(w http.ResponseWriter, r *http.Request) {
 	_, _ = w.Write(m)
 }
 
+// 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("/v1/memory", http.HandlerFunc(b.SetMemory))
 }