Go 底层锁：原子操作和sema信号量

unitiny
720 阅读9分钟

前言

在Go开发中，我们会经常用到互斥锁，读写锁。而这些锁是基于更底层的锁来实现的。今天我们就来了解下原子操作和sema信号量，看它们能做些什么事情？

原子操作

用法

原子操作只支持简单的运算操作，比如：

  • atomic.AddInt64

变量加减，举例：

var a int64 = 2
b := atomic.AddInt64(&a, 3)
fmt.Println(a, b) // 结果为 a=5 b=5

atomic.AddInt64会返回运算后的结果，目的是并发时，不需要再使用atomic.LoadInt64来读取a的值。

  • atomic.CompareAndSwapInt64

如果变量等于旧的值，则赋值为新的值，返回true；否则不赋值，返回false。举例：

var a int64 = 2
b := atomic.CompareAndSwapInt64(&a, 2, 3)
fmt.Println(a, b) // 结果为 a=3 b=true

类似于

if a == 2 {
   a = 3
   return true
} else {
   return false
}

源码解析

原子操作怎么实现的呢？追踪其AddInt64的汇编代码，发现使用了Xadd方法。

image.png

这里加了个cpu级别的内存锁，是最底层的硬件锁，来同一时间只有一个线程操作该变量的内存。

sema信号量

意义和作用

原子操作只能用于运算操作，无法满足各线程的复杂操作。因此Go实现了一个sema信号量去解决各线程并发时同步互斥问题。

sema信号量呢，是基于原子操作实现的，和操作系统的信号量差不多，也有类似PV操作。核心是一个uint32值，可看作资源的数量。

在互斥锁和读写锁中，我们都可以看到sema信号量的身影，因此研究sema信号量是很有必要的。

type Mutex struct {
   state int32
   sema  uint32 // 使用了sema信号量
}

type RWMutex struct {
   w           Mutex        
   writerSem   uint32       // 写信号量
   readerSem   uint32       // 读信号量
   readerCount atomic.Int32
   readerWait  atomic.Int32 
}

在两种锁中，都用上了sema，可它只是一个uint32的变量，能干什么用呢？

semaRoot

其实通过sema，会使用到一个叫semaRoot结构体，semaRoot结构如下：

type semaRoot struct {
   lock  mutex
   treap *sudog        // 平衡树的根节点
   nwait atomic.Uint32 // 协程等待的数量，即平衡树的节点个数
}

其中treap字段是一个sudog结构体指针，sudog是一个很重要的结构体，部分结构如下：

type sudog struct {
   g *g // 使用该结构体的协程
   next *sudog
   prev *sudog
   elem unsafe.Pointer // 协程等待的信号量地址
   waitlink *sudog // 每个信号量对应的sudog队列
   ...
}

通过sudog，我们能找到标记的协程，以及上一个和下一个sudog结构体，这些sudog又有各自标记的协程。

sudog.elem字段记录了协程对应sema信号量的地址，用于平衡树的排序标准。

小结梳理

用一张图梳理下结构关系：

image.png

sema本质是一个uint32的值，如果协程进入互斥等待状态的话，sema就会对应上一个semaRoot结构体，semaRoot里面的treap字段是sudog结构体指针，是平衡二叉树的根节点。

sudog作用是什么？

当有协程得不到锁而进入等待状态时，就会被放入sudog结构体中。sudog将等待协程封装好，并用上prev, next等字段封装为一个树节点。

平衡树的构造是什么？

这里Go的平衡树定义是：

  • 每个节点都是一个sudog结构体
  • sudog.prev可看作左节点
  • sudog.next可看作右节点
  • 树高度差不超过一层
  • 按照sudog.elem的值排序，左节点的elem值小于父节点，右节点elem值大于父节点

注意，sudog.elem的值为sema信号量的地址，而平衡树按sudog.elem的值排序，说明平衡树不同节点对应不同的sema。

sudog.elem与sudog.waitlink作用后边扩展中讲。

sema的具体实现

一切都得从sema信号量谈起。

sema > 0

sema>0，代表有资源，不需要互斥竞争同一资源了，所以处理比较简单：

  • 要获取资源，则sema减一，代表获取成功
  • 要释放资源，则sema加一，代表释放成功

将sema>0代入获取资源的源码分析：

首先调用/src/runtime/sema中的semaquire方法获取资源。

func semacquire(addr *uint32) {
   semacquire1(addr, false, 0, 0, waitReasonSemacquire)
}

然后调用semaquire1。

func semacquire1(addr *uint32, lifo bool, profile semaProfileFlags, skipframes int, reason waitReason) {
   gp := getg()
   if gp != gp.m.curg {
      throw("semacquire not on the G stack")
   }

   // Easy case.
   if cansemacquire(addr) {
      return
   }
   ...
}

semaquire1中把sema地址传入一个cansemacquire方法中，如果cansemacquire方法返回true，获取锁的过程就直接结束了。那来看下cansemacquire方法是什么：

