<commit_before>package image

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/taskcluster/slugid-go/slugid"
	"github.com/taskcluster/taskcluster-worker/engines/qemu/vm"
	"github.com/taskcluster/taskcluster-worker/runtime"
	"github.com/taskcluster/taskcluster-worker/runtime/gc"
)

// Manager loads and tracks images.
type Manager struct {
	m           sync.Mutex
	images      map[string]*image
	imageFolder string
	gc          gc.ResourceTracker
	monitor     runtime.Monitor
}

// Downloader is a function capable of downloading an image to an *os.File.
// The downloader writes the image file to the imageFile supplied, and returns
// an error if all retries etc. fails.
type Downloader func(imageFile *os.File) error

// image represents an image of which multiple instances can be created
type image struct {
	gc.DisposableResource
	imageID string
	folder  string
	machine *vm.Machine
	done    <-chan struct{}
	manager *Manager
	err     error
}

// Instance represents an instance of an image.
type Instance struct {
	m        sync.Mutex
	image    *image
	diskFile string
}

// NewManager creates a new image manager using the imageFolder for storing
// images and instances of images.
func NewManager(imageFolder string, gc gc.ResourceTracker, monitor runtime.Monitor) (*Manager, error) {
	// Ensure the image folder is created
	err := os.MkdirAll(imageFolder, 0777)
	if err != nil {
		return nil, fmt.Errorf("Failed to create imageFolder: %s, error: %s", imageFolder, err)
	}
	return &Manager{
		images:      make(map[string]*image),
		imageFolder: imageFolder,
		gc:          gc,
		monitor:     monitor,
	}, nil
}

// Instance will return an Instance of the image with imageID. If no such
// image exists in the cache, download() will be called to download it to a
// temporary filename.
//
// This method will insert the downloaded image into the cache, and ensures that
// if won't be downloaded twice, if another invocation already is downloading
// an image with the same imageID.
//
// It is the responsibility of the caller to make sure that imageID is a string
// that uniquely identifies the image. Sane patterns includes "url:<url>", or
// "taskId:<taskId>/<runId>/<artifact>". It also the callers responsibility to
// enforce any sort of access control.
func (m *Manager) Instance(imageID string, download Downloader) (*Instance, error) {
	m.m.Lock()

	// Get image from cache and insert it if not present
	img := m.images[imageID]
	if img == nil {
		imageDone := make(chan struct{})
		img = &image{
			imageID: imageID,
			folder:  filepath.Join(m.imageFolder, slugid.Nice()),
			done:    imageDone,
			manager: m,
		}
		m.images[imageID] = img
		// Start loading the image
		go img.loadImage(download, imageDone)
	}

	// Acqure the image, so we can release lock without risking the image gets
	// garbage collected.
	img.Acquire()
	m.m.Unlock() // Release lock we don't need it anymore

	// Wait for image to be done, then either return the error, or return an
	// instance of the image.
	<-img.done
	if img.err != nil {
		img.Release()
		return nil, img.err
	}
	return img.instance()
}

