一个简单的goroutine pool

开发十年，就只剩下这套Java开发体系了 >>>   

package main

import (
	"fmt"
	"runtime"
	"strings"
	"sync"
)

type GoPool struct {
	Queue    chan func() (interface{}, error)
	PoolSize int //goroutine数量

	finishCallback func(interface{}, error)
}

//初始化
func (self *GoPool) Init(pool_size, task_size int) {
	self.Queue = make(chan func() (interface{}, error), task_size)
	self.PoolSize = pool_size
}

//start
func (self *GoPool) Start() {
	//开启PoolSize个goroutine
	for i := 0; i < self.PoolSize; i++ {
		go func() {
			for {
				task, ok := <-self.Queue
				if !ok {
					break
				}
				self.finishCallback(task())
			}
		}()
	}
}

//stop
func (self *GoPool) Stop() {
	close(self.Queue)
}

//添加任务
func (self *GoPool) AddTask(task func() (interface{}, error)) {
	self.Queue <- task
}

//设置callback
func (self *GoPool) SetFinishCallback(callback func(interface{}, error)) {
	self.finishCallback = callback
}

func main() {
	//打印CPU核数
	fmt.Println("cpu nums:", runtime.NumCPU())
	//设置最多使用CPU数量
	runtime.GOMAXPROCS(runtime.NumCPU())

	arr := []string{
		"ab/cd/ef",
		"11/22/33",
		"rr/qq/33",
		"vv/hh/mm",
	}

	pool := new(GoPool)
	pool.Init(3, len(arr))

	//使用WaitGroup控制同步
	var wg sync.WaitGroup
	for i := range arr {
		wg.Add(1)
		str := arr[i]
		pool.AddTask(func() (interface{}, error) {
			return split(str)
		})
	}

	//设置callback
	pool.SetFinishCallback(func(res interface{}, err error) {
		fmt.Println("res:", res)
		if err != nil {
			fmt.Println("err:", err)
		}
		wg.Done()
	})

	pool.Start()

	wg.Wait()

	pool.Stop()

	fmt.Println("finished...")
}

func split(str string) ([]string, error) {
	return strings.Split(str, "/"), nil
}