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