Skip to content

Commit

Permalink
feat: cancel copy for upload
Browse files Browse the repository at this point in the history
  • Loading branch information
xhofe committed Jun 25, 2022
1 parent 935416d commit ee2bc99
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 3 deletions.
6 changes: 4 additions & 2 deletions drivers/local/driver.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ package local
import (
"context"
"github.com/alist-org/alist/v3/internal/errs"
"io"
"io/ioutil"
"os"
"path/filepath"
Expand Down Expand Up @@ -144,8 +143,11 @@ func (d *Driver) Put(ctx context.Context, dstDir model.Obj, stream model.FileStr
}
defer func() {
_ = out.Close()
if errors.Is(err, context.Canceled) {
_ = os.Remove(fullPath)
}
}()
_, err = io.Copy(out, stream)
err = utils.CopyWithCtx(ctx, out, stream)
if err != nil {
return errors.Wrapf(err, "error while copy file %s", fullPath)
}
Expand Down
4 changes: 3 additions & 1 deletion pkg/utils/ctx.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
package utils

import "context"
import (
"context"
)

func IsCanceled(ctx context.Context) bool {
select {
Expand Down
35 changes: 35 additions & 0 deletions pkg/utils/io.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
package utils

import (
"context"
"io"
)

// here is some syntaxic sugar inspired by the Tomas Senart's video,
// it allows me to inline the Reader interface
type readerFunc func(p []byte) (n int, err error)

func (rf readerFunc) Read(p []byte) (n int, err error) { return rf(p) }

// CopyWithCtx slightly modified function signature:
// - context has been added in order to propagate cancelation
// - I do not return the number of bytes written, has it is not useful in my use case
func CopyWithCtx(ctx context.Context, out io.Writer, in io.Reader) error {
// Copy will call the Reader and Writer interface multiple time, in order
// to copy by chunk (avoiding loading the whole file in memory).
// I insert the ability to cancel before read time as it is the earliest
// possible in the call process.
_, err := io.Copy(out, readerFunc(func(p []byte) (int, error) {
// golang non-blocking channel: https://gobyexample.com/non-blocking-channel-operations
select {
// if context has been canceled
case <-ctx.Done():
// stop process and propagate "context canceled" error
return 0, ctx.Err()
default:
// otherwise just run default io.Reader implementation
return in.Read(p)
}
}))
return err
}

0 comments on commit ee2bc99

Please sign in to comment.