Upgradable Read Write Lock for Go

In this blog post, we'll explore the implementation of an upgradable read-write lock in Go. We will talk about why we needed it by giving concrete examples from the real-world use case and also discuss potential pitfalls during the blog post.

Why do we need an upgradable read-write lock?

In Go, even though the guidelines say to avoid locks, when building a Redis® server that should be concurrently accessed by multiple connections, we usually had to fall back to locks in order to get maximum performance. And, this time standard Go library wasn't enough.

In the Upstash Redis® server implementation, we are guarding the keyspace with locks. While searching for how we can scale it more, there were a couple of ideas. One is partitioning the locks, and the other is replacing the sync.Mutex with a sync.RWMutex. Partitioning locks could be a blog for another day, we will focus on RWMutex for this blog post.

sync.RWMutex is the read-write lock in Go library. It allows multiple readers to have the lock at the same time when there is no writer around. It allows the single writer to have the lock ensuring that there are no readers in the critical section.

Since all of our commands work in memory, they are pretty quick and don't keep long under the lock. But there are commands that have a tendency to take long. These commands are usually the ones that their complexity depends on how large the data is like SUNION,SDIFF, ZDIFF etc.

Since they are read operations, it is mostly ok. It means that other readers can still continue under RWMutex. The writers will be blocked. We realized another problem. All of these commands have STORE versions like SUNIONSTORE, SDIFFSTORE, ZDIFFSTORE, etc. They need to get the write lock. And if that takes long both reads and writes need to wait. The solution comes after realizing that the part taking long is the reading and preparing of the result and not the STORE.

First naive wrong attempt

Let’s assume that we are doing this for SUNIONSTORE. We may try to unlock for the reader first, then lock for the write access as follows:

mutex := sync.RWMutex()
 
mutex.RLock()
//Calculate the UNION of two sets under read lock but don't assign it anywhere
mutex.RUnlock()
 
mutex.Lock()
//Assign it to the in-memory store
mutex.Unlock()

This does not work because the SUNIONSTORE is not atomic anymore. Once we release the read lock another operation can get in and modify the set. After that, this code will try to store a stale set in the memory, overriding the change made in between.

Second naive wrong attempt

Ok, if releasing the read-lock is a problem. Let’s not release it and take the write lock instead as follows:

mutex := sync.RWMutex()
 
mutex.RLock()
//Calculate UNION of two sets under read lock but don't assign it anywhere
 
//Upgrade to the write lock without giving up the read lock
mutex.Lock()
 
//Assign it to the in-memory store
mutex.Unlock()
 
 
mutex.RUnlock()

If you do this the code will hang at mutex.Lock. You can't get a write-lock as long as readers are holding the read lock.

Correct solution

What we need is an upgradeable read-write lock. There can be various different APIs. I will first show what we ended up with.

 
mutex.UpgradableRLock()
//Calculate the UNION of two sets under read lock but don't assign it anywhere
 
//Upgrade to the write lock without giving up the read lock
mutex.UpgradeWLock()
//Assign it to the in-memory store
 
mutex.UpgradableRUnlock()

We need to mention at this point that this idea is well-known in the literature and there are implementations of upgrading a readlock even in the standard libraries of other languages. Examples: JAVA StampedLock .NET ReaderWriterLockSlim

Our API is very close to the .NET ReaderWriterLockSlim. At this point, you might say that this usage looks very similar to the second attempt. How can it get away with deadlock? The trick is that it has one more mode upgradable-read together with read and write. Semantics are as follows:

• Both read and upgradable-read locks can be held together at the same time.
• write-lock prevents both read and upgradable-read from entering the critical section.
• There can be multiple goroutines holding the read lock but there can be only a single upgradable-read.

The last point is important. If we don't have that guarantee, we will again end up with a deadlock. If we allow let’s say two upgradable-read locks can be acquired, when they try to upgrade to write, one of them has to give up the read-lock. Therefore, the atomicity will be broken again.

The implementation

To tell you the truth, we ended up with the API after realizing that it is trivial to modify GO sync.RWMutex into this API.

 
mutex.UpgradableRLock()
//Calculate the UNION of two sets under read lock but don't assign it anywhere
 
