在 go 中实现非阻塞的 channel

发表于 | 阅读次数

今天在写用于限流的漏桶算法，使用了 channel 来模拟“桶”。大致方法是这样：

  1. 定时从 channel 中读出数据（模拟桶漏水）
  2. 收到请求时如果 channel 没有满，就写入
  3. 收到请求时如果 channel 满了，就抛弃请求并返回“超出 QPS 限制”（不要阻塞）

众所周知，channel 是阻塞的，但有办法模拟出第 3 点的要求，代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
var bucket = make(chan struct{})

func main() {
    go leak()
    for i := 0; i < 10; i++ {
        time.Sleep(500 * time.Millisecond)
        select {
        case bucket <- struct{}{}:
            fmt.Println("add data to bucket")
        default: // 注意这里
            fmt.Println("bucket is full")
        }
    }
}

func leak() {
    tick := time.Tick(time.Second)
    for range tick {
        <-bucket
    }
}

运行结果如下（或点击这里修改运行程序）

1
2
3
4
5
6
7
8
9
10
bucket is full
bucket is full
add data to bucket
bucket is full
add data to bucket
bucket is full
add data to bucket
bucket is full
add data to bucket
bucket is full

关键在于 select 中的 default，当其他分支阻塞时会执行 default 分支，这样就实现了“非阻塞 channel”的需求。