From 2c36099717836324c10ec8adfb8f5e8febfc949b Mon Sep 17 00:00:00 2001 From: xibz Date: Mon, 17 Dec 2018 14:29:53 -0800 Subject: [PATCH] add support for rate limiting Adds support for rate limiting for inbound and outbound data streams of firecracker. This also introduces a TokenBuilder which can be passed into the NewRateLimiter factory function to return a new rate limiter. --- example_test.go | 61 ++++++++++++++++++++++++++++++++++++++++++++ machine.go | 13 ++++++++++ rate_limiter.go | 58 +++++++++++++++++++++++++++++++++++++++++ rate_limiter_test.go | 54 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 186 insertions(+) create mode 100644 rate_limiter.go create mode 100644 rate_limiter_test.go diff --git a/example_test.go b/example_test.go index 997659cd..a0d3d688 100644 --- a/example_test.go +++ b/example_test.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "os" + "time" "github.com/firecracker-microvm/firecracker-go-sdk" models "github.com/firecracker-microvm/firecracker-go-sdk/client/models" @@ -151,3 +152,63 @@ func ExampleDrivesBuilder_DriveOpt() { panic(err) } } + +func ExampleNetworkInterfaces_RateLimiting() { + // construct the limitations of the bandwidth for firecracker + bandwidthBuilder := firecracker.TokenBucketBuilder{}. + WithInitialSize(1024 * 1024). // Initial token amount + WithBucketSize(1024 * 1024). // Max number of tokens + WithRefillDuration(30 * time.Second) // Refill rate + + // construct the limitations of the number of operations per duration for firecracker + opsBuilder := firecracker.TokenBucketBuilder{}. + WithInitialSize(5). + WithBucketSize(5). + WithRefillDuration(5 * time.Second) + + // create the inbound rate limiter + inbound := firecracker.NewRateLimiter(bandwidthBuilder.Build(), opsBuilder.Build()) + + bandwidthBuilder = bandwidthBuilder.WithBucketSize(1024 * 1024 * 10) + opsBuilder = opsBuilder. + WithBucketSize(100). + WithInitialSize(100) + // create the outbound rate limiter + outbound := firecracker.NewRateLimiter(bandwidthBuilder.Build(), opsBuilder.Build()) + + networkIfaces := []firecracker.NetworkInterface{ + { + MacAddress: "01-23-45-67-89-AB-CD-EF", + HostDevName: "tap-name", + InRateLimiter: inbound, + OutRateLimiter: outbound, + }, + } + + cfg := firecracker.Config{ + SocketPath: "/path/to/socket", + KernelImagePath: "/path/to/kernel", + Drives: firecracker.NewDrivesBuilder("/path/to/rootfs").Build(), + MachineCfg: models.MachineConfiguration{ + VcpuCount: 1, + }, + NetworkInterfaces: networkIfaces, + } + + ctx := context.Background() + m, err := firecracker.NewMachine(ctx, cfg) + if err != nil { + panic(fmt.Errorf("failed to create new machine: %v", err)) + } + + defer os.Remove(cfg.SocketPath) + + if err := m.Start(ctx); err != nil { + panic(fmt.Errorf("failed to initialize machine: %v", err)) + } + + // wait for VMM to execute + if err := m.Wait(ctx); err != nil { + panic(err) + } +} diff --git a/machine.go b/machine.go index c19f0a7f..db6912e4 100644 --- a/machine.go +++ b/machine.go @@ -161,6 +161,11 @@ type NetworkInterface struct { HostDevName string // AllowMMDS makes the Firecracker MMDS available on this network interface. AllowMDDS bool + + // InRateLimiter limits the incoming bytes. + InRateLimiter *models.RateLimiter + // OutRateLimiter limits the outgoing bytes. + OutRateLimiter *models.RateLimiter } // VsockDevice represents a vsock connection between the host and the guest @@ -475,6 +480,14 @@ func (m *Machine) createNetworkInterface(ctx context.Context, iface NetworkInter AllowMmdsRequests: iface.AllowMDDS, } + if iface.InRateLimiter != nil { + ifaceCfg.RxRateLimiter = iface.InRateLimiter + } + + if iface.OutRateLimiter != nil { + ifaceCfg.TxRateLimiter = iface.OutRateLimiter + } + resp, err := m.client.PutGuestNetworkInterfaceByID(ctx, ifaceID, &ifaceCfg) if err == nil { m.logger.Printf("PutGuestNetworkInterfaceByID: %s", resp.Error()) diff --git a/rate_limiter.go b/rate_limiter.go new file mode 100644 index 00000000..f61d457a --- /dev/null +++ b/rate_limiter.go @@ -0,0 +1,58 @@ +package firecracker + +import ( + "time" + + models "github.com/firecracker-microvm/firecracker-go-sdk/client/models" +) + +// RateLimiterOpt represents a functional option for rate limiting construction +type RateLimiterOpt func(*models.RateLimiter) + +// NewRateLimiter will construct a new RateLimiter with given parameters. +func NewRateLimiter(bandwidth, ops models.TokenBucket, opts ...RateLimiterOpt) *models.RateLimiter { + limiter := &models.RateLimiter{ + Bandwidth: &bandwidth, + Ops: &ops, + } + + for _, opt := range opts { + opt(limiter) + } + + return limiter +} + +// TokenBucketBuilder is a builder that allows of building components of the +// models.RateLimiter structure. +type TokenBucketBuilder struct { + bucket models.TokenBucket +} + +// WithBucketSize will set the bucket size for the given token bucket +func (b TokenBucketBuilder) WithBucketSize(size int64) TokenBucketBuilder { + b.bucket.Size = &size + + return b +} + +// WithRefillDuration will set the given refill duration of the bucket fill rate. +func (b TokenBucketBuilder) WithRefillDuration(dur time.Duration) TokenBucketBuilder { + ms := dur.Nanoseconds() + ms /= int64(time.Millisecond) + b.bucket.RefillTime = &ms + + return b +} + +// WithInitialSize will set the initial token bucket size +func (b TokenBucketBuilder) WithInitialSize(size int64) TokenBucketBuilder { + b.bucket.OneTimeBurst = &size + + return b +} + +// Build will return a new token bucket +func (b TokenBucketBuilder) Build() models.TokenBucket { + return b.bucket +} diff --git a/rate_limiter_test.go b/rate_limiter_test.go new file mode 100644 index 00000000..ef24b780 --- /dev/null +++ b/rate_limiter_test.go @@ -0,0 +1,54 @@ +package firecracker_test + +import ( + "testing" + "time" + + "github.com/firecracker-microvm/firecracker-go-sdk" + models "github.com/firecracker-microvm/firecracker-go-sdk/client/models" + "github.com/stretchr/testify/assert" +) + +func TestRateLimiter(t *testing.T) { + bucket := firecracker.TokenBucketBuilder{}. + WithRefillDuration(1 * time.Hour). + WithBucketSize(100). + WithInitialSize(100). + Build() + + expectedBucket := models.TokenBucket{ + OneTimeBurst: firecracker.Int64(100), + RefillTime: firecracker.Int64(3600000), + Size: firecracker.Int64(100), + } + + assert.Equal(t, expectedBucket, bucket) +} + +func TestRateLimiter_RefillTime(t *testing.T) { + cases := []struct { + Name string + Dur time.Duration + ExpectedMilliseconds int64 + }{ + { + Name: "one hour", + Dur: 1 * time.Hour, + ExpectedMilliseconds: 3600000, + }, + { + Name: "zero", + ExpectedMilliseconds: 0, + }, + } + + for _, c := range cases { + t.Run(c.Name, func(t *testing.T) { + bucket := firecracker.TokenBucketBuilder{}. + WithRefillDuration(c.Dur). + Build() + + assert.Equal(t, &c.ExpectedMilliseconds, bucket.RefillTime) + }) + } +}