OpenList/drivers/teldrive/copy.go

package teldrive

import (
	"fmt"
	"net/http"

	"github.com/OpenListTeam/OpenList/v4/drivers/base"
	"github.com/OpenListTeam/OpenList/v4/internal/model"
	"github.com/OpenListTeam/OpenList/v4/pkg/utils"
	"github.com/go-resty/resty/v2"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
	"golang.org/x/sync/semaphore"
)

func NewCopyManager(ctx context.Context, concurrent int, d *Teldrive) *CopyManager {
	g, ctx := errgroup.WithContext(ctx)

	return &CopyManager{
		TaskChan: make(chan CopyTask, concurrent*2),
		Sem:      semaphore.NewWeighted(int64(concurrent)),
		G:        g,
		Ctx:      ctx,
		d:        d,
	}
}

func (cm *CopyManager) startWorkers() {
	workerCount := cap(cm.TaskChan) / 2
	for i := 0; i < workerCount; i++ {
		cm.G.Go(func() error {
			return cm.worker()
		})
	}
}

func (cm *CopyManager) worker() error {
	for {
		select {
		case task, ok := <-cm.TaskChan:
			if !ok {
				return nil
			}

			if err := cm.Sem.Acquire(cm.Ctx, 1); err != nil {
				return err
			}

			var err error

			err = cm.processFile(task)

			cm.Sem.Release(1)

			if err != nil {
				return fmt.Errorf("task processing failed: %w", err)
			}

		case <-cm.Ctx.Done():
			return cm.Ctx.Err()
		}
	}
}

func (cm *CopyManager) generateTasks(ctx context.Context, srcObj, dstDir model.Obj) error {
	if srcObj.IsDir() {
		return cm.generateFolderTasks(ctx, srcObj, dstDir)
	} else {
		// add single file task directly
		select {
		case cm.TaskChan <- CopyTask{SrcObj: srcObj, DstDir: dstDir}:
			return nil
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (cm *CopyManager) generateFolderTasks(ctx context.Context, srcDir, dstDir model.Obj) error {
	objs, err := cm.d.List(ctx, srcDir, model.ListArgs{})
	if err != nil {
		return fmt.Errorf("failed to list directory %s: %w", srcDir.GetPath(), err)
	}

	err = cm.d.MakeDir(cm.Ctx, dstDir, srcDir.GetName())
	if err != nil || len(objs) == 0 {
		return err
	}
	newDstDir := &model.Object{
		ID:       dstDir.GetID(),
		Path:     dstDir.GetPath() + "/" + srcDir.GetName(),
		Name:     srcDir.GetName(),
		IsFolder: true,
	}

	for _, file := range objs {
		if utils.IsCanceled(ctx) {
			return ctx.Err()
		}

		srcFile := &model.Object{
			ID:       file.GetID(),
			Path:     srcDir.GetPath() + "/" + file.GetName(),
			Name:     file.GetName(),
			IsFolder: file.IsDir(),
		}

		// 递归生成任务
		if err := cm.generateTasks(ctx, srcFile, newDstDir); err != nil {
			return err
		}
	}

	return nil
}

func (cm *CopyManager) processFile(task CopyTask) error {
	return cm.copySingleFile(cm.Ctx, task.SrcObj, task.DstDir)
}

func (cm *CopyManager) copySingleFile(ctx context.Context, srcObj, dstDir model.Obj) error {
	// `override copy mode` should delete the existing file
	if obj, err := cm.d.getFile(dstDir.GetPath(), srcObj.GetName(), srcObj.IsDir()); err == nil {
		if err := cm.d.Remove(ctx, obj); err != nil {
			return fmt.Errorf("failed to remove existing file: %w", err)
		}
	}

	// Do copy
	return cm.d.request(http.MethodPost, "/api/files/{id}/copy", func(req *resty.Request) {
		req.SetPathParam("id", srcObj.GetID())
		req.SetBody(base.Json{
			"newName":     srcObj.GetName(),
			"destination": dstDir.GetPath(),
		})
	}, nil)
}