Revert "[cachetools] Re-introduce UploadWriter; use in UploadFromReader" (#9398)

Reverts #9346

This change appears to have introduced a race condition (exposed in
`firecracker_test. TestFirecrackerRunWithDockerMirror`). [Recent test
run](https://app.buildbuddy.io/invocation/bd6e59bc-ced5-481c-81cb-291ef09a8f9a?target=%2F%2Fenterprise%2Fserver%2Fremote_execution%2Fcontainers%2Ffirecracker%3Afirecracker_test&targetStatus=6)
on `master`.

With this revert applied, I have not yet provoked a race (16 test runs).

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",

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 {
@@ -249,58 +246,85 @@ func ComputeFileDigest(fullFilePath, instanceName string, digestFunction repb.Di
 	return computeDigest(f, instanceName, digestFunction)
 }
 
-func uploadFromReader(ctx context.Context, bsClient bspb.ByteStreamClient, rn *digest.CASResourceName, in io.Reader) (*repb.Digest, int64, error) {
+func uploadFromReader(ctx context.Context, bsClient bspb.ByteStreamClient, r *digest.CASResourceName, in io.Reader) (*repb.Digest, int64, error) {
 	if bsClient == nil {
 		return nil, 0, status.FailedPreconditionError("ByteStreamClient not configured")
 	}
-	if rn.IsEmpty() {
-		return rn.GetDigest(), 0, nil
+	if r.IsEmpty() {
+		return r.GetDigest(), 0, nil
 	}
-
-	w, err := NewUploadWriter(ctx, bsClient, rn)
+	stream, err := bsClient.Write(ctx)
 	if err != nil {
 		return nil, 0, err
 	}
-	defer w.Close()
 
-	rc := io.NopCloser(in)
-	if rn.GetCompressor() == repb.Compressor_ZSTD {
+	var rc io.ReadCloser = io.NopCloser(in)
+	if r.GetCompressor() == repb.Compressor_ZSTD {
 		rbuf := make([]byte, 0, uploadBufSizeBytes)
 		cbuf := make([]byte, 0, uploadBufSizeBytes)
-		reader, err := compression.NewZstdCompressingReader(rc, rbuf[:uploadBufSizeBytes], cbuf[:uploadBufSizeBytes])
+		reader, err := compression.NewZstdCompressingReader(io.NopCloser(in), rbuf[:uploadBufSizeBytes], cbuf[:uploadBufSizeBytes])
 		if err != nil {
 			return nil, 0, status.InternalErrorf("Failed to compress blob: %s", err)
 		}
 		rc = reader
 	}
 	defer rc.Close()
 
-	if _, err := w.ReadFrom(rc); err != nil {
-		return nil, 0, err
+	buf := make([]byte, uploadBufSizeBytes)
+	bytesUploaded := int64(0)
+	sender := rpcutil.NewSender[*bspb.WriteRequest](ctx, stream)
+	resourceName := r.NewUploadString()
+	for {
+		n, err := ioutil.ReadTryFillBuffer(rc, buf)
+		if err != nil && err != io.EOF {
+			return nil, bytesUploaded, err
+		}
+		readDone := err == io.EOF
+
+		req := &bspb.WriteRequest{
+			Data:         buf[:n],
+			ResourceName: resourceName,
+			WriteOffset:  bytesUploaded,
+			FinishWrite:  readDone,
+		}
+
+		err = sender.SendWithTimeoutCause(req, *casRPCTimeout, status.DeadlineExceededError("Timed out sending Write request"))
+		if err != nil {
+			// If the blob already exists in the CAS, the server will respond EOF.
+			// It is safe to stop sending writes.
+			if err == io.EOF {
+				break
+			}
+			return nil, bytesUploaded, err
+		}
+		bytesUploaded += int64(len(req.Data))
+		if readDone {
+			break
+		}
+
 	}
-	if err := w.Commit(); err != nil {
-		return nil, 0, err
+	rsp, err := stream.CloseAndRecv()
+	if err != nil {
+		return nil, bytesUploaded, err
 	}
 
-	bytesUploaded := w.GetBytesUploaded()
-	committedSize := w.GetCommittedSize()
-
-	if rn.GetCompressor() == repb.Compressor_IDENTITY {
+	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 committedSize != rn.GetDigest().GetSizeBytes() {
-			return nil, bytesUploaded, status.DataLossErrorf("Remote size (%d) != uploaded size: (%d)", committedSize, rn.GetDigest().GetSizeBytes())
+		if remoteSize != r.GetDigest().GetSizeBytes() {
+			return nil, bytesUploaded, status.DataLossErrorf("Remote size (%d) != uploaded size: (%d)", remoteSize, r.GetDigest().GetSizeBytes())
 		}
 	} else {
 		// -1 is returned if the blob already exists, otherwise the
-		// remote (committed) size should agree with what we uploaded.
-		if committedSize != bytesUploaded && committedSize != -1 {
-			return nil, bytesUploaded, status.DataLossErrorf("Remote size (%d) != uploaded size: (%d)", committedSize, rn.GetDigest().GetSizeBytes())
+		// remoteSize should agree with what we uploaded.
+		if remoteSize != bytesUploaded && remoteSize != -1 {
+			return nil, bytesUploaded, status.DataLossErrorf("Remote size (%d) != uploaded size: (%d)", remoteSize, r.GetDigest().GetSizeBytes())
 		}
 	}
 
-	return rn.GetDigest(), bytesUploaded, nil
+	return r.GetDigest(), bytesUploaded, nil
 }
 
 type uploadRetryResult = struct {
@@ -1077,167 +1101,3 @@ func maybeSetCompressor(rn *digest.CASResourceName) {
 		rn.SetCompressor(repb.Compressor_ZSTD)
 	}
 }
-
-type UploadWriter struct {
-	ctx          context.Context
-	stream       bspb.ByteStream_WriteClient
-	sender       rpcutil.Sender[*bspb.WriteRequest]
-	uploadString string
-
-	bytesUploaded int64
-	finished      bool
-	committedSize int64
-	closed        bool
-
-	buf           []byte
-	bytesBuffered int
-}
-
-// Write copies the input bytes to an internal buffer and may send some or all of the bytes to the CAS.
-// Bytes are not guaranteed to be uploaded to the CAS until a call to Commit() succeeds.
-// Returning EOF indicates that the blob already exists in the CAS and no further writes are necessary.
-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")
-	}
-	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
-			}
-		}
-		written += n
-		p = p[n:]
-	}
-	return written, nil
-}
-
-func (uw *UploadWriter) flush(finish bool) error {
-	if uw.finished {
-		return nil
-	}
-	if 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 {
-		// If the blob already exists in the CAS, the server will respond EOF
-		// to indicate no further writes are needed.
-		if err == io.EOF {
-			uw.finished = true
-		}
-		return err
-	}
-	uw.bytesUploaded += int64(uw.bytesBuffered)
-	uw.bytesBuffered = 0
-	uw.finished = finish
-	return nil
-}
-
-// ReadFrom reads all the bytes from the input Reader until encountering EOF,
-// copies them to an internal buffer, and may send some or all of the bytes to the CAS.
-// The bytes are not guaranteed uploaded to the CAS until Commit is called and returns successfully.
-// ReadFrom returns the number of bytes read from the input reader.
-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)
-	for {
-		n, err := ioutil.ReadTryFillBuffer(r, uw.buf[uw.bytesBuffered:])
-		uw.bytesBuffered += n
-		bytesRead += int64(n)
-		if err != nil && err != io.EOF {
-			return bytesRead, err
-		}
-		readDone := err == io.EOF
-
-		if uw.bytesBuffered == uploadBufSizeBytes {
-			err := uw.flush(false /* finish */)
-			if err != nil && err != io.EOF {
-				return bytesRead, err
-			}
-			// If the blob already exists in the CAS, the server can respond with an EOF.
-			// Stop sending bytes to the server.
-			if err == io.EOF {
-				break
-			}
-		}
-		if readDone {
-			break
-		}
-	}
-	return bytesRead, nil
-}
-
-// Commit sends any bytes remaining in the internal buffer to the CAS
-// and tells the server that the stream is done sending writes.
-func (uw *UploadWriter) Commit() error {
-	if uw.closed {
-		return status.FailedPreconditionError("UploadWriteCloser already closed, cannot commit")
-	}
-	err := uw.flush(true /* finish */)
-	// If the blob already exists in the CAS, the server can respond with an EOF.
-	// The blob exists, it is safe to finish committing.
-	if err != nil && err != io.EOF {
-		return err
-	}
-	rsp, err := uw.stream.CloseAndRecv()
-	if err != nil {
-		return err
-	}
-	uw.committedSize = rsp.GetCommittedSize()
-	return nil
-}
-
-// Close closes the underlying stream and returns an internal buffer to the pool.
-// It is expected (and safe) to call Close even if Commit fails.
-func (uw *UploadWriter) Close() error {
-	if uw.closed {
-		return status.FailedPreconditionError("UploadWriteCloser already closed, cannot close again")
-	}
-	uw.closed = true
-	uploadBufPool.Put(uw.buf)
-	return uw.stream.CloseSend()
-}
-
-func (uw *UploadWriter) GetCommittedSize() int64 {
-	return uw.committedSize
-}
-
-func (uw *UploadWriter) GetBytesUploaded() int64 {
-	return uw.bytesUploaded
-}
-
-// Assert that UploadWriter implements CommittedWriteCloser
-var _ interfaces.CommittedWriteCloser = (*UploadWriter)(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) {
-	if bsClient == nil {
-		return nil, status.FailedPreconditionError("ByteStreamClient not configured")
-	}
-	stream, err := bsClient.Write(ctx)
-	if err != nil {
-		return nil, err
-	}
-	sender := rpcutil.NewSender[*bspb.WriteRequest](ctx, stream)
-	return &UploadWriter{
-		ctx:          ctx,
-		stream:       stream,
-		sender:       sender,
-		uploadString: r.NewUploadString(),
-		buf:          uploadBufPool.Get(),
-	}, nil
-}