diff --git a/internal/offline_download/tool/add.go b/internal/offline_download/tool/add.go index 9ad8d055..68712122 100644 --- a/internal/offline_download/tool/add.go +++ b/internal/offline_download/tool/add.go @@ -16,6 +16,7 @@ const ( DeleteOnUploadSucceed DeletePolicy = "delete_on_upload_succeed" DeleteOnUploadFailed DeletePolicy = "delete_on_upload_failed" DeleteNever DeletePolicy = "delete_never" + DeleteAlways DeletePolicy = "delete_always" ) type AddURLArgs struct { diff --git a/internal/offline_download/tool/download.go b/internal/offline_download/tool/download.go index 36ab6c82..f4ff164c 100644 --- a/internal/offline_download/tool/download.go +++ b/internal/offline_download/tool/download.go @@ -125,10 +125,11 @@ func (t *DownloadTask) Complete() error { for i, _ := range files { file := files[i] TransferTaskManager.Add(&TransferTask{ - file: file, - dstDirPath: t.DstDirPath, - wg: &wg, - tempDir: t.TempDir, + file: file, + dstDirPath: t.DstDirPath, + wg: &wg, + tempDir: t.TempDir, + deletePolicy: t.DeletePolicy, }) } return nil diff --git a/internal/offline_download/tool/transfer.go b/internal/offline_download/tool/transfer.go index c39e4ba0..0744b333 100644 --- a/internal/offline_download/tool/transfer.go +++ b/internal/offline_download/tool/transfer.go @@ -9,16 +9,18 @@ import ( "github.com/pkg/errors" log "github.com/sirupsen/logrus" "github.com/xhofe/tache" + "os" "path/filepath" "sync" ) type TransferTask struct { tache.Base - file File - dstDirPath string - wg *sync.WaitGroup - tempDir string + file File + dstDirPath string + wg *sync.WaitGroup + tempDir string + deletePolicy DeletePolicy } func (t *TransferTask) Run() error { @@ -61,6 +63,24 @@ func (t *TransferTask) GetStatus() string { return "transferring" } +func (t *TransferTask) OnSucceeded() { + if t.deletePolicy == DeleteOnUploadSucceed || t.deletePolicy == DeleteAlways { + err := os.Remove(t.file.Path) + if err != nil { + log.Errorf("failed to delete file %s, error: %s", t.file.Path, err.Error()) + } + } +} + +func (t *TransferTask) OnFailed() { + if t.deletePolicy == DeleteOnUploadFailed || t.deletePolicy == DeleteAlways { + err := os.Remove(t.file.Path) + if err != nil { + log.Errorf("failed to delete file %s, error: %s", t.file.Path, err.Error()) + } + } +} + var ( TransferTaskManager *tache.Manager[*TransferTask] )