diff --git a/balloon.go b/balloon.go new file mode 100644 index 00000000..c189f809 --- /dev/null +++ b/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 +} diff --git a/balloon_test.go b/balloon_test.go new file mode 100644 index 00000000..5bf0d3ba --- /dev/null +++ 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) + } +} diff --git a/firecracker.go b/firecracker.go index af14e54d..3f15c352 100644 --- a/firecracker.go +++ b/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) +} diff --git a/handlers.go b/handlers.go index 5ee1719f..214d57f6 100644 --- a/handlers.go +++ b/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, diff --git a/machine.go b/machine.go index be40eedb..1b600eb5 100644 --- a/machine.go +++ b/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 +} diff --git a/machine_test.go b/machine_test.go index 6ce69db6..b51e49bb 100644 --- a/machine_test.go +++ b/machine_test.go @@ -66,6 +66,12 @@ var ( testDataBin = filepath.Join(testDataPath, "bin") testRootfs = filepath.Join(testDataPath, "root-drive.img") + + testBalloonMemory = int64(10) + testBalloonNewMemory = int64(6) + testBalloonDeflateOnOom = true + testStatsPollingIntervals = int64(1) + testNewStatsPollingIntervals = int64(6) ) func envOrDefault(k, empty string) string { @@ -369,9 +375,12 @@ func TestMicroVMExecution(t *testing.T) { t.Run("TestAttachVsock", func(t *testing.T) { testAttachVsock(ctx, t, m) }) t.Run("SetMetadata", func(t *testing.T) { testSetMetadata(ctx, t, m) }) t.Run("UpdateMetadata", func(t *testing.T) { testUpdateMetadata(ctx, t, m) }) - t.Run("GetMetadata", func(t *testing.T) { testGetMetadata(ctx, t, m) }) // Should be after testSetMetadata and testUpdateMetadata + t.Run("GetMetadata", func(t *testing.T) { testGetMetadata(ctx, t, m) }) // Should be after testSetMetadata and testUpdateMetadata + t.Run("TestCreateBalloon", func(t *testing.T) { testCreateBalloon(ctx, t, m) }) // should be before the microVM is started(only pre-boot) t.Run("TestStartInstance", func(t *testing.T) { testStartInstance(ctx, t, m) }) t.Run("TestGetInstanceInfo", func(t *testing.T) { testGetInstanceInfo(ctx, t, m) }) + t.Run("TestGetBalloonConfig", func(t *testing.T) { testGetBalloonConfig(ctx, t, m) }) // should be after testCreateBalloon + t.Run("TestGetBalloonStats", func(t *testing.T) { testGetBalloonStats(ctx, t, m) }) // Let the VMM start and stabilize... timer := time.NewTimer(5 * time.Second) @@ -379,6 +388,8 @@ func TestMicroVMExecution(t *testing.T) { case <-timer.C: t.Run("TestUpdateGuestDrive", func(t *testing.T) { testUpdateGuestDrive(ctx, t, m) }) t.Run("TestUpdateGuestNetworkInterface", func(t *testing.T) { testUpdateGuestNetworkInterface(ctx, t, m) }) + t.Run("TestUpdateBalloon", func(t *testing.T) { testUpdateBalloon(ctx, t, m) }) + t.Run("TestUpdateBalloonStats", func(t *testing.T) { testUpdateBalloonStats(ctx, t, m) }) t.Run("TestShutdown", func(t *testing.T) { testShutdown(ctx, t, m) }) case <-exitchannel: // if we've already exited, there's no use waiting for the timer @@ -933,7 +944,6 @@ func TestLogFiles(t *testing.T) { if _, err := os.Stat(stdoutPath); os.IsNotExist(err) { t.Errorf("expected log file to be present") - } if _, err := os.Stat(stderrPath); os.IsNotExist(err) { @@ -1716,3 +1726,48 @@ func TestCreateSnapshot(t *testing.T) { }) } } + +func testCreateBalloon(ctx context.Context, t *testing.T, m *Machine) { + if err := m.CreateBalloon(ctx, testBalloonMemory, testBalloonDeflateOnOom, testStatsPollingIntervals); err != nil { + t.Errorf("Create balloon device failed from testAttachBalloon: %s", err) + } +} + +func testGetBalloonConfig(ctx context.Context, t *testing.T, m *Machine) { + expectedBalloonConfig := models.Balloon{ + AmountMib: &testBalloonMemory, + DeflateOnOom: &testBalloonDeflateOnOom, + StatsPollingIntervals: testStatsPollingIntervals, + } + + balloonConfig, err := m.GetBalloonConfig(ctx) + if err != nil { + t.Errorf("failed to get config: %s", err) + } + assert.Equal(t, expectedBalloonConfig, balloonConfig) +} + +func testUpdateBalloon(ctx context.Context, t *testing.T, m *Machine) { + if err := m.UpdateBalloon(ctx, testBalloonNewMemory); err != nil { + t.Errorf("Updating balloon device failed from testUpdateBalloon: %s", err) + } +} + +func testGetBalloonStats(ctx context.Context, t *testing.T, m *Machine) { + expectedBalloonStats := models.BalloonStats{ + TargetMib: &testBalloonMemory, + } + + balloonStat, err := m.GetBalloonStats(ctx) + if err != nil { + t.Errorf("failed to get balloon statistics: %s", err) + } + + assert.Equal(t, expectedBalloonStats.TargetMib, balloonStat.TargetMib) +} + +func testUpdateBalloonStats(ctx context.Context, t *testing.T, m *Machine) { + if err := m.UpdateBalloonStats(ctx, testNewStatsPollingIntervals); err != nil { + t.Errorf("Updating balloon staistics failed from testUpdateBalloonStats: %s", err) + } +}