golden/go/storage/gcsclient.go

package storage

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"time"

	gstorage "cloud.google.com/go/storage"
	ttlcache "github.com/patrickmn/go-cache"
	"go.skia.org/infra/go/gcs"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/golden/go/types"
	"google.golang.org/api/option"
)

// GCSClientOptions is used to define input parameters to the GCSClient.
type GCSClientOptions struct {
	// KnownHashesGCSPath is the bucket and path for storing the list of known digests.
	KnownHashesGCSPath string

	// If DryRun is true, don't actually write the files (e.g. running locally)
	Dryrun bool
}

// GCSClient provides an abstraction around read/writes to Google storage.
type GCSClient interface {
	// WriteKnownDigests writes the given list of digests to GCS as newline separated strings.
	WriteKnownDigests(ctx context.Context, digests types.DigestSlice) error

	// LoadKnownDigests loads the digests that have previously been written
	// to GS via WriteKnownDigests. The digests should be copied to the
	// provided writer 'w'.
	LoadKnownDigests(ctx context.Context, w io.Writer) error

	// Options returns the options that were used to initialize the client
	Options() GCSClientOptions
}

const (
	// Used to cache the digests in a time-to-live (TTL) cache.
	digestsCacheKey = "digestsCacheKey"

	// We re-index the tile (and thus re-compute the known digests) no
	// faster than once per minute.
	digestsCacheFreshness = time.Minute
)

// ClientImpl implements the GCSClient interface.
type ClientImpl struct {
	storageClient *gstorage.Client
	options       GCSClientOptions

	digestsCache *ttlcache.Cache
}

// NewGCSClient creates a new instance of ClientImpl. The various
// output paths are set in GCSClientOptions.
func NewGCSClient(ctx context.Context, client *http.Client, options GCSClientOptions) (*ClientImpl, error) {
	storageClient, err := gstorage.NewClient(ctx, option.WithHTTPClient(client))
	if err != nil {
		return nil, err
	}

	return &ClientImpl{
		storageClient: storageClient,
		options:       options,
		digestsCache:  ttlcache.New(digestsCacheFreshness, digestsCacheFreshness),
	}, nil
}

// Options implements the GCSClient interface.
func (g *ClientImpl) Options() GCSClientOptions {
	return g.options
}

// WriteKnownDigests fulfills the GCSClient interface.
func (g *ClientImpl) WriteKnownDigests(ctx context.Context, digests types.DigestSlice) error {
	if g.options.Dryrun {
		sklog.Infof("dryrun: Writing %d digests", len(digests))
		return nil
	}
	writeFn := func(w *gstorage.Writer) error {
		for _, digest := range digests {
			if _, err := w.Write([]byte(digest + "\n")); err != nil {
				return fmt.Errorf("Error writing digests: %s", err)
			}
		}
		return nil
	}
	// Clear the read cache when the write completes or fails
	defer g.digestsCache.Delete(digestsCacheKey)

	return g.writeToPath(ctx, g.options.KnownHashesGCSPath, "text/plain", writeFn)
}

// LoadKnownDigests fulfills the GCSClient interface.
func (g *ClientImpl) LoadKnownDigests(ctx context.Context, w io.Writer) error {
	if cachedBytes, ok := g.digestsCache.Get(digestsCacheKey); ok {
		b := cachedBytes.([]byte)
		n, err := w.Write(b)
		return skerr.Wrapf(err, "copying %d cached bytes to writer", n)
	}

	bucketName, storagePath := gcs.SplitGSPath(g.options.KnownHashesGCSPath)

	target := g.storageClient.Bucket(bucketName).Object(storagePath)

	// If the item doesn't exist this will return gstorage.ErrObjectNotExist
	_, err := target.Attrs(ctx)
	if err != nil {
		// We simply assume an empty hashes file if the object was not found.
		if err == gstorage.ErrObjectNotExist {
			sklog.Warningf("No known digests found - maybe %s is a wrong path?", g.options.KnownHashesGCSPath)
			return nil
		}
		return err
	}

	// Copy the content to the output writer.
	reader, err := target.NewReader(ctx)
	if err != nil {
		return skerr.Wrapf(err, "opening %s for reading", g.options.KnownHashesGCSPath)
	}
	defer util.Close(reader)

	b, err := ioutil.ReadAll(reader)
	if err != nil {
		return skerr.Wrapf(err, "reading digests from GCS")
	}
	g.digestsCache.SetDefault(digestsCacheKey, b)

	n, err := w.Write(b)
	return skerr.Wrapf(err, "writing %d bytes from cache to writer", n)
}

// removeForTestingOnly removes the given file. Should only be used for testing.
func (g *ClientImpl) removeForTestingOnly(ctx context.Context, targetPath string) error {
	bucketName, storagePath := gcs.SplitGSPath(targetPath)
	target := g.storageClient.Bucket(bucketName).Object(storagePath)
	return target.Delete(ctx)
}

// writeToPath is a generic function that allows to write data to the given
// target path in GCS. The actual writing is done in the passed write function.
func (g *ClientImpl) writeToPath(ctx context.Context, targetPath, contentType string, wrtFn func(w *gstorage.Writer) error) error {
	bucketName, storagePath := gcs.SplitGSPath(targetPath)

	// Only write the known digests if a target path was given.
	if (bucketName == "") || (storagePath == "") {
		return nil
	}

	target := g.storageClient.Bucket(bucketName).Object(storagePath)
	writer := target.NewWriter(ctx)
	writer.ObjectAttrs.ContentType = contentType
	writer.ObjectAttrs.ACL = []gstorage.ACLRule{{Entity: gstorage.AllUsers, Role: gstorage.RoleReader}}

	// Write the actual data.
	if err := wrtFn(writer); err != nil {
		return skerr.Wrapf(err, "writing data to %s", targetPath)
	}

	if err := writer.Close(); err != nil {
		return skerr.Wrapf(err, "closing writer for %s", targetPath)
	}

	return nil
}

// Ensure ClientImpl fulfills the GCSClient interface.
var _ GCSClient = (*ClientImpl)(nil)