深入理解 go Mutex

var sum int // 和
var mu sync.Mutex // 互斥锁
// add 将 sum 加 1
func add() {
    // 获取锁，只能有一个协程获取到锁，
    // 其他协程需要阻塞等待锁释放才能获取到锁。
  mu.Lock()
  // 临界区域
  sum++
  mu.Unlock()
}
func TestMutex(t *testing.T) {
  // 启动 1000 个协程
  var wg sync.WaitGroup
  wg.Add(1000)
  for i := 0; i < 1000; i++ {
    go func() {
      // 每个协程里面调用 add()
      add()
      wg.Done()
    }()
  }
  // 等待所有协程执行完毕
  wg.Wait()
  // 最终 sum 的值应该是 1000
  assert.Equal(t, 1000, sum)
}

// Set is used to store a new key/value pair exclusively for this context.
// It also lazy initializes  c.Keys if it was not used previously.
func (c *Context) Set(key string, value any) {
    // 获取锁
  c.mu.Lock()
    // 如果 Keys 还没初始化，则进行初始化
  if c.Keys == nil {
    c.Keys = make(map[string]any)
  }
    // 设置键值对
  c.Keys[key] = value
    // 释放锁
  c.mu.Unlock()
}

// Get returns the value for the given key, ie: (value, true).
// If the value does not exist it returns (nil, false)
func (c *Context) Get(key string) (value any, exists bool) {
    // 获取锁
  c.mu.RLock()
    // 读取 key
  value, exists = c.Keys[key]
    // 释放锁
  c.mu.RUnlock()
  return
}

var (
  // 保护 allPools 和 oldPools 的互斥锁。
  allPoolsMu Mutex
  // allPools is the set of pools that have non-empty primary
  // caches. Protected by either 1) allPoolsMu and pinning or 2)
  // STW.
  allPools []*Pool
  // oldPools is the set of pools that may have non-empty victim
  // caches. Protected by STW.
  oldPools []*Pool
)

func (p *Pool) pinSlow() (*poolLocal, int) {
  // Retry under the mutex.
  // Can not lock the mutex while pinned.
  runtime_procUnpin()
  allPoolsMu.Lock() // 获取锁
  defer allPoolsMu.Unlock() // 函数返回的时候释放锁
  pid := runtime_procPin()
  // poolCleanup won't be called while we are pinned.
  s := p.localSize
  l := p.local
  if uintptr(pid) < s {
    return indexLocal(l, pid), pid
  }
  if p.local == nil {
    allPools = append(allPools, p) // 全局变量修改
  }
  // If GOMAXPROCS changes between GCs, we re-allocate the array and lose the old one.
  size := runtime.GOMAXPROCS(0)
  local := make([]poolLocal, size)
  atomic.StorePointer(&p.local, unsafe.Pointer(&local[0])) // store-release
  runtime_StoreReluintptr(&p.localSize, uintptr(size))     // store-release
  return &local[pid], pid
}

// Take takes the first element out of q if not empty.
func (q *Queue) Take() (any, bool) {
  // 获取互斥锁（只能有一个协程获取到锁）
  q.lock.Lock()
  // 函数返回的时候释放互斥锁（获取到锁的协程释放锁之后，其他协程才能进行抢占锁）
  defer q.lock.Unlock()
  // 下面的代码只有抢占到（也就是互斥锁锁定的代码块）
  if q.count == 0 {
    return nil, false
  }
  element := q.elements[q.head]
  q.head = (q.head + 1) % len(q.elements)
  q.count--
  return element, true
}

import (
  "fmt"
  "sync"
  "testing"
)
var mu sync.Mutex
var sum int
// 在锁的保护下对 sum 做读写操作
func test() {
  mu.Lock()
  sum++
  mu.Unlock()
}
func TestMutex(t *testing.T) {
  var wg sync.WaitGroup
  wg.Add(1000)
  for i := 0; i < 500; i++ {
    go func() {
      test()
      wg.Done()
    }()
    // 位于临界区外，也依然是可以对 sum 做读写操作的。
    sum++
  }
  wg.Wait()
  fmt.Println(sum)
}

type Mutex struct {
  state int32 // 状态字段
  sema  uint32 // 信号量
}

const (
  mutexLocked      = 1 << iota // 表示有 goroutine 拥有锁
  mutexWoken                   // 唤醒（就是第 2 位）
  mutexStarving                // 饥饿（第 3 位）
  mutexWaiterShift = iota      // 表示第 4 位开始，表示等待者的数量
  starvationThresholdNs = 1e6  // 1ms 进入饥饿模式的等待时间阈值
)

new |= mutexLocked

atomic.CompareAndSwapInt32(&m.state, old, new)

