golden/go/diffstore/imgloader.go - buildbot - Git at Google

 package diffstore

 import (
 	"bytes"
 	"context"
 	"crypto/md5"
 	"errors"
 	"fmt"
 	"image"
 	"io"
 	"net/http"
 	"os"
 	"path/filepath"
 	"sync"

 	"cloud.google.com/go/storage"
 	"go.skia.org/infra/go/fileutil"
 	"go.skia.org/infra/go/rtcache"
 	"go.skia.org/infra/go/sklog"
 	"go.skia.org/infra/go/util"
 	"go.skia.org/infra/golden/go/diff"
 	"google.golang.org/api/option"
 )

 const (
 	// MAX_URI_GET_TRIES is the number of tries we do to load an image.
 	MAX_URI_GET_TRIES = 4

 	// Number of concurrent workers downloading images.
 	N_IMG_WORKERS = 10
 )

 // ImageLoader facilitates to continously download images and cache them in RAM.
 type ImageLoader struct {
 	// storageClient is the Google client to local content from GCS.
 	storageClient *storage.Client

 	// localImgDir is the local directory where images should be written to.
 	localImgDir string

 	// gsBucketNames is the list of GCS bucket where images are stored.
 	gsBucketNames []string

 	// gsImageBaseDir is the GCS directory (prefix) where images are stored.
 	gsImageBaseDir string

 	// imageCache caches and calculates images.
 	imageCache rtcache.ReadThroughCache

 	// failureStore persists failures in retrieving images.
 	failureStore *failureStore

 	wg sync.WaitGroup

 	// mapper contains various functions for creating image IDs and paths.
 	mapper DiffStoreMapper
 }

 // Creates a new instance of ImageLoader.
 func NewImgLoader(client *http.Client, baseDir, imgDir string, gsBucketNames []string, gsImageBaseDir string, maxCacheSize int, mapper DiffStoreMapper) (*ImageLoader, error) {
 	storageClient, err := storage.NewClient(context.Background(), option.WithHTTPClient(client))
 	if err != nil {
 		return nil, err
 	}

 	fStore, err := newFailureStore(filepath.Join(baseDir, FAILUREDB_NAME))
 	if err != nil {
 		return nil, err
 	}

 	ret := &ImageLoader{
 		storageClient:  storageClient,
 		localImgDir:    imgDir,
 		gsBucketNames:  gsBucketNames,
 		gsImageBaseDir: gsImageBaseDir,
 		failureStore:   fStore,
 		mapper:         mapper,
 	}

 	// Set up the work queues that balance the load.
 	if ret.imageCache, err = rtcache.New(ret.imageLoadWorker, maxCacheSize, N_IMG_WORKERS); err != nil {
 		return nil, err
 	}
 	return ret, nil
 }

 // Warm makes sure the images are cached. It does not return a result to avoid
 // deserialization and unnecessary memory allocation.
 // The synchronous flag determines whether the call is blocking or not.
 // It workes in sync with Get, any image that is scheduled be retrieved by get
 // will not be fetched again.
 func (il *ImageLoader) Warm(priority int64, images []string, synchronous bool) {
 	var localWg sync.WaitGroup
 	il.wg.Add(len(images))
 	localWg.Add(len(images))
 	for _, id := range images {
 		go func(id string) {
 			defer il.wg.Done()
 			defer localWg.Done()
 			if err := il.imageCache.Warm(priority, id); err != nil {
 				sklog.Errorf("Unable to retrive image %s. Got error: %s", id, err)
 			}
 		}(id)
 	}

 	// Wait for the all images to be loaded.
 	if synchronous {
 		localWg.Wait()
 	}
 }

 // sync waits until all pending go routines have terminated.
 func (il *ImageLoader) Sync() {
 	il.wg.Wait()
 }

 // errResult is a helper type to capture error information in Get(...).
 type errResult struct {
 	err error
 	id  string
 }

 // Get returns the images identified by digests and returns it as an NRGBA image.
 // Priority determines the order in which multiple concurrent calls are processed.
 // func (il *ImageLoader) Get(priority int64, digests []string) (*image.NRGBA, error) {
 // 	// img, err := il.imageCache.Get(priority, digests)
 // 	// if err != nil {
 // 	// 	return nil, err
 // 	// }
 // 	// return img.(*image.NRGBA), nil
 // 	return nil, nil
 // }

 func (il *ImageLoader) Get(priority int64, images []string) ([]*image.NRGBA, error) {
 	// Parallel load the requested images.
 	result := make([]*image.NRGBA, len(images))
 	errCh := make(chan errResult, len(images))
 	var wg sync.WaitGroup
 	wg.Add(len(images))
 	for idx, id := range images {
 		go func(idx int, id string) {
 			defer wg.Done()
 			img, err := il.imageCache.Get(priority, id)
 			if err != nil {
 				errCh <- errResult{err: err, id: id}
 			} else {
 				result[idx] = img.(*image.NRGBA)
 			}
 		}(idx, id)
 	}
 	wg.Wait()
 	if len(errCh) > 0 {
 		close(errCh)
 		var msg bytes.Buffer
 		for errRet := range errCh {
 			_, _ = msg.WriteString(errRet.err.Error())
 			_, _ = msg.WriteString("\n")

 			// This captures the edge case when the error is cached in the image loader.
 			util.LogErr(il.failureStore.addDigestFailureIfNew(diff.NewDigestFailure(errRet.id, diff.OTHER)))
 		}
 		return nil, errors.New(msg.String())
 	}

 	return result, nil
 }

 // IsOnDisk returns true if the image that corresponds to the given imageID is in the disk cache.
 func (il *ImageLoader) IsOnDisk(imageID string) bool {
 	localRelPath, _, _ := il.mapper.ImagePaths(imageID)
 	return fileutil.FileExists(filepath.Join(il.localImgDir, localRelPath))
 }

 // PurgeImages removes the images that correspond to the given images.
 func (il *ImageLoader) PurgeImages(images []string, purgeGCS bool) error {
 	for _, id := range images {
 		localRelPath, bucket, gsRelPath := il.mapper.ImagePaths(id)
 		localPath := filepath.Join(il.localImgDir, localRelPath)
 		if fileutil.FileExists(localPath) {
 			if err := os.Remove(localPath); err != nil {
 				sklog.Errorf("Unable to remove image %s. Got error: %s", localPath, err)
 			}
 		}

 		if purgeGCS {
 			il.removeImg(bucket, gsRelPath)
 		}
 	}
 	return nil
 }

 // imageLoadWorker implements the rtcache.ReadThroughFunc signature.
 // It loads an image file either from disk or from Google storage.
 func (il *ImageLoader) imageLoadWorker(priority int64, imageID string) (interface{}, error) {
 	// Check if the image is in the disk cache.
 	localRelPath, bucket, gsRelPath := il.mapper.ImagePaths(imageID)
 	localPath := filepath.Join(il.localImgDir, localRelPath)
 	if fileutil.FileExists(localPath) {
 		img, err := loadImg(localPath)
 		if err != nil {
 			util.LogErr(il.failureStore.addDigestFailure(diff.NewDigestFailure(imageID, diff.CORRUPTED)))
 			return nil, err
 		}
 		util.LogErr(il.failureStore.purgeDigestFailures([]string{imageID}))
 		return img, err
 	}

 	// Download the image
 	imgBytes, err := il.downloadImg(bucket, gsRelPath)
 	if err != nil {
 		util.LogErr(il.failureStore.addDigestFailure(diff.NewDigestFailure(imageID, diff.HTTP)))
 		return nil, err
 	}

 	// Decode it and return it.
 	img, err := decodeImg(bytes.NewBuffer(imgBytes))
 	if err != nil {
 		util.LogErr(il.failureStore.addDigestFailure(diff.NewDigestFailure(imageID, diff.CORRUPTED)))
 		return nil, err
 	}

 	// Save the file to disk.
 	il.saveImgInfoAsync(imageID, imgBytes)
 	return img, nil
 }

 func (il *ImageLoader) saveImgInfoAsync(imageID string, imgBytes []byte) {
 	il.wg.Add(1)
 	go func() {
 		defer il.wg.Done()
 		localRelPath, _, _ := il.mapper.ImagePaths(imageID)
 		if err := saveFilePath(filepath.Join(il.localImgDir, localRelPath), bytes.NewBuffer(imgBytes)); err != nil {
 			sklog.Error(err)
 		}
 	}()
 }

 // downloadImg retrieves the given image from Google storage. If bucket is not empty
 // the gsPath is considered absolute within the bucket, otherwise it is relative to gsImageBaseDir.
 func (il *ImageLoader) downloadImg(bucket, gsPath string) ([]byte, error) {
 	var bucketNames []string
 	var objLocation string
 	if bucket != "" {
 		objLocation = gsPath
 		bucketNames = []string{bucket}
 	} else {
 		objLocation = filepath.Join(il.gsImageBaseDir, gsPath)
 		bucketNames = il.gsBucketNames
 	}

 	var err error
 	var imgData []byte
 	for _, bucketName := range bucketNames {
 		imgData, err = il.downloadImgFromBucket(objLocation, bucketName)
 		if err == nil {
 			return imgData, nil
 		}
 	}
 	return nil, fmt.Errorf("Failed finding image %s in buckets %v. Last error: %s", gsPath, bucketNames, err)
 }

 // downloadImgFromBucket retrieves the given image from the given Google storage bucket.
 // It returns storage.ErrObjectNotExist if the given image does not exist in the bucket.
 func (il *ImageLoader) downloadImgFromBucket(objLocation, bucketName string) ([]byte, error) {
 	ctx := context.Background()

 	// Retrieve the attributes.
 	attrs, err := il.storageClient.Bucket(bucketName).Object(objLocation).Attrs(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("Unable to retrieve attributes for %s/%s: %.80s", bucketName, objLocation, err)
 	}

 	var buf *bytes.Buffer
 	for i := 0; i < MAX_URI_GET_TRIES; i++ {
 		err = func() error {
 			reader, err := il.storageClient.Bucket(bucketName).Object(objLocation).NewReader(ctx)
 			if err != nil {
 				return fmt.Errorf("New reader failed for %s/%s: %.80s", bucketName, objLocation, err)
 			}
 			defer util.Close(reader)

 			size := reader.Size()
 			buf = bytes.NewBuffer(make([]byte, 0, size))
 			md5Hash := md5.New()
 			multiOut := io.MultiWriter(md5Hash, buf)

 			if _, err = io.Copy(multiOut, reader); err != nil {
 				return err
 			}

 			// Check the MD5.
 			if !bytes.Equal(md5Hash.Sum(nil), attrs.MD5) {
 				return fmt.Errorf("MD5 hash for digest %s incorrect.", objLocation)
 			}

 			return nil
 		}()

 		if err == nil {
 			break
 		}
 		sklog.Errorf("Error fetching file for path %s: %s", objLocation, err)
 	}

 	if err != nil {
 		sklog.Errorf("Failed fetching file after %d attempts", MAX_URI_GET_TRIES)
 		return nil, err
 	}

 	sklog.Infof("Done downloading image for: %s. Length: %d", objLocation, buf.Len())
 	return buf.Bytes(), err
 }

 // removeImg removes the image that corresponds to the given relative path from GCS.
 func (il *ImageLoader) removeImg(bucket, gsRelPath string) {
 	// If the bucket is not empty then look there otherwise use the default buckets.
 	objLocation := filepath.Join(il.gsImageBaseDir, gsRelPath)
 	bucketNames := il.gsBucketNames
 	if bucket != "" {
 		bucketNames = []string{bucket}
 	}

 	ctx := context.Background()
 	for _, bucketName := range bucketNames {
 		// Retrieve the attributes to test if the file exists.
 		_, err := il.storageClient.Bucket(bucketName).Object(objLocation).Attrs(ctx)
 		if err != nil {
 			// We ignore the error because it most likely indicates that the requested object
 			// does not exist. Currently the Attrs(...) call does not return ErrObjectNotExist
 			// as documented.
 			continue
 		}

 		// Log an error and continue to the next bucket if we cannot delete the existing file.
 		if err := il.storageClient.Bucket(bucketName).Object(objLocation).Delete(ctx); err != nil {
 			sklog.Errorf("Unable to delete existing object at %s. Got error: %s", objLocation, err)
 			continue
 		}
 	}
 }
	package diffstore

	import (
	"bytes"
	"context"
	"crypto/md5"
	"errors"
	"fmt"
	"image"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"sync"

	"cloud.google.com/go/storage"
	"go.skia.org/infra/go/fileutil"
	"go.skia.org/infra/go/rtcache"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/golden/go/diff"
	"google.golang.org/api/option"
	)

	const (
	// MAX_URI_GET_TRIES is the number of tries we do to load an image.
	MAX_URI_GET_TRIES = 4

	// Number of concurrent workers downloading images.
	N_IMG_WORKERS = 10
	)

	// ImageLoader facilitates to continously download images and cache them in RAM.
	type ImageLoader struct {
	// storageClient is the Google client to local content from GCS.
	storageClient *storage.Client

	// localImgDir is the local directory where images should be written to.
	localImgDir string

	// gsBucketNames is the list of GCS bucket where images are stored.
	gsBucketNames []string

	// gsImageBaseDir is the GCS directory (prefix) where images are stored.
	gsImageBaseDir string

	// imageCache caches and calculates images.
	imageCache rtcache.ReadThroughCache

	// failureStore persists failures in retrieving images.
	failureStore *failureStore

	wg sync.WaitGroup

	// mapper contains various functions for creating image IDs and paths.
	mapper DiffStoreMapper
	}

	// Creates a new instance of ImageLoader.
	func NewImgLoader(client http.Client, baseDir, imgDir string, gsBucketNames []string, gsImageBaseDir string, maxCacheSize int, mapper DiffStoreMapper) (ImageLoader, error) {
	storageClient, err := storage.NewClient(context.Background(), option.WithHTTPClient(client))
	if err != nil {
	return nil, err
	}

	fStore, err := newFailureStore(filepath.Join(baseDir, FAILUREDB_NAME))
	if err != nil {
	return nil, err
	}

	ret := &ImageLoader{
	storageClient: storageClient,
	localImgDir: imgDir,
	gsBucketNames: gsBucketNames,
	gsImageBaseDir: gsImageBaseDir,
	failureStore: fStore,
	mapper: mapper,
	}

	// Set up the work queues that balance the load.
	if ret.imageCache, err = rtcache.New(ret.imageLoadWorker, maxCacheSize, N_IMG_WORKERS); err != nil {
	return nil, err
	}
	return ret, nil
	}

	// Warm makes sure the images are cached. It does not return a result to avoid
	// deserialization and unnecessary memory allocation.
	// The synchronous flag determines whether the call is blocking or not.
	// It workes in sync with Get, any image that is scheduled be retrieved by get
	// will not be fetched again.
	func (il *ImageLoader) Warm(priority int64, images []string, synchronous bool) {
	var localWg sync.WaitGroup
	il.wg.Add(len(images))
	localWg.Add(len(images))
	for _, id := range images {
	go func(id string) {
	defer il.wg.Done()
	defer localWg.Done()
	if err := il.imageCache.Warm(priority, id); err != nil {
	sklog.Errorf("Unable to retrive image %s. Got error: %s", id, err)
	}
	}(id)
	}

	// Wait for the all images to be loaded.
	if synchronous {
	localWg.Wait()
	}
	}

	// sync waits until all pending go routines have terminated.
	func (il *ImageLoader) Sync() {
	il.wg.Wait()
	}

	// errResult is a helper type to capture error information in Get(...).
	type errResult struct {
	err error
	id string
	}

	// Get returns the images identified by digests and returns it as an NRGBA image.
	// Priority determines the order in which multiple concurrent calls are processed.
	// func (il ImageLoader) Get(priority int64, digests []string) (image.NRGBA, error) {
	// // img, err := il.imageCache.Get(priority, digests)
	// // if err != nil {
	// // return nil, err
	// // }
	// // return img.(*image.NRGBA), nil
	// return nil, nil
	// }

	func (il ImageLoader) Get(priority int64, images []string) ([]image.NRGBA, error) {
	// Parallel load the requested images.
	result := make([]*image.NRGBA, len(images))
	errCh := make(chan errResult, len(images))
	var wg sync.WaitGroup
	wg.Add(len(images))
	for idx, id := range images {
	go func(idx int, id string) {
	defer wg.Done()
	img, err := il.imageCache.Get(priority, id)
	if err != nil {
	errCh <- errResult{err: err, id: id}
	} else {
	result[idx] = img.(*image.NRGBA)
	}
	}(idx, id)
	}
	wg.Wait()
	if len(errCh) > 0 {
	close(errCh)
	var msg bytes.Buffer
	for errRet := range errCh {
	_, _ = msg.WriteString(errRet.err.Error())
	_, _ = msg.WriteString("\n")

	// This captures the edge case when the error is cached in the image loader.
	util.LogErr(il.failureStore.addDigestFailureIfNew(diff.NewDigestFailure(errRet.id, diff.OTHER)))
	}
	return nil, errors.New(msg.String())
	}

	return result, nil
	}

	// IsOnDisk returns true if the image that corresponds to the given imageID is in the disk cache.
	func (il *ImageLoader) IsOnDisk(imageID string) bool {
	localRelPath, _, _ := il.mapper.ImagePaths(imageID)
	return fileutil.FileExists(filepath.Join(il.localImgDir, localRelPath))
	}

	// PurgeImages removes the images that correspond to the given images.
	func (il *ImageLoader) PurgeImages(images []string, purgeGCS bool) error {
	for _, id := range images {
	localRelPath, bucket, gsRelPath := il.mapper.ImagePaths(id)
	localPath := filepath.Join(il.localImgDir, localRelPath)
	if fileutil.FileExists(localPath) {
	if err := os.Remove(localPath); err != nil {
	sklog.Errorf("Unable to remove image %s. Got error: %s", localPath, err)
	}
	}

	if purgeGCS {
	il.removeImg(bucket, gsRelPath)
	}
	}
	return nil
	}

	// imageLoadWorker implements the rtcache.ReadThroughFunc signature.
	// It loads an image file either from disk or from Google storage.
	func (il *ImageLoader) imageLoadWorker(priority int64, imageID string) (interface{}, error) {
	// Check if the image is in the disk cache.
	localRelPath, bucket, gsRelPath := il.mapper.ImagePaths(imageID)
	localPath := filepath.Join(il.localImgDir, localRelPath)
	if fileutil.FileExists(localPath) {
	img, err := loadImg(localPath)
	if err != nil {
	util.LogErr(il.failureStore.addDigestFailure(diff.NewDigestFailure(imageID, diff.CORRUPTED)))
	return nil, err
	}
	util.LogErr(il.failureStore.purgeDigestFailures([]string{imageID}))
	return img, err
	}

	// Download the image
	imgBytes, err := il.downloadImg(bucket, gsRelPath)
	if err != nil {
	util.LogErr(il.failureStore.addDigestFailure(diff.NewDigestFailure(imageID, diff.HTTP)))
	return nil, err
	}

	// Decode it and return it.
	img, err := decodeImg(bytes.NewBuffer(imgBytes))
	if err != nil {
	util.LogErr(il.failureStore.addDigestFailure(diff.NewDigestFailure(imageID, diff.CORRUPTED)))
	return nil, err
	}

	// Save the file to disk.
	il.saveImgInfoAsync(imageID, imgBytes)
	return img, nil
	}

	func (il *ImageLoader) saveImgInfoAsync(imageID string, imgBytes []byte) {
	il.wg.Add(1)
	go func() {
	defer il.wg.Done()
	localRelPath, _, _ := il.mapper.ImagePaths(imageID)
	if err := saveFilePath(filepath.Join(il.localImgDir, localRelPath), bytes.NewBuffer(imgBytes)); err != nil {
	sklog.Error(err)
	}
	}()
	}

	// downloadImg retrieves the given image from Google storage. If bucket is not empty
	// the gsPath is considered absolute within the bucket, otherwise it is relative to gsImageBaseDir.
	func (il *ImageLoader) downloadImg(bucket, gsPath string) ([]byte, error) {
	var bucketNames []string
	var objLocation string
	if bucket != "" {
	objLocation = gsPath
	bucketNames = []string{bucket}
	} else {
	objLocation = filepath.Join(il.gsImageBaseDir, gsPath)
	bucketNames = il.gsBucketNames
	}

	var err error
	var imgData []byte
	for _, bucketName := range bucketNames {
	imgData, err = il.downloadImgFromBucket(objLocation, bucketName)
	if err == nil {
	return imgData, nil
	}
	}
	return nil, fmt.Errorf("Failed finding image %s in buckets %v. Last error: %s", gsPath, bucketNames, err)
	}

	// downloadImgFromBucket retrieves the given image from the given Google storage bucket.
	// It returns storage.ErrObjectNotExist if the given image does not exist in the bucket.
	func (il *ImageLoader) downloadImgFromBucket(objLocation, bucketName string) ([]byte, error) {
	ctx := context.Background()

	// Retrieve the attributes.
	attrs, err := il.storageClient.Bucket(bucketName).Object(objLocation).Attrs(ctx)
	if err != nil {
	return nil, fmt.Errorf("Unable to retrieve attributes for %s/%s: %.80s", bucketName, objLocation, err)
	}

	var buf *bytes.Buffer
	for i := 0; i < MAX_URI_GET_TRIES; i++ {
	err = func() error {
	reader, err := il.storageClient.Bucket(bucketName).Object(objLocation).NewReader(ctx)
	if err != nil {
	return fmt.Errorf("New reader failed for %s/%s: %.80s", bucketName, objLocation, err)
	}
	defer util.Close(reader)

	size := reader.Size()
	buf = bytes.NewBuffer(make([]byte, 0, size))
	md5Hash := md5.New()
	multiOut := io.MultiWriter(md5Hash, buf)

	if _, err = io.Copy(multiOut, reader); err != nil {
	return err
	}

	// Check the MD5.
	if !bytes.Equal(md5Hash.Sum(nil), attrs.MD5) {
	return fmt.Errorf("MD5 hash for digest %s incorrect.", objLocation)
	}

	return nil
	}()

	if err == nil {
	break
	}
	sklog.Errorf("Error fetching file for path %s: %s", objLocation, err)
	}

	if err != nil {
	sklog.Errorf("Failed fetching file after %d attempts", MAX_URI_GET_TRIES)
	return nil, err
	}

	sklog.Infof("Done downloading image for: %s. Length: %d", objLocation, buf.Len())
	return buf.Bytes(), err
	}

	// removeImg removes the image that corresponds to the given relative path from GCS.
	func (il *ImageLoader) removeImg(bucket, gsRelPath string) {
	// If the bucket is not empty then look there otherwise use the default buckets.
	objLocation := filepath.Join(il.gsImageBaseDir, gsRelPath)
	bucketNames := il.gsBucketNames
	if bucket != "" {
	bucketNames = []string{bucket}
	}

	ctx := context.Background()
	for _, bucketName := range bucketNames {
	// Retrieve the attributes to test if the file exists.
	_, err := il.storageClient.Bucket(bucketName).Object(objLocation).Attrs(ctx)
	if err != nil {
	// We ignore the error because it most likely indicates that the requested object
	// does not exist. Currently the Attrs(...) call does not return ErrObjectNotExist
	// as documented.
	continue
	}

	// Log an error and continue to the next bucket if we cannot delete the existing file.
	if err := il.storageClient.Bucket(bucketName).Object(objLocation).Delete(ctx); err != nil {
	sklog.Errorf("Unable to delete existing object at %s. Got error: %s", objLocation, err)
	continue
	}
	}
	}