GO 扩展库： semaphore 实现原理与使用

Go的semaphore包提供了一个加权信号量的实现，用于控制对资源的并发访问，确保总权重不超过预设值。以下从源码出发解析其实现原理和并发控制机制：

---

1. 核心结构

semaphore中主要定义了如下两个结构体：

type Weighted struct {
	size    int64
	cur     int64
	mu      sync.Mutex
	waiters list.List
}

type waiter struct {
	n     int64
	ready chan<- struct{} // Closed when semaphore acquired.
}

1. Weighted结构体：

   • size：信号量总容量（最大权重和）。
   • cur：当前已分配的权重总和。
   • mu：互斥锁，保护对cur和waiters的并发访问。
   • waiters：双向链表（由container/list包实现），用于存储等待队列中的请求。

2. waiter结构体：

   • n：请求的权重值。
   • ready：只写通道，当请求成功获取信号量时被关闭，通知等待的协程。

---

2. 关键方法解析

1. Acquire：获取信号量

以下是Acquire的示例代码：

func (s *Weighted) Acquire(ctx context.Context, n int64) error {
	done := ctx.Done()
	s.mu.Lock()
	select {
	case <-done:
		s.mu.Unlock()
		return ctx.Err()
	default:
	}
	if s.size-s.cur >= n && s.waiters.Len() == 0 {
		s.cur += n
		s.mu.Unlock()
		return nil
	}

	if n > s.size {
		s.mu.Unlock()
		<-done
		return ctx.Err()
	}

	ready := make(chan struct{})
	w := waiter{n: n, ready: ready}
	elem := s.waiters.PushBack(w)
	s.mu.Unlock()

	select {
	case <-done:
		s.mu.Lock()
		select {
		case <-ready:
			s.cur -= n
			s.notifyWaiters()
		default:
			isFront := s.waiters.Front() == elem
			s.waiters.Remove(elem)
			if isFront && s.size > s.cur {
				s.notifyWaiters()
			}
		}
		s.mu.Unlock()
		return ctx.Err()

	case <-ready:
		select {
		case <-done:
			s.Release(n)
			return ctx.Err()
		default:
		}
		return nil
	}
}

• 流程： 从源码可知，Acquire方法主要包含以下步骤：

  1. 检查上下文是否已取消：若取消则直接返回错误。
  2. 尝试直接获取 ：
     • 若当前剩余资源足够（size - cur >= n）且无等待者，立即分配资源（增加cur）并返回。
  3. 处理超限请求 ：若请求的n超过总容量size，直接等待上下文取消（注定失败）。
  4. 加入等待队列 ：
     • 创建waiter并加入队列尾部。
     • 等待两个事件：ready通道关闭（资源就绪）或上下文取消。
  5. 处理取消或完成 ：
     • 若上下文取消，从队列中移除请求。若该请求是队首且释放后有余量，触发唤醒后续等待者。
     • 若成功获取，再次检查上下文是否取消（避免在获取后但未返回时被取消）。

• 并发控制：

  • 通过cur累加请求的权重，确保总和不超过size。
  • 无可用资源时，协程阻塞并加入队列，等待资源释放。

2. 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()
}

• 流程：

  1. 减少cur（释放资源）。
  2. 调用notifyWaiters，按顺序唤醒队列中可满足的请求。

• notifyWaiters逻辑：

  • 遍历等待队列，从队首开始检查每个请求。
  • 若当前资源足够（size - cur >= waiter.n）：
    • 分配资源（cur += waiter.n）。
    • 从队列中移除该请求，关闭其ready通道以唤醒协程。
  • 若资源不足，停止处理后续请求（防止小请求插队导致大请求饥饿）。

3. TryAcquire：非阻塞获取

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
}

• 直接检查是否有足够资源且无等待者，若有则立即分配，否则返回失败。

---

3. 并发控制机制

1. 权重累加与检查：

   • 通过cur记录已分配的权重总和，确保cur + n <= size时才允许分配。
   • 若资源不足，协程进入等待队列，按FIFO顺序唤醒，保证公平性。

2. 等待队列管理：

   • 新请求加入队列尾部，唤醒时从队首开始处理。
   • 严格按顺序尝试分配资源，避免大请求被小请求"饿死"。

3. 上下文感知：

   • 在等待过程中监听上下文取消事件，及时清理队列中的无效请求。
   • 若在等待期间被取消，释放可能已分配的资源并触发后续唤醒。

---

4. 示例场景

下面是一段通过semaphore做并发控制的示例：

func main() {
	// 创建信号量，允许最多同时运行 3 个 Goroutines
	sem := semaphore.NewWeighted(3)

	var wg sync.WaitGroup

	// 启动 10 个 Goroutines
	for i := 0; i < 10; i++ {
		wg.Add(1)

		go func(id int) {
			defer wg.Done()

			// 请求信号量
			if err := sem.Acquire(context.Background(), 1); err != nil {
				fmt.Printf("Goroutine %d failed to acquire semaphore: %v\n", id, err)
				return
			}

			// 模拟任务处理
			fmt.Printf("Goroutine %d is running\n", id)
			time.Sleep(2 * time.Second)

			// 释放信号量
			sem.Release(1)
			fmt.Printf("Goroutine %d is done\n", id)
		}(i)
	}

	// 等待所有 Goroutines 完成
	wg.Wait()
}

---

5. 总结

semaphore实现通过维护当前权重值cur和等待队列，结合互斥锁和通道同步，确保资源分配的权重总和不超过预设值。等待队列的FIFO处理保证公平性，避免饥饿问题，同时通过上下文集成支持超时和取消操作。