Golang 如何控制固定线程数

在 golang 里面大家应该都知道是 go 关键字调用一个函数,就开启了个线程。 如果一个循环,每个里面都调用 go 函数,那就会可能启动很多个线程, 如果想控制固定线程数来处理业务,该怎么写呢?

写过 golang 多线程的应该都知道,多线程之间可以通过管道通信协作,作为通信的数据。 我们可以通过 sync.WaitGroup 和管道配合来达到控制多线程的线程数。

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
package main

import (
  "fmt"
  "sync"
  "time"
)

var MAX_THREADS = 5

func main() {
  jobs := make(chan int, 10)

  var wg sync.WaitGroup
    // 设置需要多少个线程阻塞
  wg.Add(MAX_THREADS)
    fmt.Println(time.Now())

    // 根据线程数控制启动多少个消费者线程
  for n := 0; n < MAX_THREADS; n++ {
      go worker(jobs, &wg)
  }

    // 生产者
  for i := 0; i < 10; i++ {
      jobs <- i
  }
    close(jobs)

    // 等待所有线程执行完毕的阻塞方法
  wg.Wait()
  fmt.Println(time.Now())
}
// 消费者
func worker(jobs <-chan int, wg *sync.WaitGroup) {
  for job := range jobs {
      fmt.Println(job)
      time.Sleep(1 * time.Second)
    }
    // 消费完毕则调用 Done,减少需要等待的线程
  wg.Done()
}

上面这个代码控制了 5 (MAX_THREADS) 个线程,你可以修改成 2 和 10 看下效果。

sync.WaitGroup 是一个线程控制组, 它能够一直等到所有的goroutine执行完成,并且阻塞主线程的执行,直到所有的goroutine 执行完成。

它一共有 3 个方法,Add(),Done(),Wait()。

Add - 设置需要阻塞的线程数 Done - 代表单线程已经执行完毕 Wait - 等待所有线程执行完毕的阻塞方法

通过它就可以控制多少个消费者同时在处理了,而不是生成无限多个消费者。

评论