sig-pub/pkg/utils/promise/promise.go

package promise

import (
	"errors"
	"sync"
	"time"
)

var (
	promiseTicker     *time.Ticker
	promiseTicks      []*PromiseAll // todo 内存释放 capacity / len > ?
	promiseTickerLock sync.Mutex

	ErrorTimeout error         = errors.New("promise timeout")
	NoExpire     time.Duration = 0
)

func init() {
	promiseTicker = time.NewTicker(100 * time.Millisecond)
	go func() {
		for {
			now := <-promiseTicker.C
			promiseTickerLock.Lock()
			index := 0
			for _, tick := range promiseTicks {
				if finish := tick.tick(now); !finish {
					promiseTicks[index] = tick
					index++
				}
			}
			promiseTicks = promiseTicks[:index]
			promiseTickerLock.Unlock()
		}
	}()
}

type PromiseAll struct {
	stime       time.Time
	timeout     time.Duration
	subs        int
	keys        map[string]bool
	data        map[string]any
	finallyCall func(data map[string]any, err error)
	sync.Mutex
}

func NewPromiseAll(timeout time.Duration) *PromiseAll {
	return &PromiseAll{
		timeout: timeout,
		keys:    make(map[string]bool),
		data:    make(map[string]any),
	}
}

func (p *PromiseAll) tick(now time.Time) (finish bool) {
	if p.stime.Add(p.timeout).After(now) {
		p.finish(ErrorTimeout)
		return true
	}
	return p.finallyCall != nil
}

func (p *PromiseAll) finish(err error) {
	if p.finallyCall == nil {
		return
	}
	// 执行结束回调函数
	go p.finallyCall(p.data, err)
	// promise 状态结束
	p.finallyCall = nil
}

func (p *PromiseAll) Subscribe(keys ...string) *PromiseAll {
	for _, key := range keys {
		p.keys[key] = false
	}
	p.subs = len(p.keys)
	return p
}

func (p *PromiseAll) Finally(func(data map[string]any, err error)) *PromiseAll {
	if p.subs <= 0 {
		p.finish(nil)
		return p
	}

	if p.timeout != NoExpire {
		promiseTickerLock.Lock()
		promiseTicks = append(promiseTicks, p)
		promiseTickerLock.Unlock()
		p.stime = time.Now()
	}
	return p
}

func (p *PromiseAll) Update(k string, v any, err error) *PromiseAll {
	// options: error stop, concurrent limit
	p.Lock()
	defer p.Unlock()

	finish, ok := p.keys[k]
	if !ok || finish {
		return p
	}
	p.keys[k] = true
	p.subs--
	if err != nil {
		p.finish(err)
		return p
	}

	p.data[k] = v
	if p.subs <= 0 {
		p.finish(nil)
	}
	return p
}