semaphore包提供了带权重的信号量的实现。常用来限制最大并发数。

使用场景

我们使用semaphore.Weighted来限制统一时刻最多有n个线程可以同时工作，当maxWorkers个线程处于运行状态时，Acquire会被阻塞，直到其中的一个工作线程执行完成。

var (
   maxWorkers = runtime.GOMAXPROCS(0)
   sem        = semaphore.NewWeighted(int64(maxWorkers))
   out        = make([]int, 32)
)

func semaphoreDemo() {
   ctx := context.TODO()
   for i := range out {
      // When maxWorkers goroutines are in flight, Acquire blocks until one of the
      // workers finishes.
      if err := sem.Acquire(ctx, 1); err != nil {
         log.Printf("Failed to acquire semaphore: %v", err)
         break
      }

      go func(i int) {
         defer sem.Release(1)
         out[i] = collatzSteps(i + 1)
      }(i)
   }
}

源码解析

type Weighted struct {
	size    int64        // 资源上限
	cur     int64        // 计数器，存储当前已使用资源
	mu      sync.Mutex   // 用来控制并发
	waiters list.List    // 存储等待获取资源的goroutine
}

NewWeighted

NewWeighted用来创建新的信号量，通过参数n来指定初始值size。

func NewWeighted(n int64) *Weighted {
	w := &Weighted{size: n}
	return w
}

Acquire

Acquire用来获取指定权重n的资源。为了避免长时间获取不到资源导致线程阻塞，semaphore引入了上下文，可以为信号量的获取设置超时时间。

如果当前可用资源充足且没有在等待资源分配的goroutine，直接分配资源并返回；
如果获取的资源权重n>Weighted.size，说明永远不可能满足该goutine的需求，直接返回，
现存资源不够，需要将当前goroutine加入到等待队列中，阻塞当前goroutine，直到资源可用或者ctx执行完成。

func (s *Weighted) Acquire(ctx context.Context, n int64) error {
   s.mu.Lock()
   // 可用资源充足且没有在等待的goroutine
   if (s.size-s.cur) >= n && s.waiters.Len() == 0 {
      s.cur = s.cur + n
      s.mu.Unlock()
      return nil
   }
   // 永远不可能满足
   if s.size < n {
      s.mu.Unlock()
      <-ctx.Done()
      return ctx.Err()
   }
    
   // 创建一个channel等待唤醒信号
   ready := make(chan struct{})
   w := waiter{n: n, ready: ready}
   elem := s.waiters.PushBack(w)
   s.mu.Unlock()

   select {
   case <-ctx.Done():
      err := ctx.Err()
      s.mu.Lock()
      select {
      case <-ready:
         err = nil
      default:
         // 被唤醒后如果还有剩余资源，通知其他waiters
         isFront := s.waiters.Front() == elem
         s.waiters.Remove(elem)
         if isFront && s.size > s.cur {
            s.notifyWaiters()
         }
      }
      s.mu.Unlock()
      return err

   case <-ready:
      return nil
   }
}

semaphore.Weighted采用FIFO的方式对资源进行分配，如果可分配资源不满足当前队列头部等待的gouroutine所需，则会继续阻塞，直到资源满足。这样是为了避免资源需求较大的goroutine被饿死。

func (s *Weighted) notifyWaiters() {
   for {
      next := s.waiters.Front()
      if next == nil {
         break // No more waiters blocked.
      }

      w := next.Value.(waiter)
      if s.size-s.cur < w.n {
         break
      }

      s.cur += w.n
      s.waiters.Remove(next)
      close(w.ready)
   }
}

Release

释放占用资源。

func (s *Weighted) Release(n int64) {
   s.mu.Lock()
   s.cur -= n
   if s.cur < 0 {
      s.mu.Unlock()
      panic("semaphore: released more than held")
   }
   s.notifyWaiters()
   s.mu.Unlock()
}

TryAcquire

使用TryAcquire可以非阻塞地获取资源，满足条件就返回true，不满足直接返回false，不会导致线程阻塞。

func (s *Weighted) TryAcquire(n int64) bool {
   s.mu.Lock()
   success := s.size-s.cur >= n && s.waiters.Len() == 0
   if success {
      s.cur += n
   }
   s.mu.Unlock()
   return success
}

参考资料

draveness.me/golang/docs…

pkg.go.dev/golang.org/…

go并发编程-Semaphore