set firelens mem_buf_limit

yinyic · Oct 22, 2024 · 86594d1 · 86594d1
1 parent 79f17a5
commit 86594d1
Show file tree

Hide file tree

Showing 8 changed files with 237 additions and 85 deletions.
diff --git a/agent/api/container/container.go b/agent/api/container/container.go
@@ -1272,6 +1272,25 @@ func (c *Container) GetNetworkModeFromHostConfig() string {
 	return hostConfig.NetworkMode.NetworkName()
 }
 
+// GetMemoryReservationFromHostConfig returns the container memory reservation
+func (c *Container) GetMemoryReservationFromHostConfig() int64 {
+	c.lock.RLock()
+	defer c.lock.RUnlock()
+
+	if c.DockerConfig.HostConfig == nil {
+		return 0
+	}
+
+	hostConfig := &dockercontainer.HostConfig{}
+	err := json.Unmarshal([]byte(*c.DockerConfig.HostConfig), hostConfig)
+	if err != nil {
+		seelog.Warnf("Encountered error when trying to get memory reservation for container %s: %v", c.RuntimeID, err)
+		return 0
+	}
+
+	return hostConfig.MemoryReservation
+}
+
 // GetHostConfig returns the container's host config.
 func (c *Container) GetHostConfig() *string {
 	c.lock.RLock()

diff --git a/agent/api/container/container_test.go b/agent/api/container/container_test.go
@@ -751,6 +751,48 @@ func TestGetNetworkModeFromHostConfig(t *testing.T) {
 	}
 }
 