func cansemacquire(addr *uint32) bool {
   for {
      v := atomic.Load(addr)
      if v == 0 {
         return false
      }
      if atomic.Cas(addr, v, v-1) {
         return true
      }
   }
}

cansemacquire会循环判断，由于sema大于0，进入了原子操作atomic.Cas(addr, v, v-1)。

atomic.Cas和atomic.CompareAndSwapInt64方法一样，都是判断addr是否等于v，是则将addr减一，并返回true。addr为sema的地址，所以相当于sema减一了。

所以当sema>0时，表示有资源，协程直接将sema减一，然后结束。

将sema>0代入释放资源的源码分析：

首先调用/src/runtime/sema中的semrelease方法释放锁。

func semrelease(addr *uint32) {
   semrelease1(addr, false, 0)
}

调用了semrelease1方法。

func semrelease1(addr *uint32, handoff bool, skipframes int) {
   root := semtable.rootFor(addr)
   atomic.Xadd(addr, 1)

   // Easy case: no waiters?
   // This check must happen after the xadd, to avoid a missed wakeup
   // (see loop in semacquire).
   if root.nwait.Load() == 0 {
      return
   }
   ...
}

semrelease1调用了atomic.Xadd(addr, 1)，将addr加一。

若这时root.nwait等于0，表示没有等待的协程了（平衡树中没有节点了），则semrelease1只需要释放锁，而不需要做唤醒协程的任务了，于是提前return回家。

sema = 0

sema=0，代表资源紧张，需要互斥竞争同一资源了，协程若没竞争到资源就进入等待状态了，这时semaRoot上场了。

将sema=0代入获取资源的源码分析：

当走到semaquire -> semaquire1 -> cansemacquire，由于sema=0，直接返回false，继续 semaquire1。

func semacquire1(addr *uint32, lifo bool, profile semaProfileFlags, skipframes int, reason waitReason) {
   gp := getg()
   if gp != gp.m.curg {
      throw("semacquire not on the G stack")
   }

   // Easy case.
   if cansemacquire(addr) {
      return 
   }

   // Harder case:
   // increment waiter count
   // try cansemacquire one more time, return if succeeded
   // enqueue itself as a waiter
   // sleep
   // (waiter descriptor is dequeued by signaler)
   s := acquireSudog()
   root := semtable.rootFor(addr)
   ...
   for {
      lockWithRank(&root.lock, lockRankRoot)
      // Add ourselves to nwait to disable "easy case" in semrelease.
      root.nwait.Add(1)
      // Check cansemacquire to avoid missed wakeup.
      if cansemacquire(addr) {
         root.nwait.Add(-1)
         unlock(&root.lock)
         break
      }
      // Any semrelease after the cansemacquire knows we're waiting
      // (we set nwait above), so go to sleep.
      root.queue(addr, s, lifo) // 入队
      goparkunlock(&root.lock, reason, traceEvGoBlockSync, 4+skipframes)
      if s.ticket != 0 || cansemacquire(addr) {
         break
      }
   }
   if s.releasetime > 0 {
      blockevent(s.releasetime-t0, 3+skipframes)
   }
   releaseSudog(s)
}

流程如下：

  • 调用acquireSudog，获取到一个sudog。

  • 使用semtable.rootFor(addr)，semaRoot登场，用于协程入队平衡树中。

  • 进入一个for循环，它很倔，没拿到资源就不退出，要么休眠等待，要么继续循环。

    • 该协程要进入等待状态了，因此将root.nwait加一，表示平衡树中要加多一个等待协程的位置。
    • 在休眠前不甘心，又调用cansemacquire看一下是否有资源。
    • 还是没资源，于是调用root.queue，将自己这个协程放入sudog中，sudog又放入平衡树中。
    • 然后调用goparkunlock，主动挂起，休眠等待。

  • 拿到资源后，退出循环。因为该协程不是等待协程了，所以调用releaseSudog释放sudog。

将sema=0代入释放资源的源码分析：

当走到semrelease -> semrelease1时，将sema加一，释放了资源。

func semrelease1(addr *uint32, handoff bool, skipframes int) {
   root := semtable.rootFor(addr)
   atomic.Xadd(addr, 1)

   // Easy case: no waiters?
   // This check must happen after the xadd, to avoid a missed wakeup
   // (see loop in semacquire).
   if root.nwait.Load() == 0 {
      return
   }

   // Harder case: search for a waiter and wake it.
   lockWithRank(&root.lock, lockRankRoot)
   if root.nwait.Load() == 0 {
      // The count is already consumed by another goroutine,
      // so no need to wake up another goroutine.
      unlock(&root.lock)
      return
   }
   s, t0 := root.dequeue(addr)
   if s != nil {
      root.nwait.Add(-1)
   }
   unlock(&root.lock)
   if s != nil { // May be slow or even yield, so unlock first
      ...
      readyWithTime(s, 5+skipframes) // 唤醒等待协程
      ...
   }
}

