go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/resultdb/internal/artifactcontent/rbecas.go

// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package artifactcontent

import (
	"context"
	"flag"
	"fmt"
	"io"
	"strings"

	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	"go.opentelemetry.io/otel/attribute"
	"golang.org/x/sync/errgroup"
	"google.golang.org/genproto/googleapis/bytestream"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/status"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/grpc/appstatus"
	"go.chromium.org/luci/grpc/grpcmon"
	"go.chromium.org/luci/resultdb/internal/tracing"
	"go.chromium.org/luci/server/auth"
	"go.chromium.org/luci/server/router"
)

// RBEConn creates a gRPC connection to RBE authenticated as self.
func RBEConn(ctx context.Context) (*grpc.ClientConn, error) {
	creds, err := auth.GetPerRPCCredentials(ctx, auth.AsSelf, auth.WithScopes(auth.CloudOAuthScopes...))
	if err != nil {
		return nil, err
	}

	return grpc.Dial(
		"remotebuildexecution.googleapis.com:443",
		grpc.WithTransportCredentials(credentials.NewTLS(nil)),
		grpc.WithPerRPCCredentials(creds),
		grpc.WithStatsHandler(&grpcmon.ClientRPCStatsMonitor{}),
		grpc.WithUnaryInterceptor(otelgrpc.UnaryClientInterceptor()),
		grpc.WithStreamInterceptor(otelgrpc.StreamClientInterceptor()),
	)
}

// RegisterRBEInstanceFlag registers -artifact-rbe-instance flag.
func RegisterRBEInstanceFlag(fs *flag.FlagSet, target *string) {
	fs.StringVar(
		target,
		"artifact-rbe-instance",
		"",
		"Name of the RBE instance to use for artifact storage",
	)
}

// handleRBECASContent serves artifact content stored in RBE-CAS.
func (r *contentRequest) handleRBECASContent(c *router.Context, hash string) {
	// Protocol:
	// https://github.com/bazelbuild/remote-apis/blob/7802003e00901b4e740fe0ebec1243c221e02ae2/build/bazel/remote/execution/v2/remote_execution.proto#L229-L233
	// https://github.com/googleapis/googleapis/blob/c8e291e6a4d60771219205b653715d5aeec3e96b/google/bytestream/bytestream.proto#L50-L53

	// Start a reading stream.
	stream, err := r.ReadCASBlob(c.Request.Context(), &bytestream.ReadRequest{
		ResourceName: resourceName(r.RBECASInstanceName, hash, r.size.Int64),
		ReadLimit:    r.limit,
	})
	if err != nil {
		if status.Code(err) == codes.NotFound {
			// Do not lose the original error message.
			logging.Warningf(c.Request.Context(), "RBE-CAS responded: %s", err)
			err = appstatus.Errorf(codes.NotFound, "artifact content no longer exists")
		}
		r.sendError(c.Request.Context(), err)
		return
	}

	// Forward the blob to the client.
	wroteHeader := false
	for {
		_, readSpan := tracing.Start(c.Request.Context(), "resultdb.readChunk")
		chunk, err := stream.Recv()
		if err == nil {
			readSpan.SetAttributes(attribute.Int("size", len(chunk.Data)))
		}
		tracing.End(readSpan, err)

		switch {
		case err == io.EOF:
			// We are done.
			return

		case err != nil:
			if wroteHeader {
				// The response was already partially written, so it is too late to
				// write headers. Write at least something indicating the incomplete
				// response.
				fmt.Fprintf(c.Writer, "\nResultDB: internal error while writing the response!\n")
				logging.Errorf(c.Request.Context(), "Failed to read from RBE-CAS in the middle of response: %s", err)
			} else {
				if status.Code(err) == codes.NotFound {
					// Sometimes RBE-CAS doesn't report NotFound until the read
					// of the first chunk, so duplicate the NotFound handling
					// above.
					// Do not lose the original error message.
					logging.Warningf(c.Request.Context(), "RBE-CAS responded: %s", err)
					err = appstatus.Errorf(codes.NotFound, "artifact content no longer exists")
				}
				r.sendError(c.Request.Context(), err)
			}
			return

		default:
			// Forward the chunk.
			if !wroteHeader {
				r.writeContentHeaders()
				wroteHeader = true
			}

			_, writeSpan := tracing.Start(c.Request.Context(), "resultdb.writeChunk",
				attribute.Int("size", len(chunk.Data)),
			)
			_, err := c.Writer.Write(chunk.Data)
			tracing.End(writeSpan, err)
			if err != nil {
				logging.Warningf(c.Request.Context(), "Failed to write a response chunk: %s", err)
				return
			}
		}
	}
}

func resourceName(instance, hash string, size int64) string {
	return fmt.Sprintf("%s/blobs/%s/%d", instance, strings.TrimPrefix(hash, "sha256:"), size)
}

// Reader reads the artifact content from RBE-CAS.
type Reader struct {
	// RBEInstance is the name of the RBE instance where the artifact is stored.
	// Example: "projects/luci-resultdb/instances/artifacts".
	RBEInstance string
	// Hash is the hash of the artifact content stored in RBE-CAS.
	Hash string
	// Size is the content size in bytes.
	Size int64
}

// DownloadRBECASContent calls f for the downloaded artifact content.
func (r *Reader) DownloadRBECASContent(ctx context.Context, bs bytestream.ByteStreamClient, f func(context.Context, io.Reader) error) error {
	stream, err := bs.Read(ctx, &bytestream.ReadRequest{
		ResourceName: resourceName(r.RBEInstance, r.Hash, r.Size),
	})
	if err != nil {
		if status.Code(err) == codes.NotFound {
			logging.Warningf(ctx, "RBE-CAS responded: %s", err)
		}
		return err
	}

	pr, pw := io.Pipe()
	eg, ctx := errgroup.WithContext(ctx)
	defer eg.Wait()
	eg.Go(func() error {
		defer pr.Close()
		return f(ctx, pr)
	})

	eg.Go(func() error {
		defer pw.Close()
		for {
			_, readSpan := tracing.Start(ctx, "resultdb.readChunk")
			chunk, err := stream.Recv()
			if err == nil {
				readSpan.SetAttributes(attribute.Int("size", len(chunk.Data)))
			}
			tracing.End(readSpan, err)

			switch {
			case err == io.EOF:
				// We are done.
				return nil

			case err != nil:
				return err

			default:
				if _, err := pw.Write(chunk.Data); err != nil {
					if err == io.ErrClosedPipe {
						// If f() exits early, return nil here to see what error that f() returns.
						return nil
					}
					return errors.Annotate(err, "write to pipe").Err()
				}
			}
		}
	})

	return eg.Wait()
}