status-go/services/wallet/async/async.go

package async

import (
	"context"
	"sync"
	"time"
)

type Command func(context.Context) error

type Commander interface {
	Command() Command
}

// SingleShotCommand runs once.
type SingleShotCommand struct {
	Interval time.Duration
	Init     func(context.Context) error
	Runable  func(context.Context) error
}

func (c SingleShotCommand) Run(ctx context.Context) error {
	timer := time.NewTimer(c.Interval)
	if c.Init != nil {
		err := c.Init(ctx)
		if err != nil {
			return err
		}
	}
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-timer.C:
			_ = c.Runable(ctx)
		}
	}
}

// FiniteCommand terminates when error is nil.
type FiniteCommand struct {
	Interval time.Duration
	Runable  func(context.Context) error
}

func (c FiniteCommand) Run(ctx context.Context) error {
	err := c.Runable(ctx)
	if err == nil {
		return nil
	}
	ticker := time.NewTicker(c.Interval)
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			err := c.Runable(ctx)
			if err == nil {
				return nil
			}
		}
	}
}

// InfiniteCommand runs until context is closed.
type InfiniteCommand struct {
	Interval time.Duration
	Runable  func(context.Context) error
}

func (c InfiniteCommand) Run(ctx context.Context) error {
	_ = c.Runable(ctx)
	ticker := time.NewTicker(c.Interval)
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			_ = c.Runable(ctx)
		}
	}
}

func NewGroup(parent context.Context) *Group {
	ctx, cancel := context.WithCancel(parent)
	return &Group{
		ctx:    ctx,
		cancel: cancel,
	}
}

type Group struct {
	ctx    context.Context
	cancel func()
	wg     sync.WaitGroup
}

func (g *Group) Add(cmd Command) {
	g.wg.Add(1)
	go func() {
		_ = cmd(g.ctx)
		g.wg.Done()
	}()
}

func (g *Group) Stop() {
	g.cancel()
}

func (g *Group) Wait() {
	g.wg.Wait()
}

func (g *Group) WaitAsync() <-chan struct{} {
	ch := make(chan struct{})
	go func() {
		g.Wait()
		close(ch)
	}()
	return ch
}

func NewAtomicGroup(parent context.Context) *AtomicGroup {
	ctx, cancel := context.WithCancel(parent)
	ag := &AtomicGroup{ctx: ctx, cancel: cancel}
	ag.done = ag.onFinish
	return ag
}

// AtomicGroup terminates as soon as first goroutine terminates..
type AtomicGroup struct {
	ctx    context.Context
	cancel func()
	done   func()
	wg     sync.WaitGroup

	mu    sync.Mutex
	error error
}

// Go spawns function in a goroutine and stores results or errors.
func (d *AtomicGroup) Add(cmd Command) {
	d.wg.Add(1)
	go func() {
		defer d.done()
		err := cmd(d.ctx)
		d.mu.Lock()
		defer d.mu.Unlock()
		if err != nil {
			// do not overwrite original error by context errors
			if d.error != nil {
				return
			}
			d.error = err
			d.cancel()
			return
		}
	}()
}

// Wait for all downloaders to finish.
func (d *AtomicGroup) Wait() {
	d.wg.Wait()
	if d.Error() == nil {
		d.mu.Lock()
		defer d.mu.Unlock()
		d.cancel()
	}
}

func (d *AtomicGroup) WaitAsync() <-chan struct{} {
	ch := make(chan struct{})
	go func() {
		d.Wait()
		close(ch)
	}()
	return ch
}

// Error stores an error that was reported by any of the downloader. Should be called after Wait.
func (d *AtomicGroup) Error() error {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.error
}

func (d *AtomicGroup) Stop() {
	d.cancel()
}

func (d *AtomicGroup) onFinish() {
	d.wg.Done()
}

func NewQueuedAtomicGroup(parent context.Context, limit uint32) *QueuedAtomicGroup {
	qag := &QueuedAtomicGroup{NewAtomicGroup(parent), limit, 0, []Command{}, sync.Mutex{}}
	baseDoneFunc := qag.done // save original done function
	qag.AtomicGroup.done = func() {
		baseDoneFunc()
		qag.onFinish()
	}
	return qag
}

type QueuedAtomicGroup struct {
	*AtomicGroup
	limit       uint32
	count       uint32
	pendingCmds []Command
	mu          sync.Mutex
}

func (d *QueuedAtomicGroup) Add(cmd Command) {

	d.mu.Lock()
	if d.limit > 0 && d.count >= d.limit {
		d.pendingCmds = append(d.pendingCmds, cmd)
		d.mu.Unlock()
		return
	}

	d.mu.Unlock()
	d.run(cmd)
}

func (d *QueuedAtomicGroup) run(cmd Command) {
	d.mu.Lock()
	d.count++
	d.mu.Unlock()
	d.AtomicGroup.Add(cmd)
}

func (d *QueuedAtomicGroup) onFinish() {
	d.mu.Lock()
	d.count--

	if d.count < d.limit && len(d.pendingCmds) > 0 {
		cmd := d.pendingCmds[0]
		d.pendingCmds = d.pendingCmds[1:]
		d.mu.Unlock()
		d.run(cmd)
		return
	}

	d.mu.Unlock()
}