// Lock 获取锁。
// 如果锁已在使用中，则调用 goroutine 将阻塞，直到互斥量可用。
func (m *Mutex) Lock() {
  // Fast path: grab unlocked mutex.
  // 上锁成功则直接返回
  if atomic.CompareAndSwapInt32(&m.state, 0, mutexLocked) {
    return
  }
  // Slow path (outlined so that the fast path can be inlined)
  // 没有上锁成功，这个时候需要做的事情就有点多了。
  m.lockSlow()
}

// 获取 mutex 锁
func (m *Mutex) lockSlow() {
  var waitStartTime int64 // 当前协程开始等待的时间
  starving := false       // 当前协程是否是饥饿模式
  awoke := false          // 唤醒标志（是否当前协程就是被唤醒的协程）
  iter := 0               // 自旋次数（超过一定次数如果还没能获得锁，就进入等待）
  old := m.state          // 旧的状态，每次 for 循环会重新获取当前的状态字段
  for {
    // 自旋：目的是让正在运行中的 goroutine 尽快获取到锁。
    // 两种情况不会自旋：
    // 1. 饥饿模式：在饥饿模式下，锁会直接交给等待队列中的 goroutine，所以不会自旋。
    // 2. 锁被释放了：另外如果运行到这里的时候，发现锁已经被释放了，也就不需要自旋了。
    if old&(mutexLocked|mutexStarving) == mutexLocked && runtime_canSpin(iter) {
      // 设置 mutexWoken 标识
      // 如果自旋是有意义的，则会进入到这里，尝试设置 mutexWoken 标识。
      // 设置成功在持有锁的 goroutine 获取锁的时候不会唤醒等待队列中的 goroutine，下一个获取锁的就是当前 goroutine。
      if !awoke && old&mutexWoken == 0 && old>>mutexWaiterShift != 0 &&
        atomic.CompareAndSwapInt32(&m.state, old, old|mutexWoken) {
        // 各个判断的含义：
        // !awoke 已经被唤醒过一次了，说明当前协程是被从等待队列中唤醒的协程/又或者已经成功设置 mutexWoken 标识了，不需要再唤醒了。
        // old&mutexWoken == 0 如果不等于 0 说明有 goroutine 被唤醒了，不会尝试设置 mutexWoken 标识
        // old>>mutexWaiterShift != 0 如果等待队列为空，当前 goroutine 就是下一个抢占锁的 goroutine
        // 前面的判断都通过了，才会进行 CAS 操作尝试设置 mutexWoken 标识
        awoke = true
      }
      runtime_doSpin() // 自旋
      iter++           // 自旋次数 +1（超过一定次数会停止自旋）
      old = m.state    // 再次获取锁的最新状态，之后会检查是否锁被释放了
      continue         // 继续下一次检查
    }
    new := old
    // 饥饿模式下，新到达的 goroutines 必须排队。
    // 不是饥饿状态，直接竞争锁。
    if old&mutexStarving == 0 {
      new |= mutexLocked
    }
    // 进入等待队列的两种情况：
    // 1. 锁依然被占用。
    // 2. 进入了饥饿模式。
    if old&(mutexLocked|mutexStarving) != 0 {
      new += 1 << mutexWaiterShift // 等待者数量 +1
    }
     // 已经等待超过了 1ms，且锁被其他协程占用，则进入饥饿模式
    if starving && old&mutexLocked != 0 {
      new |= mutexStarving
    }
    // 唤醒之后，需要重置唤醒标志。
    // 不管有没有获取到锁，都是要清除这个标识的：
    // 获取到锁肯定要清除，如果获取到锁，需要让其他运行中的 goroutine 来抢占锁；
    // 如果没有获取到锁，goroutine 会阻塞，这个时候是需要持有锁的 goroutine 来唤醒的，如果有 mutexWoken 标识，持有锁的 goroutine 唤醒不了。
    if awoke {
      if new&mutexWoken == 0 {
        throw("sync: inconsistent mutex state")
      }
      new &^= mutexWoken // 重置唤醒标志
    }
    // 成功设置新状态
    if atomic.CompareAndSwapInt32(&m.state, old, new) {
      // 原来锁的状态已释放，并且不是饥饿状态，正常请求到了锁，返回
      if old&(mutexLocked|mutexStarving) == 0 { // 这意味着当前的 goroutine 成功获取了锁
        break
      }
      // 如果已经被唤醒过，会被加入到等待队列头。
      queueLifo := waitStartTime != 0
      if waitStartTime == 0 {
        waitStartTime = runtime_nanotime()
      }
      // 阻塞等待
      // queueLifo 为 true，表示加入到队列头。否则，加入到队列尾。
      // (首次加入队列加入到队尾，不是首次加入则加入队头，这样等待最久的 goroutine 优先能够获取到锁。)
      runtime_SemacquireMutex(&m.sema, queueLifo, 1)
      // 从等待队列中唤醒，检查锁是否应该进入饥饿模式。
      starving = starving || runtime_nanotime()-waitStartTime > starvationThresholdNs
      // 获取当前的锁最新状态
      old = m.state
      // 如果锁已经处于饥饿状态，直接抢到锁，返回。
      // 饥饿模式下，被唤醒的协程可以直接获取到锁。
      // 新来的 goroutine 都需要进入队列等待。
      if old&mutexStarving != 0 {
        // 如果这个 goroutine 被唤醒并且 Mutex 处于饥饿模式，P 的所有权已经移交给我们，
        // 但 Mutex 处于不一致的状态：mutexLocked 未设置，我们仍然被视为等待者。修复这个问题。
        if old&(mutexLocked|mutexWoken) != 0 || old>>mutexWaiterShift == 0 {
          throw("sync: inconsistent mutex state")
        }
        // 加锁，并且减少等待者数量。
        // 实际上是两步操作合成了一步：
        // 1. m.state = m.state + 1 （获取锁）
        // 2. m.state = m.state - 1<<mutexWaiterShift（waiter - 1）
        delta := int32(mutexLocked - 1<<mutexWaiterShift)
        // 清除饥饿状态的两种情况：
        // 1. 如果不需要进入饥饿模式（当前被唤醒的 goroutine 的等待时间小于 1ms）
        // 2. 原来的等待者数量为 1，说明是最后一个被唤醒的 goroutine。
        if !starving || old>>mutexWaiterShift == 1 {
          // 退出饥饿模式
          delta -= mutexStarving
        }
        // 原子操作，设置新状态。
        atomic.AddInt32(&m.state, delta)
        break
      }
      // 设置唤醒标记，重新抢占锁（会与那些运行中的 goroutine 一起竞争锁）
      awoke = true
      iter = 0
    } else {
      // CAS 更新状态失败，获取最新状态，然后重试
      old = m.state
    }
  }
}

