pikpak: use pooled memory for the gcid hash buffer and spool unknown-size uploads to disk
When the source can't be re-opened, the whole input is read once to compute its gcid before upload and held back for the upload proper. For inputs at or below --pikpak-hash-memory-limit this used a bytes.Buffer, a fresh heap allocation of up to the limit (and more while growing) per file. Hold the data in a buffer from the global memory pool instead so the pages are reused and released on cleanup. Inputs of unknown size were also always held in memory regardless of their length, as only sizes above the limit chose the temp file. Spool unknown sizes to the temp file so a large stream can't exhaust memory.
This commit is contained in:
@@ -1,7 +1,6 @@
|
|||||||
package pikpak
|
package pikpak
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"bytes"
|
|
||||||
"context"
|
"context"
|
||||||
"crypto/md5"
|
"crypto/md5"
|
||||||
"crypto/sha1"
|
"crypto/sha1"
|
||||||
@@ -23,6 +22,7 @@ import (
|
|||||||
"github.com/rclone/rclone/fs"
|
"github.com/rclone/rclone/fs"
|
||||||
"github.com/rclone/rclone/fs/config/configmap"
|
"github.com/rclone/rclone/fs/config/configmap"
|
||||||
"github.com/rclone/rclone/fs/fserrors"
|
"github.com/rclone/rclone/fs/fserrors"
|
||||||
|
"github.com/rclone/rclone/lib/multipart"
|
||||||
"github.com/rclone/rclone/lib/rest"
|
"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() {}
|
cleanup = func() {}
|
||||||
|
|
||||||
// don't cache small files on disk to reduce wear of the disk
|
// 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
|
var tempFile *os.File
|
||||||
|
|
||||||
// create the cache 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
|
// replace the already read source with a reader of our cached file
|
||||||
out = tempFile
|
out = tempFile
|
||||||
} else {
|
} else {
|
||||||
buf := &bytes.Buffer{}
|
rw := multipart.NewRW()
|
||||||
teeReader := io.TeeReader(in, buf)
|
cleanup = func() {
|
||||||
|
_ = rw.Close()
|
||||||
|
}
|
||||||
|
teeReader := io.TeeReader(in, rw)
|
||||||
|
|
||||||
if gcid, err = calcGcid(teeReader, size); err != nil {
|
if gcid, err = calcGcid(teeReader, size); err != nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
out = buf
|
out = rw
|
||||||
}
|
}
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -224,7 +224,7 @@ Fill in for rclone to use a non root folder as its starting point.
|
|||||||
Advanced: true,
|
Advanced: true,
|
||||||
}, {
|
}, {
|
||||||
Name: "hash_memory_limit",
|
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),
|
Default: fs.SizeSuffix(10 * 1024 * 1024),
|
||||||
Advanced: true,
|
Advanced: true,
|
||||||
}, {
|
}, {
|
||||||
|
|||||||
Reference in New Issue
Block a user