some enhancement and unit tests passing

danehlim · Jun 6, 2024 · f375a3c · f375a3c
1 parent 7172f3b
commit f375a3c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 12 deletions.
diff --git a/agent/stats/container.go b/agent/stats/container.go
@@ -173,15 +173,16 @@ func (container *StatsContainer) processStatsStream() error {
 				return err
 			}
 
-			isFirstStatAfterContainerRestart := container.syncContainerRestartAggregationData(rawStat)
+			isFirstStatAfterContainerRestart := apiContainer != nil && apiContainer.RestartPolicyEnabled() &&
+				container.syncContainerRestartAggregationData(rawStat)
 			err = container.statsQueue.AddContainerStat(rawStat, getNonDockerContainerStats(apiContainer),
 				&container.restartAggregationData.LastStatBeforeLastRestart, container.hasRestartedBefore())
 			if err != nil {
 				seelog.Warnf("Container [%s]: error converting stats for container: %v", dockerID, err)
 				continue
 			}
 			if isFirstStatAfterContainerRestart {
-				err = container.saveRestartAggregationData()
+				err = container.saveRestartAggregationData(apiContainer)
 				if err != nil {
 					logger.Error("Failed to update container's stats restart aggregation data in database",
 						logger.Fields{
@@ -227,15 +228,9 @@ func (container *StatsContainer) syncContainerRestartAggregationData(stat *types
 }
 
 // saveRestartAggregationData is used to save a container's restart aggregation data in Agent's local state database.
-func (container *StatsContainer) saveRestartAggregationData() error {
-	dockerID := container.containerMetadata.DockerID
-	agentContainer, err := container.resolver.ResolveContainer(dockerID)
-	if err != nil {
-		return errors.Wrap(err, "failed to resolve container via container metadata resolver")
-	}
-
-	agentContainer.Container.SetRestartAggregationDataForStats(container.restartAggregationData)
-	err = container.dataClient.SaveContainer(agentContainer.Container)
+func (container *StatsContainer) saveRestartAggregationData(apiContainer *apicontainer.Container) error {
+	apiContainer.SetRestartAggregationDataForStats(container.restartAggregationData)
+	err := container.dataClient.SaveContainer(apiContainer)
 	if err != nil {
 		return errors.Wrap(err, "failed to save data for container")
 	}

diff --git a/agent/stats/container_test.go b/agent/stats/container_test.go
@@ -87,7 +87,6 @@ func TestContainerStatsCollection(t *testing.T) {
 	errC := make(chan error)
 	numStats := 8
 	mockDockerClient.EXPECT().Stats(ctx, dockerID, dockerclient.StatsInactivityTimeout).Return(statChan, errC)
-	mockDockerClient.EXPECT().DescribeContainer(ctx, dockerID).Times(numStats)
 	go metricSenderFunc(statChan, numStats, nil)()
 
 	container := &StatsContainer{