// Unlock 释放互斥锁。
// 如果 m 在进入 Unlock 时未被锁定，则会出现运行时错误。
func (m *Mutex) Unlock() {
  // Fast path: drop lock bit.
  // unlock 成功
  // unLock 操作实际上是将 state 减去 1。
  new := atomic.AddInt32(&m.state, -mutexLocked)
  if new != 0 { // 等待队列为空的时候直接返回了
    // 唤醒一个等待锁的 goroutine
    m.unlockSlow(new)
  }
}

// unlockSlow 唤醒下一个等待锁的协程。
func (m *Mutex) unlockSlow(new int32) {
  // 如果未加锁，则会抛出错误。
  if (new+mutexLocked)&mutexLocked == 0 {
    fatal("sync: unlock of unlocked mutex")
  }
  // 下面的操作是唤醒一个在等待锁的协程。
  // 存在两种情况：
  // 1. 正常模式：
  //   a. 不需要唤醒：没有等待者、锁已经被抢占、有其他运行中的协程在尝试获取锁、已经进入了饥饿模式
  //   b. 需要唤醒：其他情况
  // 2. 饥饿模式：唤醒等待队列头部的那个协程
  if new&mutexStarving == 0 {
    // 不是饥饿模式
    old := new
    // 自旋
    for {
      // 下面几种情况不需要唤醒：
      // 1. 没有等待者了（没得唤醒）
      // 2. 锁已经被占用（只能有一个 goroutine 持有锁）
      // 3. 有其他运行中的协程已经被唤醒（运行中的 goroutine 通过自旋先抢占到了锁）
      // 4. 饥饿模式（饥饿模式下，所有新的 goroutine 都要排队，饥饿模式会直接唤醒等待队列头部的 gorutine）
      if old>>mutexWaiterShift == 0 || old&(mutexLocked|mutexWoken|mutexStarving) != 0 {
        return
      }
      // 获取到唤醒等待者的权力，开始唤醒一个等待者。
      // 下面这一行实际上是两个操作：
      // 1. waiter 数量 - 1
      // 2. 设置 mutexWoken 标志
      new = (old - 1<<mutexWaiterShift) | mutexWoken
      if atomic.CompareAndSwapInt32(&m.state, old, new) {
        // 正常模式下唤醒了一个 goroutine
        //（第二个参数为 false，表示当前的 goroutine 在释放信号量后还会继续执行直到用完时间片）
        runtime_Semrelease(&m.sema, false, 1)
        return
      }
      // 唤醒失败，进行下一次尝试。
      old = m.state
    }
  } else {
    // 饥饿模式：将互斥锁的所有权移交给下一个等待者，并放弃我们的时间片，以便下一个等待者可以立即开始运行。
    // 注意：如果“mutexLocked”未设置，等待者在唤醒后会将其设置。
    // 但是，如果设置了“mutexStarving”，则仍然认为互斥锁已被锁定，因此新到来的goroutine不会获取它。
    //
    // 当前的 goroutine 放弃 CPU 时间片，让给阻塞在 sema 的 goroutine。
    runtime_Semrelease(&m.sema, true, 1)
  }
}