并发控制

package main
// demo 参考地址https://studygolang.com/articles/25950

import (
	"github.com/siddontang/go-log/log"
	"math/rand"
	"sync"
	"time"
)

func task() int {
	//取[n,m]的数 语法入下:
	rand.Seed(time.Now().UnixNano())

	//num := rand.Intn(m-n+1) + n
	num := rand.Intn(5-1+1) + 1
	return num
}


// WaitGroup 一个异步结构体
type WaitGroup struct {
	workChan chan int
	wg       sync.WaitGroup
}

// NewPool 生成一个工作池, coreNum 限制
func NewPool(coreNum int) *WaitGroup {
	ch := make(chan int, coreNum)
	return &WaitGroup{
		workChan: ch,
		wg:       sync.WaitGroup{},
	}
}

// Add 添加
func (ap *WaitGroup) Add(num int) {
	// 不向channel里面写大量数据，channel中传递的都是数据的拷贝，可能会影响性能
	// channel只控制并发的数量
	ap.workChan <- num
	ap.wg.Add(1)
}

// Done 完结
func (ap *WaitGroup) Done() {
	select {
	case <-ap.workChan:
		ap.wg.Done()
	}
}

// Wait 等待
func (ap *WaitGroup) Wait() {
	ap.wg.Wait()
}

func testFunc(i int, wg *WaitGroup) {
	defer wg.Done()
	t := task()
	log.Infof("Index %d, duration %d, begin!", i, t)
	time.Sleep(time.Duration(t) * time.Second)
	log.Infof("Index %d, duration %d, finish!", i, t)
}

func main() {
	work := NewPool(4)

	for i := 0; i < 20; i++ {
		work.Add(1)
		go testFunc(i, work)
	}

	log.Info("waiting...")
	work.Wait()
	log.Info("done")
}

安全退出

func stopSignal(pid int, r *sql.Db) {
	sigs := make(chan os.Signal, 1)
	log.Errorf("register signal notify, pid: %d", pid)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
	for {
		s := <-sigs
		log.Error("recv signal: ", s)
		switch s {
		case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT:
			// close kafka
			r.Close()

			log.Errorf("pid: %d exit", pid)
			os.Exit(0)
		}
	}
}

func main() {
    stopSignal(os.Getpid(), r)
}