golang实现简单的goroutine worker池

客官°小女子只卖身不卖艺 2024-04-18 11:05 33阅读 0赞

实现了简单的worker池，可以执行一些后台常驻类的任务，并且可以实现worker数量的实时控制。

利用两个worker数组，alive和idle，分别保存活跃的worker和空闲的worker，并且空闲的worker可以从暂停状态快速转变为活跃状态。由于每次在增加worker数量的时候没有创建goroutine的情况，所以效率会稍微高一些。

type taskFunc func() error
    
    type workerManager struct {
    	alive     []*worker
    	idle      []*worker
    	lck       *sync.Mutex
    	count     int
    	NewWorker func(int) *worker
    	MaxAlive  int
    	MaxIdle   int
    }
    
    func (m *workerManager) SetWorkerNum(num int) (err error) {
    	if num > m.MaxAlive {
    		err = errors.New("worker num can not bigger than MaxAlive")
    		return
    	}
    	m.lck.Lock()
    	defer m.lck.Unlock()
    	if len(m.alive) == num {
    		return
    	}
    	//新增活跃的worker
    	if len(m.alive) < num {
    		need := num - len(m.alive)
    		for ; need > 0 && len(m.idle) > 0; need-- {
    			m.setAlive(m.idle[len(m.idle)-1])
    			m.delIdle(1)
    		}
    		for ; need > 0; need-- {
    			m.count++
    			w := m.NewWorker(m.count)
    			m.setAlive(w)
    		}
    		return
    	}
    	//减少活跃的worker
    	del := len(m.alive) - num
    	if del > len(m.alive) {
    		del = len(m.alive)
    	}
    	for ; del > 0; del-- {
    		if len(m.idle) < m.MaxIdle {
    			m.setIdle(m.alive[len(m.alive)-1])
    		} else {
    			m.destroyWorker(m.alive[len(m.alive)-1])
    		}
    		m.delAlive(1)
    	}
    
    	return
    }
    
    func (m *workerManager) destroyWorker(w *worker) {
    	if w != nil {
    		w.destroy <- struct{}{}
    	}
    }
    
    func (m *workerManager) delAlive(num int) {
    	rem := len(m.alive) - num
    	m.alive = m.alive[:rem]
    }
    
    func (m *workerManager) setAlive(w *worker) {
    	w.Start()
    	m.alive = append(m.alive, w)
    }
    
    func (m *workerManager) delIdle(num int) {
    	rem := len(m.idle) - num
    	m.idle = m.idle[:rem]
    }
    
    func (m *workerManager) setIdle(w *worker) {
    	w.Pause()
    	m.idle = append(m.idle, w)
    }
    
    type worker struct {
    	name    int
    	start   chan struct{}
    	pause   chan struct{}
    	destroy chan struct{}
    	Task    taskFunc
    	once    sync.Once
    }
    
    func (w *worker) Pause() {
    	w.pause <- struct{}{}
    }
    
    func NewWorker(symbol int) *worker {
    	w := new(worker)
    	w.name = symbol
    	w.start = make(chan struct{})
    	w.pause = make(chan struct{})
    	w.destroy = make(chan struct{})
    	return w
    }
    
    func (w *worker) Start() {
    	w.once.Do(func() {
    		go func() {
    			<-w.start
    			fmt.Println("start", w.name)
    			for {
    				select {
    				case <-w.pause:
    					<-w.start
    				case <-w.start:
    				case <-w.destroy:
    					goto END
    				default:
    				}
    				//...........task
    				fmt.Println("i am ", w.name)
    				time.Sleep(time.Second)
    			}
    		END:
    			fmt.Println("i am finish", w.name)
    		}()
    	})
    	w.start <- struct{}{}
    }
    
    func TestWorker(t *testing.T) {
    	manager := &workerManager{
    		lck:       new(sync.Mutex),
    		NewWorker: NewWorker,
    		MaxIdle:   2,
    		MaxAlive:  5,
    	}
    	manager.SetWorkerNum(3)
    	fmt.Println("alive len:", len(manager.alive), "idle len", len(manager.idle))
    	//fmt.Println("goroutine num:", runtime.NumGoroutine())
    	time.Sleep(time.Second * 5)
    	manager.SetWorkerNum(4)
    	fmt.Println("alive len:", len(manager.alive), "idle len", len(manager.idle))
    	//fmt.Println("goroutine num:", runtime.NumGoroutine())
    	time.Sleep(time.Second * 5)
    	manager.SetWorkerNum(1)
    	fmt.Println("alive len:", len(manager.alive), "idle len", len(manager.idle))
    	//fmt.Println("goroutine num:", runtime.NumGoroutine())
    	time.Sleep(time.Second * 5)
    	manager.SetWorkerNum(4)
    	fmt.Println("alive len:", len(manager.alive), "idle len", len(manager.idle))
    	//fmt.Println("goroutine num:", runtime.NumGoroutine())
    	time.Sleep(time.Second * 5)
    
    	manager.SetWorkerNum(5)
    	fmt.Println("alive len:", len(manager.alive), "idle len", len(manager.idle))
    	//fmt.Println("goroutine num:", runtime.NumGoroutine())
    	time.Sleep(time.Second * 5)
    
    	manager.SetWorkerNum(1)
    	fmt.Println("alive len:", len(manager.alive), "idle len", len(manager.idle))
    	//fmt.Println("goroutine num:", runtime.NumGoroutine())
    	time.Sleep(time.Second * 5)
    }