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

     1  // Copyright 2018 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 split
     6  
     7  import (
     8  	"sync"
     9  	"sync/atomic"
    10  )
    11  
    12  // RWMutex is a scalable reader/writer mutual exclusion lock. The lock
    13  // can be held by an arbitrary number of readers or a single writer.
    14  // The zero value for a RWMutex is an unlocked mutex.
    15  //
    16  // In contrast with sync.RWMutex, this lock attempts to scale to any
    17  // number of cores simultaneously acquiring read locks. However, this
    18  // makes obtaining the lock in write mode more expensive.
    19  type RWMutex struct {
    20  	readLocks *Value
    21  	writeLock sync.Mutex
    22  	initLock  sync.Mutex
    23  	init      uint32
    24  }
    25  
    26  // doInit performs lazily initialization on the first use of m.
    27  func (m *RWMutex) doInit() {
    28  	// Acquire the initialization lock to protect against
    29  	// concurrent initialization.
    30  	m.initLock.Lock()
    31  	defer m.initLock.Unlock()
    32  	if atomic.LoadUint32(&m.init) != 0 {
    33  		// Another goroutine initialized the mutex while we
    34  		// were waiting on the shard lock.
    35  		return
    36  	}
    37  	m.readLocks = New(func(*sync.Mutex) {
    38  		// Block creating new shards while the write lock is
    39  		// held.
    40  		m.writeLock.Lock()
    41  		m.writeLock.Unlock()
    42  	})
    43  	atomic.StoreUint32(&m.init, 1)
    44  }
    45  
    46  // Lock acquires m in writer mode. This blocks all readers and
    47  // writers.
    48  func (m *RWMutex) Lock() {
    49  	if atomic.LoadUint32(&m.init) == 0 {
    50  		m.doInit()
    51  	}
    52  	// Block other writers and creation of new shards.
    53  	m.writeLock.Lock()
    54  	// Acquire all read locks.
    55  	m.readLocks.Range(func(s *sync.Mutex) {
    56  		s.Lock()
    57  	})
    58  }
    59  
    60  // Unlock releases m from writer mode. The mutex must currently be
    61  // held in writer mode.
    62  func (m *RWMutex) Unlock() {
    63  	m.readLocks.Range(func(s *sync.Mutex) {
    64  		s.Unlock()
    65  	})
    66  	m.writeLock.Unlock()
    67  }
    68  
    69  // RWMutexRUnlocker is a token used to unlock an RWMutex in read mode.
    70  type RWMutexRUnlocker struct {
    71  	shard *sync.Mutex
    72  }
    73  
    74  // RLock acquires m in read mode. This blocks other goroutines from
    75  // acquiring it in write mode, but does not generally block them from
    76  // acquiring it in read mode. The caller must used the returned
    77  // RWMutexRUnlocker to release the lock.
    78  func (m *RWMutex) RLock() RWMutexRUnlocker {
    79  	if atomic.LoadUint32(&m.init) == 0 {
    80  		m.doInit()
    81  	}
    82  	shard := m.readLocks.Get().(*sync.Mutex)
    83  	shard.Lock()
    84  	return RWMutexRUnlocker{shard}
    85  }
    86  
    87  // RUnlock releases an RWMutex from read mode.
    88  func (c RWMutexRUnlocker) RUnlock() {
    89  	c.shard.Unlock()
    90  }
    91  
    92  func Example_rwMutex() {
    93  	var m RWMutex
    94  
    95  	var wg sync.WaitGroup
    96  	for i := 0; i < 64; i++ {
    97  		wg.Add(1)
    98  		go func() {
    99  			m.RLock().RUnlock()
   100  			wg.Done()
   101  		}()
   102  	}
   103  	wg.Wait()
   104  }