func (img *image) loadImage(download Downloader, done chan<- struct{}) {
	imageFilePath := filepath.Join(img.manager.imageFolder, slugid.Nice()+".tar.zst")
	var imageFile *os.File

	// Create image folder
	err := os.Mkdir(img.folder, 0777)
	if err != nil {
		goto cleanup
	}

	// Create image file
	imageFile, err = os.Create(imageFilePath)
	if err != nil {
		goto cleanup
	}

	// Download image to tempory file
	err = download(imageFile)
	if err != nil {
		goto cleanup
	}

	// close image file
	err = imageFile.Close()
	imageFile = nil // don't close twice
	if err != nil {
		goto cleanup<commit_msg>Wrapping built-in errors to better error messaging<commit_after>package image

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/pkg/errors"
	"github.com/taskcluster/slugid-go/slugid"
	"github.com/taskcluster/taskcluster-worker/engines/qemu/vm"
	"github.com/taskcluster/taskcluster-worker/runtime"
	"github.com/taskcluster/taskcluster-worker/runtime/gc"
)

// Manager loads and tracks images.
type Manager struct {
	m           sync.Mutex
	images      map[string]*image
	imageFolder string
	gc          gc.ResourceTracker
	monitor     runtime.Monitor
}

// Downloader is a function capable of downloading an image to an *os.File.
// The downloader writes the image file to the imageFile supplied, and returns
// an error if all retries etc. fails.
type Downloader func(imageFile *os.File) error

// image represents an image of which multiple instances can be created
type image struct {
	gc.DisposableResource
	imageID string
	folder  string
	machine *vm.Machine
	done    <-chan struct{}
	manager *Manager
	err     error
}

// Instance represents an instance of an image.
type Instance struct {
	m        sync.Mutex
	image    *image
	diskFile string
}

// NewManager creates a new image manager using the imageFolder for storing
// images and instances of images.
func NewManager(imageFolder string, gc gc.ResourceTracker, monitor runtime.Monitor) (*Manager, error) {
	// Ensure the image folder is created
	err := os.MkdirAll(imageFolder, 0777)
	if err != nil {
		return nil, fmt.Errorf("Failed to create imageFolder: %s, error: %s", imageFolder, err)
	}
	return &Manager{
		images:      make(map[string]*image),
		imageFolder: imageFolder,
		gc:          gc,
		monitor:     monitor,
	}, nil
}

// Instance will return an Instance of the image with imageID. If no such
// image exists in the cache, download() will be called to download it to a
// temporary filename.
//
// This method will insert the downloaded image into the cache, and ensures that
// if won't be downloaded twice, if another invocation already is downloading
// an image with the same imageID.
//
// It is the responsibility of the caller to make sure that imageID is a string
// that uniquely identifies the image. Sane patterns includes "url:<url>", or
// "taskId:<taskId>/<runId>/<artifact>". It also the callers responsibility to
// enforce any sort of access control.
func (m *Manager) Instance(imageID string, download Downloader) (*Instance, error) {
	m.m.Lock()

	// Get image from cache and insert it if not present
	img := m.images[imageID]
	if img == nil {
		imageDone := make(chan struct{})
		img = &image{
			imageID: imageID,
			folder:  filepath.Join(m.imageFolder, slugid.Nice()),
			done:    imageDone,
			manager: m,
		}
		m.images[imageID] = img
		// Start loading the image
		go img.loadImage(download, imageDone)
	}

	// Acqure the image, so we can release lock without risking the image gets
	// garbage collected.
	img.Acquire()
	m.m.Unlock() // Release lock we don't need it anymore

	// Wait for image to be done, then either return the error, or return an
	// instance of the image.
	<-img.done
	if img.err != nil {
		img.Release()
		return nil, img.err
	}
	return img.instance()
}

func (img *image) loadImage(download Downloader, done chan<- struct{}) {
	imageFilePath := filepath.Join(img.manager.imageFolder, slugid.Nice()+".tar.zst")
	var imageFile *os.File

	// Create image folder
	err := os.Mkdir(img.folder, 0777)
	if err != nil {
		err = errors.Wrap(err, "failed to create image folder")
		goto cleanup
	}

	// Create image file
	imageFile, err = os.Create(imageFilePath)
	if err != nil {
		err = errors.Wrap(err, "failed to create image file")
		goto cleanup
	}

	// Download image to tempory file
	err = download(imageFile)
	if err != nil {
		goto cleanup
	}

	// close image file
	err = imageFile.Close()
	imageFile = nil // don't close twice
	if err != nil {
		err = errors.Wrap(err, "failed to close image file")
		goto cleanup
{'fraction_non_alphanumeric': 0.10256744428515574, 'fraction_numerical': 0.0020852339371823275, 'mean_word_length': 3.41542002301496, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 32, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4095705', 'n_tokens_mistral': 2448, 'n_tokens_neox': 2342, 'n_words': 883}