Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
Upgradable read -> write locks in Go
package main
import (
"fmt"
"runtime"
"sync"
)
type UpgradableLock struct {
uglMutex sync.RWMutex
upgrades chan func()
}
func NewUpgradableLock() *UpgradableLock {
return &UpgradableLock{sync.RWMutex{}, make(chan func(), 1)}
}
func (ugl *UpgradableLock) Lock() {
ugl.uglMutex.Lock()
select {
case v, _ := <-ugl.upgrades:
v()
default:
}
}
// MaybeUpgrade() ensures that f() will be the next function called with a
// write lock, if it returns true. If it returns false, there is some other
// function that will be called first, and the caller must determine how to
// proceed. If the upgrade attempt was successful, also fires off a lock-
// unlock that will run f() if there are no other writers waiting. Does not
// release the read lock presumably held by the caller.
func (ugl *UpgradableLock) MaybeUpgrade(f func()) bool {
select {
case ugl.upgrades <- f:
go func() {
ugl.Lock()
ugl.Unlock()
}()
return true
default:
return false
}
}
func (ugl *UpgradableLock) Unlock() {
ugl.uglMutex.Unlock()
}
func (ugl *UpgradableLock) RLock() {
ugl.uglMutex.RLock()
}
func (ugl *UpgradableLock) RUnlock() {
ugl.uglMutex.RUnlock()
}
func main() {
runtime.GOMAXPROCS(16)
var wg sync.WaitGroup
a := NewUpgradableLock()
b := 0
for i := 0; i < 1000; i++ {
wg.Add(4)
go func() {
a.Lock()
b += 1
a.Unlock()
wg.Done()
}()
go func() {
ug := false
for !ug {
a.RLock()
ug = a.MaybeUpgrade(func() {
b += 1
wg.Done()
})
a.RUnlock()
//Otherwise this loop seems to fail to yield, preventing channel from
//being drained
runtime.Gosched()
}
wg.Done()
}()
go func() {
a.RLock()
a.RUnlock()
wg.Done()
}()
}
wg.Wait()
a.RLock()
fmt.Println(b)
a.RUnlock()
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment