feat: add unified worker pool for recursive downloads
This commit is contained in:
+88
-139
@@ -20,6 +20,7 @@ import (
|
||||
|
||||
"codeberg.org/petrbalvin/goget/internal/core"
|
||||
"codeberg.org/petrbalvin/goget/internal/output"
|
||||
"codeberg.org/petrbalvin/goget/internal/pool"
|
||||
"codeberg.org/petrbalvin/goget/internal/protocol"
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
@@ -862,56 +863,113 @@ func downloadRecursiveSequential(ctx context.Context, req *core.DownloadRequest)
|
||||
return sftpDownloadDirStandalone(ctx, c, req)
|
||||
}
|
||||
|
||||
// sftpTask represents a unit of work for the concurrent recursive
|
||||
// download: either a directory to list or a file to download.
|
||||
type sftpTask struct {
|
||||
req *core.DownloadRequest
|
||||
isDir bool
|
||||
size uint64
|
||||
}
|
||||
|
||||
// downloadRecursiveParallel downloads a remote directory using a pool
|
||||
// of N SFTP connections. All workers share a single semaphore so the
|
||||
// total in-flight goroutines never exceed N, regardless of directory
|
||||
// depth.
|
||||
// of N SFTP connections. All workers share a single pool so the total
|
||||
// in-flight goroutines never exceed N, regardless of directory depth.
|
||||
func downloadRecursiveParallel(ctx context.Context, req *core.DownloadRequest) (*core.DownloadResult, error) {
|
||||
parallel := req.RecursiveParallel
|
||||
if parallel < 1 {
|
||||
parallel = 1
|
||||
}
|
||||
pool, err := newSFTPClientPool(ctx, req, parallel)
|
||||
cpool, err := newSFTPClientPool(ctx, req, parallel)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer pool.close()
|
||||
defer cpool.close()
|
||||
|
||||
var (
|
||||
statsMu sync.Mutex
|
||||
totalBytes int64
|
||||
totalChunks int
|
||||
firstErr error
|
||||
)
|
||||
|
||||
sem := make(chan struct{}, parallel)
|
||||
var wg sync.WaitGroup
|
||||
|
||||
setErr := func(err error) {
|
||||
statsMu.Lock()
|
||||
if firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
statsMu.Unlock()
|
||||
}
|
||||
|
||||
wg.Add(1)
|
||||
sem <- struct{}{}
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
defer func() { <-sem }()
|
||||
if ctx.Err() != nil {
|
||||
p := pool.New(ctx, parallel, func(taskCtx context.Context, t sftpTask) {
|
||||
if taskCtx.Err() != nil {
|
||||
return
|
||||
}
|
||||
c := pool.acquire()
|
||||
defer pool.release(c)
|
||||
sftpDownloadDirWorker(ctx, c, req, &wg, sem, setErr, &statsMu, &totalBytes, &totalChunks, pool)
|
||||
}()
|
||||
|
||||
wg.Wait()
|
||||
if firstErr != nil {
|
||||
return nil, firstErr
|
||||
c := cpool.acquire()
|
||||
defer cpool.release(c)
|
||||
|
||||
if !t.isDir {
|
||||
if err := sftpDownloadFile(taskCtx, c, t.req, t.size); err != nil {
|
||||
pool.Get[sftpTask](taskCtx).AddError(fmt.Errorf("sftp download %s: %w", t.req.URL.Path, err))
|
||||
return
|
||||
}
|
||||
statsMu.Lock()
|
||||
totalBytes += int64(t.size)
|
||||
totalChunks++
|
||||
statsMu.Unlock()
|
||||
return
|
||||
}
|
||||
|
||||
// Directory: list and submit child tasks.
|
||||
remotePath := t.req.URL.Path
|
||||
if remotePath == "" {
|
||||
remotePath = "."
|
||||
}
|
||||
|
||||
absPath, err := c.realpath(remotePath)
|
||||
if err != nil {
|
||||
absPath = remotePath
|
||||
}
|
||||
|
||||
outputDir := t.req.Output
|
||||
if outputDir == "" {
|
||||
outputDir = filepath.Base(absPath)
|
||||
if outputDir == "" || outputDir == "/" {
|
||||
outputDir = "download"
|
||||
}
|
||||
}
|
||||
|
||||
entries, err := sftpListDir(taskCtx, c, absPath)
|
||||
if err != nil {
|
||||
pool.Get[sftpTask](taskCtx).AddError(fmt.Errorf("sftp list %s: %w", absPath, err))
|
||||
return
|
||||
}
|
||||
|
||||
if err := os.MkdirAll(outputDir, 0755); err != nil {
|
||||
pool.Get[sftpTask](taskCtx).AddError(fmt.Errorf("mkdir %s: %w", outputDir, err))
|
||||
return
|
||||
}
|
||||
|
||||
pp := pool.Get[sftpTask](taskCtx)
|
||||
for _, e := range entries {
|
||||
name := e.Name
|
||||
if name == "." || name == ".." {
|
||||
continue
|
||||
}
|
||||
entryIsDir := e.Attr != nil && sftpIsDir(e.Attr.Mode)
|
||||
remoteChild := absPath + "/" + name
|
||||
localChild := filepath.Join(outputDir, name)
|
||||
subReq := *t.req
|
||||
subReq.URL = &url.URL{Scheme: t.req.URL.Scheme, User: t.req.URL.User, Host: t.req.URL.Host, Path: remoteChild}
|
||||
subReq.Output = localChild
|
||||
if entryIsDir {
|
||||
pp.Submit(sftpTask{req: &subReq, isDir: true})
|
||||
} else {
|
||||
subReq.Recursive = false
|
||||
subReq.RecursiveParallel = 0
|
||||
pp.Submit(sftpTask{req: &subReq, isDir: false, size: e.Attr.Size})
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
p.Submit(sftpTask{req: req, isDir: true})
|
||||
|
||||
if errs := p.Wait(); len(errs) > 0 {
|
||||
return nil, errs[0]
|
||||
}
|
||||
statsMu.Lock()
|
||||
defer statsMu.Unlock()
|
||||
return &core.DownloadResult{
|
||||
BytesDownloaded: totalBytes,
|
||||
Duration: 0,
|
||||
@@ -958,115 +1016,6 @@ func sftpDownloadDirStandalone(ctx context.Context, c *client, req *core.Downloa
|
||||
}, nil
|
||||
}
|
||||
|
||||
// sftpDownloadDirWorker is the worker-pool variant. It lists the
|
||||
// directory, then dispatches file downloads and subdirectory recursion
|
||||
// through the shared semaphore. Subdirectory workers reuse the same
|
||||
// pool to acquire clients. The return value is always nil; results
|
||||
// and errors are routed through the shared fields.
|
||||
func sftpDownloadDirWorker(
|
||||
ctx context.Context,
|
||||
c *client,
|
||||
req *core.DownloadRequest,
|
||||
wg *sync.WaitGroup,
|
||||
sem chan struct{},
|
||||
setErr func(error),
|
||||
statsMu *sync.Mutex,
|
||||
totalBytes *int64,
|
||||
totalChunks *int,
|
||||
pool *sftpClientPool,
|
||||
) {
|
||||
remotePath := req.URL.Path
|
||||
if remotePath == "" {
|
||||
remotePath = "."
|
||||
}
|
||||
|
||||
absPath, err := c.realpath(remotePath)
|
||||
if err != nil {
|
||||
absPath = remotePath
|
||||
}
|
||||
|
||||
outputDir := req.Output
|
||||
if outputDir == "" {
|
||||
outputDir = filepath.Base(absPath)
|
||||
if outputDir == "" || outputDir == "/" {
|
||||
outputDir = "download"
|
||||
}
|
||||
}
|
||||
|
||||
entries, err := sftpListDir(ctx, c, absPath)
|
||||
if err != nil {
|
||||
setErr(fmt.Errorf("sftp list %s: %w", absPath, err))
|
||||
return
|
||||
}
|
||||
|
||||
if err := os.MkdirAll(outputDir, 0755); err != nil {
|
||||
setErr(fmt.Errorf("mkdir %s: %w", outputDir, err))
|
||||
return
|
||||
}
|
||||
|
||||
for i, e := range entries {
|
||||
name := e.Name
|
||||
if name == "." || name == ".." {
|
||||
continue
|
||||
}
|
||||
isDir := e.Attr != nil && sftpIsDir(e.Attr.Mode)
|
||||
remoteChild := absPath + "/" + name
|
||||
localChild := filepath.Join(outputDir, name)
|
||||
_ = i
|
||||
|
||||
if isDir {
|
||||
wg.Add(1)
|
||||
select {
|
||||
case sem <- struct{}{}:
|
||||
go func(remoteChild, localChild string) {
|
||||
defer wg.Done()
|
||||
defer func() { <-sem }()
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
cc := pool.acquire()
|
||||
defer pool.release(cc)
|
||||
subReq := *req
|
||||
subReq.URL = &url.URL{Scheme: req.URL.Scheme, User: req.URL.User, Host: req.URL.Host, Path: remoteChild}
|
||||
subReq.Output = localChild
|
||||
sftpDownloadDirWorker(ctx, cc, &subReq, wg, sem, setErr, statsMu, totalBytes, totalChunks, pool)
|
||||
}(remoteChild, localChild)
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
} else {
|
||||
wg.Add(1)
|
||||
select {
|
||||
case sem <- struct{}{}:
|
||||
go func(remoteChild, localChild string, size uint64) {
|
||||
defer wg.Done()
|
||||
defer func() { <-sem }()
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
cc := pool.acquire()
|
||||
defer pool.release(cc)
|
||||
subReq := *req
|
||||
subReq.URL = &url.URL{Scheme: req.URL.Scheme, User: req.URL.User, Host: req.URL.Host, Path: remoteChild}
|
||||
subReq.Output = localChild
|
||||
subReq.Recursive = false
|
||||
subReq.RecursiveParallel = 0
|
||||
if err := sftpDownloadFile(ctx, cc, &subReq, size); err != nil {
|
||||
setErr(fmt.Errorf("sftp download %s: %w", remoteChild, err))
|
||||
return
|
||||
}
|
||||
statsMu.Lock()
|
||||
*totalBytes += int64(size)
|
||||
*totalChunks++
|
||||
statsMu.Unlock()
|
||||
}(remoteChild, localChild, e.Attr.Size)
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// sftpWalkDir recursively walks one remote directory on `c`,
|
||||
// downloading each file and recursing into subdirectories. Stats are
|
||||
// accumulated inline. Used by the standalone mode.
|
||||
|
||||
Reference in New Issue
Block a user