Revert "[cachetools] Buffer writes in UploadWriter" (#9352)

Reverts #9294

Author: dan-stowell

server/remote_cache/cachetools/BUILD

@@ -13,7 +13,6 @@ go_library(
         "//server/interfaces",
         "//server/metrics",
         "//server/remote_cache/digest",
-        "//server/util/bytebufferpool",
         "//server/util/compression",
         "//server/util/ioutil",
         "//server/util/log",
@@ -40,12 +39,10 @@ go_test(
         "//server/interfaces",
         "//server/remote_cache/digest",
         "//server/testutil/testcache",
-        "//server/testutil/testdigest",
         "//server/testutil/testenv",
         "//server/util/random",
         "//server/util/status",
         "//server/util/testing/flags",
-        "@com_github_google_go_cmp//cmp",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
         "@org_golang_google_genproto_googleapis_bytestream//:bytestream",

server/remote_cache/cachetools/cachetools.go

@@ -18,7 +18,6 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/server/interfaces"
 	"github.com/buildbuddy-io/buildbuddy/server/metrics"
 	"github.com/buildbuddy-io/buildbuddy/server/remote_cache/digest"
-	"github.com/buildbuddy-io/buildbuddy/server/util/bytebufferpool"
 	"github.com/buildbuddy-io/buildbuddy/server/util/compression"
 	"github.com/buildbuddy-io/buildbuddy/server/util/ioutil"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
@@ -50,8 +49,6 @@ var (
 	acRPCTimeout                = flag.Duration("cache.client.ac_rpc_timeout", 15*time.Second, "Maximum time a single Action Cache RPC can take.")
 	filecacheTreeSalt           = flag.String("cache.filecache_tree_salt", "20250304", "A salt to invalidate filecache tree hashes, if/when needed.")
 	requestCachedSubtreeDigests = flag.Bool("cache.request_cached_subtree_digests", true, "If true, GetTree requests will set send_cached_subtree_digests.")
-
-	uploadBufPool = bytebufferpool.FixedSize(uploadBufSizeBytes)
 )
 
 func retryOptions(name string) *retry.Options {
@@ -258,29 +255,19 @@ func uploadFromReader(ctx context.Context, bsClient bspb.ByteStreamClient, r *di
 	}
 
 	if r.GetCompressor() == repb.Compressor_IDENTITY {
-		w, err := NewUploadWriter(ctx, bsClient, r)
+		w, err := newUploadWriteCloser(ctx, bsClient, r)
 		if err != nil {
 			return nil, 0, err
 		}
-		defer w.Close()
-		bytesRead, err := w.ReadFrom(in)
+		bytesUploaded, err := w.ReadFrom(in)
 		if err != nil {
+			w.Close()
 			return nil, 0, err
 		}
-		if err := w.Commit(); err != nil {
-			return nil, 0, err
-		}
-		// Either the write succeeded or was short-circuited, but in
-		// either case, the remoteSize for uncompressed uploads should
-		// match the file size.
-		committedSize, err := w.GetCommittedSize()
-		if err != nil {
+		if err := w.Close(); err != nil {
 			return nil, 0, err
 		}
-		if committedSize != r.GetDigest().GetSizeBytes() {
-			return nil, 0, status.DataLossErrorf("Remote size (%d) != uploaded size: (%d)", committedSize, r.GetDigest().GetSizeBytes())
-		}
-		return r.GetDigest(), bytesRead, nil
+		return r.GetDigest(), bytesUploaded, nil
 	}
 
 	stream, err := bsClient.Write(ctx)
@@ -336,15 +323,24 @@ func uploadFromReader(ctx context.Context, bsClient bspb.ByteStreamClient, r *di
 		return nil, 0, err
 	}
 
-	committedSize := rsp.GetCommittedSize()
-	// -1 is returned if the blob already exists, otherwise the
-	// remoteSize should agree with what we uploaded.
-	if committedSize != bytesUploaded && committedSize != -1 {
-		return nil, 0, status.DataLossErrorf("Remote size (%d) != uploaded size: (%d)", committedSize, r.GetDigest().GetSizeBytes())
+	remoteSize := rsp.GetCommittedSize()
+	if r.GetCompressor() == repb.Compressor_IDENTITY {
+		// Either the write succeeded or was short-circuited, but in
+		// either case, the remoteSize for uncompressed uploads should
+		// match the file size.
+		if remoteSize != r.GetDigest().GetSizeBytes() {
+			return nil, 0, status.DataLossErrorf("Remote size (%d) != uploaded size: (%d)", remoteSize, r.GetDigest().GetSizeBytes())
+		}
+	} else {
+		// -1 is returned if the blob already exists, otherwise the
+		// remoteSize should agree with what we uploaded.
+		if remoteSize != bytesUploaded && remoteSize != -1 {
+			return nil, 0, status.DataLossErrorf("Remote size (%d) != uploaded size: (%d)", remoteSize, r.GetDigest().GetSizeBytes())
+		}
 	}
 
 	// At this point, a remote size of -1 means nothing new was uploaded.
-	if committedSize == -1 {
+	if remoteSize == -1 {
 		bytesUploaded = 0
 	}
 
@@ -1117,130 +1113,91 @@ func maybeSetCompressor(rn *digest.CASResourceName) {
 	}
 }
 
-type UploadWriter struct {
+type uploadWriteCloser struct {
 	ctx          context.Context
 	stream       bspb.ByteStream_WriteClient
 	sender       rpcutil.Sender[*bspb.WriteRequest]
+	resource     *digest.CASResourceName
 	uploadString string
 
 	bytesUploaded int64
-	finished      bool
-	committedSize int64
-	closed        bool
-
-	buf           []byte
-	bytesBuffered int
 }
 
-// Assert that UploadWriter implements CommittedWriteCloser
-var _ interfaces.CommittedWriteCloser = (*UploadWriter)(nil)
-
-func (uw *UploadWriter) Write(p []byte) (int, error) {
-	if uw.finished || uw.closed {
-		return 0, status.FailedPreconditionError("Cannot write to UploadWriter after it is finished or closed")
-	}
+func (cwc *uploadWriteCloser) Write(p []byte) (int, error) {
 	written := 0
 	for len(p) > 0 {
-		n := copy(uw.buf[uw.bytesBuffered:], p)
-		uw.bytesBuffered += n
-		if uw.bytesBuffered == uploadBufSizeBytes {
-			if err := uw.flush(false /* finish */); err != nil {
-				return written, err
+		n := min(len(p), uploadBufSizeBytes)
+		req := &bspb.WriteRequest{
+			Data:         p[:n],
+			ResourceName: cwc.uploadString,
+			WriteOffset:  cwc.bytesUploaded + int64(written),
+			FinishWrite:  false,
+		}
+
+		err := cwc.sender.SendWithTimeoutCause(req, *casRPCTimeout, status.DeadlineExceededError("Timed out sending Write request"))
+		if err != nil {
+			if err == io.EOF {
+				break
 			}
+			cwc.bytesUploaded += int64(written)
+			return written, err
 		}
 		written += n
 		p = p[n:]
 	}
+	cwc.bytesUploaded += int64(written)
 	return written, nil
 }
 
-func (uw *UploadWriter) flush(finish bool) error {
-	if uw.finished || uw.closed {
-		return status.FailedPreconditionError("UploadWriteCloser already finished or closed, cannot flush")
-	}
-	req := &bspb.WriteRequest{
-		Data:         uw.buf[:uw.bytesBuffered],
-		ResourceName: uw.uploadString,
-		WriteOffset:  uw.bytesUploaded,
-		FinishWrite:  finish,
-	}
-	err := uw.sender.SendWithTimeoutCause(req, *casRPCTimeout, status.DeadlineExceededError("Timed out sending Write request"))
-	if err != nil {
-		return err
-	}
-	uw.bytesUploaded += int64(uw.bytesBuffered)
-	uw.bytesBuffered = 0
-	uw.finished = finish
-	return nil
-}
-
-func (uw *UploadWriter) ReadFrom(r io.Reader) (int64, error) {
-	if uw.finished || uw.closed {
-		return 0, status.FailedPreconditionError("UploadWriter cannot ReadFrom after it is finished or closed")
-	}
-	bytesRead := int64(0)
+func (cwc *uploadWriteCloser) ReadFrom(r io.Reader) (int64, error) {
+	buf := make([]byte, uploadBufSizeBytes)
+	var bytesUploaded int64
 	for {
-		n, err := ioutil.ReadTryFillBuffer(r, uw.buf[uw.bytesBuffered:])
-		uw.bytesBuffered += n
-		bytesRead += int64(n)
+		n, err := ioutil.ReadTryFillBuffer(r, buf)
 		done := err == io.EOF
 		if err != nil && !done {
-			return bytesRead, err
+			return bytesUploaded, err
 		}
-		if uw.bytesBuffered == uploadBufSizeBytes {
-			if err := uw.flush(false /* finish */); err != nil {
-				return bytesRead, err
-			}
+		written, err := cwc.Write(buf[:n])
+		if err != nil {
+			return bytesUploaded + int64(written), err
 		}
+		bytesUploaded += int64(written)
 		if done {
 			break
 		}
 	}
-	return bytesRead, nil
+	return bytesUploaded, nil
 }
 
-func (uw *UploadWriter) Commit() error {
-	if uw.closed {
-		return status.FailedPreconditionError("UploadWriteCloser already closed, cannot commit")
-	}
-	if uw.finished {
-		return status.FailedPreconditionError("UploadWriteCloser already committed, cannot commit again")
-	}
-	if err := uw.flush(true /* finish */); err != nil {
-		return err
+func (cwc *uploadWriteCloser) Close() error {
+	req := &bspb.WriteRequest{
+		ResourceName: cwc.uploadString,
+		WriteOffset:  cwc.bytesUploaded,
+		FinishWrite:  true,
 	}
-	rsp, err := uw.stream.CloseAndRecv()
+
+	err := cwc.sender.SendWithTimeoutCause(req, *casRPCTimeout, status.DeadlineExceededError("Timed out sending Write request"))
 	if err != nil {
 		return err
 	}
-	uw.committedSize = rsp.GetCommittedSize()
-	return nil
-}
 
-func (uw *UploadWriter) Close() error {
-	if uw.closed {
-		return status.FailedPreconditionError("UploadWriteCloser already closed, cannot close again")
+	rsp, err := cwc.stream.CloseAndRecv()
+	if err != nil {
+		return err
 	}
-	uw.closed = true
-	uploadBufPool.Put(uw.buf)
-	return uw.stream.CloseSend()
-}
 
-func (uw *UploadWriter) GetCommittedSize() (int64, error) {
-	if !uw.finished {
-		return 0, status.FailedPreconditionError("UploadWriter not committed, cannot get committed size")
+	remoteSize := rsp.GetCommittedSize()
+	// Either the write succeeded or was short-circuited, but in
+	// either case, the remoteSize for uncompressed uploads should
+	// match the file size.
+	if remoteSize != cwc.resource.GetDigest().GetSizeBytes() {
+		return status.DataLossErrorf("Remote size (%d) != uploaded size: (%d)", remoteSize, cwc.resource.GetDigest().GetSizeBytes())
 	}
-	return uw.committedSize, nil
-}
-
-func (uw *UploadWriter) GetBytesUploaded() int64 {
-	return uw.bytesUploaded
+	return nil
 }
 
-// NewUploadWriter returns an UploadWriter that writes to the CAS for the specific resource name.
-// The blob is guaranteed to be written to the CAS only if all Write(...) calls and the Close() call succeed.
-// The caller is responsible for checking data integrity using GetCommittedSize() and GetBytesUploaded().
-func NewUploadWriter(ctx context.Context, bsClient bspb.ByteStreamClient, r *digest.CASResourceName) (*UploadWriter, error) {
+func newUploadWriteCloser(ctx context.Context, bsClient bspb.ByteStreamClient, r *digest.CASResourceName) (*uploadWriteCloser, error) {
 	if bsClient == nil {
 		return nil, status.FailedPreconditionError("ByteStreamClient not configured")
 	}
@@ -1252,11 +1209,11 @@ func NewUploadWriter(ctx context.Context, bsClient bspb.ByteStreamClient, r *dig
 		return nil, err
 	}
 	sender := rpcutil.NewSender[*bspb.WriteRequest](ctx, stream)
-	return &UploadWriter{
+	return &uploadWriteCloser{
 		ctx:          ctx,
 		stream:       stream,
 		sender:       sender,
+		resource:     r,
 		uploadString: r.NewUploadString(),
-		buf:          uploadBufPool.Get(),
 	}, nil
 }