github.com/egonelbre/exp@v0.0.0-20240430123955-ed1d3aa93911/sync2/biasedmutex.go (about)

     1  package sync2
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  // BiasedMutex is a RWMutex that allows tuning progress making on both
     8  // writing and reading side. Tries to let through multiple readers and then
     9  // multiple writers.
    10  //
    11  // BiasedMutex controls the flow by keeping track on how many reads / writes
    12  // there have been. Once there have been RLock()-s called over `ReaderThreshold`,
    13  // then RLock will wait in cased there are any writes pending.
    14  // BiasedMutex will continue letting through at most `WriterThreshold` writers.
    15  // Once writers have been exhausted the variables are reset.
    16  type BiasedMutex struct {
    17  	readerThreshold int32
    18  	writerThreshold int32
    19  
    20  	mu      sync.Mutex
    21  	readers sync.Cond
    22  	writers sync.Cond
    23  
    24  	reading        int32
    25  	writing        int32
    26  	writersWaiting int32
    27  
    28  	readCount  int32
    29  	writeCount int32
    30  }
    31  
    32  func NewBiasedMutex() *BiasedMutex {
    33  	m := &BiasedMutex{}
    34  	m.readers.L = &m.mu
    35  	m.writers.L = &m.mu
    36  
    37  	m.readerThreshold = 4
    38  	m.writerThreshold = 4
    39  
    40  	return m
    41  }
    42  
    43  func (m *BiasedMutex) SetReaderThreshold(x int) { m.readerThreshold = int32(x) }
    44  func (m *BiasedMutex) SetWriterThreshold(x int) { m.writerThreshold = int32(x) }
    45  
    46  func (m *BiasedMutex) RLock() {
    47  	m.mu.Lock()
    48  	for m.writing+m.writersWaiting > 0 {
    49  		if m.writing == 0 && m.readCount < m.readerThreshold {
    50  			break
    51  		}
    52  		m.readers.Wait()
    53  	}
    54  	m.readCount++
    55  	m.reading++
    56  	m.mu.Unlock()
    57  }
    58  
    59  func (m *BiasedMutex) RUnlock() {
    60  	m.mu.Lock()
    61  	m.reading--
    62  	if m.reading == 0 && m.writersWaiting > 0 {
    63  		m.writers.Signal()
    64  	}
    65  	m.mu.Unlock()
    66  }
    67  
    68  func (m *BiasedMutex) Lock() {
    69  	m.mu.Lock()
    70  	m.writersWaiting++
    71  	// wait for other readers or writers to finish
    72  	for m.reading > 0 || m.writing > 0 {
    73  		m.writers.Wait()
    74  	}
    75  	// mark us as writing
    76  	m.writersWaiting--
    77  	m.writing++
    78  	m.mu.Unlock()
    79  }
    80  
    81  func (m *BiasedMutex) Unlock() {
    82  	m.mu.Lock()
    83  	m.writing--
    84  	m.writeCount++
    85  	if m.writeCount >= m.writerThreshold {
    86  		m.readers.Broadcast()
    87  		m.readCount = 0
    88  		m.writeCount = 0
    89  	}
    90  	if m.writersWaiting > 0 {
    91  		m.writers.Signal()
    92  	} else {
    93  		m.readCount = 0
    94  		m.readers.Broadcast()
    95  	}
    96  	m.mu.Unlock()
    97  }