diff --git a/backend/pikpak/helper.go b/backend/pikpak/helper.go index 9bc8ea041..ee0327553 100644 --- a/backend/pikpak/helper.go +++ b/backend/pikpak/helper.go @@ -1,7 +1,6 @@ package pikpak import ( - "bytes" "context" "crypto/md5" "crypto/sha1" @@ -23,6 +22,7 @@ import ( "github.com/rclone/rclone/fs" "github.com/rclone/rclone/fs/config/configmap" "github.com/rclone/rclone/fs/fserrors" + "github.com/rclone/rclone/lib/multipart" "github.com/rclone/rclone/lib/rest" ) @@ -307,7 +307,9 @@ func readGcid(in io.Reader, size, threshold int64) (gcid string, out io.Reader, cleanup = func() {} // don't cache small files on disk to reduce wear of the disk - if size > threshold { + // + // files of unknown size go to disk as they may not fit in memory + if size < 0 || size > threshold { var tempFile *os.File // create the cache file @@ -340,13 +342,16 @@ func readGcid(in io.Reader, size, threshold int64) (gcid string, out io.Reader, // replace the already read source with a reader of our cached file out = tempFile } else { - buf := &bytes.Buffer{} - teeReader := io.TeeReader(in, buf) + rw := multipart.NewRW() + cleanup = func() { + _ = rw.Close() + } + teeReader := io.TeeReader(in, rw) if gcid, err = calcGcid(teeReader, size); err != nil { return } - out = buf + out = rw } return } diff --git a/backend/pikpak/helper_internal_test.go b/backend/pikpak/helper_internal_test.go new file mode 100644 index 000000000..1384a756d --- /dev/null +++ b/backend/pikpak/helper_internal_test.go @@ -0,0 +1,46 @@ +package pikpak + +import ( + "bytes" + "io" + "os" + "testing" + + "github.com/rclone/rclone/lib/pool" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestReadGcid checks readGcid hashes and replays its input from a pooled +// buffer for small known sizes, and from a temp file for large or unknown +// sizes, releasing the buffer on cleanup. +func TestReadGcid(t *testing.T) { + content := bytes.Repeat([]byte("pikpak"), 1024) + want, err := calcGcid(bytes.NewReader(content), int64(len(content))) + require.NoError(t, err) + + for _, test := range []struct { + name string + size int64 + threshold int64 + wantFile bool + }{ + {"memory", int64(len(content)), int64(len(content)), false}, + {"file", int64(len(content)), int64(len(content)) - 1, true}, + {"unknown size", -1, int64(len(content)), true}, + } { + t.Run(test.name, func(t *testing.T) { + inUse := pool.Global().InUse() + gcid, out, cleanup, err := readGcid(bytes.NewReader(content), test.size, test.threshold) + require.NoError(t, err) + assert.Equal(t, want, gcid) + _, isFile := out.(*os.File) + assert.Equal(t, test.wantFile, isFile, "expected temp file %v", test.wantFile) + got, err := io.ReadAll(out) + require.NoError(t, err) + assert.Equal(t, content, got) + cleanup() + assert.Equal(t, inUse, pool.Global().InUse(), "buffer should be returned to the pool") + }) + } +} diff --git a/backend/pikpak/pikpak.go b/backend/pikpak/pikpak.go index df5c3f499..a765efa06 100644 --- a/backend/pikpak/pikpak.go +++ b/backend/pikpak/pikpak.go @@ -224,7 +224,7 @@ Fill in for rclone to use a non root folder as its starting point. Advanced: true, }, { Name: "hash_memory_limit", - Help: "Files bigger than this will be cached on disk to calculate hash if required.", + Help: "Files bigger than this will be cached on disk to calculate hash if required.\n\nFiles of unknown size are always cached on disk.", Default: fs.SizeSuffix(10 * 1024 * 1024), Advanced: true, }, {