feat(server): Add configurable memory balloon for MicroVMs

docs/reference/configuration.md

@@ -241,6 +241,45 @@ pools:
           size: 2000
           refill_time: 1000
     #
+    # Firecracker memory balloon device. Requires a guest kernel with
+    # CONFIG_VIRTIO_BALLOON (and CONFIG_PAGE_REPORTING for free page
+    # reporting). The device can only be attached before boot, so the
+    # options apply to MicroVMs created after the change.
+    #
+    # Default: {} (no balloon device)
+    #
+    balloon:
+      #
+      # Initial target size of the balloon in MiB. 0 attaches the device
+      # without taking memory from the guest; the target can be changed later
+      # through the Firecracker API (PATCH /balloon).
+      #
+      # Default: 0
+      #
+      amount_mib: 0
+      #
+      # Let the guest take pages back from the balloon instead of running
+      # into an out-of-memory state.
+      #
+      # Default: false
+      #
+      deflate_on_oom: true
+      #
+      # Interval in seconds between balloon statistics updates. 0 disables
+      # the statistics.
+      #
+      # Default: 0
+      #
+      stats_polling_interval_s: 0
+      #
+      # Let the guest continually report memory it no longer uses, so the
+      # host can reclaim it. Without it, a MicroVM keeps its peak memory
+      # usage on the host until it exits.
+      #
+      # Default: false
+      #
+      free_page_reporting: true
+    #
     # Metadata to pass to the Firecracker VM via MMDS.
     #
     # Default: {}

server/balloon.go

@@ -0,0 +1,79 @@
+package server
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+
+	"github.com/firecracker-microvm/firecracker-go-sdk"
+)
+
+const createBalloonHandlerName = "fireactions.CreateBalloon"
+
+// balloonRequest is the body of Firecracker's PUT /balloon. The SDK's balloon
+// model predates free page reporting, so the request is sent directly to the
+// API socket instead.
+type balloonRequest struct {
+	AmountMib             int64 `json:"amount_mib"`
+	DeflateOnOom          bool  `json:"deflate_on_oom"`
+	StatsPollingIntervalS int64 `json:"stats_polling_interval_s,omitempty"`
+	FreePageReporting     bool  `json:"free_page_reporting,omitempty"`
+}
+
+func (c *FirecrackerBalloonConfig) toRequest() balloonRequest {
+	return balloonRequest{
+		AmountMib:             c.AmountMib,
+		DeflateOnOom:          c.DeflateOnOom,
+		StatsPollingIntervalS: c.StatsPollingIntervalS,
+		FreePageReporting:     c.FreePageReporting,
+	}
+}
+
+// newCreateBalloonHandler returns a handler that attaches a balloon device to
+// the MicroVM. It must run before the instance starts.
+func newCreateBalloonHandler(socketPath string, config *FirecrackerBalloonConfig) firecracker.Handler {
+	return firecracker.Handler{
+		Name: createBalloonHandlerName,
+		Fn: func(ctx context.Context, _ *firecracker.Machine) error {
+			return putBalloon(ctx, socketPath, config.toRequest())
+		},
+	}
+}
+
+func putBalloon(ctx context.Context, socketPath string, balloon balloonRequest) error {
+	body, err := json.Marshal(balloon)
+	if err != nil {
+		return fmt.Errorf("encoding balloon: %w", err)
+	}
+
+	client := &http.Client{Transport: &http.Transport{
+		DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
+			return (&net.Dialer{}).DialContext(ctx, "unix", socketPath)
+		},
+	}}
+	defer client.CloseIdleConnections()
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPut, "http://localhost/balloon", bytes.NewReader(body))
+	if err != nil {
+		return fmt.Errorf("creating balloon request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Accept", "application/json")
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return fmt.Errorf("creating balloon: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusNoContent {
+		msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
+		return fmt.Errorf("creating balloon: %s: %s", resp.Status, bytes.TrimSpace(msg))
+	}
+
+	return nil
+}

server/balloon_test.go

