Go学习笔记 - 控制协程数量

之前因为需要并发同步单张数据表的脚本，因为设置协程数量太多连接数过高导致服务宕机 … 得不偿失，所以想怎么控制一下协程数量

Channel

利用channel阻塞，上代码：

package main

import (
	"fmt"
	"sync"
	"time"
)

func main() {
	// 允许最大协程数
	ch := make(chan int, 5)
	wg := &sync.WaitGroup{}
	for i := 0; i < 10; i++ {
		wg.Add(1)
		ch <- i // ch长度等于5时，阻塞等待
		go func(i int) {
			defer wg.Done()
			fmt.Println("记录执行 ", i)
			time.Sleep(time.Second * 3)
			fmt.Println("执行完毕 ", <-ch)
		}(i)
	}
	wg.Wait()
}

原理就是利用channel通道的特性，当通道满的时候继续向通道丢数据会阻塞代码执行。

封装一下:

package main

import (
	"fmt"
	"sync"
	"time"
)

type Worker struct {
	ch chan interface{}
	wg *sync.WaitGroup
}

func main() {
	// 允许最大协程数
	worker := NewWorker(5)
	for i := 0; i < 10; i++ {
		worker.Run(i)
		go func(i int) {
			fmt.Println("记录执行 ", i)
			time.Sleep(time.Second * 3)
			key := worker.Done()
			fmt.Println("执行完毕 ", key)
		}(i)
	}
	worker.Wait()
}

func NewWorker(size int) *Worker {
	return &Worker{
		ch: make(chan interface{}, size),
		wg: &sync.WaitGroup{},
	}
}

func (w *Worker) Run(key interface{}) {
	w.ch <- key
	w.wg.Add(1)
}

func (w *Worker) Done() interface{} {
	key := <-w.ch
	w.wg.Done()
	return key
}

func (w *Worker) Wait() {
	w.wg.Wait()
}

Goroutine Pool

感觉Duck不必 … 日后再谈