An upgradable read write lock for go

README.md

And upgradable read write lock for go

UpgradableRWMutex is an enhanced version of the standard sync.RWMutex. It has the all methods sync.RWMutex with exact same semantics. It gives more methods to give upgradable-read feature.

The new semantics for upgradable-read are as follows:

• Multiple goroutines can get read-lock together with a single upgradable-read-lock.
• Only one goroutine can have a write-lock and no read-lock/upgradable-read-lock can be acquired in this state.
• There can be only a single goroutine keeping the upgrade-read-lock.
• UpgradableRWMutex is not reentrant.

See related blog here for details:https://upstash.com/blog/upgradable-rwlock-for-go

Usage

mutex.UpgradableRLock()
defer mutex.UpgradableRUnlock()
// read-lock acquired section. We can return here safely if an error occurs
mutex.UpgradeWLock()
// critical section with exclusive right access

upgradablerw_mutex.go

// Based on source code copyright by The Go Authors.
//
// Copyright (c) 2009 The Go Authors. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//   * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//   * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//   * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package main

import (
	"sync"
	"sync/atomic"
	_ "unsafe"
)

// UpgradableRWMutex is an enhanced version of the standard sync.RWMutex.
// It has the all methods sync.RWMutex with exact same semantics.
// It gives more methods to give upgradable-read feature.
//
// The new semantics for upgradable-read are as follows:
// Multiple goroutines can get read-lock together with a single upgradable-read-lock.
// Only one goroutine can have a write-lock and no read-lock/upgradable-read-lock can be acquired in this state.
// There can be only a single goroutine keeping the upgrade-read-lock.
// UpgradableRWMutex is not reentrant.
//
// Usage of the UpgradableRWMutex:
//
//	mutex.UpgradableRLock()
//	defer mutex.UpgradableRUnlock()
//	// read-lock acquired section. We can return here safely if an error occurs
//	mutex.UpgradeWLock()
//	// critical section with exclusive right access
type UpgradableRWMutex struct {
	w           sync.Mutex   // held if there are pending writers
	writerSem   uint32       // semaphore for writers to wait for completing readers
	readerSem   uint32       // semaphore for readers to wait for completing writers
	readerCount atomic.Int32 // number of pending readers
	// number of departing readers. A negative number.
	// Number of readers left while under write lock(while write lock waiting for readers to leave)
	readerWait atomic.Int32
	// Keep track if an upgradeable read-lock is upgraded to write-lock or not. Always accessed under w locked
	upgraded           bool
	upgradableReadMode bool
}

//go:linkname semaphoreAcquire sync.runtime_Semacquire
func semaphoreAcquire(s *uint32)

//go:linkname semaphoreRelease sync.runtime_Semrelease
func semaphoreRelease(s *uint32, handoff bool, skipframes int)

const rwmutexMaxReaders = 1 << 30

// RLock is same as sync.RWMutex.RLock
func (rw *UpgradableRWMutex) RLock() {
	if rw.readerCount.Add(1) < 0 {
		// A writer is pending, wait for it.
		semaphoreAcquire(&rw.readerSem)
	}
}

// TryRLock is same as sync.RWMutex
func (rw *UpgradableRWMutex) TryRLock() bool {
	for {
		c := rw.readerCount.Load()
		if c < 0 {
			return false
		}
		if rw.readerCount.CompareAndSwap(c, c+1) {
			return true
		}
	}
}

// RUnlock is same as sync.RWMutex
func (rw *UpgradableRWMutex) RUnlock() {
	if r := rw.readerCount.Add(-1); r < 0 {
		// Outlined slow-path to allow the fast-path to be inlined
		rw.rUnlockSlow(r)
	}
}

func (rw *UpgradableRWMutex) rUnlockSlow(r int32) {
	if r+1 == 0 || r+1 == -rwmutexMaxReaders {
		panic("sync: RUnlock of unlocked UpgradableRWMutex")
	}
	// A writer is pending.
	if rw.readerWait.Add(-1) == 0 {
		// The last reader unblocks the writer.
		semaphoreRelease(&rw.writerSem, false, 1)
	}
}

// Lock is same as sync.RWMutex
func (rw *UpgradableRWMutex) Lock() {
	// First, resolve competition with other writers.
	rw.w.Lock()
	// 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)
	}
}

// TryLock is same as sync.RWMutex
func (rw *UpgradableRWMutex) TryLock() bool {
	if !rw.w.TryLock() {
		return false
	}
	if !rw.readerCount.CompareAndSwap(0, -rwmutexMaxReaders) {
		rw.w.Unlock()
		return false
	}
	return true
}

