kratos-utils/common/global_wait_group.go

package common

/**

import中必须有git.ikuban.com/server/kratos-utils/common

import (
	"context"
	"fmt"
	"time"

	"git.ikuban.com/server/kratos-utils/common"
)

func main() {
	ctx := context.Background()
	common.GlobalWaitGroup.SetContext(ctx)
	test()
	go test()
	common.GlobalWaitGroup.Wait()
}

func test() {
	ok := common.GlobalWaitGroup.Add(1)
	if !ok {
		return
	}
	common.GlobalWaitGroup.SetTimeout(time.Now().Unix() + 12)
	go func() {
		fmt.Println("任务开始")
		time.Sleep(time.Second * 11)
		fmt.Println("任务完成")
		common.GlobalWaitGroup.Done()
	}()
}


*/

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

var GlobalWaitGroup *globalWaitGroup

func init() {
	GlobalWaitGroup = NewGlobalWaitGroup(context.Background())
}

type globalWaitGroup struct {
	wg      sync.WaitGroup
	ctx     context.Context
	timeout int64 //超时时间， 时间戳
	isStop  bool
	lock    sync.Mutex
	ok      chan bool
}

func NewGlobalWaitGroup(ctx context.Context) *globalWaitGroup {
	this := &globalWaitGroup{
		wg:     sync.WaitGroup{},
		ctx:    ctx,
		isStop: false,
		ok:     make(chan bool, 1),
	}
	go func() {
		<-this.ctx.Done()
		this.isStop = true
	}()

	return this
}

func (this *globalWaitGroup) SetContext(ctx context.Context) {
	this.ctx, _ = context.WithCancel(ctx)
}

func (this *globalWaitGroup) Stop() {
	this.isStop = true
	this.Wait()
}

func (this *globalWaitGroup) Wait() {
	go func() {
		this.wait()
	}()
	diff := this.timeout - time.Now().Unix()
	if diff <= 0 {
		diff = 10
	}
	select {
	case <-this.ok:
		fmt.Println("没有任务而结束")
		return
	case <-time.After(time.Second * time.Duration(diff)):
		fmt.Println("超时结束")
		return
	}
}

func (this *globalWaitGroup) wait() {
	this.wg.Wait()
	this.ok <- true
}

func (this *globalWaitGroup) Done() {
	this.wg.Done()
}

func (this *globalWaitGroup) SetTimeout(timeout int64) {
	if timeout <= 0 {
		return
	}
	this.lock.Lock()
	if this.timeout < timeout {
		this.timeout = timeout
	}
	this.lock.Unlock()
}

func (this *globalWaitGroup) Add(delta int) bool {
	if this.isStop {
		fmt.Println("已结束,不能在添加任务")
		return false
	}
	this.wg.Add(delta)
	return true
}