-
Notifications
You must be signed in to change notification settings - Fork 151
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #804 from forta-network/caner/forta-1216-remove-im…
…ages-less-aggressively Implement image cleanup and separate it from bot teardown
- Loading branch information
Showing
16 changed files
with
335 additions
and
31 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,111 @@ | ||
package containers | ||
|
||
import ( | ||
"context" | ||
"fmt" | ||
"strings" | ||
"time" | ||
|
||
"github.com/forta-network/forta-node/clients" | ||
"github.com/forta-network/forta-node/clients/docker" | ||
"github.com/forta-network/forta-node/services/components/registry" | ||
"github.com/sirupsen/logrus" | ||
) | ||
|
||
const imageCleanupInterval = time.Hour * 1 | ||
|
||
// ImageCleanup deals with image cleanup. | ||
type ImageCleanup interface { | ||
Do(context.Context) error | ||
} | ||
|
||
type imageCleanup struct { | ||
client clients.DockerClient | ||
botRegistry registry.BotRegistry | ||
lastCleanup time.Time | ||
exclusionList []string | ||
} | ||
|
||
// NewImageCleanup creates new. | ||
func NewImageCleanup(client clients.DockerClient, botRegistry registry.BotRegistry, excludeImages ...string) *imageCleanup { | ||
return &imageCleanup{ | ||
client: client, | ||
botRegistry: botRegistry, | ||
exclusionList: excludeImages, | ||
} | ||
} | ||
|
||
// Do does the image cleanup by finding all unused Disco images and removing them. | ||
// The logic executes only after an interval. | ||
func (ic *imageCleanup) Do(ctx context.Context) error { | ||
if time.Since(ic.lastCleanup) < imageCleanupInterval { | ||
return nil | ||
} | ||
|
||
containers, err := ic.client.GetContainers(ctx) | ||
if err != nil { | ||
return fmt.Errorf("failed to get containers during image cleanup: %v", err) | ||
} | ||
|
||
// we list the digest references as the main references for all images | ||
// because we pull by digest references | ||
images, err := ic.client.ListDigestReferences(ctx) | ||
if err != nil { | ||
return fmt.Errorf("failed to list images during image cleanup: %v", err) | ||
} | ||
|
||
heartbeatBot, err := ic.botRegistry.LoadHeartbeatBot() | ||
if err != nil { | ||
return fmt.Errorf("failed to load the heartbeat bot during cleanup: %v", err) | ||
} | ||
|
||
logrus.WithField("image", heartbeatBot.Image).Debug("cleanup: loaded heartbeat bot image reference") | ||
|
||
for _, image := range images { | ||
logger := logrus.WithField("image", image) | ||
|
||
if ic.isExcluded(image) || image == heartbeatBot.Image { | ||
logger.Debug("image is excluded - skipping cleanup") | ||
continue | ||
} | ||
|
||
if ic.isImageInUse(containers, image) { | ||
logger.Debug("image is in use - skipping cleanup") | ||
continue | ||
} | ||
|
||
if err := ic.client.RemoveImage(ctx, image); err != nil { | ||
logger.WithError(err).Warn("failed to cleanup unused disco image") | ||
} else { | ||
logger.Info("successfully cleaned up unused image") | ||
} | ||
} | ||
|
||
ic.lastCleanup = time.Now() | ||
return nil | ||
} | ||
|
||
func (ic *imageCleanup) isExcluded(ref string) bool { | ||
// needs to be a Disco image | ||
if !strings.Contains(ref, "bafybei") { | ||
return true | ||
} | ||
|
||
for _, excluded := range ic.exclusionList { | ||
// expecting the ref to include the excluded ref because | ||
// we specify it and it can be a subset of the full reference | ||
if strings.Contains(ref, excluded) { | ||
return true | ||
} | ||
} | ||
return false | ||
} | ||
|
||
func (lc *imageCleanup) isImageInUse(containers docker.ContainerList, image string) bool { | ||
for _, container := range containers { | ||
if container.Image == image { | ||
return true | ||
} | ||
} | ||
return false | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,112 @@ | ||
package containers | ||
|
||
import ( | ||
"context" | ||
"errors" | ||
"testing" | ||
"time" | ||
|
||
"github.com/forta-network/forta-node/clients/docker" | ||
mock_clients "github.com/forta-network/forta-node/clients/mocks" | ||
"github.com/forta-network/forta-node/config" | ||
mock_registry "github.com/forta-network/forta-node/services/components/registry/mocks" | ||
"github.com/golang/mock/gomock" | ||
"github.com/sirupsen/logrus" | ||
"github.com/stretchr/testify/require" | ||
"github.com/stretchr/testify/suite" | ||
) | ||
|
||
const ( | ||
testCleanupImage1 = "bafybei-testcleanupimage1" | ||
testCleanupImage2 = "bafybei-testcleanupimage2" | ||
testCleanupImage3 = "bafybei-testcleanupimage3" | ||
testHeartbeatBotImage = "bafybei-heartbeatbot" | ||
) | ||
|
||
type ImageCleanupTestSuite struct { | ||
r *require.Assertions | ||
|
||
client *mock_clients.MockDockerClient | ||
botRegistry *mock_registry.MockBotRegistry | ||
|
||
imageCleanup *imageCleanup | ||
|
||
suite.Suite | ||
} | ||
|
||
func TestImageCleanupTestSuite(t *testing.T) { | ||
suite.Run(t, &ImageCleanupTestSuite{}) | ||
} | ||
|
||
func (s *ImageCleanupTestSuite) SetupTest() { | ||
logrus.SetLevel(logrus.DebugLevel) | ||
|
||
s.r = s.Require() | ||
|
||
ctrl := gomock.NewController(s.T()) | ||
s.client = mock_clients.NewMockDockerClient(ctrl) | ||
s.botRegistry = mock_registry.NewMockBotRegistry(ctrl) | ||
|
||
s.imageCleanup = NewImageCleanup(s.client, s.botRegistry) | ||
} | ||
|
||
func (s *ImageCleanupTestSuite) TestIntervalSkip() { | ||
s.imageCleanup.lastCleanup = time.Now() // very close cleanup time | ||
|
||
// no calls expected | ||
|
||
s.r.NoError(s.imageCleanup.Do(context.Background())) | ||
} | ||
|
||
func (s *ImageCleanupTestSuite) TestContainerListError() { | ||
s.client.EXPECT().GetContainers(gomock.Any()).Return(nil, errors.New("test error")) | ||
|
||
s.r.Error(s.imageCleanup.Do(context.Background())) | ||
} | ||
|
||
func (s *ImageCleanupTestSuite) TestImagesListError() { | ||
s.client.EXPECT().GetContainers(gomock.Any()).Return(docker.ContainerList{}, nil) | ||
s.client.EXPECT().ListDigestReferences(gomock.Any()).Return(nil, errors.New("test error")) | ||
|
||
s.r.Error(s.imageCleanup.Do(context.Background())) | ||
} | ||
|
||
func (s *ImageCleanupTestSuite) TestHeartbeatBotError() { | ||
s.client.EXPECT().GetContainers(gomock.Any()).Return(docker.ContainerList{}, nil) | ||
s.client.EXPECT().ListDigestReferences(gomock.Any()).Return([]string{testCleanupImage1}, nil) | ||
s.botRegistry.EXPECT().LoadHeartbeatBot().Return(nil, errors.New("test error")) | ||
|
||
s.r.Error(s.imageCleanup.Do(context.Background())) | ||
} | ||
|
||
func (s *ImageCleanupTestSuite) TestRemoveImageError() { | ||
initialLastCleanup := s.imageCleanup.lastCleanup | ||
|
||
s.client.EXPECT().GetContainers(gomock.Any()).Return(docker.ContainerList{}, nil) | ||
s.client.EXPECT().ListDigestReferences(gomock.Any()).Return([]string{testCleanupImage1}, nil) | ||
s.botRegistry.EXPECT().LoadHeartbeatBot().Return(&config.AgentConfig{Image: testHeartbeatBotImage}, nil) | ||
s.client.EXPECT().RemoveImage(gomock.Any(), testCleanupImage1).Return(errors.New("test error")) | ||
|
||
// no error and mutated last cleanup timestamp: removal errors do not affect this | ||
s.r.NoError(s.imageCleanup.Do(context.Background())) | ||
s.r.NotEqual(initialLastCleanup, s.imageCleanup.lastCleanup) | ||
} | ||
|
||
func (s *ImageCleanupTestSuite) TestCleanupSuccess() { | ||
initialLastCleanup := s.imageCleanup.lastCleanup | ||
|
||
s.imageCleanup.exclusionList = []string{testCleanupImage1} // excluding image | ||
s.client.EXPECT().GetContainers(gomock.Any()).Return(docker.ContainerList{ | ||
{ | ||
Image: testCleanupImage2, // image in use by container | ||
}, | ||
}, nil) | ||
s.client.EXPECT().ListDigestReferences(gomock.Any()).Return( | ||
[]string{testCleanupImage1, testCleanupImage2, testCleanupImage3, testHeartbeatBotImage}, nil, | ||
) | ||
s.botRegistry.EXPECT().LoadHeartbeatBot().Return(&config.AgentConfig{Image: testHeartbeatBotImage}, nil) | ||
s.client.EXPECT().RemoveImage(gomock.Any(), testCleanupImage3).Return(nil) // only removes image 3 | ||
|
||
s.r.NoError(s.imageCleanup.Do(context.Background())) | ||
s.r.NotEqual(initialLastCleanup, s.imageCleanup.lastCleanup) | ||
} |
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
Oops, something went wrong.