// Unlock is same as sync.RWMutex
func (rw *UpgradableRWMutex) Unlock() {
	// Announce to readers there is no active writer.
	r := rw.readerCount.Add(rwmutexMaxReaders)
	if r >= rwmutexMaxReaders {
		panic("sync: Unlock of unlocked UpgradableRWMutex")
	}
	// Unblock blocked readers, if any.
	for i := 0; i < int(r); i++ {
		semaphoreRelease(&rw.readerSem, false, 0)
	}
	// Allow other writers to proceed.
	rw.w.Unlock()
}

// UpgradeWLock upgrade the read lock to the write lock
func (rw *UpgradableRWMutex) UpgradeWLock() {
	if !rw.upgradableReadMode {
		panic("sync: Upgrade outside of upgradableReadLock not allowed")
	}
	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)
	}
}

// UpgradableRUnlock unlocks either the write-lock if it is upgraded
// or unlock just the upgradeableRead-lock if not upgraded
func (rw *UpgradableRWMutex) UpgradableRUnlock() {
	rw.upgradableReadMode = false
	if rw.upgraded {
		rw.upgraded = false
		rw.Unlock()
	} else {
		rw.w.Unlock()
	}
}

// UpgradableRLock acquires an upgradable-read-lock which can be later upgraded to write-lock,
// Example usage:
//
//	mutex.UpgradableRLock()
//	defer mutex.UpgradableRUnlock()
//	// read-lock acquired section. We can return here safely if an error occurs
//	mutex.UpgradeWLock()
//	// critical section with exclusive right access
func (rw *UpgradableRWMutex) UpgradableRLock() {
	// First, resolve competition with other writers.
	// Disallow writers to acquire the lock
	rw.w.Lock()
	if rw.readerCount.Load() < 0 {
		panic("reader count can not be negative. We have the write lock")
	}
	rw.upgradableReadMode = true
}

upgradablerw_mutex_test.go

package main

import (
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func Test_UpgradableLock(t *testing.T) {
	t.Run("can't get upgradableReadLock while writeLock held", func(t *testing.T) {
		m := UpgradableRWMutex{}
		m.Lock()
		i := atomic.Int32{}
		go func() {
			m.UpgradableRLock()
			defer m.UpgradableRUnlock()
			i.Add(1)
		}()
		time.Sleep(time.Millisecond * 10)
		assert.Equal(t, int32(0), i.Load())
		m.Unlock()
	})

	t.Run("can't get write lock while upgradableReadLock upgraded to write", func(t *testing.T) {
		m := UpgradableRWMutex{}
		m.UpgradableRLock()

		m.UpgradeWLock()
		i := atomic.Int32{}
		go func() {
			m.Lock()
			defer m.Unlock()
			i.Add(1)
		}()
		time.Sleep(time.Millisecond * 10)
		assert.Equal(t, int32(0), i.Load())
		m.UpgradableRUnlock()
	})

	t.Run("upgradable read lock allows other readers before upgrade", func(t *testing.T) {
		m := UpgradableRWMutex{}
		m.UpgradableRLock()

		waitGroup := sync.WaitGroup{}
		for i := 0; i < 10; i++ {
			waitGroup.Add(1)
			go func() {
				m.RLock()
				defer m.RLock()
				waitGroup.Done()
			}()
		}
		waitGroup.Wait()
		m.UpgradableRUnlock()
	})

	t.Run("upgradable read lock prevents other readers after upgrade", func(t *testing.T) {
		m := UpgradableRWMutex{}
		m.UpgradableRLock()

		m.UpgradeWLock()
		i := atomic.Int32{}
		go func() {
			m.RLock()
			defer m.RUnlock()
			i.Add(1)
		}()
		time.Sleep(time.Millisecond * 10)
		assert.Equal(t, int32(0), i.Load())
		m.UpgradableRUnlock()
	})

	t.Run("get upgradable read-lock while there are other readers", func(t *testing.T) {
		m := UpgradableRWMutex{}

		m.RLock()
		defer m.RUnlock()

		m.RLock()
		defer m.RUnlock()

		m.UpgradableRLock()
		defer m.UpgradableRUnlock()
	})

	// this is to avoid a failing upgrade under read-lock. Once an upgradeable-lock is acquired, it must be able to
	// upgrade to write-lock without causing deadlock
	t.Run("prevent getting a second upgradable-lock even after upgrade", func(t *testing.T) {
		m := UpgradableRWMutex{}

		m.RLock()
		defer m.RUnlock()

		m.RLock()
		defer m.RUnlock()

		m.UpgradableRLock()
		i := atomic.Int32{}
		go func() {
			m.UpgradableRLock()
			defer m.UpgradableRUnlock()
			i.Add(1)
		}()
		time.Sleep(time.Millisecond * 10)
		assert.Equal(t, int32(0), i.Load())
		m.UpgradableRUnlock()
	})
}