+func TestGetMemoryReservationFromHostConfig(t *testing.T) {
+	getContainer := func(hostConfig *string) *Container {
+		c := &Container{
+			Name: "c",
+		}
+		c.DockerConfig.HostConfig = hostConfig
+		return c
+	}
+
+	getStrPtr := func(s string) *string {
+		return &s
+	}
+
+	testCases := []struct {
+		name           string
+		container      *Container
+		expectedOutput int64
+	}{
+		{
+			name:           "happy case",
+			container:      getContainer(getStrPtr("{\"MemoryReservation\":50}")),
+			expectedOutput: 50,
+		},
+		{
+			name:           "invalid case",
+			container:      getContainer(getStrPtr("invalid")),
+			expectedOutput: 0,
+		},
+		{
+			name:           "nil case",
+			container:      getContainer(nil),
+			expectedOutput: 0,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Equal(t, tc.expectedOutput, tc.container.GetMemoryReservationFromHostConfig())
+		})
+	}
+}
+
 func TestShouldCreateWithEnvfiles(t *testing.T) {
 	cases := []struct {
 		in  Container

diff --git a/agent/api/task/task.go b/agent/api/task/task.go
@@ -1246,7 +1246,7 @@ func (task *Task) initializeFirelensResource(config *config.Config, resourceFiel
 
 	for _, container := range task.Containers {
 		firelensConfig := container.GetFirelensConfig()
-		if firelensConfig != nil {
+		if firelensConfig != nil { // is Firelens container
 			var ec2InstanceID string
 			if container.Environment != nil && container.Environment[awsExecutionEnvKey] == ec2ExecutionEnv {
 				ec2InstanceID = resourceFields.EC2InstanceID
@@ -1262,7 +1262,7 @@ func (task *Task) initializeFirelensResource(config *config.Config, resourceFiel
 			}
 			firelensResource, err := firelens.NewFirelensResource(config.Cluster, task.Arn, task.Family+":"+task.Version,
 				ec2InstanceID, config.DataDir, firelensConfig.Type, config.AWSRegion, networkMode, firelensConfig.Options, containerToLogOptions,
-				credentialsManager, task.ExecutionCredentialsID)
+				credentialsManager, task.ExecutionCredentialsID, container.GetMemoryReservationFromHostConfig())
 			if err != nil {
 				return errors.Wrap(err, "unable to initialize firelens resource")
 			}

diff --git a/agent/taskresource/firelens/firelens_unimplemented.go b/agent/taskresource/firelens/firelens_unimplemented.go
@@ -53,7 +53,7 @@ type FirelensResource struct{}
 // NewFirelensResource returns a new FirelensResource.
 func NewFirelensResource(cluster, taskARN, taskDefinition, ec2InstanceID, dataDir, firelensConfigType, region, networkMode string,
 	firelensOptions map[string]string, containerToLogOptions map[string]map[string]string, credentialsManager credentials.Manager,
-	executionCredentialsID string) (*FirelensResource, error) {
+	executionCredentialsID string, containerMemoryReservation int64) (*FirelensResource, error) {
 	return nil, errors.New("not implemented")
 }
 

diff --git a/agent/taskresource/firelens/firelens_unix.go b/agent/taskresource/firelens/firelens_unix.go
@@ -67,22 +67,23 @@ const (
 // FirelensResource models fluentd/fluentbit firelens container related resources as a task resource.
 type FirelensResource struct {
 	// Fields that are specific to firelens resource. They are only set at initialization so are not protected by lock.
-	cluster                string
-	taskARN                string
-	taskDefinition         string
-	ec2InstanceID          string
-	resourceDir            string
-	firelensConfigType     string
-	region                 string
-	ecsMetadataEnabled     bool
-	containerToLogOptions  map[string]map[string]string
-	credentialsManager     credentials.Manager
-	executionCredentialsID string
-	externalConfigType     string
-	externalConfigValue    string
-	networkMode            string
-	ioutil                 ioutilwrapper.IOUtil
-	s3ClientCreator        factory.S3ClientCreator
+	cluster                    string
+	taskARN                    string
+	taskDefinition             string
+	ec2InstanceID              string
+	resourceDir                string
+	firelensConfigType         string
+	region                     string
+	ecsMetadataEnabled         bool
+	containerToLogOptions      map[string]map[string]string
+	credentialsManager         credentials.Manager
+	executionCredentialsID     string
+	externalConfigType         string
+	externalConfigValue        string
+	networkMode                string
+	ioutil                     ioutilwrapper.IOUtil
+	s3ClientCreator            factory.S3ClientCreator
+	containerMemoryReservation int64
 
 	// Fields for the common functionality of task resource. Access to these fields are protected by lock.
 	createdAtUnsafe     time.Time
@@ -98,20 +99,21 @@ type FirelensResource struct {
 // NewFirelensResource returns a new FirelensResource.
 func NewFirelensResource(cluster, taskARN, taskDefinition, ec2InstanceID, dataDir, firelensConfigType, region, networkMode string,
 	firelensOptions map[string]string, containerToLogOptions map[string]map[string]string, credentialsManager credentials.Manager,
-	executionCredentialsID string) (*FirelensResource, error) {
+	executionCredentialsID string, containerMemoryReservation int64) (*FirelensResource, error) {
 	firelensResource := &FirelensResource{
-		cluster:                cluster,
-		taskARN:                taskARN,
-		taskDefinition:         taskDefinition,
-		ec2InstanceID:          ec2InstanceID,
-		firelensConfigType:     firelensConfigType,
-		region:                 region,
-		networkMode:            networkMode,
-		containerToLogOptions:  containerToLogOptions,
-		ioutil:                 ioutilwrapper.NewIOUtil(),
-		s3ClientCreator:        factory.NewS3ClientCreator(),
-		executionCredentialsID: executionCredentialsID,
-		credentialsManager:     credentialsManager,
+		cluster:                    cluster,
+		taskARN:                    taskARN,
+		taskDefinition:             taskDefinition,
+		ec2InstanceID:              ec2InstanceID,
+		firelensConfigType:         firelensConfigType,
+		region:                     region,
+		networkMode:                networkMode,
+		containerToLogOptions:      containerToLogOptions,
+		ioutil:                     ioutilwrapper.NewIOUtil(),
+		s3ClientCreator:            factory.NewS3ClientCreator(),
+		executionCredentialsID:     executionCredentialsID,
+		credentialsManager:         credentialsManager,
+		containerMemoryReservation: containerMemoryReservation,
 	}
 
 	fields := strings.Split(taskARN, "/")