//Upgrade to the write lock without giving up the read lock
mutex.UpgradeWLock()
//Assign it to the in-memory store
 
mutex.UpgradableRUnlock()

To understand our implementation, we need to take a look sync.RWMutex especially Lock and Rlock

 
type RWMutex struct {
	w           Mutex        // held if there are pending writers
	// .... skipped
}
 
func (rw *RWMutex) Lock() {
	// First, resolve competition with other writers.
	rw.w.Lock()
	// Announce to readers that there is a pending writer.
	r := rw.readerCount.Add(-rwmutexMaxReaders) + rwmutexMaxReaders
	// Wait for active readers.
	if r != 0 && rw.readerWait.Add(r) != 0 {
		semaphoreAcquire(&rw.writerSem)
	}
}
 
func (rw *RWMutex) RLock() {
	if rw.readerCount.Add(1) < 0 {
		// A writer is pending, wait for it.
		semaphoreAcquire(&rw.readerSem)
	}
}

The RLock is really simple. It just increments a reader count and sleep for a writer to leave the critical section if necessary.

The Lock gets the underlying w lock first. At this point, no other writes can get in, but it has nothing to do with reads yet. Then it announces to the readers that they can not get in anymore. That means if we just split that method into two, the first half is basically an UpgradebleRLock and the second part is UpgradeWLock. In other words, the second part upgrades the upgradable-read to write.

This gives us everything we need already and all we did is to separate a method into two.

 
func (rw *UpgradableRWMutex) UpgradableRLock() {
	// First, resolve competition with other writers.
	// Disallow writers to acquire the lock
	rw.w.Lock()	
}
 
func (rw *UpgradableRWMutex) UpgradeWLock() {
	// Announce to readers there is a pending writer.
	r := rw.readerCount.Add(-rwmutexMaxReaders) + rwmutexMaxReaders
	// Wait for active readers.
	if r != 0 && rw.readerWait.Add(r) != 0 {
		semaphoreAcquire(&rw.writerSem)
	}
}
 

We still need to figure out how to unlock though, a.k.a UpgradableRUnlock. This method should call the sync.RWMutex.Unlock() if it is upgraded to the write. It should call just the rw.w.Unlock() if it is not upgraded. Because all we did was just rw.w.Lock() anyway so far.

We will add a bool to keep track if the lock is upgraded or not. And we don't need to guard this bool for concurrent access because we will access it only under rw.w.Lock()

 
type UpgradableRWMutex struct {
	//... skipped. It’s the same as sync.RWMutex
 
	// Keep track if an upgradeable read-lock is upgraded to a write-lock or not. Always accessed under w locked
	upgraded           bool
 
}
 
func (rw *UpgradableRWMutex) UpgradableRLock() {
	// First, resolve competition with other writers.
	// Disallow writers to acquire the lock
	rw.w.Lock()	
}
 
func (rw *UpgradableRWMutex) UpgradeWLock() {
	rw.upgraded = true
	// Announce to readers there is a pending writer.
	r := rw.readerCount.Add(-rwmutexMaxReaders) + rwmutexMaxReaders
	// Wait for active readers.
	if r != 0 && rw.readerWait.Add(r) != 0 {
		semaphoreAcquire(&rw.writerSem)
	}
}
 
func (rw *UpgradableRWMutex) UpgradableRUnlock() {
	rw.upgradableReadMode = false
	if rw.upgraded {
		rw.upgraded = false
		rw.Unlock()
	} else {
		rw.w.Unlock()
	}
}

Conclusion

In this blog post, we've examined the need and the implementation of an upgradable read-write lock in Go, which can be a useful tool in concurrent applications. We discussed the details of the implementation and how this construct optimizes shared resource access. We also highlighted common pitfalls and naive approaches in the quest for an upgradable read-write lock.

For the full implementation code and deeper exploration, follow this link. Thank you for joining us on this journey. Stay tuned for more technical insights and implementations as we continue to improve the Upstash Redis® server.