siso/reapi/cache.go

// Copyright 2023 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package reapi

import (
	"cmp"
	"context"
	"fmt"
	"io"

	bpb "google.golang.org/genproto/googleapis/bytestream"

	"go.chromium.org/build/hashigo/digest"
	rpb "go.chromium.org/build/remote-apis/build/bazel/remote/execution/v2"

	"go.chromium.org/build/siso/blob"
	"go.chromium.org/build/siso/o11y/clog"
	"go.chromium.org/build/siso/reapi/bytestreamio"
)

// CacheStore provides a thin wrapper around REAPI client that gets and uploads blobs and action results.
type CacheStore struct {
	client *Client
}

// CacheStore returns cache store of the client.
func (c *Client) CacheStore() CacheStore {
	return CacheStore{
		client: c,
	}
}

func (c CacheStore) String() string {
	return fmt.Sprintf("cachestore:reapi addr:%s instance:%s", c.client.opt.Address, c.client.opt.Instance)
}

// GetActionResult gets action result for the action identified by the digest.
func (c CacheStore) GetActionResult(ctx context.Context, d digest.Digest) (*rpb.ActionResult, error) {
	return c.client.GetActionResult(ctx, d)
}

// SetActionResult sets action result for the action identified by the digest.
// If a failing action is provided, caching will be skipped.
func (c CacheStore) SetActionResult(ctx context.Context, d digest.Digest, ar *rpb.ActionResult) error {
	// Intentionally no-op. Setting the action result could allow a malicious
	// actor to perform arbitrary code execution.
	// May reconsider in the future for trusted workers, but for now, do nothing.

	// Intentionally return no error, as it is intended for this function to be
	// called.
	return nil
}

// GetContent gets contents for the fname identified by the digest.
func (c CacheStore) GetContent(ctx context.Context, d digest.Digest, fname string) ([]byte, error) {
	return c.client.Get(ctx, d, fname)
}

// SetContent sets contents for the fname identified by the digest.
func (c CacheStore) SetContent(ctx context.Context, d digest.Digest, fname string, content []byte) error {
	data := blob.FromBytes(c.client.digestFn, fname, content)
	if d != data.Digest() {
		return fmt.Errorf("digest mismatch: d=%s content=%s", d, data.Digest())
	}
	ds := blob.NewStore()
	ds.Set(data)
	_, err := c.client.UploadAll(ctx, ds)
	return err
}

// HasContent returns whether digest is in cache store.
func (c CacheStore) HasContent(ctx context.Context, d digest.Digest) bool {
	missing, err := c.client.Missing(ctx, []digest.Digest{d})
	if err != nil {
		clog.Warningf(ctx, "failed to call missing [%s]: %v", d, err)
		return false
	}
	return len(missing) != 1
}

// Source returns digest source for fname identified by the digest.
func (c CacheStore) Source(_ context.Context, d digest.Digest, fname string) blob.Source {
	return digestSource{
		c:     c.client,
		d:     d,
		fname: fname,
	}
}

type digestSourceReader struct {
	r      io.ReadCloser
	n      int
	size   int64
	c      *Client
	cancel context.CancelFunc
}

func (r *digestSourceReader) Read(buf []byte) (int, error) {
	n, err := r.r.Read(buf)
	r.n += n
	if err == io.EOF && int64(r.n) != r.size {
		err = io.ErrUnexpectedEOF
	}
	return n, err
}

func (r *digestSourceReader) Close() error {
	// On a partial read the consumer gave up early, so skip the EOF check
	// and cancel directly instead of draining a possibly large residual.
	var eofErr error
	if int64(r.n) == r.size {
		eofErr = expectEOF(r.r)
	}
	closeErr := r.r.Close()
	r.c.m.ReadDone(r.n, cmp.Or(closeErr, eofErr))
	if r.cancel != nil {
		r.cancel()
	}
	return closeErr
}

type digestSource struct {
	c     *Client
	d     digest.Digest
	fname string
}

func (s digestSource) Open(ctx context.Context) (io.ReadCloser, error) {
	// The context passed to bytestreamio.Open is used for the gRPC stream.
	// We need to cancel this context to close the stream when we are done reading
	// or if an error occurs during setup, otherwise we leak the stream until the
	// parent context (step context) is canceled.
	ctx, cancel := context.WithCancel(ctx)
	r, err := bytestreamio.Open(ctx, bpb.NewByteStreamClient(s.c.casDataConn), s.c.resourceName(s.d))
	if err != nil {
		cancel()
		s.c.m.ReadDone(0, err)
		return nil, err
	}
	rd, err := s.c.newDecoder(r, s.d)
	if err != nil {
		cancel()
		s.c.m.ReadDone(0, err)
		return nil, err
	}
	return &digestSourceReader{r: rd, size: s.d.SizeBytes, c: s.c, cancel: cancel}, err
}

// ReadAll returns the full blob content in a single CAS fetch, retrying on
// retriable errors. blob.DataToBytes uses this instead of Open+read.
func (s digestSource) ReadAll(ctx context.Context) ([]byte, error) {
	return s.c.Get(ctx, s.d, s.fname)
}

func (s digestSource) String() string {
	return fmt.Sprintf("digest-data-in-cas:%s for %s", s.d, s.fname)
}