golang RWMutex读写锁实现读共享写独占的功能示例

// rwmutex 是一个读写互斥的锁
// 将允许多个 goroutine 持有读锁，但写锁只会有一个持有
// rwmutex 使用了 sync.RWMutex 来辅助写锁互斥
type rwmutex struct {
rLock      mutex    // 用于保护设置 readers, readerPass, writer
readers    muintptr // 休眠等待的 goroutine 读锁队列，等到写锁占有结束后将对应被唤起。
readerPass uint32   // 读锁队列需要跳过的 goroutine 数量，当在写锁结束后会唤起读锁队列里的 goroutine，但有的可能已不在队列里了，这部分需跳过。
wLock  mutex    // 用于 writer 之间的互斥锁
writer muintptr // 等待读完成的 writer
readerCount uint32 // 正在执行读操作的 goroutine数量
readerWait  uint32 // 等待读锁释放的数量。当写锁占有后，前面还有部分读锁在继续着，需要等它们释放才能继续进行。
}

func (rw *rwmutex) Lock() {
    // 用于多个写锁之间的的竞争
    lock(&rw.wLock)
    m := getg().m
    // 将读锁数量 readerCount 置为负数，用于判断当前是否处于写锁占有状态，
    // rw.readerCount < 0 则表示当前正在进行写锁占有.
    r := int32(atomic.Xadd(&rw.readerCount, -rwmutexMaxReaders)) + rwmutexMaxReaders
    // 前面还有读锁在进行着，需要等待释放完才能继续
    lock(&rw.rLock)
    if r != 0 && atomic.Xadd(&rw.readerWait, r) != 0 {
        systemstack(func() {
            rw.writer.set(m)
            unlock(&rw.rLock)
            notesleep(&m.park)
            noteclear(&m.park)
        })
    } else {
        unlock(&rw.rLock)
    }
}

func (rw *rwmutex) Rlock() {
    acquirem()
    if int32(atomic.Xadd(&rw.readerCount, 1)) < 0 {
        // 读锁数量 readerCount + 1 后小于 0，表示当前正被写锁占有，
        // 等待写锁释放
        systemstack(func() {
            lock(&rw.rLock)
            if rw.readerPass > 0 {
                rw.readerPass -= 1
                unlock(&rw.rLock)
            } else {
                // 等待写锁唤起
                m := getg().m
                m.schedlink = rw.readers
                rw.readers.set(m)
                unlock(&rw.rLock)
                notesleep(&m.park)
                noteclear(&m.park)
            }
        })
    }
}

func (rw *rwmutex) Unlock() {
    // 将原来被写锁置为负数的 readerCount 重新恢复回来.
    r := int32(atomic.Xadd(&rw.readerCount, rwmutexMaxReaders))
    if r >= rwmutexMaxReaders {
        throw("unlock of unlocked rwmutex")
    }
    // 唤起之前等待的读锁.
    lock(&rw.rLock)
    for rw.readers.ptr() != nil {
        reader := rw.readers.ptr()
        rw.readers = reader.schedlink
        reader.schedlink.set(nil)
        notewakeup(&reader.park)
        r -= 1
    }
    // 如果 r > 0, 说明读锁队列里有的 goroutine 已不在队列里了，这部分需跳过
    rw.readerPass += uint32(r)
    unlock(&rw.rLock)
    // 解除写锁
    unlock(&rw.wLock)
}

func (rw *rwmutex) RUnlock() {
    // 如果释放后，readerCount < 0，表示当前写锁正在占有
    if r := int32(atomic.Xadd(&rw.readerCount, -1)); r < 0 {
        if r+1 == 0 || r+1 == -rwmutexMaxReaders {
            throw("runlock of unlocked rwmutex")
        }
        // readerWait == 0，表示前面的读锁都释放完了，
        // 需要唤起写锁
        if atomic.Xadd(&rw.readerWait, -1) == 0 {
            // The last reader unblocks the writer.
            lock(&rw.rLock)
            w := rw.writer.ptr()
            if w != nil {
                notewakeup(&w.park)
            }
            unlock(&rw.rLock)
        }
    }
    releasem(getg().m)
}