新增全局waitgroup

common/global_wait_group.go

@@ -0,0 +1,125 @@
+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 = ctx
+}
+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
+}