status-im/status-go

package async

import (
    "context"
    "sync"
    "time"

    "github.com/ethereum/go-ethereum/log"
)

type Command func(context.Context) error

type Commander interface {
    Command(inteval ...time.Duration) Command
}

type Runner interface {
    Run(context.Context) error
}

// 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 with error.
type AtomicGroup struct {
    ctx    context.Context
    cancel func()
    done   func()
    wg     sync.WaitGroup

    mu    sync.Mutex
    error error
}

type AtomicGroupKey string

func (d *AtomicGroup) SetName(name string) {
    d.ctx = context.WithValue(d.ctx, AtomicGroupKey("name"), name)
}

func (d *AtomicGroup) Name() string {
    val := d.ctx.Value(AtomicGroupKey("name"))
    if val != nil {
        return val.(string)
    }
    return ""
}

// 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 {
                log.Info("async.Command failed", "error", err, "d.error", d.error, "group", d.Name())
                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()
}

func NewErrorCounter(maxErrors int, msg string) *ErrorCounter {
    return &ErrorCounter{maxErrors: maxErrors, msg: msg}
}

type ErrorCounter struct {
    cnt       int
    maxErrors int
    err       error
    msg       string
}

// Returns false in case of counter overflow
func (ec *ErrorCounter) SetError(err error) bool {
    log.Debug("ErrorCounter setError", "msg", ec.msg, "err", err, "cnt", ec.cnt)

    ec.cnt++

    // do not overwrite the first error
    if ec.err == nil {
        ec.err = err
    }

    if ec.cnt >= ec.maxErrors {
        log.Error("ErrorCounter overflow", "msg", ec.msg)
        return false
    }

    return true
}

func (ec *ErrorCounter) Error() error {
    return ec.err
}

func (ec *ErrorCounter) MaxErrors() int {
    return ec.maxErrors
}

type FiniteCommandWithErrorCounter struct {
    FiniteCommand
    *ErrorCounter
}

func (c FiniteCommandWithErrorCounter) Run(ctx context.Context) error {
    f := func(ctx context.Context) (quit bool, err error) {
        err = c.Runable(ctx)
        if err == nil {
            return true, err
        }

        if c.ErrorCounter.SetError(err) {
            return false, err
        }
        return true, err
    }

    quit, err := f(ctx)
    if quit {
        return err
    }

    ticker := time.NewTicker(c.Interval)
    for {
        select {
        case <-ctx.Done():
            return ctx.Err()
        case <-ticker.C:
            quit, err := f(ctx)
            if quit {
                return err
            }
        }
    }
}