123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960 |
- package entity
- import (
- "context"
- "sync"
- "time"
- )
- //并发限制
- type reqLimit struct {
- doPool chan struct{}
- waitPool chan struct{}
- }
- var (
- ReqLimitInit *reqLimit
- ReqLimitLock = &sync.Mutex{}
- )
- func NewLimit(num int) {
- doPool := make(chan struct{}, num)
- for i := 0; i < num; i++ {
- doPool <- struct{}{}
- }
- waitPool := make(chan struct{}, num)
- for i := 0; i < num; i++ {
- waitPool <- struct{}{}
- }
- ReqLimitInit = &reqLimit{
- doPool: doPool,
- waitPool: waitPool,
- }
- }
- // -2 等待池已满
- // -1 超时
- // 1:可以执行查询
- func (r *reqLimit) Limit(ctx context.Context) int {
- ctx, cancel := context.WithTimeout(ctx, 30*time.Second) //30秒
- defer cancel()
- select {
- case <-r.waitPool:
- defer func() {
- r.waitPool <- struct{}{}
- }()
- select {
- case <-r.doPool:
- return 1
- case <-ctx.Done(): //超时
- return -1
- }
- default:
- return -2
- }
- }
- func (r *reqLimit) Release() {
- r.doPool <- struct{}{}
- }
|