chore: 初始化 Cloudreve 源码仓库
This commit is contained in:
@@ -0,0 +1,705 @@
|
||||
package workflows
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/cloudreve/Cloudreve/v4/application/dependency"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent/task"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory/types"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/encrypt"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/fs"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/manager"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/manager/entitysource"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/hashid"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/logging"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/queue"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/util"
|
||||
"github.com/gofrs/uuid"
|
||||
)
|
||||
|
||||
type (
|
||||
CreateArchiveTask struct {
|
||||
*queue.DBTask
|
||||
|
||||
l logging.Logger
|
||||
state *CreateArchiveTaskState
|
||||
progress queue.Progresses
|
||||
node cluster.Node
|
||||
}
|
||||
|
||||
CreateArchiveTaskPhase string
|
||||
|
||||
CreateArchiveTaskState struct {
|
||||
Uris []string `json:"uris,omitempty"`
|
||||
Dst string `json:"dst,omitempty"`
|
||||
TempPath string `json:"temp_path,omitempty"`
|
||||
ArchiveFile string `json:"archive_file,omitempty"`
|
||||
Phase CreateArchiveTaskPhase `json:"phase,omitempty"`
|
||||
SlaveUploadTaskID int `json:"slave__upload_task_id,omitempty"`
|
||||
SlaveArchiveTaskID int `json:"slave__archive_task_id,omitempty"`
|
||||
SlaveCompressState *SlaveCreateArchiveTaskState `json:"slave_compress_state,omitempty"`
|
||||
Failed int `json:"failed,omitempty"`
|
||||
NodeState `json:",inline"`
|
||||
}
|
||||
)
|
||||
|
||||
const (
|
||||
CreateArchiveTaskPhaseNotStarted CreateArchiveTaskPhase = "not_started"
|
||||
CreateArchiveTaskPhaseCompressFiles CreateArchiveTaskPhase = "compress_files"
|
||||
CreateArchiveTaskPhaseUploadArchive CreateArchiveTaskPhase = "upload_archive"
|
||||
|
||||
CreateArchiveTaskPhaseAwaitSlaveCompressing CreateArchiveTaskPhase = "await_slave_compressing"
|
||||
CreateArchiveTaskPhaseCreateAndAwaitSlaveUploading CreateArchiveTaskPhase = "await_slave_uploading"
|
||||
CreateArchiveTaskPhaseCompleteUpload CreateArchiveTaskPhase = "complete_upload"
|
||||
|
||||
ProgressTypeArchiveCount = "archive_count"
|
||||
ProgressTypeArchiveSize = "archive_size"
|
||||
ProgressTypeUpload = "upload"
|
||||
ProgressTypeUploadCount = "upload_count"
|
||||
)
|
||||
|
||||
func init() {
|
||||
queue.RegisterResumableTaskFactory(queue.CreateArchiveTaskType, NewCreateArchiveTaskFromModel)
|
||||
}
|
||||
|
||||
// NewCreateArchiveTask creates a new CreateArchiveTask
|
||||
func NewCreateArchiveTask(ctx context.Context, src []string, dst string) (queue.Task, error) {
|
||||
state := &CreateArchiveTaskState{
|
||||
Uris: src,
|
||||
Dst: dst,
|
||||
NodeState: NodeState{},
|
||||
}
|
||||
stateBytes, err := json.Marshal(state)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal state: %w", err)
|
||||
}
|
||||
|
||||
t := &CreateArchiveTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: &ent.Task{
|
||||
Type: queue.CreateArchiveTaskType,
|
||||
CorrelationID: logging.CorrelationID(ctx),
|
||||
PrivateState: string(stateBytes),
|
||||
PublicState: &types.TaskPublicState{},
|
||||
},
|
||||
DirectOwner: inventory.UserFromContext(ctx),
|
||||
},
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
func NewCreateArchiveTaskFromModel(task *ent.Task) queue.Task {
|
||||
return &CreateArchiveTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: task,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (m *CreateArchiveTask) Do(ctx context.Context) (task.Status, error) {
|
||||
dep := dependency.FromContext(ctx)
|
||||
m.l = dep.Logger()
|
||||
|
||||
m.Lock()
|
||||
if m.progress == nil {
|
||||
m.progress = make(queue.Progresses)
|
||||
}
|
||||
m.Unlock()
|
||||
|
||||
// unmarshal state
|
||||
state := &CreateArchiveTaskState{}
|
||||
if err := json.Unmarshal([]byte(m.State()), state); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to unmarshal state: %w", err)
|
||||
}
|
||||
m.state = state
|
||||
|
||||
// select node
|
||||
node, err := allocateNode(ctx, dep, &m.state.NodeState, types.NodeCapabilityCreateArchive)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to allocate node: %w", err)
|
||||
}
|
||||
m.node = node
|
||||
|
||||
next := task.StatusCompleted
|
||||
|
||||
if m.node.IsMaster() {
|
||||
// Initialize temp folder
|
||||
// Compress files
|
||||
// Upload files to dst
|
||||
switch m.state.Phase {
|
||||
case CreateArchiveTaskPhaseNotStarted, "":
|
||||
next, err = m.initializeTempFolder(ctx, dep)
|
||||
case CreateArchiveTaskPhaseCompressFiles:
|
||||
next, err = m.createArchiveFile(ctx, dep)
|
||||
case CreateArchiveTaskPhaseUploadArchive:
|
||||
next, err = m.uploadArchive(ctx, dep)
|
||||
default:
|
||||
next, err = task.StatusError, fmt.Errorf("unknown phase %q: %w", m.state.Phase, queue.CriticalErr)
|
||||
}
|
||||
} else {
|
||||
// Listing all files and send to slave node for compressing
|
||||
// Await compressing and send to slave for uploading
|
||||
// Await uploading and complete upload
|
||||
switch m.state.Phase {
|
||||
case CreateArchiveTaskPhaseNotStarted, "":
|
||||
next, err = m.listEntitiesAndSendToSlave(ctx, dep)
|
||||
case CreateArchiveTaskPhaseAwaitSlaveCompressing:
|
||||
next, err = m.awaitSlaveCompressing(ctx, dep)
|
||||
case CreateArchiveTaskPhaseCreateAndAwaitSlaveUploading:
|
||||
next, err = m.createAndAwaitSlaveUploading(ctx, dep)
|
||||
case CreateArchiveTaskPhaseCompleteUpload:
|
||||
next, err = m.completeUpload(ctx, dep)
|
||||
default:
|
||||
next, err = task.StatusError, fmt.Errorf("unknown phase %q: %w", m.state.Phase, queue.CriticalErr)
|
||||
}
|
||||
}
|
||||
|
||||
newStateStr, marshalErr := json.Marshal(m.state)
|
||||
if marshalErr != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to marshal state: %w", marshalErr)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.Task.PrivateState = string(newStateStr)
|
||||
m.Unlock()
|
||||
return next, err
|
||||
}
|
||||
|
||||
func (m *CreateArchiveTask) Cleanup(ctx context.Context) error {
|
||||
if m.state.SlaveCompressState != nil && m.state.SlaveCompressState.TempPath != "" && m.node != nil {
|
||||
if err := m.node.CleanupFolders(context.Background(), m.state.SlaveCompressState.TempPath); err != nil {
|
||||
m.l.Warning("Failed to cleanup slave temp folder %s: %s", m.state.SlaveCompressState.TempPath, err)
|
||||
}
|
||||
}
|
||||
|
||||
if m.state.TempPath != "" {
|
||||
time.Sleep(time.Duration(1) * time.Second)
|
||||
return os.RemoveAll(m.state.TempPath)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *CreateArchiveTask) initializeTempFolder(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
tempPath, err := prepareTempFolder(ctx, dep, m)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to prepare temp folder: %w", err)
|
||||
}
|
||||
|
||||
m.state.TempPath = tempPath
|
||||
m.state.Phase = CreateArchiveTaskPhaseCompressFiles
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *CreateArchiveTask) listEntitiesAndSendToSlave(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
uris, err := fs.NewUriFromStrings(m.state.Uris...)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create uri from strings: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
payload := &SlaveCreateArchiveTaskState{
|
||||
Entities: make([]SlaveCreateArchiveEntity, 0, len(uris)),
|
||||
Policies: make(map[int]*ent.StoragePolicy),
|
||||
}
|
||||
|
||||
user := inventory.UserFromContext(ctx)
|
||||
fm := manager.NewFileManager(dep, user)
|
||||
storagePolicyClient := dep.StoragePolicyClient()
|
||||
masterKey, _ := dep.MasterEncryptKeyVault(ctx).GetMasterKey(ctx)
|
||||
|
||||
failed, err := fm.CreateArchive(ctx, uris, io.Discard,
|
||||
fs.WithDryRun(func(name string, e fs.Entity) {
|
||||
entityModel, err := decryptEntityKeyIfNeeded(masterKey, e.Model())
|
||||
if err != nil {
|
||||
m.l.Warning("Failed to decrypt entity key for %q: %s", name, err)
|
||||
return
|
||||
}
|
||||
|
||||
payload.Entities = append(payload.Entities, SlaveCreateArchiveEntity{
|
||||
Entity: entityModel,
|
||||
Path: name,
|
||||
})
|
||||
if _, ok := payload.Policies[e.PolicyID()]; !ok {
|
||||
policy, err := storagePolicyClient.GetPolicyByID(ctx, e.PolicyID())
|
||||
if err != nil {
|
||||
m.l.Warning("Failed to get policy %d: %s", e.PolicyID(), err)
|
||||
} else {
|
||||
payload.Policies[e.PolicyID()] = policy
|
||||
}
|
||||
}
|
||||
}),
|
||||
fs.WithMaxArchiveSize(user.Edges.Group.Settings.CompressSize),
|
||||
)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to compress files: %w", err)
|
||||
}
|
||||
|
||||
m.state.Failed = failed
|
||||
payloadStr, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to marshal payload: %w", err)
|
||||
}
|
||||
|
||||
taskId, err := m.node.CreateTask(ctx, queue.SlaveCreateArchiveTaskType, string(payloadStr))
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create slave task: %w", err)
|
||||
}
|
||||
|
||||
m.state.Phase = CreateArchiveTaskPhaseAwaitSlaveCompressing
|
||||
m.state.SlaveArchiveTaskID = taskId
|
||||
m.ResumeAfter((10 * time.Second))
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *CreateArchiveTask) awaitSlaveCompressing(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
t, err := m.node.GetTask(ctx, m.state.SlaveArchiveTaskID, false)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get slave task: %w", err)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.state.NodeState.progress = t.Progress
|
||||
m.Unlock()
|
||||
|
||||
m.state.SlaveCompressState = &SlaveCreateArchiveTaskState{}
|
||||
if err := json.Unmarshal([]byte(t.PrivateState), m.state.SlaveCompressState); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to unmarshal slave compress state: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
if t.Status == task.StatusError {
|
||||
return task.StatusError, fmt.Errorf("slave task failed: %s (%w)", t.Error, queue.CriticalErr)
|
||||
}
|
||||
|
||||
if t.Status == task.StatusCanceled {
|
||||
return task.StatusError, fmt.Errorf("slave task canceled (%w)", queue.CriticalErr)
|
||||
}
|
||||
|
||||
if t.Status == task.StatusCompleted {
|
||||
m.state.Phase = CreateArchiveTaskPhaseCreateAndAwaitSlaveUploading
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
m.l.Info("Slave task %d is still compressing, resume after 30s.", m.state.SlaveArchiveTaskID)
|
||||
m.ResumeAfter((time.Second * 30))
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *CreateArchiveTask) createAndAwaitSlaveUploading(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
u := inventory.UserFromContext(ctx)
|
||||
|
||||
if m.state.SlaveUploadTaskID == 0 {
|
||||
dst, err := fs.NewUriFromString(m.state.Dst)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to parse dst uri %q: %s (%w)", m.state.Dst, err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
// Create slave upload task
|
||||
payload := &SlaveUploadTaskState{
|
||||
Files: []SlaveUploadEntity{
|
||||
{
|
||||
Size: m.state.SlaveCompressState.CompressedSize,
|
||||
Uri: dst,
|
||||
Src: m.state.SlaveCompressState.ZipFilePath,
|
||||
},
|
||||
},
|
||||
MaxParallel: dep.SettingProvider().MaxParallelTransfer(ctx),
|
||||
UserID: u.ID,
|
||||
}
|
||||
|
||||
payloadStr, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to marshal payload: %w", err)
|
||||
}
|
||||
|
||||
taskId, err := m.node.CreateTask(ctx, queue.SlaveUploadTaskType, string(payloadStr))
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create slave task: %w", err)
|
||||
}
|
||||
|
||||
m.state.NodeState.progress = nil
|
||||
m.state.SlaveUploadTaskID = taskId
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
m.l.Info("Checking slave upload task %d...", m.state.SlaveUploadTaskID)
|
||||
t, err := m.node.GetTask(ctx, m.state.SlaveUploadTaskID, true)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get slave task: %w", err)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.state.NodeState.progress = t.Progress
|
||||
m.Unlock()
|
||||
|
||||
if t.Status == task.StatusError {
|
||||
return task.StatusError, fmt.Errorf("slave task failed: %s (%w)", t.Error, queue.CriticalErr)
|
||||
}
|
||||
|
||||
if t.Status == task.StatusCanceled {
|
||||
return task.StatusError, fmt.Errorf("slave task canceled (%w)", queue.CriticalErr)
|
||||
}
|
||||
|
||||
if t.Status == task.StatusCompleted {
|
||||
m.state.Phase = CreateArchiveTaskPhaseCompleteUpload
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
m.l.Info("Slave task %d is still uploading, resume after 30s.", m.state.SlaveUploadTaskID)
|
||||
m.ResumeAfter(time.Second * 30)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *CreateArchiveTask) completeUpload(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
return task.StatusCompleted, nil
|
||||
}
|
||||
|
||||
func (m *CreateArchiveTask) createArchiveFile(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
uris, err := fs.NewUriFromStrings(m.state.Uris...)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create uri from strings: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
user := inventory.UserFromContext(ctx)
|
||||
fm := manager.NewFileManager(dep, user)
|
||||
|
||||
// Create temp zip file
|
||||
fileName := fmt.Sprintf("%s.zip", uuid.Must(uuid.NewV4()))
|
||||
zipFilePath := filepath.Join(
|
||||
m.state.TempPath,
|
||||
fileName,
|
||||
)
|
||||
zipFile, err := util.CreatNestedFile(zipFilePath)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create zip file: %w", err)
|
||||
}
|
||||
|
||||
defer zipFile.Close()
|
||||
|
||||
// Start compressing
|
||||
m.Lock()
|
||||
m.progress[ProgressTypeArchiveCount] = &queue.Progress{}
|
||||
m.progress[ProgressTypeArchiveSize] = &queue.Progress{}
|
||||
m.Unlock()
|
||||
failed, err := fm.CreateArchive(ctx, uris, zipFile,
|
||||
fs.WithArchiveCompression(true),
|
||||
fs.WithMaxArchiveSize(user.Edges.Group.Settings.CompressSize),
|
||||
fs.WithProgressFunc(func(current, diff int64, total int64) {
|
||||
atomic.AddInt64(&m.progress[ProgressTypeArchiveSize].Current, diff)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeArchiveCount].Current, 1)
|
||||
}),
|
||||
)
|
||||
if err != nil {
|
||||
zipFile.Close()
|
||||
_ = os.Remove(zipFilePath)
|
||||
return task.StatusError, fmt.Errorf("failed to compress files: %w", err)
|
||||
}
|
||||
|
||||
m.state.Failed = failed
|
||||
m.Lock()
|
||||
delete(m.progress, ProgressTypeArchiveSize)
|
||||
delete(m.progress, ProgressTypeArchiveCount)
|
||||
m.Unlock()
|
||||
|
||||
m.state.Phase = CreateArchiveTaskPhaseUploadArchive
|
||||
m.state.ArchiveFile = fileName
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *CreateArchiveTask) uploadArchive(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
fm := manager.NewFileManager(dep, inventory.UserFromContext(ctx))
|
||||
zipFilePath := filepath.Join(
|
||||
m.state.TempPath,
|
||||
m.state.ArchiveFile,
|
||||
)
|
||||
|
||||
m.l.Info("Uploading archive file %s to %s...", zipFilePath, m.state.Dst)
|
||||
|
||||
uri, err := fs.NewUriFromString(m.state.Dst)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf(
|
||||
"failed to parse dst uri %q: %s (%w)",
|
||||
m.state.Dst,
|
||||
err,
|
||||
queue.CriticalErr,
|
||||
)
|
||||
}
|
||||
|
||||
file, err := os.Open(zipFilePath)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to open compressed archive %q: %s", m.state.ArchiveFile, err)
|
||||
}
|
||||
defer file.Close()
|
||||
fi, err := file.Stat()
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get file info: %w", err)
|
||||
}
|
||||
size := fi.Size()
|
||||
|
||||
m.Lock()
|
||||
m.progress[ProgressTypeUpload] = &queue.Progress{}
|
||||
m.Unlock()
|
||||
fileData := &fs.UploadRequest{
|
||||
Props: &fs.UploadProps{
|
||||
Uri: uri,
|
||||
Size: size,
|
||||
},
|
||||
ProgressFunc: func(current, diff int64, total int64) {
|
||||
atomic.StoreInt64(&m.progress[ProgressTypeUpload].Current, current)
|
||||
atomic.StoreInt64(&m.progress[ProgressTypeUpload].Total, total)
|
||||
},
|
||||
File: file,
|
||||
Seeker: file,
|
||||
}
|
||||
|
||||
_, err = fm.Update(ctx, fileData)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to upload archive file: %w", err)
|
||||
}
|
||||
|
||||
return task.StatusCompleted, nil
|
||||
}
|
||||
|
||||
func (m *CreateArchiveTask) Progress(ctx context.Context) queue.Progresses {
|
||||
m.Lock()
|
||||
defer m.Unlock()
|
||||
|
||||
if m.state.NodeState.progress != nil {
|
||||
merged := make(queue.Progresses)
|
||||
for k, v := range m.progress {
|
||||
merged[k] = v
|
||||
}
|
||||
|
||||
for k, v := range m.state.NodeState.progress {
|
||||
merged[k] = v
|
||||
}
|
||||
|
||||
return merged
|
||||
}
|
||||
return m.progress
|
||||
}
|
||||
|
||||
func (m *CreateArchiveTask) Summarize(hasher hashid.Encoder) *queue.Summary {
|
||||
// unmarshal state
|
||||
if m.state == nil {
|
||||
if err := json.Unmarshal([]byte(m.State()), &m.state); err != nil {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
failed := m.state.Failed
|
||||
if m.state.SlaveCompressState != nil {
|
||||
failed = m.state.SlaveCompressState.Failed
|
||||
}
|
||||
|
||||
return &queue.Summary{
|
||||
NodeID: m.state.NodeID,
|
||||
Phase: string(m.state.Phase),
|
||||
Props: map[string]any{
|
||||
SummaryKeySrcMultiple: m.state.Uris,
|
||||
SummaryKeyDst: m.state.Dst,
|
||||
SummaryKeyFailed: failed,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
type (
|
||||
SlaveCreateArchiveEntity struct {
|
||||
Entity *ent.Entity `json:"entity"`
|
||||
Path string `json:"path"`
|
||||
}
|
||||
SlaveCreateArchiveTaskState struct {
|
||||
Entities []SlaveCreateArchiveEntity `json:"entities"`
|
||||
Policies map[int]*ent.StoragePolicy `json:"policies"`
|
||||
CompressedSize int64 `json:"compressed_size"`
|
||||
TempPath string `json:"temp_path"`
|
||||
ZipFilePath string `json:"zip_file_path"`
|
||||
Failed int `json:"failed"`
|
||||
}
|
||||
SlaveCreateArchiveTask struct {
|
||||
*queue.InMemoryTask
|
||||
|
||||
mu sync.RWMutex
|
||||
progress queue.Progresses
|
||||
l logging.Logger
|
||||
state *SlaveCreateArchiveTaskState
|
||||
}
|
||||
)
|
||||
|
||||
// NewSlaveCreateArchiveTask creates a new SlaveCreateArchiveTask from raw private state
|
||||
func NewSlaveCreateArchiveTask(ctx context.Context, props *types.SlaveTaskProps, id int, state string) queue.Task {
|
||||
return &SlaveCreateArchiveTask{
|
||||
InMemoryTask: &queue.InMemoryTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: &ent.Task{
|
||||
ID: id,
|
||||
CorrelationID: logging.CorrelationID(ctx),
|
||||
PublicState: &types.TaskPublicState{
|
||||
SlaveTaskProps: props,
|
||||
},
|
||||
PrivateState: state,
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
progress: make(queue.Progresses),
|
||||
}
|
||||
}
|
||||
|
||||
func (t *SlaveCreateArchiveTask) Do(ctx context.Context) (task.Status, error) {
|
||||
ctx = prepareSlaveTaskCtx(ctx, t.Model().PublicState.SlaveTaskProps)
|
||||
dep := dependency.FromContext(ctx)
|
||||
t.l = dep.Logger()
|
||||
fm := manager.NewFileManager(dep, nil)
|
||||
|
||||
// unmarshal state
|
||||
state := &SlaveCreateArchiveTaskState{}
|
||||
if err := json.Unmarshal([]byte(t.State()), state); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to unmarshal state: %w", err)
|
||||
}
|
||||
|
||||
t.state = state
|
||||
|
||||
totalFiles := int64(0)
|
||||
totalFileSize := int64(0)
|
||||
for _, e := range t.state.Entities {
|
||||
totalFiles++
|
||||
totalFileSize += e.Entity.Size
|
||||
}
|
||||
|
||||
t.Lock()
|
||||
t.progress[ProgressTypeArchiveCount] = &queue.Progress{Total: totalFiles}
|
||||
t.progress[ProgressTypeArchiveSize] = &queue.Progress{Total: totalFileSize}
|
||||
t.Unlock()
|
||||
|
||||
// 3. Create temp workspace
|
||||
tempPath, err := prepareTempFolder(ctx, dep, t)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to prepare temp folder: %w", err)
|
||||
}
|
||||
t.state.TempPath = tempPath
|
||||
|
||||
// 2. Create archive file
|
||||
fileName := fmt.Sprintf("%s.zip", uuid.Must(uuid.NewV4()))
|
||||
zipFilePath := filepath.Join(
|
||||
t.state.TempPath,
|
||||
fileName,
|
||||
)
|
||||
zipFile, err := util.CreatNestedFile(zipFilePath)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create zip file: %w", err)
|
||||
}
|
||||
|
||||
defer zipFile.Close()
|
||||
|
||||
zipWriter := zip.NewWriter(zipFile)
|
||||
defer zipWriter.Close()
|
||||
|
||||
// 3. Download each entity and write into zip file
|
||||
for _, e := range t.state.Entities {
|
||||
policy, ok := t.state.Policies[e.Entity.StoragePolicyEntities]
|
||||
if !ok {
|
||||
state.Failed++
|
||||
t.l.Warning("Policy not found for entity %d, skipping...", e.Entity.ID)
|
||||
continue
|
||||
}
|
||||
|
||||
entity := fs.NewEntity(e.Entity)
|
||||
es, err := fm.GetEntitySource(ctx, 0,
|
||||
fs.WithEntity(entity),
|
||||
fs.WithPolicy(fm.CastStoragePolicyOnSlave(ctx, policy)),
|
||||
)
|
||||
if err != nil {
|
||||
state.Failed++
|
||||
t.l.Warning("Failed to get entity source for entity %d: %s, skipping...", e.Entity.ID, err)
|
||||
continue
|
||||
}
|
||||
|
||||
// Write to zip file
|
||||
header := &zip.FileHeader{
|
||||
Name: e.Path,
|
||||
Modified: entity.UpdatedAt(),
|
||||
UncompressedSize64: uint64(entity.Size()),
|
||||
Method: zip.Deflate,
|
||||
}
|
||||
|
||||
writer, err := zipWriter.CreateHeader(header)
|
||||
if err != nil {
|
||||
es.Close()
|
||||
state.Failed++
|
||||
t.l.Warning("Failed to create zip header for %s: %s, skipping...", e.Path, err)
|
||||
continue
|
||||
}
|
||||
|
||||
es.Apply(entitysource.WithContext(ctx))
|
||||
_, err = io.Copy(writer, es)
|
||||
es.Close()
|
||||
if err != nil {
|
||||
state.Failed++
|
||||
t.l.Warning("Failed to write entity %d to zip file: %s, skipping...", e.Entity.ID, err)
|
||||
}
|
||||
|
||||
atomic.AddInt64(&t.progress[ProgressTypeArchiveSize].Current, entity.Size())
|
||||
atomic.AddInt64(&t.progress[ProgressTypeArchiveCount].Current, 1)
|
||||
}
|
||||
|
||||
zipWriter.Close()
|
||||
stat, err := zipFile.Stat()
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get compressed file info: %w", err)
|
||||
}
|
||||
|
||||
t.state.CompressedSize = stat.Size()
|
||||
t.state.ZipFilePath = zipFilePath
|
||||
// Clear unused fields to save space
|
||||
t.state.Entities = nil
|
||||
t.state.Policies = nil
|
||||
|
||||
newStateStr, marshalErr := json.Marshal(t.state)
|
||||
if marshalErr != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to marshal state: %w", marshalErr)
|
||||
}
|
||||
|
||||
t.Lock()
|
||||
t.Task.PrivateState = string(newStateStr)
|
||||
t.Unlock()
|
||||
return task.StatusCompleted, nil
|
||||
}
|
||||
|
||||
func (m *SlaveCreateArchiveTask) Progress(ctx context.Context) queue.Progresses {
|
||||
m.Lock()
|
||||
defer m.Unlock()
|
||||
|
||||
return m.progress
|
||||
}
|
||||
|
||||
func decryptEntityKeyIfNeeded(masterKey []byte, entity *ent.Entity) (*ent.Entity, error) {
|
||||
if entity.Props == nil || entity.Props.EncryptMetadata == nil || entity.Props.EncryptMetadata.KeyPlainText != nil {
|
||||
return entity, nil
|
||||
}
|
||||
|
||||
decryptedKey, err := encrypt.DecryptWithMasterKey(masterKey, entity.Props.EncryptMetadata.Key)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to decrypt entity key: %w", err)
|
||||
}
|
||||
|
||||
entity.Props.EncryptMetadata.KeyPlainText = decryptedKey
|
||||
entity.Props.EncryptMetadata.Key = nil
|
||||
return entity, nil
|
||||
}
|
||||
@@ -0,0 +1,847 @@
|
||||
package workflows
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/cloudreve/Cloudreve/v4/application/dependency"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent/task"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory/types"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/fs"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/fs/dbfs"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/manager"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/hashid"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/logging"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/queue"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/util"
|
||||
"github.com/gofrs/uuid"
|
||||
"github.com/mholt/archives"
|
||||
)
|
||||
|
||||
type (
|
||||
ExtractArchiveTask struct {
|
||||
*queue.DBTask
|
||||
|
||||
l logging.Logger
|
||||
state *ExtractArchiveTaskState
|
||||
progress queue.Progresses
|
||||
node cluster.Node
|
||||
}
|
||||
ExtractArchiveTaskPhase string
|
||||
ExtractArchiveTaskState struct {
|
||||
Uri string `json:"uri,omitempty"`
|
||||
Encoding string `json:"encoding,omitempty"`
|
||||
Dst string `json:"dst,omitempty"`
|
||||
TempPath string `json:"temp_path,omitempty"`
|
||||
TempZipFilePath string `json:"temp_zip_file_path,omitempty"`
|
||||
ProcessedCursor string `json:"processed_cursor,omitempty"`
|
||||
SlaveTaskID int `json:"slave_task_id,omitempty"`
|
||||
Password string `json:"password,omitempty"`
|
||||
FileMask []string `json:"file_mask,omitempty"`
|
||||
NodeState `json:",inline"`
|
||||
Phase ExtractArchiveTaskPhase `json:"phase,omitempty"`
|
||||
}
|
||||
)
|
||||
|
||||
const (
|
||||
ExtractArchivePhaseNotStarted ExtractArchiveTaskPhase = ""
|
||||
ExtractArchivePhaseDownloadZip ExtractArchiveTaskPhase = "download_zip"
|
||||
ExtractArchivePhaseAwaitSlaveComplete ExtractArchiveTaskPhase = "await_slave_complete"
|
||||
|
||||
ProgressTypeExtractCount = "extract_count"
|
||||
ProgressTypeExtractSize = "extract_size"
|
||||
ProgressTypeDownload = "download"
|
||||
|
||||
SummaryKeySrc = "src"
|
||||
SummaryKeySrcPhysical = "src_physical"
|
||||
SummaryKeyDst = "dst"
|
||||
)
|
||||
|
||||
func init() {
|
||||
queue.RegisterResumableTaskFactory(queue.ExtractArchiveTaskType, NewExtractArchiveTaskFromModel)
|
||||
}
|
||||
|
||||
// NewExtractArchiveTask creates a new ExtractArchiveTask
|
||||
func NewExtractArchiveTask(ctx context.Context, src, dst, encoding, password string, mask []string) (queue.Task, error) {
|
||||
state := &ExtractArchiveTaskState{
|
||||
Uri: src,
|
||||
Dst: dst,
|
||||
Encoding: encoding,
|
||||
NodeState: NodeState{},
|
||||
Password: password,
|
||||
FileMask: mask,
|
||||
}
|
||||
stateBytes, err := json.Marshal(state)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal state: %w", err)
|
||||
}
|
||||
|
||||
t := &ExtractArchiveTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: &ent.Task{
|
||||
Type: queue.ExtractArchiveTaskType,
|
||||
CorrelationID: logging.CorrelationID(ctx),
|
||||
PrivateState: string(stateBytes),
|
||||
PublicState: &types.TaskPublicState{},
|
||||
},
|
||||
DirectOwner: inventory.UserFromContext(ctx),
|
||||
},
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
func NewExtractArchiveTaskFromModel(task *ent.Task) queue.Task {
|
||||
return &ExtractArchiveTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: task,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (m *ExtractArchiveTask) Do(ctx context.Context) (task.Status, error) {
|
||||
dep := dependency.FromContext(ctx)
|
||||
m.l = dep.Logger()
|
||||
|
||||
m.Lock()
|
||||
if m.progress == nil {
|
||||
m.progress = make(queue.Progresses)
|
||||
}
|
||||
m.Unlock()
|
||||
|
||||
// unmarshal state
|
||||
state := &ExtractArchiveTaskState{}
|
||||
if err := json.Unmarshal([]byte(m.State()), state); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to unmarshal state: %w", err)
|
||||
}
|
||||
m.state = state
|
||||
|
||||
// select node
|
||||
node, err := allocateNode(ctx, dep, &m.state.NodeState, types.NodeCapabilityExtractArchive)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to allocate node: %w", err)
|
||||
}
|
||||
m.node = node
|
||||
|
||||
next := task.StatusCompleted
|
||||
|
||||
if node.IsMaster() {
|
||||
switch m.state.Phase {
|
||||
case ExtractArchivePhaseNotStarted:
|
||||
next, err = m.masterExtractArchive(ctx, dep)
|
||||
case ExtractArchivePhaseDownloadZip:
|
||||
next, err = m.masterDownloadZip(ctx, dep)
|
||||
default:
|
||||
next, err = task.StatusError, fmt.Errorf("unknown phase %q: %w", m.state.Phase, queue.CriticalErr)
|
||||
}
|
||||
} else {
|
||||
switch m.state.Phase {
|
||||
case ExtractArchivePhaseNotStarted:
|
||||
next, err = m.createSlaveExtractTask(ctx, dep)
|
||||
case ExtractArchivePhaseAwaitSlaveComplete:
|
||||
next, err = m.awaitSlaveExtractComplete(ctx, dep)
|
||||
default:
|
||||
next, err = task.StatusError, fmt.Errorf("unknown phase %q: %w", m.state.Phase, queue.CriticalErr)
|
||||
}
|
||||
}
|
||||
|
||||
newStateStr, marshalErr := json.Marshal(m.state)
|
||||
if marshalErr != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to marshal state: %w", marshalErr)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.Task.PrivateState = string(newStateStr)
|
||||
m.Unlock()
|
||||
return next, err
|
||||
}
|
||||
|
||||
func (m *ExtractArchiveTask) createSlaveExtractTask(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
uri, err := fs.NewUriFromString(m.state.Uri)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to parse src uri: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
user := inventory.UserFromContext(ctx)
|
||||
fm := manager.NewFileManager(dep, user)
|
||||
|
||||
// Get entity source to extract
|
||||
archiveFile, err := fm.Get(ctx, uri, dbfs.WithFileEntities(), dbfs.WithRequiredCapabilities(dbfs.NavigatorCapabilityDownloadFile), dbfs.WithNotRoot())
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get archive file: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
// Validate file size
|
||||
if user.Edges.Group.Settings.DecompressSize > 0 && archiveFile.Size() > user.Edges.Group.Settings.DecompressSize {
|
||||
return task.StatusError,
|
||||
fmt.Errorf("file size %d exceeds the limit %d (%w)", archiveFile.Size(), user.Edges.Group.Settings.DecompressSize, queue.CriticalErr)
|
||||
}
|
||||
|
||||
// Create slave task
|
||||
storagePolicyClient := dep.StoragePolicyClient()
|
||||
policy, err := storagePolicyClient.GetPolicyByID(ctx, archiveFile.PrimaryEntity().PolicyID())
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get policy: %w", err)
|
||||
}
|
||||
|
||||
masterKey, _ := dep.MasterEncryptKeyVault(ctx).GetMasterKey(ctx)
|
||||
entityModel, err := decryptEntityKeyIfNeeded(masterKey, archiveFile.PrimaryEntity().Model())
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to decrypt entity key for archive file %q: %s", archiveFile.DisplayName(), err)
|
||||
}
|
||||
|
||||
payload := &SlaveExtractArchiveTaskState{
|
||||
FileName: archiveFile.DisplayName(),
|
||||
Entity: entityModel,
|
||||
Policy: policy,
|
||||
Encoding: m.state.Encoding,
|
||||
Dst: m.state.Dst,
|
||||
UserID: user.ID,
|
||||
Password: m.state.Password,
|
||||
FileMask: m.state.FileMask,
|
||||
}
|
||||
|
||||
payloadStr, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to marshal payload: %w", err)
|
||||
}
|
||||
|
||||
taskId, err := m.node.CreateTask(ctx, queue.SlaveExtractArchiveType, string(payloadStr))
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create slave task: %w", err)
|
||||
}
|
||||
|
||||
m.state.Phase = ExtractArchivePhaseAwaitSlaveComplete
|
||||
m.state.SlaveTaskID = taskId
|
||||
m.ResumeAfter((10 * time.Second))
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *ExtractArchiveTask) awaitSlaveExtractComplete(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
t, err := m.node.GetTask(ctx, m.state.SlaveTaskID, true)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get slave task: %w", err)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.state.NodeState.progress = t.Progress
|
||||
m.Unlock()
|
||||
|
||||
if t.Status == task.StatusError {
|
||||
return task.StatusError, fmt.Errorf("slave task failed: %s (%w)", t.Error, queue.CriticalErr)
|
||||
}
|
||||
|
||||
if t.Status == task.StatusCanceled {
|
||||
return task.StatusError, fmt.Errorf("slave task canceled (%w)", queue.CriticalErr)
|
||||
}
|
||||
|
||||
if t.Status == task.StatusCompleted {
|
||||
return task.StatusCompleted, nil
|
||||
}
|
||||
|
||||
m.l.Info("Slave task %d is still compressing, resume after 30s.", m.state.SlaveTaskID)
|
||||
m.ResumeAfter((time.Second * 30))
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *ExtractArchiveTask) masterExtractArchive(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
uri, err := fs.NewUriFromString(m.state.Uri)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to parse src uri: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
dst, err := fs.NewUriFromString(m.state.Dst)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to parse dst uri: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
user := inventory.UserFromContext(ctx)
|
||||
fm := manager.NewFileManager(dep, user)
|
||||
|
||||
// Get entity source to extract
|
||||
archiveFile, err := fm.Get(ctx, uri, dbfs.WithFileEntities(), dbfs.WithRequiredCapabilities(dbfs.NavigatorCapabilityDownloadFile), dbfs.WithNotRoot())
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get archive file: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
// Validate file size
|
||||
if user.Edges.Group.Settings.DecompressSize > 0 && archiveFile.Size() > user.Edges.Group.Settings.DecompressSize {
|
||||
return task.StatusError,
|
||||
fmt.Errorf("file size %d exceeds the limit %d (%w)", archiveFile.Size(), user.Edges.Group.Settings.DecompressSize, queue.CriticalErr)
|
||||
}
|
||||
|
||||
es, err := fm.GetEntitySource(ctx, 0, fs.WithEntity(archiveFile.PrimaryEntity()))
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get entity source: %w", err)
|
||||
}
|
||||
|
||||
defer es.Close()
|
||||
|
||||
m.l.Info("Extracting archive %q to %q", uri, m.state.Dst)
|
||||
// Identify file format
|
||||
format, readStream, err := archives.Identify(ctx, archiveFile.DisplayName(), es)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to identify archive format: %w", err)
|
||||
}
|
||||
|
||||
m.l.Info("Archive file %q format identified as %q", uri, format.Extension())
|
||||
|
||||
extractor, ok := format.(archives.Extractor)
|
||||
if !ok {
|
||||
return task.StatusError, fmt.Errorf("format not an extractor %s", format.Extension())
|
||||
}
|
||||
|
||||
formatExt := format.Extension()
|
||||
if formatExt == ".zip" || formatExt == ".7z" {
|
||||
// Zip/7Z extractor requires a Seeker+ReadAt
|
||||
if m.state.TempZipFilePath == "" && !es.IsLocal() {
|
||||
m.state.Phase = ExtractArchivePhaseDownloadZip
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
if m.state.TempZipFilePath != "" {
|
||||
// Use temp zip file path
|
||||
zipFile, err := os.Open(m.state.TempZipFilePath)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to open temp zip file: %w", err)
|
||||
}
|
||||
|
||||
defer zipFile.Close()
|
||||
readStream = zipFile
|
||||
}
|
||||
|
||||
if es.IsLocal() {
|
||||
if _, err = es.Seek(0, 0); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to seek entity source: %w", err)
|
||||
}
|
||||
|
||||
readStream = es
|
||||
}
|
||||
}
|
||||
|
||||
if zipExtractor, ok := extractor.(archives.Zip); ok {
|
||||
if m.state.Encoding != "" {
|
||||
m.l.Info("Using encoding %q for zip archive", m.state.Encoding)
|
||||
encoding, ok := manager.ZipEncodings[strings.ToLower(m.state.Encoding)]
|
||||
if !ok {
|
||||
m.l.Warning("Unknown encoding %q, fallback to default encoding", m.state.Encoding)
|
||||
} else {
|
||||
zipExtractor.TextEncoding = encoding
|
||||
extractor = zipExtractor
|
||||
}
|
||||
}
|
||||
} else if rarExtractor, ok := extractor.(archives.Rar); ok && m.state.Password != "" {
|
||||
rarExtractor.Password = m.state.Password
|
||||
extractor = rarExtractor
|
||||
} else if sevenZipExtractor, ok := extractor.(archives.SevenZip); ok && m.state.Password != "" {
|
||||
sevenZipExtractor.Password = m.state.Password
|
||||
extractor = sevenZipExtractor
|
||||
}
|
||||
|
||||
needSkipToCursor := false
|
||||
if m.state.ProcessedCursor != "" {
|
||||
needSkipToCursor = true
|
||||
}
|
||||
m.Lock()
|
||||
m.progress[ProgressTypeExtractCount] = &queue.Progress{}
|
||||
m.progress[ProgressTypeExtractSize] = &queue.Progress{}
|
||||
m.Unlock()
|
||||
|
||||
// extract and upload
|
||||
err = extractor.Extract(ctx, readStream, func(ctx context.Context, f archives.FileInfo) error {
|
||||
if needSkipToCursor && f.NameInArchive != m.state.ProcessedCursor {
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractCount].Current, 1)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractSize].Current, f.Size())
|
||||
m.l.Info("File %q already processed, skipping...", f.NameInArchive)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Found cursor, start from cursor +1
|
||||
if m.state.ProcessedCursor == f.NameInArchive {
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractCount].Current, 1)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractSize].Current, f.Size())
|
||||
needSkipToCursor = false
|
||||
return nil
|
||||
}
|
||||
|
||||
rawPath := util.FormSlash(f.NameInArchive)
|
||||
savePath := dst.JoinRaw(rawPath)
|
||||
|
||||
// If file mask is not empty, check if the path is in the mask
|
||||
if len(m.state.FileMask) > 0 && !isFileInMask(rawPath, m.state.FileMask) {
|
||||
m.l.Warning("File %q is not in the mask, skipping...", f.NameInArchive)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractCount].Current, 1)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractSize].Current, f.Size())
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if path is legit
|
||||
if !strings.HasPrefix(savePath.Path(), util.FillSlash(path.Clean(dst.Path()))) {
|
||||
m.l.Warning("Path %q is not legit, skipping...", f.NameInArchive)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractCount].Current, 1)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractSize].Current, f.Size())
|
||||
return nil
|
||||
}
|
||||
|
||||
if f.FileInfo.IsDir() {
|
||||
_, err := fm.Create(ctx, savePath, types.FileTypeFolder)
|
||||
if err != nil {
|
||||
m.l.Warning("Failed to create directory %q: %s, skipping...", rawPath, err)
|
||||
}
|
||||
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractCount].Current, 1)
|
||||
m.state.ProcessedCursor = f.NameInArchive
|
||||
return nil
|
||||
}
|
||||
|
||||
fileStream, err := f.Open()
|
||||
if err != nil {
|
||||
m.l.Warning("Failed to open file %q in archive file: %s, skipping...", rawPath, err)
|
||||
return nil
|
||||
}
|
||||
|
||||
fileData := &fs.UploadRequest{
|
||||
Props: &fs.UploadProps{
|
||||
Uri: savePath,
|
||||
Size: f.Size(),
|
||||
LastModified: func() *time.Time {
|
||||
t := f.FileInfo.ModTime().Local()
|
||||
return &t
|
||||
}(),
|
||||
},
|
||||
ProgressFunc: func(current, diff int64, total int64) {
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractSize].Current, diff)
|
||||
},
|
||||
File: fileStream,
|
||||
}
|
||||
|
||||
_, err = fm.Update(ctx, fileData, fs.WithNoEntityType())
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to upload file %q in archive file: %w", rawPath, err)
|
||||
}
|
||||
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractCount].Current, 1)
|
||||
m.state.ProcessedCursor = f.NameInArchive
|
||||
return nil
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to extract archive: %w", err)
|
||||
}
|
||||
|
||||
return task.StatusCompleted, nil
|
||||
}
|
||||
|
||||
func (m *ExtractArchiveTask) masterDownloadZip(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
uri, err := fs.NewUriFromString(m.state.Uri)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to parse src uri: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
user := inventory.UserFromContext(ctx)
|
||||
fm := manager.NewFileManager(dep, user)
|
||||
|
||||
// Get entity source to extract
|
||||
archiveFile, err := fm.Get(ctx, uri, dbfs.WithFileEntities(), dbfs.WithRequiredCapabilities(dbfs.NavigatorCapabilityDownloadFile), dbfs.WithNotRoot())
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get archive file: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
es, err := fm.GetEntitySource(ctx, 0, fs.WithEntity(archiveFile.PrimaryEntity()))
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get entity source: %w", err)
|
||||
}
|
||||
|
||||
defer es.Close()
|
||||
|
||||
// For non-local entity, we need to download the whole zip file first
|
||||
tempPath, err := prepareTempFolder(ctx, dep, m)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to prepare temp folder: %w", err)
|
||||
}
|
||||
m.state.TempPath = tempPath
|
||||
|
||||
fileName := fmt.Sprintf("%s.zip", uuid.Must(uuid.NewV4()))
|
||||
zipFilePath := filepath.Join(
|
||||
m.state.TempPath,
|
||||
fileName,
|
||||
)
|
||||
|
||||
zipFile, err := util.CreatNestedFile(zipFilePath)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create zip file: %w", err)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.progress[ProgressTypeDownload] = &queue.Progress{Total: es.Entity().Size()}
|
||||
m.Unlock()
|
||||
|
||||
defer zipFile.Close()
|
||||
if _, err := io.Copy(zipFile, util.NewCallbackReader(es, func(i int64) {
|
||||
atomic.AddInt64(&m.progress[ProgressTypeDownload].Current, i)
|
||||
})); err != nil {
|
||||
zipFile.Close()
|
||||
if err := os.Remove(zipFilePath); err != nil {
|
||||
m.l.Warning("Failed to remove temp zip file %q: %s", zipFilePath, err)
|
||||
}
|
||||
return task.StatusError, fmt.Errorf("failed to copy zip file to local temp: %w", err)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
delete(m.progress, ProgressTypeDownload)
|
||||
m.Unlock()
|
||||
m.state.TempZipFilePath = zipFilePath
|
||||
m.state.Phase = ExtractArchivePhaseNotStarted
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *ExtractArchiveTask) Summarize(hasher hashid.Encoder) *queue.Summary {
|
||||
if m.state == nil {
|
||||
if err := json.Unmarshal([]byte(m.State()), &m.state); err != nil {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
return &queue.Summary{
|
||||
NodeID: m.state.NodeID,
|
||||
Phase: string(m.state.Phase),
|
||||
Props: map[string]any{
|
||||
SummaryKeySrc: m.state.Uri,
|
||||
SummaryKeyDst: m.state.Dst,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (m *ExtractArchiveTask) Progress(ctx context.Context) queue.Progresses {
|
||||
m.Lock()
|
||||
defer m.Unlock()
|
||||
|
||||
if m.state.NodeState.progress != nil {
|
||||
merged := make(queue.Progresses)
|
||||
for k, v := range m.progress {
|
||||
merged[k] = v
|
||||
}
|
||||
|
||||
for k, v := range m.state.NodeState.progress {
|
||||
merged[k] = v
|
||||
}
|
||||
|
||||
return merged
|
||||
}
|
||||
return m.progress
|
||||
}
|
||||
|
||||
func (m *ExtractArchiveTask) Cleanup(ctx context.Context) error {
|
||||
if m.state.TempPath != "" {
|
||||
time.Sleep(time.Duration(1) * time.Second)
|
||||
return os.RemoveAll(m.state.TempPath)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type (
|
||||
SlaveExtractArchiveTask struct {
|
||||
*queue.InMemoryTask
|
||||
|
||||
l logging.Logger
|
||||
state *SlaveExtractArchiveTaskState
|
||||
progress queue.Progresses
|
||||
node cluster.Node
|
||||
}
|
||||
|
||||
SlaveExtractArchiveTaskState struct {
|
||||
FileName string `json:"file_name"`
|
||||
Entity *ent.Entity `json:"entity"`
|
||||
Policy *ent.StoragePolicy `json:"policy"`
|
||||
Encoding string `json:"encoding,omitempty"`
|
||||
Dst string `json:"dst,omitempty"`
|
||||
UserID int `json:"user_id"`
|
||||
TempPath string `json:"temp_path,omitempty"`
|
||||
TempZipFilePath string `json:"temp_zip_file_path,omitempty"`
|
||||
ProcessedCursor string `json:"processed_cursor,omitempty"`
|
||||
Password string `json:"password,omitempty"`
|
||||
FileMask []string `json:"file_mask,omitempty"`
|
||||
}
|
||||
)
|
||||
|
||||
// NewSlaveExtractArchiveTask creates a new SlaveExtractArchiveTask from raw private state
|
||||
func NewSlaveExtractArchiveTask(ctx context.Context, props *types.SlaveTaskProps, id int, state string) queue.Task {
|
||||
return &SlaveExtractArchiveTask{
|
||||
InMemoryTask: &queue.InMemoryTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: &ent.Task{
|
||||
ID: id,
|
||||
CorrelationID: logging.CorrelationID(ctx),
|
||||
PublicState: &types.TaskPublicState{
|
||||
SlaveTaskProps: props,
|
||||
},
|
||||
PrivateState: state,
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
progress: make(queue.Progresses),
|
||||
}
|
||||
}
|
||||
|
||||
func (m *SlaveExtractArchiveTask) Do(ctx context.Context) (task.Status, error) {
|
||||
ctx = prepareSlaveTaskCtx(ctx, m.Model().PublicState.SlaveTaskProps)
|
||||
dep := dependency.FromContext(ctx)
|
||||
m.l = dep.Logger()
|
||||
np, err := dep.NodePool(ctx)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get node pool: %w", err)
|
||||
}
|
||||
|
||||
m.node, err = np.Get(ctx, types.NodeCapabilityNone, 0)
|
||||
if err != nil || !m.node.IsMaster() {
|
||||
return task.StatusError, fmt.Errorf("failed to get master node: %w", err)
|
||||
}
|
||||
|
||||
fm := manager.NewFileManager(dep, nil)
|
||||
|
||||
// unmarshal state
|
||||
state := &SlaveExtractArchiveTaskState{}
|
||||
if err := json.Unmarshal([]byte(m.State()), state); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to unmarshal state: %w", err)
|
||||
}
|
||||
|
||||
m.state = state
|
||||
m.Lock()
|
||||
if m.progress == nil {
|
||||
m.progress = make(queue.Progresses)
|
||||
}
|
||||
m.progress[ProgressTypeExtractCount] = &queue.Progress{}
|
||||
m.progress[ProgressTypeExtractSize] = &queue.Progress{}
|
||||
m.Unlock()
|
||||
|
||||
dst, err := fs.NewUriFromString(m.state.Dst)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to parse dst uri: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
// 1. Get entity source
|
||||
entity := fs.NewEntity(m.state.Entity)
|
||||
es, err := fm.GetEntitySource(ctx, 0, fs.WithEntity(entity), fs.WithPolicy(fm.CastStoragePolicyOnSlave(ctx, m.state.Policy)))
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get entity source: %w", err)
|
||||
}
|
||||
|
||||
defer es.Close()
|
||||
|
||||
// 2. Identify file format
|
||||
format, readStream, err := archives.Identify(ctx, m.state.FileName, es)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to identify archive format: %w", err)
|
||||
}
|
||||
m.l.Info("Archive file %q format identified as %q", m.state.FileName, format.Extension())
|
||||
|
||||
extractor, ok := format.(archives.Extractor)
|
||||
if !ok {
|
||||
return task.StatusError, fmt.Errorf("format not an extractor %q", format.Extension())
|
||||
}
|
||||
|
||||
formatExt := format.Extension()
|
||||
if formatExt == ".zip" || formatExt == ".7z" {
|
||||
if _, err = es.Seek(0, 0); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to seek entity source: %w", err)
|
||||
}
|
||||
|
||||
if m.state.TempZipFilePath == "" && !es.IsLocal() {
|
||||
tempPath, err := prepareTempFolder(ctx, dep, m)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to prepare temp folder: %w", err)
|
||||
}
|
||||
m.state.TempPath = tempPath
|
||||
|
||||
fileName := fmt.Sprintf("%s.zip", uuid.Must(uuid.NewV4()))
|
||||
zipFilePath := filepath.Join(
|
||||
m.state.TempPath,
|
||||
fileName,
|
||||
)
|
||||
zipFile, err := util.CreatNestedFile(zipFilePath)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create zip file: %w", err)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.progress[ProgressTypeDownload] = &queue.Progress{Total: es.Entity().Size()}
|
||||
m.Unlock()
|
||||
|
||||
defer zipFile.Close()
|
||||
if _, err := io.Copy(zipFile, util.NewCallbackReader(es, func(i int64) {
|
||||
atomic.AddInt64(&m.progress[ProgressTypeDownload].Current, i)
|
||||
})); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to copy zip file to local temp: %w", err)
|
||||
}
|
||||
|
||||
zipFile.Close()
|
||||
m.state.TempZipFilePath = zipFilePath
|
||||
}
|
||||
|
||||
if es.IsLocal() {
|
||||
readStream = es
|
||||
} else if m.state.TempZipFilePath != "" {
|
||||
// Use temp zip file path
|
||||
zipFile, err := os.Open(m.state.TempZipFilePath)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to open temp zip file: %w", err)
|
||||
}
|
||||
|
||||
defer zipFile.Close()
|
||||
readStream = zipFile
|
||||
}
|
||||
|
||||
if es.IsLocal() {
|
||||
readStream = es
|
||||
}
|
||||
}
|
||||
|
||||
if zipExtractor, ok := extractor.(archives.Zip); ok {
|
||||
if m.state.Encoding != "" {
|
||||
m.l.Info("Using encoding %q for zip archive", m.state.Encoding)
|
||||
encoding, ok := manager.ZipEncodings[strings.ToLower(m.state.Encoding)]
|
||||
if !ok {
|
||||
m.l.Warning("Unknown encoding %q, fallback to default encoding", m.state.Encoding)
|
||||
} else {
|
||||
zipExtractor.TextEncoding = encoding
|
||||
extractor = zipExtractor
|
||||
}
|
||||
}
|
||||
} else if rarExtractor, ok := extractor.(archives.Rar); ok && m.state.Password != "" {
|
||||
rarExtractor.Password = m.state.Password
|
||||
extractor = rarExtractor
|
||||
} else if sevenZipExtractor, ok := extractor.(archives.SevenZip); ok && m.state.Password != "" {
|
||||
sevenZipExtractor.Password = m.state.Password
|
||||
extractor = sevenZipExtractor
|
||||
}
|
||||
|
||||
needSkipToCursor := false
|
||||
if m.state.ProcessedCursor != "" {
|
||||
needSkipToCursor = true
|
||||
}
|
||||
|
||||
// 3. Extract and upload
|
||||
err = extractor.Extract(ctx, readStream, func(ctx context.Context, f archives.FileInfo) error {
|
||||
if needSkipToCursor && f.NameInArchive != m.state.ProcessedCursor {
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractCount].Current, 1)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractSize].Current, f.Size())
|
||||
m.l.Info("File %q already processed, skipping...", f.NameInArchive)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Found cursor, start from cursor +1
|
||||
if m.state.ProcessedCursor == f.NameInArchive {
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractCount].Current, 1)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractSize].Current, f.Size())
|
||||
needSkipToCursor = false
|
||||
return nil
|
||||
}
|
||||
|
||||
rawPath := util.FormSlash(f.NameInArchive)
|
||||
savePath := dst.JoinRaw(rawPath)
|
||||
|
||||
// If file mask is not empty, check if the path is in the mask
|
||||
if len(m.state.FileMask) > 0 && !isFileInMask(rawPath, m.state.FileMask) {
|
||||
m.l.Debug("File %q is not in the mask, skipping...", f.NameInArchive)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if path is legit
|
||||
if !strings.HasPrefix(savePath.Path(), util.FillSlash(path.Clean(dst.Path()))) {
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractCount].Current, 1)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractSize].Current, f.Size())
|
||||
m.l.Warning("Path %q is not legit, skipping...", f.NameInArchive)
|
||||
return nil
|
||||
}
|
||||
|
||||
if f.FileInfo.IsDir() {
|
||||
_, err := fm.Create(ctx, savePath, types.FileTypeFolder, fs.WithNode(m.node), fs.WithStatelessUserID(m.state.UserID))
|
||||
if err != nil {
|
||||
m.l.Warning("Failed to create directory %q: %s, skipping...", rawPath, err)
|
||||
}
|
||||
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractCount].Current, 1)
|
||||
m.state.ProcessedCursor = f.NameInArchive
|
||||
return nil
|
||||
}
|
||||
|
||||
fileStream, err := f.Open()
|
||||
if err != nil {
|
||||
m.l.Warning("Failed to open file %q in archive file: %s, skipping...", rawPath, err)
|
||||
return nil
|
||||
}
|
||||
|
||||
fileData := &fs.UploadRequest{
|
||||
Props: &fs.UploadProps{
|
||||
Uri: savePath,
|
||||
Size: f.Size(),
|
||||
LastModified: func() *time.Time {
|
||||
t := f.FileInfo.ModTime().Local()
|
||||
return &t
|
||||
}(),
|
||||
},
|
||||
ProgressFunc: func(current, diff int64, total int64) {
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractSize].Current, diff)
|
||||
},
|
||||
File: fileStream,
|
||||
}
|
||||
|
||||
_, err = fm.Update(ctx, fileData, fs.WithNode(m.node), fs.WithStatelessUserID(m.state.UserID), fs.WithNoEntityType())
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to upload file %q in archive file: %w", rawPath, err)
|
||||
}
|
||||
|
||||
atomic.AddInt64(&m.progress[ProgressTypeExtractCount].Current, 1)
|
||||
m.state.ProcessedCursor = f.NameInArchive
|
||||
return nil
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to extract archive: %w", err)
|
||||
}
|
||||
|
||||
return task.StatusCompleted, nil
|
||||
}
|
||||
|
||||
func (m *SlaveExtractArchiveTask) Cleanup(ctx context.Context) error {
|
||||
if m.state.TempPath != "" {
|
||||
time.Sleep(time.Duration(1) * time.Second)
|
||||
return os.RemoveAll(m.state.TempPath)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *SlaveExtractArchiveTask) Progress(ctx context.Context) queue.Progresses {
|
||||
m.Lock()
|
||||
defer m.Unlock()
|
||||
return m.progress
|
||||
}
|
||||
|
||||
func isFileInMask(path string, mask []string) bool {
|
||||
if len(mask) == 0 {
|
||||
return true
|
||||
}
|
||||
|
||||
for _, m := range mask {
|
||||
if path == m || strings.HasPrefix(path, m+"/") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
package workflows
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"runtime"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/cloudreve/Cloudreve/v4/application/dependency"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent/task"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory/types"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/fs"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/manager"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/hashid"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/logging"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/queue"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/serializer"
|
||||
)
|
||||
|
||||
type (
|
||||
ImportTask struct {
|
||||
*queue.DBTask
|
||||
|
||||
l logging.Logger
|
||||
state *ImportTaskState
|
||||
progress queue.Progresses
|
||||
}
|
||||
ImportTaskState struct {
|
||||
PolicyID int `json:"policy_id"`
|
||||
Src string `json:"src"`
|
||||
Recursive bool `json:"is_recursive"`
|
||||
Dst string `json:"dst"`
|
||||
Phase ImportTaskPhase `json:"phase"`
|
||||
Failed int `json:"failed,omitempty"`
|
||||
ExtractMediaMeta bool `json:"extract_media_meta"`
|
||||
}
|
||||
ImportTaskPhase string
|
||||
)
|
||||
|
||||
const (
|
||||
ProgressTypeImported = "imported"
|
||||
ProgressTypeIndexed = "indexed"
|
||||
|
||||
// ImportBatchSize is the number of files to process in each batch
|
||||
// to control memory usage during large imports.
|
||||
ImportBatchSize = 100
|
||||
)
|
||||
|
||||
func init() {
|
||||
queue.RegisterResumableTaskFactory(queue.ImportTaskType, NewImportTaskFromModel)
|
||||
}
|
||||
|
||||
func NewImportTask(ctx context.Context, u *ent.User, src string, recursive bool, dst string, policyID int) (queue.Task, error) {
|
||||
state := &ImportTaskState{
|
||||
Src: src,
|
||||
Recursive: recursive,
|
||||
Dst: dst,
|
||||
PolicyID: policyID,
|
||||
}
|
||||
stateBytes, err := json.Marshal(state)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal state: %w", err)
|
||||
}
|
||||
|
||||
t := &ImportTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: &ent.Task{
|
||||
Type: queue.ImportTaskType,
|
||||
CorrelationID: logging.CorrelationID(ctx),
|
||||
PrivateState: string(stateBytes),
|
||||
PublicState: &types.TaskPublicState{},
|
||||
},
|
||||
DirectOwner: u,
|
||||
},
|
||||
}
|
||||
|
||||
return t, nil
|
||||
}
|
||||
|
||||
func NewImportTaskFromModel(task *ent.Task) queue.Task {
|
||||
return &ImportTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: task,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (m *ImportTask) Do(ctx context.Context) (task.Status, error) {
|
||||
dep := dependency.FromContext(ctx)
|
||||
m.l = dep.Logger()
|
||||
|
||||
m.Lock()
|
||||
if m.progress == nil {
|
||||
m.progress = make(queue.Progresses)
|
||||
}
|
||||
m.progress[ProgressTypeIndexed] = &queue.Progress{}
|
||||
m.Unlock()
|
||||
|
||||
// unmarshal state
|
||||
state := &ImportTaskState{}
|
||||
if err := json.Unmarshal([]byte(m.State()), state); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to unmarshal state: %w", err)
|
||||
}
|
||||
m.state = state
|
||||
|
||||
next, err := m.processImport(ctx, dep)
|
||||
|
||||
newStateStr, marshalErr := json.Marshal(m.state)
|
||||
if marshalErr != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to marshal state: %w", marshalErr)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.Task.PrivateState = string(newStateStr)
|
||||
m.Unlock()
|
||||
return next, err
|
||||
}
|
||||
|
||||
func (m *ImportTask) processImport(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
user := inventory.UserFromContext(ctx)
|
||||
|
||||
dst, err := fs.NewUriFromString(m.state.Dst)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to parse dst: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
// Use a temporary file manager just for listing physical files
|
||||
listFm := manager.NewFileManager(dep, user)
|
||||
physicalFiles, err := listFm.ListPhysical(ctx, m.state.Src, m.state.PolicyID, m.state.Recursive,
|
||||
func(i int) {
|
||||
atomic.AddInt64(&m.progress[ProgressTypeIndexed].Current, int64(i))
|
||||
})
|
||||
listFm.Recycle()
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to list physical files: %w", err)
|
||||
}
|
||||
|
||||
m.l.Info("Importing %d physical files", len(physicalFiles))
|
||||
|
||||
m.Lock()
|
||||
m.progress[ProgressTypeImported] = &queue.Progress{
|
||||
Total: int64(len(physicalFiles)),
|
||||
}
|
||||
delete(m.progress, ProgressTypeIndexed)
|
||||
m.Unlock()
|
||||
|
||||
failed := 0
|
||||
totalFiles := len(physicalFiles)
|
||||
|
||||
// Process files in batches to control memory usage
|
||||
for batchStart := 0; batchStart < totalFiles; batchStart += ImportBatchSize {
|
||||
batchEnd := min(batchStart+ImportBatchSize, totalFiles)
|
||||
|
||||
batch := physicalFiles[batchStart:batchEnd]
|
||||
batchFailed := m.processBatch(ctx, dep, user, dst, batch)
|
||||
failed += batchFailed
|
||||
|
||||
// Clear batch elements to allow GC of individual items
|
||||
for i := batchStart; i < batchEnd; i++ {
|
||||
physicalFiles[i] = fs.PhysicalObject{}
|
||||
}
|
||||
|
||||
// Run GC after each batch to free memory
|
||||
runtime.GC()
|
||||
}
|
||||
|
||||
// Clear the entire slice to allow GC
|
||||
physicalFiles = nil
|
||||
runtime.GC()
|
||||
|
||||
m.state.Failed = failed
|
||||
return task.StatusCompleted, nil
|
||||
}
|
||||
|
||||
// processBatch processes a batch of physical files with a fresh file manager.
|
||||
func (m *ImportTask) processBatch(ctx context.Context, dep dependency.Dep, user *ent.User, dst *fs.URI, batch []fs.PhysicalObject) int {
|
||||
fm := manager.NewFileManager(dep, user)
|
||||
defer fm.Recycle()
|
||||
|
||||
failed := 0
|
||||
for _, physicalFile := range batch {
|
||||
if physicalFile.IsDir {
|
||||
m.l.Info("Creating folder %s", physicalFile.RelativePath)
|
||||
_, err := fm.Create(ctx, dst.Join(physicalFile.RelativePath), types.FileTypeFolder)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeImported].Current, 1)
|
||||
if err != nil {
|
||||
m.l.Warning("Failed to create folder %s: %s", physicalFile.RelativePath, err)
|
||||
failed++
|
||||
}
|
||||
} else {
|
||||
m.l.Info("Importing file %s", physicalFile.RelativePath)
|
||||
err := fm.ImportPhysical(ctx, dst, m.state.PolicyID, physicalFile, m.state.ExtractMediaMeta)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeImported].Current, 1)
|
||||
if err != nil {
|
||||
var appErr serializer.AppError
|
||||
if errors.As(err, &appErr) && appErr.Code == serializer.CodeObjectExist {
|
||||
m.l.Info("File %s already exists, skipping", physicalFile.RelativePath)
|
||||
continue
|
||||
}
|
||||
m.l.Error("Failed to import file %s: %s, skipping", physicalFile.RelativePath, err)
|
||||
failed++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return failed
|
||||
}
|
||||
|
||||
func (m *ImportTask) Progress(ctx context.Context) queue.Progresses {
|
||||
m.Lock()
|
||||
defer m.Unlock()
|
||||
return m.progress
|
||||
}
|
||||
|
||||
func (m *ImportTask) Summarize(hasher hashid.Encoder) *queue.Summary {
|
||||
// unmarshal state
|
||||
if m.state == nil {
|
||||
if err := json.Unmarshal([]byte(m.State()), &m.state); err != nil {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
return &queue.Summary{
|
||||
Phase: string(m.state.Phase),
|
||||
Props: map[string]any{
|
||||
SummaryKeyDst: m.state.Dst,
|
||||
SummaryKeySrcStr: m.state.Src,
|
||||
SummaryKeyFailed: m.state.Failed,
|
||||
SummaryKeySrcDstPolicyID: hashid.EncodePolicyID(hasher, m.state.PolicyID),
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
package workflows
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"slices"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/cloudreve/Cloudreve/v4/application/dependency"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent/task"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory/types"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/fs/dbfs"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/manager"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/hashid"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/logging"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/queue"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/searcher"
|
||||
)
|
||||
|
||||
type (
|
||||
RebuildIndexTask struct {
|
||||
*queue.DBTask
|
||||
|
||||
l logging.Logger
|
||||
state *RebuildIndexTaskState
|
||||
progress queue.Progresses
|
||||
}
|
||||
RebuildIndexTaskPhase string
|
||||
RebuildIndexTaskState struct {
|
||||
Phase RebuildIndexTaskPhase `json:"phase"`
|
||||
Total int `json:"total"`
|
||||
Indexed int `json:"indexed"`
|
||||
LastFileID int `json:"last_file_id"`
|
||||
Failed int `json:"failed"`
|
||||
FilteredStoragePolicy []int `json:"filtered_storage_policy"`
|
||||
}
|
||||
)
|
||||
|
||||
const (
|
||||
RebuildIndexPhaseNuke RebuildIndexTaskPhase = "nuke"
|
||||
RebuildIndexPhaseIndex RebuildIndexTaskPhase = "index"
|
||||
|
||||
RebuildIndexBatchSize = 1000
|
||||
RebuildIndexConcurrent = 4
|
||||
|
||||
ProgressTypeRebuildIndex = "rebuild_index"
|
||||
)
|
||||
|
||||
func init() {
|
||||
queue.RegisterResumableTaskFactory(queue.FullTextRebuildTaskType, NewRebuildIndexTaskFromModel)
|
||||
}
|
||||
|
||||
func NewRebuildIndexTask(ctx context.Context, u *ent.User, filteredStoragePolicy []int) (queue.Task, error) {
|
||||
state := &RebuildIndexTaskState{
|
||||
Phase: RebuildIndexPhaseNuke,
|
||||
FilteredStoragePolicy: filteredStoragePolicy,
|
||||
}
|
||||
stateBytes, err := json.Marshal(state)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal state: %w", err)
|
||||
}
|
||||
|
||||
return &RebuildIndexTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: &ent.Task{
|
||||
Type: queue.FullTextRebuildTaskType,
|
||||
CorrelationID: logging.CorrelationID(ctx),
|
||||
PrivateState: string(stateBytes),
|
||||
PublicState: &types.TaskPublicState{},
|
||||
},
|
||||
DirectOwner: u,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func NewRebuildIndexTaskFromModel(t *ent.Task) queue.Task {
|
||||
return &RebuildIndexTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: t,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (m *RebuildIndexTask) Do(ctx context.Context) (task.Status, error) {
|
||||
dep := dependency.FromContext(ctx)
|
||||
m.l = dep.Logger()
|
||||
|
||||
m.Lock()
|
||||
if m.progress == nil {
|
||||
m.progress = make(queue.Progresses)
|
||||
}
|
||||
m.progress[ProgressTypeRebuildIndex] = &queue.Progress{}
|
||||
m.Unlock()
|
||||
|
||||
state := &RebuildIndexTaskState{}
|
||||
if err := json.Unmarshal([]byte(m.State()), state); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to unmarshal state: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
m.state = state
|
||||
|
||||
var (
|
||||
next = task.StatusCompleted
|
||||
err error
|
||||
)
|
||||
switch m.state.Phase {
|
||||
case RebuildIndexPhaseNuke, "":
|
||||
next, err = m.nuke(ctx, dep)
|
||||
case RebuildIndexPhaseIndex:
|
||||
next, err = m.index(ctx, dep)
|
||||
default:
|
||||
next, err = task.StatusError, fmt.Errorf("unknown phase %q: %w", m.state.Phase, queue.CriticalErr)
|
||||
}
|
||||
|
||||
newStateStr, marshalErr := json.Marshal(m.state)
|
||||
if marshalErr != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to marshal state: %w", marshalErr)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.Task.PrivateState = string(newStateStr)
|
||||
m.Unlock()
|
||||
return next, err
|
||||
}
|
||||
|
||||
// nuke deletes all existing index documents and ensures a fresh index exists,
|
||||
// then counts total indexable files for progress tracking.
|
||||
func (m *RebuildIndexTask) nuke(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
indexer := dep.SearchIndexer(ctx)
|
||||
|
||||
m.l.Info("Deleting all existing index documents...")
|
||||
if err := indexer.DeleteAll(ctx); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to delete all index documents: %w", err)
|
||||
}
|
||||
|
||||
if err := dep.FileClient().DeleteAllMetadataByName(ctx, dbfs.FullTextIndexKey); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to delete all metadata by name: %w", err)
|
||||
}
|
||||
|
||||
m.l.Info("Ensuring index exists with correct configuration...")
|
||||
if err := indexer.EnsureIndex(ctx); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to ensure index: %w", err)
|
||||
}
|
||||
|
||||
// Count total indexable files
|
||||
total, err := dep.FileClient().CountIndexableFiles(ctx)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to count indexable files: %w", err)
|
||||
}
|
||||
|
||||
m.state.Total = total
|
||||
m.state.Phase = RebuildIndexPhaseIndex
|
||||
m.state.LastFileID = 0
|
||||
m.state.Indexed = 0
|
||||
|
||||
m.l.Info("Found %d indexable files, starting rebuild...", total)
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
// index processes a batch of files and suspends for the next batch.
|
||||
func (m *RebuildIndexTask) index(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
atomic.StoreInt64(&m.progress[ProgressTypeRebuildIndex].Total, int64(m.state.Total))
|
||||
atomic.StoreInt64(&m.progress[ProgressTypeRebuildIndex].Current, int64(m.state.Indexed))
|
||||
|
||||
files, err := dep.FileClient().ListIndexableFiles(ctx, m.state.LastFileID, RebuildIndexBatchSize)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to list indexable files after ID %d: %w", m.state.LastFileID, err)
|
||||
}
|
||||
|
||||
if len(files) == 0 {
|
||||
m.l.Info("Rebuild complete. %d files indexed, %d failed.", m.state.Indexed-m.state.Failed, m.state.Failed)
|
||||
return task.StatusCompleted, nil
|
||||
}
|
||||
|
||||
batchFailed := m.processBatch(ctx, dep, files)
|
||||
m.state.Failed += batchFailed
|
||||
m.state.Indexed += len(files)
|
||||
m.state.LastFileID = files[len(files)-1].ID
|
||||
|
||||
atomic.StoreInt64(&m.progress[ProgressTypeRebuildIndex].Current, int64(m.state.Indexed))
|
||||
|
||||
// Suspend and resume for next batch
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
// processBatch indexes a batch of files concurrently.
|
||||
func (m *RebuildIndexTask) processBatch(ctx context.Context, dep dependency.Dep, files []*ent.File) int {
|
||||
user := inventory.UserFromContext(ctx)
|
||||
|
||||
var (
|
||||
wg sync.WaitGroup
|
||||
mu sync.Mutex
|
||||
failed int
|
||||
)
|
||||
|
||||
sem := make(chan struct{}, RebuildIndexConcurrent)
|
||||
indexer := dep.SearchIndexer(ctx)
|
||||
extractor := dep.TextExtractor(ctx)
|
||||
|
||||
for _, f := range files {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return failed
|
||||
case sem <- struct{}{}:
|
||||
}
|
||||
|
||||
wg.Add(1)
|
||||
go func(f *ent.File) {
|
||||
defer func() {
|
||||
<-sem
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
if err := m.indexSingleFile(ctx, dep, user, indexer, extractor, f); err != nil {
|
||||
m.l.Warning("Failed to index file %d (%s): %s", f.ID, f.Name, err)
|
||||
mu.Lock()
|
||||
failed++
|
||||
mu.Unlock()
|
||||
}
|
||||
}(f)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
return failed
|
||||
}
|
||||
|
||||
// indexSingleFile extracts text from a single file and indexes it.
|
||||
func (m *RebuildIndexTask) indexSingleFile(
|
||||
ctx context.Context,
|
||||
dep dependency.Dep,
|
||||
user *ent.User,
|
||||
indexer searcher.SearchIndexer,
|
||||
extractor searcher.TextExtractor,
|
||||
f *ent.File,
|
||||
) error {
|
||||
fm := manager.NewFileManager(dep, user)
|
||||
defer fm.Recycle()
|
||||
|
||||
entityID := f.PrimaryEntity
|
||||
if entityID == 0 {
|
||||
// No primary entity, index with just the file name (no text content).
|
||||
m.l.Debug("No primary entity for file %d, skipping.", f.ID)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if this file type is eligible for text extraction
|
||||
var text string
|
||||
if manager.ShouldExtractText(extractor, f.Name, f.Size) {
|
||||
source, err := fm.GetEntitySource(ctx, entityID)
|
||||
if err != nil {
|
||||
// Cannot get source; index with file name only.
|
||||
m.l.Debug("Cannot get entity source for file %d: %s, skipping.", f.ID, err)
|
||||
return fmt.Errorf("cannot get entity source for file %d: %w", f.ID, err)
|
||||
}
|
||||
defer source.Close()
|
||||
|
||||
if len(m.state.FilteredStoragePolicy) > 0 {
|
||||
if !slices.Contains(m.state.FilteredStoragePolicy, source.Entity().PolicyID()) {
|
||||
m.l.Debug("Entity source for file %d is not in filtered storage policy, skipping.", f.ID)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
extracted, err := extractor.Extract(ctx, source)
|
||||
if err != nil {
|
||||
m.l.Debug("Failed to extract text for file %d: %s, skipping", f.ID, err)
|
||||
return nil
|
||||
} else {
|
||||
text = extracted
|
||||
}
|
||||
|
||||
if err := indexer.IndexFile(ctx, f.OwnerID, f.ID, entityID, f.Name, text); err != nil {
|
||||
return fmt.Errorf("failed to index file %d: %w", f.ID, err)
|
||||
}
|
||||
|
||||
if err := dep.FileClient().UpsertMetadata(ctx, f, map[string]string{
|
||||
dbfs.FullTextIndexKey: hashid.EncodeEntityID(dep.HashIDEncoder(), entityID),
|
||||
}, nil); err != nil {
|
||||
m.l.Warning("Failed to upsert metadata for file %d: %s", f.ID, err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *RebuildIndexTask) Progress(ctx context.Context) queue.Progresses {
|
||||
m.Lock()
|
||||
defer m.Unlock()
|
||||
return m.progress
|
||||
}
|
||||
|
||||
func (m *RebuildIndexTask) Summarize(hasher hashid.Encoder) *queue.Summary {
|
||||
if m.state == nil {
|
||||
if err := json.Unmarshal([]byte(m.State()), &m.state); err != nil {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
return &queue.Summary{
|
||||
Phase: string(m.state.Phase),
|
||||
Props: map[string]any{
|
||||
SummaryKeyFailed: m.state.Failed,
|
||||
SummaryKeyTotal: m.state.Total,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,684 @@
|
||||
package workflows
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/cloudreve/Cloudreve/v4/application/dependency"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent/task"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory/types"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/downloader"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/fs"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/manager"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/hashid"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/logging"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/queue"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/request"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/serializer"
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
type (
|
||||
RemoteDownloadTask struct {
|
||||
*queue.DBTask
|
||||
|
||||
l logging.Logger
|
||||
state *RemoteDownloadTaskState
|
||||
node cluster.Node
|
||||
d downloader.Downloader
|
||||
progress queue.Progresses
|
||||
}
|
||||
RemoteDownloadTaskPhase string
|
||||
RemoteDownloadTaskState struct {
|
||||
SrcFileUri string `json:"src_file_uri,omitempty"`
|
||||
SrcUri string `json:"src_uri,omitempty"`
|
||||
Dst string `json:"dst,omitempty"`
|
||||
Handle *downloader.TaskHandle `json:"handle,omitempty"`
|
||||
Status *downloader.TaskStatus `json:"status,omitempty"`
|
||||
NodeState `json:",inline"`
|
||||
Phase RemoteDownloadTaskPhase `json:"phase,omitempty"`
|
||||
SlaveUploadTaskID int `json:"slave__upload_task_id,omitempty"`
|
||||
SlaveUploadState *SlaveUploadTaskState `json:"slave_upload_state,omitempty"`
|
||||
GetTaskStatusTried int `json:"get_task_status_tried,omitempty"`
|
||||
Transferred map[int]interface{} `json:"transferred,omitempty"`
|
||||
Failed int `json:"failed,omitempty"`
|
||||
}
|
||||
)
|
||||
|
||||
const (
|
||||
RemoteDownloadTaskPhaseNotStarted RemoteDownloadTaskPhase = ""
|
||||
RemoteDownloadTaskPhaseMonitor = "monitor"
|
||||
RemoteDownloadTaskPhaseTransfer = "transfer"
|
||||
RemoteDownloadTaskPhaseAwaitSeeding = "seeding"
|
||||
|
||||
GetTaskStatusMaxTries = 5
|
||||
|
||||
SummaryKeyDownloadStatus = "download"
|
||||
SummaryKeySrcStr = "src_str"
|
||||
|
||||
ProgressTypeRelocateTransferCount = "relocate"
|
||||
ProgressTypeUploadSinglePrefix = "upload_single_"
|
||||
|
||||
SummaryKeySrcMultiple = "src_multiple"
|
||||
SummaryKeySrcDstPolicyID = "dst_policy_id"
|
||||
SummaryKeyFailed = "failed"
|
||||
SummaryKeyTotal = "total"
|
||||
)
|
||||
|
||||
func init() {
|
||||
queue.RegisterResumableTaskFactory(queue.RemoteDownloadTaskType, NewRemoteDownloadTaskFromModel)
|
||||
}
|
||||
|
||||
// NewRemoteDownloadTask creates a new RemoteDownloadTask
|
||||
func NewRemoteDownloadTask(ctx context.Context, src string, srcFile, dst string) (queue.Task, error) {
|
||||
state := &RemoteDownloadTaskState{
|
||||
SrcUri: src,
|
||||
SrcFileUri: srcFile,
|
||||
Dst: dst,
|
||||
NodeState: NodeState{},
|
||||
}
|
||||
stateBytes, err := json.Marshal(state)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal state: %w", err)
|
||||
}
|
||||
|
||||
t := &RemoteDownloadTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: &ent.Task{
|
||||
Type: queue.RemoteDownloadTaskType,
|
||||
CorrelationID: logging.CorrelationID(ctx),
|
||||
PrivateState: string(stateBytes),
|
||||
PublicState: &types.TaskPublicState{},
|
||||
},
|
||||
DirectOwner: inventory.UserFromContext(ctx),
|
||||
},
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
func NewRemoteDownloadTaskFromModel(task *ent.Task) queue.Task {
|
||||
return &RemoteDownloadTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: task,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (m *RemoteDownloadTask) Do(ctx context.Context) (task.Status, error) {
|
||||
dep := dependency.FromContext(ctx)
|
||||
m.l = dep.Logger()
|
||||
|
||||
// unmarshal state
|
||||
state := &RemoteDownloadTaskState{}
|
||||
if err := json.Unmarshal([]byte(m.State()), state); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to unmarshal state: %w", err)
|
||||
}
|
||||
m.state = state
|
||||
|
||||
// select node
|
||||
node, err := allocateNode(ctx, dep, &m.state.NodeState, types.NodeCapabilityRemoteDownload)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to allocate node: %w", err)
|
||||
}
|
||||
m.node = node
|
||||
|
||||
// create downloader instance
|
||||
if m.d == nil {
|
||||
d, err := node.CreateDownloader(ctx, dep.RequestClient(), dep.SettingProvider())
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create downloader: %w", err)
|
||||
}
|
||||
|
||||
m.d = d
|
||||
}
|
||||
|
||||
next := task.StatusCompleted
|
||||
switch m.state.Phase {
|
||||
case RemoteDownloadTaskPhaseNotStarted:
|
||||
next, err = m.createDownloadTask(ctx, dep)
|
||||
case RemoteDownloadTaskPhaseMonitor, RemoteDownloadTaskPhaseAwaitSeeding:
|
||||
next, err = m.monitor(ctx, dep)
|
||||
case RemoteDownloadTaskPhaseTransfer:
|
||||
if m.node.IsMaster() {
|
||||
next, err = m.masterTransfer(ctx, dep)
|
||||
} else {
|
||||
next, err = m.slaveTransfer(ctx, dep)
|
||||
}
|
||||
}
|
||||
|
||||
newStateStr, marshalErr := json.Marshal(m.state)
|
||||
if marshalErr != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to marshal state: %w", marshalErr)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.Task.PrivateState = string(newStateStr)
|
||||
m.Unlock()
|
||||
return next, err
|
||||
}
|
||||
|
||||
func (m *RemoteDownloadTask) createDownloadTask(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
if m.state.Handle != nil {
|
||||
m.state.Phase = RemoteDownloadTaskPhaseMonitor
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
user := inventory.UserFromContext(ctx)
|
||||
torrentUrl := m.state.SrcUri
|
||||
|
||||
// SSRF policy is owned by the node that actually performs the fetch.
|
||||
// Validate the user-supplied URL (m.state.SrcUri) against the assigned
|
||||
// node's URLValidation policy, with the operator-configured site URL
|
||||
// host(s) always added to the allowlist. We don't validate SrcFileUri
|
||||
// resolution because that's a Cloudreve-internal entity URL pointing at
|
||||
// the user's own torrent file.
|
||||
if m.state.SrcUri != "" {
|
||||
opt := buildSSRFOptions(ctx, dep, m.node.Settings(ctx))
|
||||
if err := request.ValidateExternalURL(ctx, m.state.SrcUri, opt); err != nil {
|
||||
return task.StatusError, fmt.Errorf("url rejected: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
}
|
||||
|
||||
if m.state.SrcFileUri != "" {
|
||||
// Target is a torrent file
|
||||
uri, err := fs.NewUriFromString(m.state.SrcFileUri)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to parse src file uri: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
fm := manager.NewFileManager(dep, user)
|
||||
expire := time.Now().Add(dep.SettingProvider().EntityUrlValidDuration(ctx))
|
||||
torrentUrls, _, err := fm.GetEntityUrls(ctx, []manager.GetEntityUrlArgs{
|
||||
{URI: uri},
|
||||
}, fs.WithUrlExpire(&expire))
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get torrent entity urls: %w", err)
|
||||
}
|
||||
|
||||
if len(torrentUrls) == 0 {
|
||||
return task.StatusError, fmt.Errorf("no torrent urls found")
|
||||
}
|
||||
|
||||
torrentUrl = torrentUrls[0].Url
|
||||
}
|
||||
|
||||
// Create download task
|
||||
handle, err := m.d.CreateTask(ctx, torrentUrl, user.Edges.Group.Settings.RemoteDownloadOptions)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create download task: %w", err)
|
||||
}
|
||||
|
||||
m.state.Handle = handle
|
||||
m.state.Phase = RemoteDownloadTaskPhaseMonitor
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
// buildSSRFOptions composes the SSRF policy for a download: the assigned
|
||||
// node's URLValidation settings, plus the operator-configured site URL hosts
|
||||
// (always allowlisted so users can fetch files served by Cloudreve itself).
|
||||
func buildSSRFOptions(ctx context.Context, dep dependency.Dep, node *types.NodeSetting) request.SSRFOptions {
|
||||
opt := request.SSRFOptions{}
|
||||
if node != nil && node.URLValidation != nil {
|
||||
opt.Disabled = node.URLValidation.Disabled
|
||||
opt.AllowedHosts = append(opt.AllowedHosts, node.URLValidation.AllowedHosts...)
|
||||
opt.AllowedCIDRs = append(opt.AllowedCIDRs, node.URLValidation.AllowedCIDRs...)
|
||||
}
|
||||
for _, u := range dep.SettingProvider().AllSiteURLs(ctx) {
|
||||
if h := u.Hostname(); h != "" {
|
||||
opt.AllowedHosts = append(opt.AllowedHosts, h)
|
||||
}
|
||||
}
|
||||
return opt
|
||||
}
|
||||
|
||||
func (m *RemoteDownloadTask) monitor(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
resumeAfter := time.Duration(m.node.Settings(ctx).Interval) * time.Second
|
||||
|
||||
// Update task status
|
||||
status, err := m.d.Info(ctx, m.state.Handle)
|
||||
if err != nil {
|
||||
if errors.Is(err, downloader.ErrTaskNotFount) && m.state.Status != nil {
|
||||
// If task is not found, but it previously existed, consider it as canceled
|
||||
m.l.Warning("task not found, consider it as canceled")
|
||||
return task.StatusCanceled, nil
|
||||
}
|
||||
|
||||
m.state.GetTaskStatusTried++
|
||||
if m.state.GetTaskStatusTried >= GetTaskStatusMaxTries {
|
||||
return task.StatusError, fmt.Errorf("failed to get task status after %d retry: %w", m.state.GetTaskStatusTried, err)
|
||||
}
|
||||
|
||||
m.l.Warning("failed to get task info: %s, will retry.", err)
|
||||
m.ResumeAfter(resumeAfter)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
// Follow to new handle if needed
|
||||
if status.FollowedBy != nil {
|
||||
m.l.Info("Task handle updated to %v", status.FollowedBy)
|
||||
m.state.Handle = status.FollowedBy
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
if m.state.Status == nil || m.state.Status.Total != status.Total {
|
||||
m.l.Info("download size changed, re-validate files.")
|
||||
// First time to get status / total size changed, check user capacity
|
||||
if err := m.validateFiles(ctx, dep, status); err != nil {
|
||||
m.state.Status = status
|
||||
return task.StatusError, fmt.Errorf("failed to validate files: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
}
|
||||
|
||||
m.state.Status = status
|
||||
m.state.GetTaskStatusTried = 0
|
||||
|
||||
m.l.Debug("Monitor %q task state: %s", status.Name, status.State)
|
||||
switch status.State {
|
||||
case downloader.StatusSeeding:
|
||||
m.l.Info("Download task seeding")
|
||||
if m.state.Phase == RemoteDownloadTaskPhaseMonitor {
|
||||
// Not transferred
|
||||
m.state.Phase = RemoteDownloadTaskPhaseTransfer
|
||||
return task.StatusSuspending, nil
|
||||
} else if !m.node.Settings(ctx).WaitForSeeding {
|
||||
// Skip seeding
|
||||
m.l.Info("Download task seeding skipped.")
|
||||
return task.StatusCompleted, nil
|
||||
} else {
|
||||
// Still seeding
|
||||
m.ResumeAfter(resumeAfter)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
case downloader.StatusCompleted:
|
||||
m.l.Info("Download task completed")
|
||||
if m.state.Phase == RemoteDownloadTaskPhaseMonitor {
|
||||
// Not transferred
|
||||
m.state.Phase = RemoteDownloadTaskPhaseTransfer
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
// Seeding complete
|
||||
m.l.Info("Download task seeding completed")
|
||||
return task.StatusCompleted, nil
|
||||
case downloader.StatusDownloading:
|
||||
m.ResumeAfter(resumeAfter)
|
||||
return task.StatusSuspending, nil
|
||||
case downloader.StatusUnknown, downloader.StatusError:
|
||||
return task.StatusError, fmt.Errorf("download task failed with state %q (%w), errorMsg: %s", status.State, queue.CriticalErr, status.ErrorMessage)
|
||||
}
|
||||
|
||||
m.ResumeAfter(resumeAfter)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *RemoteDownloadTask) slaveTransfer(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
u := inventory.UserFromContext(ctx)
|
||||
if m.state.Transferred == nil {
|
||||
m.state.Transferred = make(map[int]interface{})
|
||||
}
|
||||
|
||||
if m.state.SlaveUploadTaskID == 0 {
|
||||
dstUri, err := fs.NewUriFromString(m.state.Dst)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to parse dst uri %q: %s (%w)", m.state.Dst, err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
// Create slave upload task
|
||||
payload := &SlaveUploadTaskState{
|
||||
Files: []SlaveUploadEntity{},
|
||||
MaxParallel: dep.SettingProvider().MaxParallelTransfer(ctx),
|
||||
UserID: u.ID,
|
||||
}
|
||||
|
||||
// Construct files to be transferred
|
||||
for _, f := range m.state.Status.Files {
|
||||
if !f.Selected {
|
||||
continue
|
||||
}
|
||||
|
||||
// Skip already transferred
|
||||
if _, ok := m.state.Transferred[f.Index]; ok {
|
||||
continue
|
||||
}
|
||||
|
||||
dst := dstUri.JoinRaw(sanitizeFileName(f.Name))
|
||||
src := path.Join(m.state.Status.SavePath, f.Name)
|
||||
payload.Files = append(payload.Files, SlaveUploadEntity{
|
||||
Src: src,
|
||||
Uri: dst,
|
||||
Size: f.Size,
|
||||
Index: f.Index,
|
||||
})
|
||||
}
|
||||
|
||||
payloadStr, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to marshal payload: %w", err)
|
||||
}
|
||||
|
||||
taskId, err := m.node.CreateTask(ctx, queue.SlaveUploadTaskType, string(payloadStr))
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to create slave task: %w", err)
|
||||
}
|
||||
|
||||
m.state.NodeState.progress = nil
|
||||
m.state.SlaveUploadTaskID = taskId
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
m.l.Info("Checking slave upload task %d...", m.state.SlaveUploadTaskID)
|
||||
t, err := m.node.GetTask(ctx, m.state.SlaveUploadTaskID, true)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get slave task: %w", err)
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.state.NodeState.progress = t.Progress
|
||||
m.Unlock()
|
||||
|
||||
m.state.SlaveUploadState = &SlaveUploadTaskState{}
|
||||
if err := json.Unmarshal([]byte(t.PrivateState), m.state.SlaveUploadState); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to unmarshal slave compress state: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
if t.Status == task.StatusError || t.Status == task.StatusCompleted {
|
||||
if len(m.state.SlaveUploadState.Transferred) < len(m.state.SlaveUploadState.Files) {
|
||||
// Not all files transferred, retry
|
||||
slaveTaskId := m.state.SlaveUploadTaskID
|
||||
m.state.SlaveUploadTaskID = 0
|
||||
for i, _ := range m.state.SlaveUploadState.Transferred {
|
||||
m.state.Transferred[m.state.SlaveUploadState.Files[i].Index] = struct{}{}
|
||||
}
|
||||
|
||||
m.l.Warning("Slave task %d failed to transfer %d files, retrying...", slaveTaskId, len(m.state.SlaveUploadState.Files)-len(m.state.SlaveUploadState.Transferred))
|
||||
return task.StatusError, fmt.Errorf(
|
||||
"slave task failed to transfer %d files, first 5 errors: %s",
|
||||
len(m.state.SlaveUploadState.Files)-len(m.state.SlaveUploadState.Transferred),
|
||||
m.state.SlaveUploadState.First5TransferErrors,
|
||||
)
|
||||
} else {
|
||||
m.state.Phase = RemoteDownloadTaskPhaseAwaitSeeding
|
||||
m.ResumeAfter(0)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
}
|
||||
|
||||
if t.Status == task.StatusCanceled {
|
||||
return task.StatusError, fmt.Errorf("slave task canceled (%w)", queue.CriticalErr)
|
||||
}
|
||||
|
||||
m.l.Info("Slave task %d is still uploading, resume after 30s.", m.state.SlaveUploadTaskID)
|
||||
m.ResumeAfter(time.Second * 30)
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *RemoteDownloadTask) masterTransfer(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
if m.state.Transferred == nil {
|
||||
m.state.Transferred = make(map[int]interface{})
|
||||
}
|
||||
|
||||
maxParallel := dep.SettingProvider().MaxParallelTransfer(ctx)
|
||||
wg := sync.WaitGroup{}
|
||||
worker := make(chan int, maxParallel)
|
||||
for i := 0; i < maxParallel; i++ {
|
||||
worker <- i
|
||||
}
|
||||
|
||||
// Sum up total count and select files
|
||||
totalCount := 0
|
||||
totalSize := int64(0)
|
||||
allFiles := make([]downloader.TaskFile, 0, len(m.state.Status.Files))
|
||||
for _, f := range m.state.Status.Files {
|
||||
if f.Selected {
|
||||
allFiles = append(allFiles, f)
|
||||
totalSize += f.Size
|
||||
totalCount++
|
||||
}
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.progress = make(queue.Progresses)
|
||||
m.progress[ProgressTypeUploadCount] = &queue.Progress{Total: int64(totalCount)}
|
||||
m.progress[ProgressTypeUpload] = &queue.Progress{Total: totalSize}
|
||||
m.Unlock()
|
||||
|
||||
dstUri, err := fs.NewUriFromString(m.state.Dst)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to parse dst uri: %s (%w)", err, queue.CriticalErr)
|
||||
}
|
||||
|
||||
user := inventory.UserFromContext(ctx)
|
||||
fm := manager.NewFileManager(dep, user)
|
||||
failed := int64(0)
|
||||
ae := serializer.NewAggregateError()
|
||||
|
||||
transferFunc := func(workerId int, file downloader.TaskFile) {
|
||||
sanitizedName := sanitizeFileName(file.Name)
|
||||
dst := dstUri.JoinRaw(sanitizedName)
|
||||
src := filepath.FromSlash(path.Join(m.state.Status.SavePath, file.Name))
|
||||
m.l.Info("Uploading file %s to %s...", src, sanitizedName, dst)
|
||||
|
||||
progressKey := fmt.Sprintf("%s%d", ProgressTypeUploadSinglePrefix, workerId)
|
||||
m.Lock()
|
||||
m.progress[progressKey] = &queue.Progress{Identifier: dst.String(), Total: file.Size}
|
||||
fileProgress := m.progress[progressKey]
|
||||
uploadProgress := m.progress[ProgressTypeUpload]
|
||||
uploadCountProgress := m.progress[ProgressTypeUploadCount]
|
||||
m.Unlock()
|
||||
|
||||
defer func() {
|
||||
atomic.AddInt64(&uploadCountProgress.Current, 1)
|
||||
worker <- workerId
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
fileStream, err := os.Open(src)
|
||||
if err != nil {
|
||||
m.l.Warning("Failed to open file %s: %s", src, err.Error())
|
||||
atomic.AddInt64(&uploadProgress.Current, file.Size)
|
||||
atomic.AddInt64(&failed, 1)
|
||||
ae.Add(file.Name, fmt.Errorf("failed to open file: %w", err))
|
||||
return
|
||||
}
|
||||
|
||||
defer fileStream.Close()
|
||||
|
||||
fileData := &fs.UploadRequest{
|
||||
Props: &fs.UploadProps{
|
||||
Uri: dst,
|
||||
Size: file.Size,
|
||||
},
|
||||
ProgressFunc: func(current, diff int64, total int64) {
|
||||
atomic.AddInt64(&fileProgress.Current, diff)
|
||||
atomic.AddInt64(&uploadProgress.Current, diff)
|
||||
},
|
||||
File: fileStream,
|
||||
}
|
||||
|
||||
_, err = fm.Update(ctx, fileData, fs.WithNoEntityType())
|
||||
if err != nil {
|
||||
m.l.Warning("Failed to upload file %s: %s", src, err.Error())
|
||||
atomic.AddInt64(&failed, 1)
|
||||
atomic.AddInt64(&uploadProgress.Current, file.Size)
|
||||
ae.Add(file.Name, fmt.Errorf("failed to upload file: %w", err))
|
||||
return
|
||||
}
|
||||
|
||||
m.Lock()
|
||||
m.state.Transferred[file.Index] = nil
|
||||
m.Unlock()
|
||||
}
|
||||
|
||||
// Start upload files
|
||||
for _, file := range allFiles {
|
||||
// Check if file is already transferred
|
||||
if _, ok := m.state.Transferred[file.Index]; ok {
|
||||
m.l.Info("File %s already transferred, skipping...", file.Name)
|
||||
m.Lock()
|
||||
atomic.AddInt64(&m.progress[ProgressTypeUpload].Current, file.Size)
|
||||
atomic.AddInt64(&m.progress[ProgressTypeUploadCount].Current, 1)
|
||||
m.Unlock()
|
||||
continue
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return task.StatusError, ctx.Err()
|
||||
case workerId := <-worker:
|
||||
wg.Add(1)
|
||||
|
||||
go transferFunc(workerId, file)
|
||||
}
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
if failed > 0 {
|
||||
m.state.Failed = int(failed)
|
||||
m.l.Error("Failed to transfer %d file(s).", failed)
|
||||
return task.StatusError, fmt.Errorf("failed to transfer %d file(s), first 5 errors: %s", failed, ae.FormatFirstN(5))
|
||||
}
|
||||
|
||||
m.l.Info("All files transferred.")
|
||||
m.state.Phase = RemoteDownloadTaskPhaseAwaitSeeding
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *RemoteDownloadTask) awaitSeeding(ctx context.Context, dep dependency.Dep) (task.Status, error) {
|
||||
return task.StatusSuspending, nil
|
||||
}
|
||||
|
||||
func (m *RemoteDownloadTask) validateFiles(ctx context.Context, dep dependency.Dep, status *downloader.TaskStatus) error {
|
||||
// Validate files
|
||||
user := inventory.UserFromContext(ctx)
|
||||
fm := manager.NewFileManager(dep, user)
|
||||
|
||||
dstUri, err := fs.NewUriFromString(m.state.Dst)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to parse dst uri: %w", err)
|
||||
}
|
||||
|
||||
selectedFiles := lo.Filter(status.Files, func(f downloader.TaskFile, _ int) bool {
|
||||
return f.Selected
|
||||
})
|
||||
if len(selectedFiles) == 0 {
|
||||
return fmt.Errorf("no selected file found in download task")
|
||||
}
|
||||
|
||||
validateArgs := lo.Map(selectedFiles, func(f downloader.TaskFile, _ int) fs.PreValidateFile {
|
||||
return fs.PreValidateFile{
|
||||
Name: sanitizeFileName(f.Name),
|
||||
Size: f.Size,
|
||||
OmitName: f.Name == "",
|
||||
}
|
||||
})
|
||||
|
||||
if err := fm.PreValidateUpload(ctx, dstUri, validateArgs...); err != nil {
|
||||
return fmt.Errorf("failed to pre-validate files: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *RemoteDownloadTask) Cleanup(ctx context.Context) error {
|
||||
if m.state.Handle != nil {
|
||||
if err := m.d.Cancel(ctx, m.state.Handle); err != nil {
|
||||
m.l.Warning("failed to cancel download task: %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
if m.state.Status != nil && m.node.IsMaster() && m.state.Status.SavePath != "" {
|
||||
if err := os.RemoveAll(m.state.Status.SavePath); err != nil {
|
||||
m.l.Warning("failed to remove download temp folder: %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetDownloadTarget sets the files to download for the task
|
||||
func (m *RemoteDownloadTask) SetDownloadTarget(ctx context.Context, args ...*downloader.SetFileToDownloadArgs) error {
|
||||
if m.state.Handle == nil {
|
||||
return fmt.Errorf("download task not created")
|
||||
}
|
||||
|
||||
return m.d.SetFilesToDownload(ctx, m.state.Handle, args...)
|
||||
}
|
||||
|
||||
// CancelDownload cancels the download task
|
||||
func (m *RemoteDownloadTask) CancelDownload(ctx context.Context) error {
|
||||
if m.state.Handle == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
return m.d.Cancel(ctx, m.state.Handle)
|
||||
}
|
||||
|
||||
func (m *RemoteDownloadTask) Summarize(hasher hashid.Encoder) *queue.Summary {
|
||||
// unmarshal state
|
||||
if m.state == nil {
|
||||
if err := json.Unmarshal([]byte(m.State()), &m.state); err != nil {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
var status *downloader.TaskStatus
|
||||
if m.state.Status != nil {
|
||||
status = &*m.state.Status
|
||||
|
||||
// Redact save path
|
||||
status.SavePath = ""
|
||||
}
|
||||
|
||||
failed := m.state.Failed
|
||||
if m.state.SlaveUploadState != nil && m.state.Phase != RemoteDownloadTaskPhaseTransfer {
|
||||
failed = len(m.state.SlaveUploadState.Files) - len(m.state.SlaveUploadState.Transferred)
|
||||
}
|
||||
|
||||
return &queue.Summary{
|
||||
Phase: string(m.state.Phase),
|
||||
NodeID: m.state.NodeID,
|
||||
Props: map[string]any{
|
||||
SummaryKeySrcStr: m.state.SrcUri,
|
||||
SummaryKeySrc: m.state.SrcFileUri,
|
||||
SummaryKeyDst: m.state.Dst,
|
||||
SummaryKeyFailed: failed,
|
||||
SummaryKeyDownloadStatus: status,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (m *RemoteDownloadTask) Progress(ctx context.Context) queue.Progresses {
|
||||
m.Lock()
|
||||
defer m.Unlock()
|
||||
|
||||
merged := make(queue.Progresses)
|
||||
for k, v := range m.progress {
|
||||
merged[k] = v
|
||||
}
|
||||
|
||||
if m.state.NodeState.progress != nil {
|
||||
for k, v := range m.state.NodeState.progress {
|
||||
merged[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
return merged
|
||||
}
|
||||
|
||||
func sanitizeFileName(name string) string {
|
||||
r := strings.NewReplacer("\\", "_", ":", "_", "*", "_", "?", "_", "\"", "_", "<", "_", ">", "_", "|", "_")
|
||||
return r.Replace(name)
|
||||
}
|
||||
@@ -0,0 +1,234 @@
|
||||
package workflows
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/cloudreve/Cloudreve/v4/application/dependency"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent"
|
||||
"github.com/cloudreve/Cloudreve/v4/ent/task"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory/types"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/fs"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/filemanager/manager"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/logging"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/queue"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/serializer"
|
||||
)
|
||||
|
||||
type (
|
||||
SlaveUploadEntity struct {
|
||||
Uri *fs.URI `json:"uri"`
|
||||
Src string `json:"src"`
|
||||
Size int64 `json:"size"`
|
||||
Index int `json:"index"`
|
||||
}
|
||||
SlaveUploadTaskState struct {
|
||||
MaxParallel int `json:"max_parallel"`
|
||||
Files []SlaveUploadEntity `json:"files"`
|
||||
Transferred map[int]interface{} `json:"transferred"`
|
||||
UserID int `json:"user_id"`
|
||||
First5TransferErrors string `json:"first_5_transfer_errors,omitempty"`
|
||||
}
|
||||
SlaveUploadTask struct {
|
||||
*queue.InMemoryTask
|
||||
|
||||
progress queue.Progresses
|
||||
l logging.Logger
|
||||
state *SlaveUploadTaskState
|
||||
node cluster.Node
|
||||
}
|
||||
)
|
||||
|
||||
// NewSlaveUploadTask creates a new SlaveUploadTask from raw private state
|
||||
func NewSlaveUploadTask(ctx context.Context, props *types.SlaveTaskProps, id int, state string) queue.Task {
|
||||
return &SlaveUploadTask{
|
||||
InMemoryTask: &queue.InMemoryTask{
|
||||
DBTask: &queue.DBTask{
|
||||
Task: &ent.Task{
|
||||
ID: id,
|
||||
CorrelationID: logging.CorrelationID(ctx),
|
||||
PublicState: &types.TaskPublicState{
|
||||
SlaveTaskProps: props,
|
||||
},
|
||||
PrivateState: state,
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
progress: make(queue.Progresses),
|
||||
}
|
||||
}
|
||||
|
||||
func (t *SlaveUploadTask) Do(ctx context.Context) (task.Status, error) {
|
||||
ctx = prepareSlaveTaskCtx(ctx, t.Model().PublicState.SlaveTaskProps)
|
||||
dep := dependency.FromContext(ctx)
|
||||
t.l = dep.Logger()
|
||||
|
||||
np, err := dep.NodePool(ctx)
|
||||
if err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to get node pool: %w", err)
|
||||
}
|
||||
|
||||
t.node, err = np.Get(ctx, types.NodeCapabilityNone, 0)
|
||||
if err != nil || !t.node.IsMaster() {
|
||||
return task.StatusError, fmt.Errorf("failed to get master node: %w", err)
|
||||
}
|
||||
|
||||
fm := manager.NewFileManager(dep, nil)
|
||||
|
||||
// unmarshal state
|
||||
state := &SlaveUploadTaskState{}
|
||||
if err := json.Unmarshal([]byte(t.State()), state); err != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to unmarshal state: %w", err)
|
||||
}
|
||||
|
||||
t.state = state
|
||||
if t.state.Transferred == nil {
|
||||
t.state.Transferred = make(map[int]interface{})
|
||||
}
|
||||
|
||||
wg := sync.WaitGroup{}
|
||||
worker := make(chan int, t.state.MaxParallel)
|
||||
for i := 0; i < t.state.MaxParallel; i++ {
|
||||
worker <- i
|
||||
}
|
||||
|
||||
// Sum up total count
|
||||
totalCount := 0
|
||||
totalSize := int64(0)
|
||||
for _, res := range state.Files {
|
||||
totalSize += res.Size
|
||||
totalCount++
|
||||
}
|
||||
t.Lock()
|
||||
t.progress[ProgressTypeUploadCount] = &queue.Progress{}
|
||||
t.progress[ProgressTypeUpload] = &queue.Progress{}
|
||||
t.Unlock()
|
||||
atomic.StoreInt64(&t.progress[ProgressTypeUploadCount].Total, int64(totalCount))
|
||||
atomic.StoreInt64(&t.progress[ProgressTypeUpload].Total, totalSize)
|
||||
ae := serializer.NewAggregateError()
|
||||
transferFunc := func(workerId, fileId int, file SlaveUploadEntity) {
|
||||
t.l.Info("Uploading file %s to %s...", file.Src, file.Uri.String())
|
||||
|
||||
progressKey := fmt.Sprintf("%s%d", ProgressTypeUploadSinglePrefix, workerId)
|
||||
t.Lock()
|
||||
t.progress[progressKey] = &queue.Progress{Identifier: file.Uri.String(), Total: file.Size}
|
||||
fileProgress := t.progress[progressKey]
|
||||
uploadProgress := t.progress[ProgressTypeUpload]
|
||||
uploadCountProgress := t.progress[ProgressTypeUploadCount]
|
||||
t.Unlock()
|
||||
|
||||
defer func() {
|
||||
atomic.AddInt64(&uploadCountProgress.Current, 1)
|
||||
worker <- workerId
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
handle, err := os.Open(filepath.FromSlash(file.Src))
|
||||
if err != nil {
|
||||
t.l.Warning("Failed to open file %s: %s", file.Src, err.Error())
|
||||
atomic.AddInt64(&fileProgress.Current, file.Size)
|
||||
ae.Add(path.Base(file.Src), fmt.Errorf("failed to open file: %w", err))
|
||||
return
|
||||
}
|
||||
|
||||
stat, err := handle.Stat()
|
||||
if err != nil {
|
||||
t.l.Warning("Failed to get file stat for %s: %s", file.Src, err.Error())
|
||||
handle.Close()
|
||||
atomic.AddInt64(&fileProgress.Current, file.Size)
|
||||
ae.Add(path.Base(file.Src), fmt.Errorf("failed to get file stat: %w", err))
|
||||
return
|
||||
}
|
||||
|
||||
fileData := &fs.UploadRequest{
|
||||
Props: &fs.UploadProps{
|
||||
Uri: file.Uri,
|
||||
Size: stat.Size(),
|
||||
},
|
||||
ProgressFunc: func(current, diff int64, total int64) {
|
||||
atomic.AddInt64(&fileProgress.Current, diff)
|
||||
atomic.AddInt64(&uploadCountProgress.Current, 1)
|
||||
atomic.StoreInt64(&fileProgress.Total, total)
|
||||
},
|
||||
File: handle,
|
||||
Seeker: handle,
|
||||
}
|
||||
|
||||
_, err = fm.Update(ctx, fileData, fs.WithNode(t.node), fs.WithStatelessUserID(t.state.UserID), fs.WithNoEntityType())
|
||||
if err != nil {
|
||||
handle.Close()
|
||||
t.l.Warning("Failed to upload file %s: %s", file.Src, err.Error())
|
||||
atomic.AddInt64(&uploadProgress.Current, file.Size)
|
||||
ae.Add(path.Base(file.Src), fmt.Errorf("failed to upload file: %w", err))
|
||||
return
|
||||
}
|
||||
|
||||
t.Lock()
|
||||
t.state.Transferred[fileId] = nil
|
||||
t.Unlock()
|
||||
handle.Close()
|
||||
}
|
||||
|
||||
// Start upload files
|
||||
for fileId, file := range t.state.Files {
|
||||
// Check if file is already transferred
|
||||
if _, ok := t.state.Transferred[fileId]; ok {
|
||||
t.l.Info("File %s already transferred, skipping...", file.Src)
|
||||
t.Lock()
|
||||
atomic.AddInt64(&t.progress[ProgressTypeUpload].Current, file.Size)
|
||||
atomic.AddInt64(&t.progress[ProgressTypeUploadCount].Current, 1)
|
||||
t.Unlock()
|
||||
continue
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return task.StatusError, ctx.Err()
|
||||
case workerId := <-worker:
|
||||
wg.Add(1)
|
||||
|
||||
go transferFunc(workerId, fileId, file)
|
||||
}
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
t.state.First5TransferErrors = ae.FormatFirstN(5)
|
||||
newStateStr, marshalErr := json.Marshal(t.state)
|
||||
if marshalErr != nil {
|
||||
return task.StatusError, fmt.Errorf("failed to marshal state: %w", marshalErr)
|
||||
}
|
||||
t.Lock()
|
||||
t.Task.PrivateState = string(newStateStr)
|
||||
t.Unlock()
|
||||
|
||||
// If all files are failed to transfer, return error
|
||||
if len(t.state.Transferred) != len(t.state.Files) {
|
||||
t.l.Warning("%d files not transferred", len(t.state.Files)-len(t.state.Transferred))
|
||||
if len(t.state.Transferred) == 0 {
|
||||
return task.StatusError, fmt.Errorf("all file failed to transfer")
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
return task.StatusCompleted, nil
|
||||
}
|
||||
|
||||
func (m *SlaveUploadTask) Progress(ctx context.Context) queue.Progresses {
|
||||
m.Lock()
|
||||
defer m.Unlock()
|
||||
|
||||
res := make(queue.Progresses)
|
||||
for k, v := range m.progress {
|
||||
res[k] = v
|
||||
}
|
||||
return res
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package workflows
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/cloudreve/Cloudreve/v4/application/dependency"
|
||||
"github.com/cloudreve/Cloudreve/v4/inventory/types"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/queue"
|
||||
"github.com/cloudreve/Cloudreve/v4/pkg/util"
|
||||
)
|
||||
|
||||
const (
|
||||
TaskTempPath = "fm_workflows"
|
||||
slaveProgressRefreshInterval = 5 * time.Second
|
||||
)
|
||||
|
||||
type NodeState struct {
|
||||
NodeID int `json:"node_id"`
|
||||
|
||||
progress queue.Progresses
|
||||
}
|
||||
|
||||
// allocateNode allocates a node for the task.
|
||||
func allocateNode(ctx context.Context, dep dependency.Dep, state *NodeState, capability types.NodeCapability) (cluster.Node, error) {
|
||||
np, err := dep.NodePool(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get node pool: %w", err)
|
||||
}
|
||||
|
||||
node, err := np.Get(ctx, capability, state.NodeID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get node: %w", err)
|
||||
}
|
||||
|
||||
state.NodeID = node.ID()
|
||||
return node, nil
|
||||
}
|
||||
|
||||
// prepareSlaveTaskCtx prepares the context for the slave task.
|
||||
func prepareSlaveTaskCtx(ctx context.Context, props *types.SlaveTaskProps) context.Context {
|
||||
ctx = context.WithValue(ctx, cluster.SlaveNodeIDCtx{}, strconv.Itoa(props.NodeID))
|
||||
ctx = context.WithValue(ctx, cluster.MasterSiteUrlCtx{}, props.MasterSiteURl)
|
||||
ctx = context.WithValue(ctx, cluster.MasterSiteVersionCtx{}, props.MasterSiteVersion)
|
||||
ctx = context.WithValue(ctx, cluster.MasterSiteIDCtx{}, props.MasterSiteID)
|
||||
return ctx
|
||||
}
|
||||
|
||||
func prepareTempFolder(ctx context.Context, dep dependency.Dep, t queue.Task) (string, error) {
|
||||
settings := dep.SettingProvider()
|
||||
tempPath := util.DataPath(path.Join(settings.TempPath(ctx), TaskTempPath, strconv.Itoa(t.ID())))
|
||||
if err := util.CreatNestedFolder(tempPath); err != nil {
|
||||
return "", fmt.Errorf("failed to create temp folder: %w", err)
|
||||
}
|
||||
|
||||
dep.Logger().Info("Temp folder created: %s", tempPath)
|
||||
return tempPath, nil
|
||||
}
|
||||
Reference in New Issue
Block a user