acp/acp.go

package acp

import (
	"context"
	"fmt"
	"sync"

	"github.com/sirupsen/logrus"
)

type Copyer struct {
	*option
	running           sync.WaitGroup
	errLock           sync.Mutex
	err               error
	eventCh           chan Event
	getDevice         func(in string) string
	getDiskUsageCache func(mountPoint string) *diskUsageCache
}

func New(ctx context.Context, opts ...Option) (*Copyer, error) {
	opt := newOption()
	for _, o := range opts {
		if o == nil {
			continue
		}
		opt = o(opt)
	}
	if err := opt.check(); err != nil {
		return nil, err
	}

	getDevice, err := getMountpointCache()
	if err != nil {
		return nil, err
	}

	c := &Copyer{
		option:    opt,
		eventCh:   make(chan Event, 128),
		getDevice: getDevice,
		getDiskUsageCache: Cache(func(mountPoint string) *diskUsageCache {
			return newDiskUsageCache(mountPoint, defaultDiskUsageFreshInterval)
		}),
	}

	// Account for both pipeline and event dispatch before either goroutine starts.
	c.running.Add(2)
	go wrap(ctx, func() { c.run(ctx) })

	return c, nil
}

func (c *Copyer) Wait() {
	c.running.Wait()
}

// WaitErr waits for the copy pipeline and returns its first error.
func (c *Copyer) WaitErr() error {
	c.Wait()
	c.errLock.Lock()
	defer c.errLock.Unlock()
	return c.err
}

func (c *Copyer) setError(err error) {
	if err == nil {
		return
	}
	c.errLock.Lock()
	defer c.errLock.Unlock()
	if c.err == nil {
		c.err = err
	}
}

func (c *Copyer) run(ctx context.Context) error {
	defer c.running.Done()
	defer close(c.eventCh)

	go wrap(ctx, func() { c.eventLoop(ctx) })

	indexed, err := c.index(ctx)
	if err != nil {
		c.setError(err)
		return err
	}

	prepared := c.prepare(ctx, indexed)
	copyed := c.copy(ctx, prepared)
	sinkFailed := c.cleanupJob(ctx, copyed)
	if c.streamSink != nil && !sinkFailed {
		if err := c.streamSink.Flush(ctx); err != nil {
			c.setError(fmt.Errorf("flush stream sink failed, %w", err))
		}
	}

	// empty pipes
	for range indexed {
	}
	for range prepared {
	}
	for range copyed {
	}

	return nil
}

func (c *Copyer) eventLoop(ctx context.Context) {
	defer c.running.Done()

	chans := make([]chan Event, len(c.eventHanders))
	for idx := range chans {
		chans[idx] = make(chan Event, 128)
	}

	var handlers sync.WaitGroup
	for idx, ch := range chans {
		handler := c.eventHanders[idx]
		events := ch

		handlers.Add(1)
		go wrap(ctx, func() {
			defer handlers.Done()

			for {
				e, ok := <-events
				if !ok {
					handler(&EventFinished{})
					return
				}
				handler(e)
			}
		})
	}

	defer func() {
		for _, ch := range chans {
			close(ch)
		}
		handlers.Wait()
	}()
	for e := range c.eventCh {
		for _, ch := range chans {
			ch <- e
		}
	}
}

func (c *Copyer) logf(l logrus.Level, format string, args ...any) {
	c.logger.Logf(l, format, args...)
}

func (c *Copyer) submit(e Event) {
	c.eventCh <- e
}

func (c *Copyer) reportError(src, dst string, err error) {
	e := &Error{Src: src, Dst: dst, Err: err}
	c.setError(fmt.Errorf("copy failed, source=%q target=%q, %w", src, dst, err))
	c.logf(logrus.ErrorLevel, e.Error())
	c.submit(&EventReportError{Error: e})
}