golang限流

1.channel限流

package main

import (
    "errors"
    "fmt"
    "time"
)

type (
    Limiter struct {
        maxCount int
        tokens   chan *Token
    }

    Token struct {
    }
)

var (
    ErrInvalidCount = errors.New("invalid count settings")
    ErrInvalidToken = errors.New("invalid token")
)

func NewLimiter(maxCount int) (*Limiter, error) {
    if maxCount < 0 {
        return nil, ErrInvalidCount
    }
    limiter := &Limiter{
        tokens: make(chan *Token, maxCount),
    }
    for i := 0; i < maxCount; i++ {
        tokenObj := &Token{}
        limiter.tokens <- tokenObj
    }
    return limiter, nil
}

func (limiter *Limiter) Get() *Token {
    return <-limiter.tokens
}

func (limiter *Limiter) Put(tokenObj *Token) error {
    if tokenObj == nil {
        return ErrInvalidToken
    }
    limiter.tokens <- tokenObj
    return nil
}

func (limiter *Limiter) Len() int {
    return len(limiter.tokens)
}

//以下是測試代碼
func main() {
    limiter, _ := NewLimiter(1000)
    input := make([]int, 0)
    for index := 0; index < 10000; index++ {
        input = append(input, index)
    }

    timeout := 2
    chs := make([]chan string, len(input))
    startTime := time.Now()
    fmt.Println("Multirun start")
    for i, _ := range input {
        chs[i] = make(chan string, 1)
        token := limiter.Get()
        go func(token *Token, ch chan string, taskId, sleepTime, timeout int) {
            defer limiter.Put(token)
            time.Sleep(time.Duration(sleepTime) * time.Second)
            ch <- fmt.Sprintf("task id %d , sleep %d second", taskId, sleepTime)
        }(token, chs[i], i, 1, timeout)
    }
    for _, ch := range chs {
        fmt.Println(<-ch)
    }
    endTime := time.Now()
    fmt.Printf("Multissh finished. Process time %s. Number of task is %d", endTime.Sub(startTime), len(input))

}