github.com/aclements/go-misc@v0.0.0-20240129233631-2f6ede80790c/go-weave/weave/mutex.go (about)

     1  // Copyright 2016 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package weave
     6  
     7  import "fmt"
     8  
     9  type Mutex struct {
    10  	locked  bool
    11  	waiters []*thread
    12  }
    13  
    14  func (m *Mutex) Lock() {
    15  	if !m.locked {
    16  		m.locked = true
    17  		return
    18  	}
    19  	this := globalSched.curThread
    20  	m.waiters = append(m.waiters, this)
    21  	this.block(m.reset)
    22  }
    23  
    24  func (m *Mutex) Unlock() {
    25  	if !m.locked {
    26  		panic("attempt to Unlock unlocked Mutex")
    27  	}
    28  	if len(m.waiters) == 0 {
    29  		m.locked = false
    30  	} else {
    31  		// Pick an arbitrary thread to wake up.
    32  		next := globalSched.Amb(len(m.waiters))
    33  		t := m.waiters[next]
    34  		m.waiters[next] = m.waiters[len(m.waiters)-1]
    35  		m.waiters = m.waiters[:len(m.waiters)-1]
    36  		t.unblock()
    37  	}
    38  	globalSched.Sched()
    39  }
    40  
    41  func (m *Mutex) reset() {
    42  	*m = Mutex{}
    43  }
    44  
    45  type RWMutex struct {
    46  	r, w             int
    47  	readers, writers []*thread
    48  }
    49  
    50  func (rw *RWMutex) Lock() {
    51  	if rw.r == 0 && rw.w == 0 {
    52  		rw.w++
    53  		return
    54  	}
    55  	this := globalSched.curThread
    56  	rw.writers = append(rw.writers, this)
    57  	this.block(rw.reset)
    58  }
    59  
    60  func (rw *RWMutex) RLock() {
    61  	if rw.w == 0 {
    62  		rw.r++
    63  		return
    64  	}
    65  	this := globalSched.curThread
    66  	rw.readers = append(rw.readers, this)
    67  	this.block(rw.reset)
    68  }
    69  
    70  func (rw *RWMutex) reset() {
    71  	*rw = RWMutex{}
    72  }
    73  
    74  func (rw *RWMutex) Unlock() {
    75  	rw.w--
    76  	rw.release()
    77  }
    78  
    79  func (rw *RWMutex) RUnlock() {
    80  	rw.r--
    81  	rw.release()
    82  }
    83  
    84  func (rw *RWMutex) release() {
    85  	if rw.w != 0 {
    86  		panic(fmt.Sprintf("bad RWMutex writer count: %d", rw.w))
    87  	}
    88  	if len(rw.readers) > 0 {
    89  		// Wake all readers.
    90  		rw.r += len(rw.readers)
    91  		for _, t := range rw.readers {
    92  			t.unblock()
    93  		}
    94  		rw.readers = rw.readers[:0]
    95  	} else if rw.r == 0 && len(rw.writers) > 0 {
    96  		// Wake one writer.
    97  		rw.w++
    98  		next := globalSched.Amb(len(rw.writers))
    99  		t := rw.writers[next]
   100  		rw.writers[next] = rw.writers[len(rw.writers)-1]
   101  		rw.writers = rw.writers[:len(rw.writers)-1]
   102  		t.unblock()
   103  	}
   104  	globalSched.Sched()
   105  }