set firelens mem_buf_limit

yinyic · Oct 22, 2024 · 8bef39f · 8bef39f
1 parent 79f17a5
commit 8bef39f
Show file tree

Hide file tree

Showing 7 changed files with 195 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 int64(hostConfig.MemoryReservation)
+}
+
 // GetHostConfig returns the container's host config.
 func (c *Container) GetHostConfig() *string {
 	c.lock.RLock()

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, "/")

diff --git a/agent/taskresource/firelens/firelens_unix_test.go b/agent/taskresource/firelens/firelens_unix_test.go
@@ -42,18 +42,19 @@ import (
 )
 
 const (
-	testCluster                = "mycluster"
-	testTaskARN                = "arn:aws:ecs:us-east-2:01234567891011:task/mycluster/3de392df-6bfa-470b-97ed-aa6f482cd7a"
-	testTaskDefinition         = "taskdefinition:1"
-	testEC2InstanceID          = "i-123456789a"
-	testDataDir                = "testdatadir"
-	testResourceDir            = "testresourcedir"
-	testTerminalResason        = "testterminalreason"
-	testTempFile               = "testtempfile"
-	testRegion                 = "us-west-2"
-	testExecutionCredentialsID = "testexecutioncredentialsid"
-	testExternalConfigType     = "testexternalconfigtype"
-	testExternalConfigValue    = "testexternalconfigvalue"
+	testCluster                    = "mycluster"
+	testTaskARN                    = "arn:aws:ecs:us-east-2:01234567891011:task/mycluster/3de392df-6bfa-470b-97ed-aa6f482cd7a"
+	testTaskDefinition             = "taskdefinition:1"
+	testEC2InstanceID              = "i-123456789a"
+	testDataDir                    = "testdatadir"
+	testResourceDir                = "testresourcedir"
+	testTerminalResason            = "testterminalreason"
+	testTempFile                   = "testtempfile"
+	testRegion                     = "us-west-2"
+	testExecutionCredentialsID     = "testexecutioncredentialsid"
+	testExternalConfigType         = "testexternalconfigtype"
+	testExternalConfigValue        = "testexternalconfigvalue"
+	testContainerMemoryReservation = 100
 )
 
 var (
@@ -105,7 +106,7 @@ func mockMkdirAllError() func() {
 
 func newMockFirelensResource(firelensConfigType, networkMode string, lopOptions map[string]string,
 	mockIOUtil *mock_ioutilwrapper.MockIOUtil, mockCredentialsManager *mock_credentials.MockManager,
-	mockS3ClientCreator *mock_factory.MockS3ClientCreator) *FirelensResource {
+	mockS3ClientCreator *mock_factory.MockS3ClientCreator, containerMemoryReservation int64) *FirelensResource {
 	return &FirelensResource{
 		cluster:            testCluster,
 		taskARN:            testTaskARN,
@@ -118,10 +119,11 @@ func newMockFirelensResource(firelensConfigType, networkMode string, lopOptions
 		containerToLogOptions: map[string]map[string]string{
 			"container": lopOptions,
 		},
-		executionCredentialsID: testExecutionCredentialsID,
-		credentialsManager:     mockCredentialsManager,
-		ioutil:                 mockIOUtil,
-		s3ClientCreator:        mockS3ClientCreator,
+		executionCredentialsID:     testExecutionCredentialsID,
+		credentialsManager:         mockCredentialsManager,
+		ioutil:                     mockIOUtil,
+		s3ClientCreator:            mockS3ClientCreator,
+		containerMemoryReservation: containerMemoryReservation,
 	}
 }
 
@@ -158,7 +160,7 @@ func TestCreateFirelensResourceFluentdBridgeMode(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	defer mockRename()()
 	gomock.InOrder(
@@ -173,7 +175,7 @@ func TestCreateFirelensResourceFluentdAWSVPCMode(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, awsvpcNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	defer mockRename()()
 	gomock.InOrder(
@@ -188,7 +190,7 @@ func TestCreateFirelensResourceFluentdDefaultMode(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, "", testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	defer mockRename()()
 	gomock.InOrder(
@@ -203,7 +205,7 @@ func TestCreateFirelensResourceFluentbit(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentbit, bridgeNetworkMode, testFluentbitOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	defer mockRename()()
 	gomock.InOrder(
@@ -218,7 +220,7 @@ func TestCreateFirelensResourceInvalidType(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 	firelensResource.firelensConfigType = "invalid"
 
 	assert.Error(t, firelensResource.Create())
@@ -230,7 +232,7 @@ func TestCreateFirelensResourceCreateConfigDirError(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	defer mockMkdirAllError()()
 
@@ -243,7 +245,7 @@ func TestCreateFirelensResourceCreateSocketDirError(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	defer mockMkdirAllError()()
 
@@ -256,7 +258,7 @@ func TestCreateFirelensResourceGenerateConfigError(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 	firelensResource.containerToLogOptions = map[string]map[string]string{
 		"container": {
 			"invalid": "invalid",
@@ -272,7 +274,7 @@ func TestCreateFirelensResourceCreateTempFileError(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	gomock.InOrder(
 		mockIOUtil.EXPECT().TempFile(testResourceDir, tempFile).Return(nil, errors.New("test error")),
@@ -287,7 +289,7 @@ func TestCreateFirelensResourceWriteConfigFileError(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	mockFile.(*mock_oswrapper.MockFile).WriteImpl = func(bytes []byte) (i int, e error) {
 		return 0, errors.New("test error")
@@ -306,7 +308,7 @@ func TestCreateFirelensResourceChmodError(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	mockFile.(*mock_oswrapper.MockFile).ChmodImpl = func(mode os.FileMode) error {
 		return errors.New("test error")
@@ -325,7 +327,7 @@ func TestCreateFirelensResourceRenameError(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	gomock.InOrder(
 		mockIOUtil.EXPECT().TempFile(testResourceDir, tempFile).Return(mockFile, nil),
@@ -347,7 +349,7 @@ func TestCreateFirelensResourceWithS3Config(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	err := firelensResource.parseOptions(testFirelensOptionsS3)
 	require.NoError(t, err)
@@ -385,7 +387,7 @@ func TestCreateFirelensResourceWithS3ConfigMissingCredentials(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	err := firelensResource.parseOptions(testFirelensOptionsS3)
 	require.NoError(t, err)
@@ -403,7 +405,7 @@ func TestCreateFirelensResourceWithS3ConfigInvalidS3ARN(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	err := firelensResource.parseOptions(testFirelensOptionsS3)
 	require.NoError(t, err)
@@ -422,7 +424,7 @@ func TestCreateFirelensResourceWithS3ConfigDownloadFailure(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	err := firelensResource.parseOptions(testFirelensOptionsS3)
 	require.NoError(t, err)
@@ -450,7 +452,7 @@ func TestCleanupFirelensResource(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	assert.NoError(t, firelensResource.Cleanup())
 }
@@ -460,7 +462,7 @@ func TestCleanupFirelensResourceError(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, mockIOUtil,
-		mockCredentialsManager, mockS3ClientCreator)
+		mockCredentialsManager, mockS3ClientCreator, testContainerMemoryReservation)
 
 	removeAll = func(path string) error {
 		return errors.New("test error")
@@ -478,7 +480,7 @@ func TestInitializeFirelensResource(t *testing.T) {
 	defer done()
 
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, nil, nil,
-		nil)
+		nil, testContainerMemoryReservation)
 	firelensResource.Initialize(&taskresource.ResourceFields{
 		ResourceFieldsCommon: &taskresource.ResourceFieldsCommon{
 			CredentialsManager: mockCredentialsManager,
@@ -494,7 +496,7 @@ func TestInitializeFirelensResource(t *testing.T) {
 
 func TestSetKnownStatus(t *testing.T) {
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, nil, nil,
-		nil)
+		nil, testContainerMemoryReservation)
 	firelensResource.appliedStatusUnsafe = resourcestatus.ResourceStatus(FirelensCreated)
 
 	firelensResource.SetKnownStatus(resourcestatus.ResourceStatus(FirelensCreated))
@@ -504,7 +506,7 @@ func TestSetKnownStatus(t *testing.T) {
 
 func TestSetKnownStatusNoAppliedStatusUpdate(t *testing.T) {
 	firelensResource := newMockFirelensResource(FirelensConfigTypeFluentd, bridgeNetworkMode, testFluentdOptions, nil, nil,
-		nil)
+		nil, testContainerMemoryReservation)
 	firelensResource.appliedStatusUnsafe = resourcestatus.ResourceStatus(FirelensCreated)
 
 	firelensResource.SetKnownStatus(resourcestatus.ResourceStatus(FirelensStatusNone))