Skip to content

Commit

Permalink
Merge pull request #804 from forta-network/caner/forta-1216-remove-im…
Browse files Browse the repository at this point in the history
…ages-less-aggressively

Implement image cleanup and separate it from bot teardown
  • Loading branch information
canercidam authored Aug 22, 2023
2 parents 550e8d2 + ca387b7 commit ec12e38
Show file tree
Hide file tree
Showing 16 changed files with 335 additions and 31 deletions.
1 change: 1 addition & 0 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ mocks:
mockgen -source services/components/lifecycle/bot_manager.go -destination services/components/lifecycle/mocks/mock_bot_manager.go
mockgen -source services/components/lifecycle/bot_monitor.go -destination services/components/lifecycle/mocks/mock_bot_monitor.go
mockgen -source services/components/containers/bot_client.go -destination services/components/containers/mocks/mock_bot_client.go
mockgen -source services/components/containers/image_cleanup.go -destination services/components/containers/mocks/mock_image_cleanup.go
mockgen -source services/components/botio/sender.go -destination services/components/botio/mocks/mock_sender.go
mockgen -source services/components/botio/bot_client.go -destination services/components/botio/mocks/mock_bot_client.go
mockgen -source services/components/botio/bot_client_factory.go -destination services/components/botio/mocks/mock_bot_client_factory.go
Expand Down
12 changes: 12 additions & 0 deletions clients/docker/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -786,6 +786,18 @@ func (d *dockerClient) EnsureLocalImages(ctx context.Context, timeoutPerPull tim
return
}

// ListDigestReferences lists all digest references of all images.
func (d *dockerClient) ListDigestReferences(ctx context.Context) (imgs []string, err error) {
imgSummaries, err := d.cli.ImageList(ctx, types.ImageListOptions{})
if err != nil {
return nil, fmt.Errorf("failed to list locally available images: %v", err)
}
for _, imgSummary := range imgSummaries {
imgs = append(imgs, imgSummary.RepoDigests...)
}
return
}

// GetContainerLogs gets the container logs.
func (d *dockerClient) GetContainerLogs(ctx context.Context, containerID, tail string, truncate int) (string, error) {
r, err := d.cli.ContainerLogs(ctx, containerID, types.ContainerLogsOptions{
Expand Down
1 change: 1 addition & 0 deletions clients/interfaces.go
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ type DockerClient interface {
HasLocalImage(ctx context.Context, ref string) (bool, error)
EnsureLocalImage(ctx context.Context, name, ref string) error
EnsureLocalImages(ctx context.Context, timeoutPerPull time.Duration, imagePulls []docker.ImagePull) []error
ListDigestReferences(ctx context.Context) ([]string, error)
GetContainerLogs(ctx context.Context, containerID, tail string, truncate int) (string, error)
GetContainerFromRemoteAddr(ctx context.Context, hostPort string) (*types.Container, error)
SetImagePullCooldown(threshold int, cooldownDuration time.Duration)
Expand Down
15 changes: 15 additions & 0 deletions clients/mocks/mock_clients.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

11 changes: 7 additions & 4 deletions services/components/components.go
Original file line number Diff line number Diff line change
Expand Up @@ -76,8 +76,9 @@ type BotLifecycleConfig struct {

// BotLifecycle contains the bot lifecycle components.
type BotLifecycle struct {
BotManager lifecycle.BotLifecycleManager
BotClient containers.BotClient
BotManager lifecycle.BotLifecycleManager
BotClient containers.BotClient
ImageCleanup containers.ImageCleanup
}

// GetBotLifecycleComponents returns the bot lifecycle management components.
Expand Down Expand Up @@ -118,9 +119,11 @@ func GetBotLifecycleComponents(ctx context.Context, botLifeConfig BotLifecycleCo
botLifeConfig.BotRegistry, botClient, lifecycleMediator,
lifecycleMetrics, botMonitor,
)
imageCleanup := containers.NewImageCleanup(dockerClient, botLifeConfig.BotRegistry)

return BotLifecycle{
BotManager: botManager,
BotClient: botClient,
BotManager: botManager,
BotClient: botClient,
ImageCleanup: imageCleanup,
}, nil
}
14 changes: 3 additions & 11 deletions services/components/containers/bot_client.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ const (
type BotClient interface {
EnsureBotImages(ctx context.Context, botConfigs []config.AgentConfig) []error
LaunchBot(ctx context.Context, botConfig config.AgentConfig) error
TearDownBot(ctx context.Context, containerName string, removeImage bool) error
TearDownBot(ctx context.Context, containerName string) error
StopBot(ctx context.Context, botConfig config.AgentConfig) error
LoadBotContainers(ctx context.Context) ([]types.Container, error)
StartWaitBotContainer(ctx context.Context, containerID string) error
Expand Down Expand Up @@ -139,8 +139,8 @@ func getServiceContainerNames() []string {
}
}

// TearDownBot tears down a bot by shutting down the docker container and removing it.
func (bc *botClient) TearDownBot(ctx context.Context, containerName string, removeImage bool) error {
// TearDownBot tears down a bot removing the docker container and network.
func (bc *botClient) TearDownBot(ctx context.Context, containerName string) error {
container, err := bc.client.GetContainerByName(ctx, containerName)
if err != nil {
return fmt.Errorf("failed to get the bot container to tear down: %v", err)
Expand Down Expand Up @@ -191,14 +191,6 @@ func (bc *botClient) TearDownBot(ctx context.Context, containerName string, remo
},
).WithError(err).Warn("failed to destroy the bot network")
}
if !removeImage {
return nil
}
if err := bc.client.RemoveImage(ctx, container.Image); err != nil {
log.WithFields(log.Fields{
"image": container.Image,
}).WithError(err).Warn("failed to remove image of the destroyed bot container")
}
return nil
}

Expand Down
3 changes: 1 addition & 2 deletions services/components/containers/bot_client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -154,9 +154,8 @@ func (s *BotClientTestSuite) TestTearDownBot() {
s.client.EXPECT().ShutdownContainer(gomock.Any(), testContainerID2, &timeout).Return(testErr)
s.client.EXPECT().RemoveContainer(gomock.Any(), testContainerID2).Return(testErr)
s.client.EXPECT().RemoveNetworkByName(gomock.Any(), botConfig.ContainerName()).Return(testErr)
s.client.EXPECT().RemoveImage(gomock.Any(), testImageRef).Return(testErr)

s.r.NoError(s.botClient.TearDownBot(context.Background(), botConfig.ContainerName(), true))
s.r.NoError(s.botClient.TearDownBot(context.Background(), botConfig.ContainerName()))
}

func (s *BotClientTestSuite) TestStopBot() {
Expand Down
111 changes: 111 additions & 0 deletions services/components/containers/image_cleanup.go
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
}
112 changes: 112 additions & 0 deletions services/components/containers/image_cleanup_test.go
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)
}
8 changes: 4 additions & 4 deletions services/components/containers/mocks/mock_bot_client.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit ec12e38

Please sign in to comment.