Skip to content

Commit

Permalink
feat: add progress for io copy
Browse files Browse the repository at this point in the history
  • Loading branch information
xhofe committed Aug 31, 2022
1 parent 8e1ed40 commit 755f4b8
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 3 deletions.
2 changes: 1 addition & 1 deletion drivers/local/driver.go
Original file line number Diff line number Diff line change
Expand Up @@ -211,7 +211,7 @@ func (d *Local) Put(ctx context.Context, dstDir model.Obj, stream model.FileStre
_ = os.Remove(fullPath)
}
}()
err = utils.CopyWithCtx(ctx, out, stream)
err = utils.CopyWithCtx(ctx, out, stream, stream.GetSize(), up)
if err != nil {
return err
}
Expand Down
11 changes: 9 additions & 2 deletions pkg/utils/io.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,13 @@ 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 {
func CopyWithCtx(ctx context.Context, out io.Writer, in io.Reader, size int64, progress func(percentage int)) 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.
var finish int64 = 0
s := size / 100
_, err := io.Copy(out, readerFunc(func(p []byte) (int, error) {
// golang non-blocking channel: https://gobyexample.com/non-blocking-channel-operations
select {
Expand All @@ -28,7 +30,12 @@ func CopyWithCtx(ctx context.Context, out io.Writer, in io.Reader) error {
return 0, ctx.Err()
default:
// otherwise just run default io.Reader implementation
return in.Read(p)
n, err := in.Read(p)
if err == nil || err == io.EOF {
finish += int64(n)
progress(int(finish / s))
}
return n, err
}
}))
return err
Expand Down

0 comments on commit 755f4b8

Please sign in to comment.