closemu.go

     1  // Copyright 2026 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 sql
     6  
     7  import (
     8  	"sync"
     9  	"sync/atomic"
    10  )
    11  
    12  // A closingMutex is an RWMutex for synchronizing close.
    13  // Unlike a sync.RWMutex, RLock takes priority over Lock.
    14  // Reads can starve out close, but reads are safely reentrant.
    15  type closingMutex struct {
    16  	// state is 2*readers+writerWaiting.
    17  	//   0 is unlocked
    18  	//   1 is unlocked and a writer needs to wake
    19  	//   >0 is read-locked
    20  	//   <0 is write-locked
    21  	state atomic.Int64
    22  	mu    sync.Mutex
    23  	read  *sync.Cond
    24  	write *sync.Cond
    25  }
    26  
    27  func (m *closingMutex) RLock() {
    28  	if m.TryRLock() {
    29  		return
    30  	}
    31  
    32  	// Wait for writer.
    33  	m.mu.Lock()
    34  	defer m.mu.Unlock()
    35  	for {
    36  		if m.TryRLock() {
    37  			return
    38  		}
    39  		m.init()
    40  		m.read.Wait()
    41  	}
    42  }
    43  
    44  func (m *closingMutex) RUnlock() {
    45  	for {
    46  		x := m.state.Load()
    47  		if x < 2 {
    48  			panic("runlock of un-rlocked mutex")
    49  		}
    50  		if m.state.CompareAndSwap(x, x-2) {
    51  			if x-2 == 1 {
    52  				// We were the last reader, and a writer is waiting.
    53  				// The lock makes sure the writer sees the broadcast.
    54  				m.mu.Lock()
    55  				defer m.mu.Unlock()
    56  				m.write.Broadcast()
    57  			}
    58  			return
    59  		}
    60  	}
    61  }
    62  
    63  func (m *closingMutex) Lock() {
    64  	m.mu.Lock()
    65  	defer m.mu.Unlock()
    66  	for {
    67  		x := m.state.Load()
    68  		if x == 0 || x == 1 {
    69  			if m.state.CompareAndSwap(x, -1) {
    70  				return
    71  			}
    72  			continue
    73  		}
    74  		// Set writer waiting bit and sleep.
    75  		if x&1 == 0 && !m.state.CompareAndSwap(x, x|1) {
    76  			continue
    77  		}
    78  		m.init()
    79  		m.write.Wait()
    80  	}
    81  }
    82  
    83  func (m *closingMutex) Unlock() {
    84  	m.mu.Lock()
    85  	defer m.mu.Unlock()
    86  	if !m.state.CompareAndSwap(-1, 0) {
    87  		panic("unlock of unlocked mutex")
    88  	}
    89  	if m.read != nil {
    90  		m.read.Broadcast()
    91  		m.write.Broadcast()
    92  	}
    93  }
    94  
    95  func (m *closingMutex) TryRLock() bool {
    96  	for {
    97  		x := m.state.Load()
    98  		// Fail if the mutex is write locked (x < 0)
    99  		// or unlocked with a writer waiting (x == 1).
   100  		//
   101  		// If the mutex is read-locked (x > 1), try to add a reader
   102  		// even if this starves out a waiting writer.
   103  		if x < 0 || x == 1 {
   104  			return false
   105  		}
   106  		if m.state.CompareAndSwap(x, x+2) {
   107  			return true
   108  		}
   109  	}
   110  }
   111  
   112  func (m *closingMutex) init() {
   113  	// Lazily create the read/write Conds.
   114  	// In the common, uncontended case, we'll never need them.
   115  	if m.read == nil {
   116  		m.read = sync.NewCond(&m.mu)
   117  		m.write = sync.NewCond(&m.mu)
   118  	}
   119  }
   120
View as plain text