localenv/collector/collector.go

package collector

import (
	"context"
	"errors"
	"io"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"localenv/utils"

	"github.com/docker/docker/client"
	"github.com/rs/zerolog/log"
)

const HostImagesStore string = "/tmp/localenv-swarm"

var ErrImageNotFound = errors.New("image not found")

func imageNameIntoFilename(name string) string {
	name = strings.Replace(name, "/", "-", -1)
	return strings.Replace(name, ":", "-", -1)
}

type collectorErr struct {
	errs []error
	sync.Mutex
}

func newCollectorError() collectorErr {
	return collectorErr{
		errs: []error{},
	}
}

func (c *collectorErr) AddError(err error) {
	c.Lock()
	defer c.Unlock()

	c.errs = append(c.errs, err)
}

func (c *collectorErr) Err() error {
	if len(c.errs) != 0 {
		return errors.Join(c.errs...)
	}

	return nil
}

type Collector struct {
	hostCLI  *client.Client
	swarmCLI *client.Client
}

func NewCollector(hostCLI, swarmCLI *client.Client) Collector {
	return Collector{
		hostCLI:  hostCLI,
		swarmCLI: swarmCLI,
	}
}

func (c Collector) saveImage(ctx context.Context, name, filepath string) error {
	log.Info().Str("image", name).Str("path", filepath).Msg("saving image into store...")

	imageSaveResponse, err := c.hostCLI.ImageSave(ctx, []string{name})
	if err != nil {
		return err
	}
	defer imageSaveResponse.Close()

	tarFile, err := os.Create(filepath)
	if err != nil {
		return err
	}
	defer tarFile.Close()

	content, err := io.ReadAll(imageSaveResponse)
	if err != nil {
		return err
	}

	if _, err := tarFile.Write(content); err != nil {
		return err
	}

	log.Info().Str("image", name).Str("path", filepath).Msg("image saved successfully")
	return nil
}

func (c Collector) loadImage(ctx context.Context, filepath string) error {
	log.Info().Str("path", filepath).Msg("loading image...")

	tarFile, err := os.Open(filepath)
	if err != nil {
		return err
	}
	defer tarFile.Close()

	loadResponse, err := c.swarmCLI.ImageLoad(ctx, tarFile, false)
	if err != nil {
		return err
	}
	defer loadResponse.Body.Close()

	log.Info().Str("path", filepath).Msg("image loaded successfully")
	return nil
}

func (c Collector) retryCheckImage(ctx context.Context, name string) error {
	fnRetry := func(ctx context.Context) error {
		ok, err := c.checkImage(ctx, name)
		if err != nil {
			log.Warn().Str("image", name).Msg("error while searching image, retrying...")
			return err
		}

		if !ok {
			log.Warn().Str("image", name).Msg("image not found, retrying...")
			return ErrImageNotFound
		}

		return nil
	}

	attempts := 30
	waitDuration := 2 * time.Second
	if err := utils.Retry(ctx, fnRetry, waitDuration, attempts); err != nil {
		return err
	}

	return nil
}

func (c Collector) checkImage(ctx context.Context, name string) (bool, error) {
	images, err := utils.FilterImagesByName(ctx, c.swarmCLI, name)
	if err != nil {
		return false, err
	}

	if len(images) != 1 {
		return false, nil
	}

	return true, nil
}

func (c Collector) DeployImages(ctx context.Context, images []string) error {
	errs := newCollectorError()

	var wg sync.WaitGroup
	wg.Add(len(images))

	for i := range images {
		go func(idx int) {
			defer wg.Done()

			ok, err := c.checkImage(ctx, images[idx])
			if err != nil {
				errs.AddError(err)
			}

			if ok {
				return
			}

			filepath := path.Join(HostImagesStore, imageNameIntoFilename(images[idx])+".tar")

			if err = c.saveImage(ctx, images[idx], filepath); err != nil {
				errs.AddError(err)
			}

			if err = c.loadImage(ctx, filepath); err != nil {
				errs.AddError(err)
			}

			if err := c.retryCheckImage(ctx, images[idx]); err != nil {
				errs.AddError(err)
			}
		}(i)
	}

	wg.Wait()

	return errs.Err()
}