若这时root.nwait!=0，说明有其它协程在等待这个资源，那当前释放资源的协程就该去唤醒等待协程。

怎么唤醒呢？

首先获取等待协程所在的sudog，使用root.dequeue(addr)方法，从平衡树中出队一个sudog，sudog里记录的协程就是等待协程。

然后使用readyWithTime(s, 5+skipframes)唤醒协程。

疑问

为什么要通过acquireSudog获取一个sudog，直接new不好吗？

在GMP模型中，每个处理器P有个sudogcache字段，会缓存一批sudog。当sudog不使用时，还会在sudogcache中缓存着，而不会被GC清理。在下次使用sudog时，就直接从缓存中获取，而不需要重新分配一个sudog空间。

获取缓存的入口就是acquireSudog方法。

treap为什么是平衡树的根节点？有什么作用？

treap实时记录平衡树根节点，如此才能有效的二分搜索，更快唤醒等待协程。

扩展一丢丢

更具体的平衡树

其实平衡树每个节点都是一个队列，对应字段sudog.waitlink。

使用同一sema的等待协程，它sudog的elem字段都等于同一个sema的地址。这些sudog会放入平衡树对应sema节点的队列中。

image.png

使用平衡树是为了在 O(logn) 时间下搜索节点。

使用队列是为了让等待协程按等待时间从大到小排序，等待时间久的放队头，这样能优先处理。

平衡树是如何插入，释放sudog的？

插入

在semaquire1中，如果拿不到资源就将sudog入队，调用root.queue(addr, s, lifo)来入队，具体代码如下：

func (root *semaRoot) queue(addr *uint32, s *sudog, lifo bool) {
   s.g = getg()
   s.elem = unsafe.Pointer(addr)
   s.next = nil
   s.prev = nil

   var last *sudog
   pt := &root.treap
   for t := *pt; t != nil; t = *pt {
      if t.elem == unsafe.Pointer(addr) {
         // Already have addr in list.
         if lifo {
            // Substitute s in t's place in treap.
            *pt = s
            s.ticket = t.ticket
            s.acquiretime = t.acquiretime
            s.parent = t.parent
            s.prev = t.prev
            s.next = t.next
            if s.prev != nil {
               s.prev.parent = s
            }
            if s.next != nil {
               s.next.parent = s
            }
            // Add t first in s's wait list.
            s.waitlink = t
            s.waittail = t.waittail
            if s.waittail == nil {
               s.waittail = t
            }
            t.parent = nil
            t.prev = nil
            t.next = nil
            t.waittail = nil
         } else {
            // Add s to end of t's wait list.
            if t.waittail == nil {
               t.waitlink = s
            } else {
               t.waittail.waitlink = s
            }
            t.waittail = s
            s.waitlink = nil
         }
         return
      }
      last = t
      
      // 二分查找
      if uintptr(unsafe.Pointer(addr)) < uintptr(t.elem) {
         pt = &t.prev
      } else {
         pt = &t.next
      }
   }

   // 维护平衡树特性
   for s.parent != nil && s.parent.ticket > s.ticket {
      if s.parent.prev == s {
         root.rotateRight(s.parent)
      } else {
         if s.parent.next != s {
            panic("semaRoot queue")
         }
         root.rotateLeft(s.parent)
      }
   }
}

这里sudog.elem = unsafe.Pointer(addr)，即记录了sema地址。

然后使用treap这个平衡树的根节点作入口，开始二分查找，查找elem = sema地址的节点。

找到对应节点后，如果该协程的sudog已经在队列中，则替换；否则将sudog加入队尾。

之后为了保持平衡树特性，进行旋转操作。

释放

在semrelease1中，为了唤醒其它等待协程，需要释放平衡树的sudog，便调用root.dequeue(addr)来释放。

func (root *semaRoot) dequeue(addr *uint32) (found *sudog, now int64) {
   ps := &root.treap
   s := *ps
   for ; s != nil; s = *ps {
      if s.elem == unsafe.Pointer(addr) {
         goto Found
      }
      if uintptr(unsafe.Pointer(addr)) < uintptr(s.elem) {
         ps = &s.prev
      } else {
         ps = &s.next
      }
   }
   return nil, 0
   
Found:   
   ...
}

先通过二分查找，找到sema对应的sudog节点（队列头），然后进入Found代码段。

Found代码中，若队列中有多个值，则删除一个sudog即可；若队列只剩该sudog，在删除后还得进行旋转平衡树操作，维持平衡树特性。

总结

sema信号量基于原子操作实现，通过semaRoot维护等待协程，合适时机唤醒对应协程，实现了协程互斥和同步功能。

千言万语不如亲自细读。有不对或疑惑的地方欢迎交流。

avatar
文章
阅读
粉丝