自实现 Goroutine pool

自实现 Goroutine pool

// limiter.go 

package limiter

import (
	"fmt"
	"go.uber.org/zap"
	"station.grab/src/common"
	"sync"
	"time"
)

const (
	// MinimaLimit is the minimal concurrency limit
	MinimaLimit = 5
)

// Job is an interface for add jobs.
type Job interface {
	Run() (resp string, err error)
}

// EasyLimiter object
type EasyLimiter struct {
	semp chan struct{} // 控制并发的chan

	wg sync.WaitGroup // waitGroup 用于等待协程执行完成, 并关闭通道\清理资源

	jobChan    chan Job         // Job 队列(实现接口即可, 解耦了任务的具体实现)
	resultChan chan interface{} // job执行结果队列
}

func NewEasyLimiter(taskCount, limit int) *EasyLimiter {
	if limit <= MinimaLimit {
		limit = MinimaLimit
	}

	c := &EasyLimiter{
		semp:       make(chan struct{}, limit),
		resultChan: make(chan interface{}, taskCount),
		jobChan:    make(chan Job, taskCount),
	}

	// 创建后马上就监听job队列
	// job队列中有数据且semp队列未满 (满了会阻塞,以此来实现并发控制), 则取出job对象, 交给单独协程处理
	go func() {
		for job := range c.jobChan {
			//c.semp <- struct{}{}

			select {
			case c.semp <- struct{}{}:

			case <-time.After(time.Millisecond * 200):
				common.ZLogger.Info("goroutine pool full, wait for 200 mis  ", zap.Int("size", len(c.semp)))
			}

			go func(ajob Job) {
				defer func() {
					c.wg.Done()
					<-c.semp
				}()
				//common.ZLogger.Info("开始执行任务")
				result, err := ajob.Run()
				//common.ZLogger.Info("完成执行任务")
				if err != nil {
					fmt.Printf("err:%v", err)
				}
				c.resultChan <- result

			}(job)
		}
		common.ZLogger.Info("task队列关闭")
	}()

	return c
}

func (c *EasyLimiter) AddJob(job Job) {
	c.wg.Add(1)
	c.jobChan <- job
	//common.ZLogger.Info("添加任务")
}

func (c *EasyLimiter) Wait() {
	// 关闭job队列 ,此时已不会再添加
	close(c.jobChan)
	c.wg.Wait()
	// 关闭result队列,以保证range方式读取chan 程序会正常向下执行
	close(c.resultChan)
}

测试

// limiter_test.go 

package limiter

import (
	"fmt"
	"go.uber.org/zap"
	"math/rand"
	"station.grab/src/common"
	"testing"
	"time"
)

func Max(a int, b int) int { //注意参数和返回值是怎么声明的

	if a > b {
		return a
	}
	return b
}
func RandomInt(n int) int {
	rand.Seed(time.Now().UnixNano())
	v := rand.Intn(n)
	if v < 1 {
		v = 1
	}
	return v
}

type SampleJob struct {
	total int
	idx   int
	key   string
}

func (s SampleJob) Run() (resp string, err error) {

	v := fmt.Sprintf("job run: %d/%d, %v", s.idx, s.total, s.key)
	common.ZLogger.Info(v)
	//time.Sleep(time.Second*time.Duration(RandomInt(5)))
	time.Sleep(time.Second * 1)

	return v, nil
}

func TestLimiter_Execute(t *testing.T) {

	fmt.Println("begin")
	total := 20
	limiter := NewEasyLimiter(total, 5)

	for i := 0; i < total; i++ {
		limiter.AddJob(&SampleJob{
			total: total,
			idx:   i,
			key:   "test",
		})
	}

	// 控制完成并退出的信号
	chanSignal := make(chan interface{})

	go func() {
		for {
			select {
			case result, ok := <-limiter.resultChan:
				common.ZLogger.Info("read from result chan ", zap.Any("result", result))
				if !ok {
					common.ZLogger.Info("result通道关闭,退出当前goroutine")
					chanSignal<- 1
					return
				}
			}
		}
	}()

	limiter.Wait()

	<-chanSignal

	fmt.Println("done")
}

func TestRandom(t *testing.T) {
	for i := 0; i < 20; i++ {
		fmt.Println(RandomInt(5))
	}

}