@@ -0,0 +1,80 @@
+package server
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"net"
+	"net/http"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func serveUnixSocket(t *testing.T, handler http.HandlerFunc) string {
+	t.Helper()
+
+	socketPath := filepath.Join(t.TempDir(), "firecracker.sock")
+	listener, err := net.Listen("unix", socketPath)
+	require.NoError(t, err)
+
+	server := &http.Server{Handler: handler}
+	go server.Serve(listener)
+	t.Cleanup(func() { server.Close() })
+
+	return socketPath
+}
+
+func TestCreateBalloonHandler(t *testing.T) {
+	var method, path string
+	var body map[string]interface{}
+	socketPath := serveUnixSocket(t, func(w http.ResponseWriter, r *http.Request) {
+		method, path = r.Method, r.URL.Path
+		data, _ := io.ReadAll(r.Body)
+		json.Unmarshal(data, &body)
+		w.WriteHeader(http.StatusNoContent)
+	})
+
+	handler := newCreateBalloonHandler(socketPath, &FirecrackerBalloonConfig{
+		AmountMib: 0, DeflateOnOom: true, StatsPollingIntervalS: 5, FreePageReporting: true,
+	})
+	require.NoError(t, handler.Fn(context.Background(), nil))
+
+	assert.Equal(t, createBalloonHandlerName, handler.Name)
+	assert.Equal(t, http.MethodPut, method)
+	assert.Equal(t, "/balloon", path)
+	assert.Equal(t, map[string]interface{}{
+		"amount_mib":               float64(0),
+		"deflate_on_oom":           true,
+		"stats_polling_interval_s": float64(5),
+		"free_page_reporting":      true,
+	}, body)
+}
+
+func TestCreateBalloonHandlerOmitsDisabledFeatures(t *testing.T) {
+	var body map[string]interface{}
+	socketPath := serveUnixSocket(t, func(w http.ResponseWriter, r *http.Request) {
+		data, _ := io.ReadAll(r.Body)
+		json.Unmarshal(data, &body)
+		w.WriteHeader(http.StatusNoContent)
+	})
+
+	handler := newCreateBalloonHandler(socketPath, &FirecrackerBalloonConfig{AmountMib: 128})
+	require.NoError(t, handler.Fn(context.Background(), nil))
+
+	assert.Equal(t, map[string]interface{}{"amount_mib": float64(128), "deflate_on_oom": false}, body)
+}
+
+func TestCreateBalloonHandlerError(t *testing.T) {
+	socketPath := serveUnixSocket(t, func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusBadRequest)
+		w.Write([]byte(`{"fault_message":"Amount of pages requested is too large."}`))
+	})
+
+	handler := newCreateBalloonHandler(socketPath, &FirecrackerBalloonConfig{AmountMib: 1 << 20})
+	err := handler.Fn(context.Background(), nil)
+	assert.ErrorContains(t, err, "400 Bad Request")
+	assert.ErrorContains(t, err, "Amount of pages requested is too large.")
+}

server/config.go

@@ -55,6 +55,7 @@ type FirecrackerConfig struct {
 	MachineConfig    FirecrackerMachineConfig           `yaml:"machine_config"`
 	NetworkInterface *FirecrackerNetworkInterfaceConfig `yaml:"network_interface"`
 	Rootfs           *FirecrackerRootfsConfig           `yaml:"rootfs"`
+	Balloon          *FirecrackerBalloonConfig          `yaml:"balloon"`
 	Metadata         map[string]interface{}             `yaml:"metadata"`
 }
 
@@ -76,6 +77,15 @@ type FirecrackerRootfsConfig struct {
 	RateLimiter *FirecrackerRateLimiterConfig `yaml:"rate_limiter"`
 }
 
