go并发控制

Posted 2021-09-03 luoguo

Get到一个并发数量控制的办法

package main
import (
    "fmt"
    "sync"
    "time"
)

type GoLimit struct {
    c  chan struct{}
    wg *sync.WaitGroup
}

func NewGoLimit(maxSize int) *GoLimit {
    return &GoLimit{
        c:  make(chan struct{}, maxSize),
        wg: new(sync.WaitGroup),
    }
}

func (g *GoLimit) Add(delta int) {
    g.wg.Add(delta)
    for i := 0; i < delta; i++ {
        g.c <- struct{}{}
    }
}

func (g *GoLimit) Done() {
    <-g.c
    g.wg.Done()
}

func (s *GoLimit) Wait() {
    s.wg.Wait()
}

var golimit = NewGoLimit(50)

func main() {
    userCount := 100
    for i := 0; i < userCount; i++ {
        go Read(i)
    }

    golimit.Wait()
}

func Read(i int) {
    defer golimit.Done()
    golimit.Add(1)

    fmt.Printf("go func: %d, time: %d\\n", i, time.Now().Unix())
    time.Sleep(time.Second)
}