mailru: buffer speedup hashing in the global memory pool and fix upload retries

With speedup enabled, files up to --mailru-speedup-max-memory are read
into memory so their hash can be tried against the server before
uploading. This used io.ReadAll, which allocates a fresh heap slice per
file and grows it by doubling, so with the default 32 MiB limit and
several transfers this churned a lot of garbage outside rclone's memory
accounting.

Buffer the file with multipart.NewRW instead, hashing it in transit,
so the memory comes from the global pool and is reused.

When the hash isn't known to the server the buffered file is uploaded
from the same buffer. Previously a low level retry of that upload
resent an already drained reader, so the retry always failed. Rewind
seekable bodies at the start of each attempt so retries resend the
whole file. Add a test which drops the connection on the first attempt
and checks the retried body is complete.

The body is sent through lib/rest, which wraps it in readers.NoCloser,
so the transport can't close the pooled buffer early; Update closes it
when it returns.

backend/mailru/mailru.go

@@ -36,6 +36,7 @@ import (
 	"github.com/rclone/rclone/fs/operations"
 
 	"github.com/rclone/rclone/lib/encoder"
+	"github.com/rclone/rclone/lib/multipart"
 	"github.com/rclone/rclone/lib/oauthutil"
 	"github.com/rclone/rclone/lib/pacer"
 	"github.com/rclone/rclone/lib/readers"
@@ -1614,7 +1615,6 @@ func (o *Object) Update(ctx context.Context, in io.Reader, src fs.ObjectInfo, op
 	}
 
 	var (
-		fileBuf  []byte
 		fileHash []byte
 		newHash  []byte
 		slowHash bool
@@ -1660,15 +1660,19 @@ func (o *Object) Update(ctx context.Context, in io.Reader, src fs.ObjectInfo, op
 
 	// Attempt to put by calculating hash in memory
 	if trySpeedup && size <= int64(o.fs.opt.SpeedupMaxMem) {
-		fileBuf, err = io.ReadAll(in)
-		if err != nil {
+		rw := multipart.NewRW()
+		defer func() {
+			_ = rw.Close()
+		}()
+		hasher := mrhash.New()
+		if _, err = io.Copy(rw, io.TeeReader(in, hasher)); err != nil {
 			return err
 		}
-		fileHash = mrhash.Sum(fileBuf)
+		fileHash = hasher.Sum(nil)
 		if o.putByHash(ctx, fileHash, src, "memory") {
 			return nil
 		}
-		wrapIn = bytes.NewReader(fileBuf)
+		wrapIn = rw
 		trySpeedup = false // speedup failed, force upload
 	}
 
@@ -1702,9 +1706,8 @@ func (o *Object) Update(ctx context.Context, in io.Reader, src fs.ObjectInfo, op
 	// Upload object data
 	if size <= mrhash.Size {
 		// Optimize upload: skip extra request if data fits in the hash buffer.
-		if fileBuf == nil {
-			fileBuf, err = io.ReadAll(wrapIn)
-		}
+		var fileBuf []byte
+		fileBuf, err = io.ReadAll(wrapIn)
 		if fileHash == nil && err == nil {
 			fileHash = mrhash.Sum(fileBuf)
 		}
@@ -1869,6 +1872,12 @@ func (o *Object) upload(ctx context.Context, in io.Reader, size int64, options .
 		strHash string
 	)
 	err = o.fs.pacer.Call(func() (bool, error) {
+		// Rewind seekable bodies so a retry resends the whole file
+		if seeker, ok := in.(io.Seeker); ok {
+			if _, err := seeker.Seek(0, io.SeekStart); err != nil {
+				return false, err
+			}
+		}
 		res, err = o.fs.srv.Call(ctx, &opts)
 		if err == nil {
 			strHash, err = readBodyWord(res)

backend/mailru/mailru_internal_test.go

@@ -0,0 +1,65 @@
+package mailru
+
+import (
+	"bytes"
+	"context"
+	"encoding/hex"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/rclone/rclone/backend/mailru/mrhash"
+	"github.com/rclone/rclone/fs"
+	"github.com/rclone/rclone/lib/multipart"
+	"github.com/rclone/rclone/lib/pacer"
+	"github.com/rclone/rclone/lib/rest"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/oauth2"
+)
+
+// TestUploadRetry checks that an upload from a seekable pooled buffer
+// whose first attempt fails with a retriable error is retried with the
+// full body rather than the drained reader.
+func TestUploadRetry(t *testing.T) {
+	content := bytes.Repeat([]byte("mailru"), 512)
+	wantHash := mrhash.Sum(content)
+	var attempts atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+		if attempts.Add(1) == 1 {
+			// Drop the connection so the client sees a retriable error
+			conn, _, err := w.(http.Hijacker).Hijack()
+			require.NoError(t, err)
+			_ = conn.Close()
+			return
+		}
+		assert.Equal(t, content, body, "retried upload should re-send the full body")
+		_, _ = w.Write([]byte(hex.EncodeToString(wantHash) + "\n"))
+	}))
+	defer server.Close()
+
+	ctx := context.Background()
+	f := &Fs{
+		srv:         rest.NewClient(server.Client()),
+		source:      oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "token"}),
+		pacer:       fs.NewPacer(ctx, pacer.NewDefault(pacer.MinSleep(time.Millisecond), pacer.MaxSleep(10*time.Millisecond))),
+		shardURL:    server.URL,
+		shardExpiry: time.Now().Add(time.Hour),
+	}
+	o := &Object{fs: f, remote: "remote"}
+
+	rw := multipart.NewRW()
+	defer func() { _ = rw.Close() }()
+	_, err := rw.Write(content)
+	require.NoError(t, err)
+
+	gotHash, err := o.upload(ctx, rw, int64(len(content)))
+	require.NoError(t, err)
+	assert.Equal(t, wantHash, gotHash)
+	assert.Equal(t, int32(2), attempts.Load(), "expected exactly one failed attempt and one retry")
+}