status-im/status-go

package transactions

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

// TaskFunc defines the task to be run. The context is canceled when Stop is
// called to early stop scheduled task.
type TaskFunc func(ctx context.Context) (done bool)

const (
    WorkNotDone = false
    WorkDone    = true
)

// ConditionalRepeater runs a task at regular intervals until the task returns
// true. It doesn't allow running task in parallel and can be triggered early
// by call to RunUntilDone.
type ConditionalRepeater struct {
    interval time.Duration
    task     TaskFunc
    // nil if not running
    ctx      context.Context
    ctxMu    sync.Mutex
    cancel   context.CancelFunc
    runNowCh chan bool
    runNowMu sync.Mutex
}

func NewConditionalRepeater(interval time.Duration, task TaskFunc) *ConditionalRepeater {
    return &ConditionalRepeater{
        interval: interval,
        task:     task,
        runNowCh: make(chan bool, 1),
    }
}

// RunUntilDone starts the task immediately and continues to run it at the defined
// interval until the task returns true. Can be called multiple times but it
// does not allow multiple concurrent executions of the task.
func (t *ConditionalRepeater) RunUntilDone() {
    t.ctxMu.Lock()
    defer func() {
        t.runNowMu.Lock()
        if len(t.runNowCh) == 0 {
            t.runNowCh <- true
        }
        t.runNowMu.Unlock()
        t.ctxMu.Unlock()
    }()

    if t.ctx != nil {
        return
    }
    t.ctx, t.cancel = context.WithCancel(context.Background())

    go func() {
        defer func() {
            t.ctxMu.Lock()
            defer t.ctxMu.Unlock()
            t.cancel()
            t.ctx = nil
        }()

        ticker := time.NewTicker(t.interval)
        defer ticker.Stop()

        for {
            select {
            // Stop was called or task returned true
            case <-t.ctx.Done():
                return
            // Scheduled execution
            case <-ticker.C:
                if t.task(t.ctx) {
                    return
                }
            // Start right away if requested
            case <-t.runNowCh:
                ticker.Reset(t.interval)
                if t.task(t.ctx) {
                    t.runNowMu.Lock()
                    if len(t.runNowCh) == 0 {
                        t.runNowMu.Unlock()
                        return
                    }
                    t.runNowMu.Unlock()
                }
            }
        }
    }()
}

// Stop forcefully stops the running task by canceling its context.
func (t *ConditionalRepeater) Stop() {
    t.ctxMu.Lock()
    defer t.ctxMu.Unlock()
    if t.ctx != nil {
        t.cancel()
    }
}

func (t *ConditionalRepeater) IsRunning() bool {
    t.ctxMu.Lock()
    defer t.ctxMu.Unlock()
    return t.ctx != nil
}