+// FirecrackerBalloonConfig configures the MicroVM's memory balloon device. A
+// nil balloon leaves the MicroVM without one.
+type FirecrackerBalloonConfig struct {
+	AmountMib             int64 `yaml:"amount_mib" validate:"gte=0"`
+	DeflateOnOom          bool  `yaml:"deflate_on_oom"`
+	StatsPollingIntervalS int64 `yaml:"stats_polling_interval_s" validate:"gte=0"`
+	FreePageReporting     bool  `yaml:"free_page_reporting"`
+}
+
 // FirecrackerRateLimiterConfig defines an IO rate limiter with independent
 // bytes/s and ops/s limits. A nil token bucket leaves that limit unlimited.
 type FirecrackerRateLimiterConfig struct {

server/config_test.go

@@ -77,6 +77,26 @@ func TestNewConfigRootfsInvalid(t *testing.T) {
 	assert.ErrorContains(t, err, "Config.Pools[0].Firecracker.Rootfs.RateLimiter.Ops.RefillTime")
 }
 
+func TestNewConfigBalloon(t *testing.T) {
+	config, err := NewConfig("testdata/config1.yaml")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	assert.Equal(t, &FirecrackerBalloonConfig{
+		AmountMib: 0, DeflateOnOom: true, StatsPollingIntervalS: 5, FreePageReporting: true,
+	}, config.Pools[0].Firecracker.Balloon)
+
+	// A pool without a balloon block gets no balloon device.
+	assert.Nil(t, config.Pools[1].Firecracker.Balloon)
+}
+
+func TestNewConfigBalloonInvalid(t *testing.T) {
+	// A negative balloon size is rejected.
+	_, err := NewConfig("testdata/config4.yaml")
+	assert.ErrorContains(t, err, "Config.Pools[0].Firecracker.Balloon.AmountMib")
+}
+
 func TestFirecrackerRateLimiterConfigToSDK(t *testing.T) {
 	var nilRateLimiter *FirecrackerRateLimiterConfig
 	assert.Nil(t, nilRateLimiter.toSDK())

server/pool.go

@@ -535,6 +535,11 @@ func (p *Pool) createMachine(ctx context.Context) error {
 
 	fcMachine.Handlers.FcInit = fcMachine.Handlers.FcInit.Append(firecracker.NewSetMetadataHandler(metadata))
 
+	if balloonConfig := p.config.Firecracker.Balloon; balloonConfig != nil {
+		socketPath := filepath.Join(p.GetDir(), fmt.Sprintf("%s.sock", runnerName))
+		fcMachine.Handlers.FcInit = fcMachine.Handlers.FcInit.Append(newCreateBalloonHandler(socketPath, balloonConfig))
+	}
+
 	vmmCtx, vmmCancel := context.WithCancel(p.ctx)
 	if err := fcMachine.Start(vmmCtx); err != nil {
 		vmmCancel()

server/testdata/config1.yaml

@@ -58,6 +58,11 @@ pools:
         ops:
           size: 2000
           refill_time: 1000
+    balloon:
+      amount_mib: 0
+      deflate_on_oom: true
+      stats_polling_interval_s: 5
+      free_page_reporting: true
     metadata:
       example1: value1
 - name: fireactions-2vcpu-4gb

server/testdata/config4.yaml

@@ -0,0 +1,30 @@
+---
+bind_address: 0.0.0.0:8080
+
+github:
+  app_private_key: |
+    -----BEGIN RSA PRIVATE KEY-----
+  app_id: 12345
+
+pools:
+- name: fireactions-2vcpu-2gb
+  replicas: 1
+  runner:
+    name: fireactions-2vcpu-2gb
+    image: ghcr.io/hostinger/fireactions/runner:ubuntu-20.04-x64-2.310.2
+    image_pull_policy: IfNotPresent
+    group_id: 1
+    organization: hostinger
+    labels:
+    - self-hosted
+  firecracker:
+    binary_path: firecracker
+    kernel_image_path: /var/lib/fireactions/vmlinux
+    kernel_args: "console=ttyS0 noapic reboot=k panic=1 pci=off nomodules rw"
+    machine_config:
+      mem_size_mib: 2048
+      vcpu_count: 2
+    balloon:
+      amount_mib: -1
+
+log_level: debug