From ee2bc99e4cc88d82aacbf4afe567044630a40798 Mon Sep 17 00:00:00 2001
From: Noah Hsu <i@nn.ci>
Date: Sat, 25 Jun 2022 15:14:03 +0800
Subject: [PATCH] feat: cancel copy for upload

---
 drivers/local/driver.go |  6 ++++--
 pkg/utils/ctx.go        |  4 +++-
 pkg/utils/io.go         | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 42 insertions(+), 3 deletions(-)
 create mode 100644 pkg/utils/io.go

diff --git a/drivers/local/driver.go b/drivers/local/driver.go
index ae21dfdfb15..016a529b1ac 100644
--- a/drivers/local/driver.go
+++ b/drivers/local/driver.go
@@ -3,7 +3,6 @@ package local
 import (
 	"context"
 	"github.com/alist-org/alist/v3/internal/errs"
-	"io"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -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)
 	}
diff --git a/pkg/utils/ctx.go b/pkg/utils/ctx.go
index d2a67f04f02..89f27f8eb3e 100644
--- a/pkg/utils/ctx.go
+++ b/pkg/utils/ctx.go
@@ -1,6 +1,8 @@
 package utils
 
-import "context"
+import (
+	"context"
+)
 
 func IsCanceled(ctx context.Context) bool {
 	select {
diff --git a/pkg/utils/io.go b/pkg/utils/io.go
new file mode 100644
index 00000000000..06c11f88907
--- /dev/null
+++ b/pkg/utils/io.go
@@ -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
+}