一文带你深入了解Golang中的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)
   }
}