RWMutex with TryLock

RWMutex.go

package sync

import (
	"sync"
	"sync/atomic"
)

type nullLocker struct {}

func (n *nullLocker) Lock() {}
func (n *nullLocker) Unlock() {}

/*
 Mostly a copy of sync.RWMutex
 With TryLock/TryRLock
*/
type RWMutex struct {
	w              sync.Mutex
	readers        int32
	readersWaiting int32
	writerSem      sync.Cond
	readerSem      sync.Cond
}

const maxReaders = 1 << 30

func NewRWMutex() *RWMutex {
	return &RWMutex{writerSem: sync.Cond{L:&nullLocker{}}, readerSem: sync.Cond{L:&nullLocker{}}}
}

func (m *RWMutex) WLock() {
	m.w.Lock()
	r := atomic.AddInt32(&m.readers, -maxReaders) + maxReaders
	if r != 0 && atomic.AddInt32(&m.readersWaiting, r) != 0 {
		m.writerSem.Wait()
	}
}

func (m *RWMutex) WUnlock() {
	atomic.AddInt32(&m.readers, maxReaders)
	m.readerSem.Broadcast()
	m.w.Unlock()
}

func (m *RWMutex) RLock() {
	if atomic.AddInt32(&m.readers, 1) < 0 {
		m.readerSem.Wait()
	}
}

func (m *RWMutex) RUnlock() {
	if atomic.AddInt32(&m.readers, -1) < 0 {
		if atomic.AddInt32(&m.readersWaiting, -1) == 0 {
			m.writerSem.Signal()
		}
	}
}

func (m *RWMutex) TryRLock() bool {
	if atomic.AddInt32(&m.readers, 1) < 0 {
		atomic.AddInt32(&m.readers, -1)
		return false
	}
	return true
}

//This will lock if there are no writers
//which means it will wait for readers to finish
func (m *RWMutex) TryWLock() bool {
	r := atomic.LoadInt32(&m.readers)
	if r < 0 {
		return false
	}
	m.WLock()
	return true
}

// This will lock for write, only if there are no readers or writers
func (m *RWMutex) TryWLockGreedy() bool {
	r := atomic.LoadInt32(&m.readers)
	if r != 0 {
		return false
	}
	m.WLock()
	return true
}

RWMutex_test.go

package sync

import (
	"fmt"
	"testing"
)

func parallelReader(m *RWMutex, num int, clocked, cunlock, cdone chan bool) {
	m.RLock()
	fmt.Println(num, " locked for read")
	clocked <- true
	<-cunlock
	m.RUnlock()
	fmt.Println(num, " unlocked read")
	cdone <- true
}

func parallelWriter(m *RWMutex, num int,clocked, cunlock, cdone chan bool) {
	m.WLock()
	fmt.Println(num, " locked for write")
	clocked <- true
	m.WUnlock()
	<-cunlock
	fmt.Println(num, " unlocked write")
	cdone <- true
}

func spawn(m *RWMutex, num int, f func(*RWMutex,int,chan bool, chan bool, chan bool)) {
	clocked := make(chan bool)
	cunlock := make(chan bool)
	cdone := make(chan bool)
	
	for i := 0; i < num; i++ {
		go f(m, i, clocked, cunlock, cdone)
	}
	
	for i := 0; i < num; i++ {
		<-clocked
	}

	for i := 0; i <num;i++ {
		cunlock <- true
	}

	for i := 0; i < num; i++ {
		<-cdone
	}
	
}

func spawnReaders(m *RWMutex, num int) {
	spawn(m, num, parallelReader)
}

func spawnWriters(m *RWMutex, num int) {
	spawn(m, num, parallelWriter)
}

func TestRWMutexParallelReads(t *testing.T) {
	mutex := NewRWMutex()
	spawnReaders(mutex, 10)
}

func TestRWMutexParallelWrites(t *testing.T) {
	mutex := NewRWMutex()
	spawnWriters(mutex, 10)
}

func TestRWMutexParallelReadWrites(t *testing.T) {
	mutex := NewRWMutex()
	for i := 0; i < 10; i++{
		go spawnReaders(mutex, 500)
		go spawnWriters(mutex, 500)
	}
}

func TestRWMutexTryRLock(t *testing.T) {
	mutex := NewRWMutex()

	//We should get the lock just fine
	r := mutex.TryRLock()
	if r != true {
		t.Error("Lock should be acquired!")
	}
	mutex.RUnlock()
	
	
	mutex.WLock()
	//When wlocked, we shouldn't get the lock
	r = mutex.TryRLock()
	if r == true {
		t.Error("Lock shouldn't be acquired!")
	}
	mutex.WUnlock()
}

func TestRWMutexTryWLock(t *testing.T) {
	mutex := NewRWMutex()
	finished := make(chan bool)

	r := mutex.TryWLock()
	if r != true {
		t.Error("WLock should be acquired!")
	}

	go func(){
		r := mutex.TryWLock()
		if r == true {
			t.Error("WLock shouldn't be acquired!")
		
		}
		finished <- true
	}()

	<-finished
}

func TestRWMutexTryWLockGreedy(t *testing.T) {
	mutex := NewRWMutex()
	finished := make(chan bool)

	r := mutex.TryWLockGreedy() 
	if r != true{
		t.Error("WLock should be acquired!")
	}
	
	go func(){
		r = mutex.TryWLockGreedy() 
		if r == true{
			t.Error("WLock shouldn't be acquired!")
		}
		finished <- true
	}()
	
	<-finished
	mutex.WUnlock()
	mutex.RLock()
	
	go func(){
		r = mutex.TryWLockGreedy() 
		if r == true{
			t.Error("WLock shouldn't be acquired!")
		}
		finished <- true
	}()
	<-finished
	mutex.RUnlock()
}