Implement memory ballooning feature (#332)

Memory ballooning support on firecracker-go-sdk.
Add wrappers for the Firecracker endpoints "/ballloon", "/balloon/statistics".

Signed-off-by: Royce Zhao <qiqinzha@amazon.com>

Author: RoyceDavison

balloon.go

@@ -0,0 +1,64 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+package firecracker
+
+import (
+	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
+)
+
+// BalloonDevice is a builder that will create a balloon used to set up
+// the firecracker microVM.
+type BalloonDevice struct {
+	balloon models.Balloon
+}
+
+type BalloonOpt func(*models.Balloon)
+
+// NewBalloonDevice will return a new BalloonDevice.
+func NewBalloonDevice(amountMib int64, deflateOnOom bool, opts ...BalloonOpt) BalloonDevice {
+	b := models.Balloon{
+		AmountMib:     &amountMib,
+		DeflateOnOom: &deflateOnOom,
+	}
+
+	for _, opt := range opts {
+		opt(&b)
+	}
+
+	return BalloonDevice{balloon: b}
+}
+
+// Build will return a new balloon
+func (b BalloonDevice) Build() models.Balloon {
+	return b.balloon
+}
+
+// WithStatsPollingIntervals is a functional option which sets the time in seconds between refreshing statistics.
+func WithStatsPollingIntervals(statsPollingIntervals int64) BalloonOpt {
+	return func(d *models.Balloon) {
+		d.StatsPollingIntervals = statsPollingIntervals
+	}
+}
+
+// UpdateAmountMiB sets the target size of the balloon
+func (b BalloonDevice) UpdateAmountMib(amountMib int64) BalloonDevice {
+	b.balloon.AmountMib = &amountMib
+	return b
+}
+
+// UpdateStatsPollingIntervals sets the time in seconds between refreshing statistics.
+// A non-zero value will enable the statistics. Defaults to 0.
+func (b BalloonDevice) UpdateStatsPollingIntervals(statsPollingIntervals int64) BalloonDevice {
+	b.balloon.StatsPollingIntervals = statsPollingIntervals
+	return b
+}

balloon_test.go

@@ -0,0 +1,58 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package firecracker
+
+import (
+	"reflect"
+	"testing"
+
+	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
+)
+
+var (
+	expectedAmountMib              = int64(6)
+	expectedDeflateOnOom          = true
+	expectedStatsPollingIntervals = int64(1)
+
+	expectedBalloon = models.Balloon{
+		AmountMib:              &expectedAmountMib,
+		DeflateOnOom:          &expectedDeflateOnOom,
+		StatsPollingIntervals: expectedStatsPollingIntervals,
+	}
+)
+
+func TestNewBalloonDevice(t *testing.T) {
+	balloon := NewBalloonDevice(expectedAmountMib, expectedDeflateOnOom, WithStatsPollingIntervals(expectedStatsPollingIntervals)).Build()
+	if e, a := expectedBalloon, balloon; !reflect.DeepEqual(e, a) {
+		t.Errorf("expected balloon %v, but received %v", e, a)
+	}
+}
+
+func TestUpdateAmountMiB(t *testing.T) {
+	BalloonDevice := NewBalloonDevice(int64(1), expectedDeflateOnOom, WithStatsPollingIntervals(expectedStatsPollingIntervals))
+	balloon := BalloonDevice.UpdateAmountMib(expectedAmountMib).Build()
+
+	if e, a := expectedBalloon, balloon; !reflect.DeepEqual(e, a) {
+		t.Errorf("expected balloon %v, but received %v", e, a)
+	}
+}
+
+func TestUpdateStatsPollingIntervals(t *testing.T) {
+	BalloonDevice := NewBalloonDevice(expectedAmountMib, expectedDeflateOnOom)
+	balloon := BalloonDevice.UpdateStatsPollingIntervals(expectedStatsPollingIntervals).Build()
+
+	if e, a := expectedBalloon, balloon; !reflect.DeepEqual(e, a) {
+		t.Errorf("expected balloon %v, but received %v", e, a)
+	}
+}

firecracker.go

@@ -397,3 +397,80 @@ func (f *Client) PatchGuestDriveByID(ctx context.Context, driveID, pathOnHost st
 
 	return f.client.Operations.PatchGuestDriveByID(params)
 }
+
+// PutBalloonOpt is a functional option to be used for the
+// PutBalloon API in setting any additional optional fields.
+type PutBalloonOpt func(*ops.PutBalloonParams)
+
+// PutBalloonOpt is a wrapper for the swagger generated client to make
+// calling of the API easier.
+func (f *Client) PutBalloon(ctx context.Context, balloon *models.Balloon, opts ...PutBalloonOpt) (*ops.PutBalloonNoContent, error) {
+	timeout, cancel := context.WithTimeout(ctx, time.Duration(f.firecrackerRequestTimeout)*time.Millisecond)
+	defer cancel()
+
+	params := ops.NewPutBalloonParamsWithContext(timeout)
+	params.SetBody(balloon)
+	for _, opt := range opts {
+		opt(params)
+	}
+
+	return f.client.Operations.PutBalloon(params)
+}
+
+// DescribeBalloonConfig is a wrapper for the swagger generated client to make
+// calling of the API easier.
+func (f *Client) DescribeBalloonConfig(ctx context.Context) (*ops.DescribeBalloonConfigOK, error) {
+	params := ops.NewDescribeBalloonConfigParams()
+	params.SetContext(ctx)
+	params.SetTimeout(time.Duration(f.firecrackerRequestTimeout) * time.Millisecond)
+
+	return f.client.Operations.DescribeBalloonConfig(params)
+}
+
+// PatchBalloonOpt is a functional option to be used for the PatchBalloon API in setting
+// any additional optional fields.
+type PatchBalloonOpt func(*ops.PatchBalloonParams)
+
+// PatchBalloon is a wrapper for the swagger generated client to make calling of the
+// API easier.
+func (f *Client) PatchBalloon(ctx context.Context, ballonUpdate *models.BalloonUpdate, opts ...PatchBalloonOpt) (*ops.PatchBalloonNoContent, error) {
+	timeout, cancel := context.WithTimeout(ctx, time.Duration(f.firecrackerRequestTimeout)*time.Millisecond)
+	defer cancel()
+
+	params := ops.NewPatchBalloonParamsWithContext(timeout)
+	params.SetBody(ballonUpdate)
+	for _, opt := range opts {
+		opt(params)
+	}
+
+	return f.client.Operations.PatchBalloon(params)
+}
+
+// DescribeBalloonStats is a wrapper for the swagger generated client to make calling of the
+// API easier.
+func (f *Client) DescribeBalloonStats(ctx context.Context) (*ops.DescribeBalloonStatsOK, error) {
+	params := ops.NewDescribeBalloonStatsParams()
+	params.SetContext(ctx)
+	params.SetTimeout(time.Duration(f.firecrackerRequestTimeout) * time.Millisecond)
+
+	return f.client.Operations.DescribeBalloonStats(params)
+}
+
+// PatchBalloonStatsIntervalOpt is a functional option to be used for the PatchBalloonStatsInterval API in setting
+// any additional optional fields.
+type PatchBalloonStatsIntervalOpt func(*ops.PatchBalloonStatsIntervalParams)
+
+// PatchBalloonStatsInterval is a wrapper for the swagger generated client to make calling of the
+// API easier.
+func (f *Client) PatchBalloonStatsInterval(ctx context.Context, balloonStatsUpdate *models.BalloonStatsUpdate, opts ...PatchBalloonStatsIntervalOpt) (*ops.PatchBalloonStatsIntervalNoContent, error) {
+	timeout, cancel := context.WithTimeout(ctx, time.Duration(f.firecrackerRequestTimeout)*time.Millisecond)
+	defer cancel()
+
+	params := ops.NewPatchBalloonStatsIntervalParamsWithContext(timeout)
+	params.SetBody(balloonStatsUpdate)
+	for _, opt := range opts {
+		opt(params)
+	}
+
+	return f.client.Operations.PatchBalloonStatsInterval(params)
+}

handlers.go

@@ -34,6 +34,7 @@ const (
 	LinkFilesToRootFSHandlerName       = "fcinit.LinkFilesToRootFS"
 	SetupNetworkHandlerName            = "fcinit.SetupNetwork"
 	SetupKernelArgsHandlerName         = "fcinit.SetupKernelArgs"
+	CreateBalloonHandlerName           = "fcint.CreateBalloon"
 
 	ValidateCfgHandlerName        = "validate.Cfg"
 	ValidateJailerCfgHandlerName  = "validate.JailerCfg"
@@ -268,6 +269,17 @@ var ConfigMmdsHandler = Handler{
 	},
 }
 
+// NewCreateBalloonHandler is a named handler that put a memory balloon into the
+// firecracker process.
+func NewCreateBalloonHandler(amountMib int64, deflateOnOom bool, StatsPollingIntervals int64) Handler {
+	return Handler{
+		Name: CreateBalloonHandlerName,
+		Fn: func(ctx context.Context, m *Machine) error {
+			return m.CreateBalloon(ctx, amountMib, deflateOnOom, StatsPollingIntervals)
+		},
+	}
+}
+
 var defaultFcInitHandlerList = HandlerList{}.Append(
 	SetupNetworkHandler,
 	SetupKernelArgsHandler,

machine.go

@@ -1084,3 +1084,79 @@ func (m *Machine) CreateSnapshot(ctx context.Context, memFilePath, snapshotPath
 	m.logger.Debug("snapshot created successfully")
 	return nil
 }
+
+// CreateBalloon creates a balloon device if one does not exist
+func (m *Machine) CreateBalloon(ctx context.Context, amountMib int64, deflateOnOom bool, statsPollingIntervals int64, opts ...PutBalloonOpt) error {
+	balloon := models.Balloon{
+		AmountMib:              &amountMib,
+		DeflateOnOom:          &deflateOnOom,
+		StatsPollingIntervals: statsPollingIntervals,
+	}
+	_, err := m.client.PutBalloon(ctx, &balloon, opts...)
+
+	if err != nil {
+		m.logger.Errorf("Create balloon device failed : %s", err)
+		return err
+	}
+
+	m.logger.Debug("Created balloon device successful")
+	return nil
+}
+
+// GetBalloonConfig gets the current balloon device configuration.
+func (m *Machine) GetBalloonConfig(ctx context.Context) (models.Balloon, error) {
+	var balloonConfig models.Balloon
+	resp, err := m.client.DescribeBalloonConfig(ctx)
+	if err != nil {
+		m.logger.Errorf("Getting balloonConfig: %s", err)
+		return balloonConfig, err
+	}
+
+	balloonConfig = *resp.Payload
+	m.logger.Debug("GetBalloonConfig successful")
+	return balloonConfig, err
+}
+
+// UpdateBalloon will update an existing balloon device, before or after machine startup
+func (m *Machine) UpdateBalloon(ctx context.Context, amountMib int64, opts ...PatchBalloonOpt) error {
+	ballonUpdate := models.BalloonUpdate{
+		AmountMib: &amountMib,
+	}
+	_, err := m.client.PatchBalloon(ctx, &ballonUpdate, opts...)
+	if err != nil {
+		m.logger.Errorf("Update balloon device failed : %s", err)
+		return err
+	}
+
+	m.logger.Debug("Update balloon device successful")
+	return nil
+}
+
+// GetBalloonStats gets the latest balloon device statistics, only if enabled pre-boot.
+func (m *Machine) GetBalloonStats(ctx context.Context) (models.BalloonStats, error) {
+	var balloonStats models.BalloonStats
+	resp, err := m.client.DescribeBalloonStats(ctx)
+	if err != nil {
+		m.logger.Errorf("Getting balloonStats: %s", err)
+		return balloonStats, err
+	}
+	balloonStats = *resp.Payload
+	m.logger.Debug("GetBalloonStats successful")
+	return balloonStats, nil
+}
+
+// UpdateBalloon will update a balloon device statistics polling interval.
+// Statistics cannot be turned on/off after boot.
+func (m *Machine) UpdateBalloonStats(ctx context.Context, statsPollingIntervals int64, opts ...PatchBalloonStatsIntervalOpt) error {
+	balloonStatsUpdate := models.BalloonStatsUpdate{
+		StatsPollingIntervals: &statsPollingIntervals,
+	}
+
+	if _, err := m.client.PatchBalloonStatsInterval(ctx, &balloonStatsUpdate, opts...); err != nil {
+		m.logger.Errorf("UpdateBalloonStats failed: %v", err)
+		return err
+	}
+
+	m.logger.Debug("UpdateBalloonStats successful")
+	return nil
+}