acp/index.go

package acp

import (
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"sync/atomic"
	"time"
)

const (
	UnexpectFileMode = os.ModeType &^ os.ModeDir
)

type counter struct {
	bytes, files int64
}

func (c *Copyer) index(ctx context.Context) (<-chan *baseJob, error) {
	if c.streamSource != nil {
		return c.indexStream(ctx), nil
	}

	jobs, err := c.walk(ctx)
	if err != nil {
		return nil, err
	}

	ch := make(chan *baseJob, 128)
	go wrap(ctx, func() {
		defer close(ch)

		for _, job := range jobs {
			select {
			case <-ctx.Done():
				return
			case ch <- job:
			}
		}
	})

	return ch, nil
}

func (c *Copyer) indexStream(ctx context.Context) <-chan *baseJob {
	ch := make(chan *baseJob, 128)
	go wrap(ctx, func() {
		defer close(ch)

		var bytes, files int64
		defer func() {
			c.submit(&EventUpdateCount{Bytes: bytes, Files: files, Finished: true})
		}()
		for {
			if c.linearTargetStopped() {
				return
			}
			request, err := c.streamSource.Next(ctx)
			if err != nil {
				if err != io.EOF {
					c.reportError("", "", fmt.Errorf("read stream source failed, %w", err))
				}
				return
			}
			if request == nil {
				c.reportError("", "", fmt.Errorf("read stream source failed, request is nil"))
				return
			}

			sourcePath := filepath.Clean(request.Source)
			info, err := os.Stat(sourcePath)
			if err != nil {
				c.reportError(sourcePath, "", fmt.Errorf("stream job get stat failed, %w", err))
				return
			}
			if !info.Mode().IsRegular() {
				c.reportError(sourcePath, "", fmt.Errorf("stream job source is not a regular file"))
				return
			}
			stat, err := newStat(sourcePath, info)
			if err != nil {
				c.reportError(sourcePath, "", fmt.Errorf("read stream job stat failed, %w", err))
				return
			}

			job := &baseJob{
				copyer:   c,
				src:      &source{base: filepath.Dir(sourcePath), path: filepath.Base(sourcePath)},
				path:     sourcePath,
				stat:     stat,
				targets:  append([]string(nil), request.Targets...),
				streamID: request.ID,
			}
			c.submit(&EventUpdateJob{job.report()})
			bytes += stat.size
			files++

			select {
			case <-ctx.Done():
				c.setError(ctx.Err())
				return
			case ch <- job:
			}
		}
	})
	return ch
}

func (c *Copyer) walk(ctx context.Context) ([]*baseJob, error) {
	done := make(chan struct{})
	var reporting sync.WaitGroup
	reporting.Add(1)
	defer func() {
		close(done)
		reporting.Wait()
	}()

	cntr := new(counter)
	go wrap(ctx, func() {
		defer reporting.Done()
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				c.submit(&EventUpdateCount{Bytes: atomic.LoadInt64(&cntr.bytes), Files: atomic.LoadInt64(&cntr.files)})
			case <-done:
				c.submit(&EventUpdateCount{Bytes: atomic.LoadInt64(&cntr.bytes), Files: atomic.LoadInt64(&cntr.files), Finished: true})
				return
			}
		}
	})

	jobs := make([]*baseJob, 0, 64)
	appendJob := func(job *baseJob) {
		if !job.stat.mode.IsRegular() {
			c.reportError(
				job.path, "",
				fmt.Errorf(
					"unexpected file mode, not regular file, mode= %s",
					job.stat.mode,
				),
			)
			return
		}

		c.submit(&EventUpdateJob{job.report()})
		jobs = append(jobs, job)
		atomic.AddInt64(&cntr.files, 1)
		atomic.AddInt64(&cntr.bytes, job.stat.size)
	}

	var walk func(src *source, dsts []string)
	walk = func(src *source, dsts []string) {
		path := src.src()

		fi, err := os.Stat(path)
		if err != nil {
			c.reportError(path, "", fmt.Errorf("walk get stat, %w", err))
			return
		}

		mode := fi.Mode()
		if mode.IsRegular() {
			targets := make([]string, 0, len(dsts))
			for _, d := range dsts {
				targets = append(targets, src.dst(d))
			}

			stat, err := newStat(path, fi)
			if err != nil {
				c.reportError(path, "", fmt.Errorf("read sys stat, %w", err))
				return
			}

			appendJob(&baseJob{
				copyer:  c,
				src:     src,
				path:    path,
				stat:    stat,
				targets: targets,
			})
			return
		}
		if mode&UnexpectFileMode != 0 {
			return
		}

		files, err := os.ReadDir(path)
		if err != nil {
			c.reportError(path, "", fmt.Errorf("walk read dir, %w", err))
			return
		}
		for _, file := range files {
			walk(src.append(file.Name()), dsts)
		}
	}

	results := make([]*baseJob, 0, 64)
	for _, j := range c.wildcardJobs {
		for _, s := range j.src {
			walk(s, j.dst)
		}

		if len(jobs) == 0 {
			continue
		}

		joined, err := c.joinJobs(jobs)
		if err != nil {
			return nil, err
		}

		results = append(results, joined...)
		jobs = jobs[:0]
	}

	for _, j := range c.accurateJobs {
		fi, err := os.Stat(j.src)
		if err != nil {
			c.reportError(j.src, "", fmt.Errorf("accurate job get stat, %w", err))
			continue
		}
		if !fi.Mode().IsRegular() {
			continue
		}

		stat, err := newStat(j.src, fi)
		if err != nil {
			c.reportError(j.src, "", fmt.Errorf("read sys stat, %w", err))
			continue
		}

		appendJob(&baseJob{
			copyer:  c,
			src:     &source{base: "/", path: j.src},
			path:    j.src,
			stat:    stat,
			targets: j.dsts,
		})
	}
	results = append(results, jobs...)

	return results, nil
}

func (c *Copyer) joinJobs(jobs []*baseJob) ([]*baseJob, error) {
	sort.Slice(jobs, func(i int, j int) bool {
		return comparePath(jobs[i].src.path, jobs[j].src.path) < 0
	})

	var last *baseJob
	filtered := make([]*baseJob, 0, len(jobs))
	for _, job := range jobs {
		if last != nil && last.src.path == job.src.path {
			c.reportError(last.path, "", fmt.Errorf("same relative path, ignored, '%s'", job.path))
			continue
		}

		filtered = append(filtered, job)
		last = job
	}

	return filtered, nil
}