go用chan实现WaitGroup并支持超时

来来来,话不多说,上代码

package main

import "fmt"
import "time"
import "sync"

type group struct {
	gc chan bool
	tk *time.Ticker
	cap int
	mutex sync.Mutex
}

func WaitGroup(timeOuteRec int) *group{
	timeout     := time.Millisecond * time.Duration(timeOuteRec)
	wg := group{
		gc   : make(chan bool),
		cap  :  0,
		tk   : time.NewTicker(timeout),
	}

	return &wg
}


func(w *group)Add(index int){
	w.mutex.Lock()
	w.cap++
	w.mutex.Unlock()

	go func(w *group,index int) {
		for i := 0;i<index;i++{
			fmt.Println("exec goruntine product")
			w.gc<- true
		}
	}(w,index)
}

func(w *group)Done(){
	<-w.gc
	fmt.Println("exec goruntine consumer")
	w.mutex.Lock()
	w.cap--
	w.mutex.Unlock()
}

func(w *group)Wait(){
	defer w.tk.Stop()
	for  {
		select {
		case <-w.tk.C:
			fmt.Println("time out exec over")
			return;
		default:
			w.mutex.Lock()
			if w.cap == 0 {
				fmt.Println("all goruntine exec over")
				return;
			}
			w.mutex.Unlock()
		}
	}
}

func main() {
	fmt.Println("start...")
	wg := WaitGroup(10)
	wg.Add(1)
	closer := func(wg *group) {
		wg.Done()
	}
	go closer(wg)

	wg.Wait()
	fmt.Println("exec return")
}

猜你喜欢

转载自blog.csdn.net/u012460314/article/details/107932605