Go语言限制协程并发数的两种方案详解

package golimit

type GoLimit struct {
    ch chan int
}

func NewGoLimit(max int) *GoLimit {
    return &GoLimit{ch: make(chan int, max)}
}

func (g *GoLimit) Add() {
    g.ch <- 1 // 写入数据，若满则阻塞
}

func (g *GoLimit) Done() {
    <-g.ch    // 读取数据，腾出空间
}

package main

import (
	"golimit"
	"log"
	"time"
)

func main() {
	log.Println("开始测试...")
	// 设置最大允许并发数为2
	g := golimit.NewGoLimit(2)
	
	for i := 0; i < 10; i++ {
		// 尝试增加一个协程，若已达上限将阻塞在此行
		g.Add()
		
		go func(id int) {
			defer g.Done() // 协程结束后释放信号量
			// 模拟耗时操作
			time.Sleep(time.Second * 2)
			log.Printf("Task %d done", id)
		}(i)
	}
	
	log.Println("所有协程已启动，主函数等待...")
	// 简单等待以观察效果
	time.Sleep(time.Second * 5)
	log.Println("测试结束")
}

package golimit

import (
	"sync"
)

type GoLimit struct {
	max       uint             // 最大并发数
	count     uint             // 当前并发数
	isAddLock bool             // 是否处于“禁止新增”状态
	zeroChan  chan interface{} // 为零时的广播信道
	addLock   sync.Mutex       // 控制新增请求的锁
	dataLock  sync.Mutex       // 保护计数数据的锁
}

func NewGoLimit(max uint) *GoLimit {
	return &GoLimit{
		max:     max,
		count:   0,
		isAddLock: false,
	}
}

// 增加并发数，若达到最大值则阻塞
func (g *GoLimit) Add() {
	g.addLock.Lock()
	g.dataLock.Lock()
	
	g.count += 1
	if g.count < g.max {
		// 未达到上限，解锁让其他请求进入
		g.addLock.Unlock()
	} else {
		// 已达到上限，不解锁 addLock，迫使其他 Add 调用阻塞
		g.isAddLock = true
	}
	g.dataLock.Unlock()
}

// 减少并发数，可能解除新增限制
func (g *GoLimit) Done() {
	g.dataLock.Lock()
	g.count -= 1
	
	// 如果之前加上了锁，现在数量小于最大值了，则解锁
	if g.isAddLock && g.count < g.max {
		g.isAddLock = false
		g.addLock.Unlock()
	}
	
	// 处理归零广播
	if g.count == 0 && g.zeroChan != nil {
		close(g.zeroChan)
		g.zeroChan = nil
	}
	g.dataLock.Unlock()
}

// 运行时动态更新最大并发数
func (g *GoLimit) SetMax(n uint) {
	g.dataLock.Lock()
	g.max = n
	
	// 如果之前是被限制的，且现在的新上限允许更多并发，则解锁
	if g.isAddLock && g.count < g.max {
		g.isAddLock = false
		g.addLock.Unlock()
	}
	// 如果之前没被限制，但现在达到了新的（更小的）上限，则重新上锁
	if !g.isAddLock && g.count >= g.max {
		g.isAddLock = true
		g.addLock.Lock()
	}
	g.dataLock.Unlock()
}

// 阻塞等待所有并发协程执行完毕
func (g *GoLimit) WaitZero() {
	g.dataLock.Lock()
	if g.count == 0 {
		g.dataLock.Unlock()
		return
	}
	if g.zeroChan == nil {
		g.zeroChan = make(chan interface{})
	}
	c := g.zeroChan // 复制引用以防同时close造成panic
	g.dataLock.Unlock()
	<-c // 等待信道被关闭
}

func (g *GoLimit) Count() uint {
	return g.count
}

func (g *GoLimit) Max() uint {
	return g.max
}

package main

import (
	"github.com/zh-five/golimit"
	"log"
	"time"
)

func main() {
	log.Println("开始测试...")
	g := golimit.NewGoLimit(2)
	
	for i := 0; i < 10; i++ {
		g.Add()
		go func(id int) {
			defer g.Done()
			// 此处可插入 SetMax(3) 演示运行时动态调大并发限制
			time.Sleep(time.Second * 2)
			log.Printf("Task %d completed", id)
		}(i)
	}
	
	log.Println("协程分发完毕")
	
	// 阻塞主程序直到所有子协程完成
	g.WaitZero()
	log.Println("全部测试结束")
}