chore: Use atomics where appropriate

This commit is contained in:
Francis Lavoie 2026-04-13 00:37:11 -04:00
parent 0722cf6fd8
commit cef419186f
No known key found for this signature in database
5 changed files with 40 additions and 36 deletions

View File

@ -766,7 +766,7 @@ func Validate(cfg *Config) error {
// code is emitted.
func exitProcess(ctx context.Context, logger *zap.Logger) {
// let the rest of the program know we're quitting; only do it once
if !atomic.CompareAndSwapInt32(exiting, 0, 1) {
if !exiting.CompareAndSwap(false, true) {
return
}
@ -845,11 +845,11 @@ func exitProcess(ctx context.Context, logger *zap.Logger) {
}()
}
var exiting = new(int32) // accessed atomically
var exiting atomic.Bool
// Exiting returns true if the process is exiting.
// EXPERIMENTAL API: subject to change or removal.
func Exiting() bool { return atomic.LoadInt32(exiting) == 1 }
func Exiting() bool { return exiting.Load() }
// OnExit registers a callback to invoke during process exit.
// This registration is PROCESS-GLOBAL, meaning that each

View File

@ -120,8 +120,8 @@ func listenReusable(ctx context.Context, lnKey string, network, address string,
// re-wrapped in a new fakeCloseListener each time the listener
// is reused. This type is atomic and values must not be copied.
type fakeCloseListener struct {
closed int32 // accessed atomically; belongs to this struct only
*sharedListener // embedded, so we also become a net.Listener
closed atomic.Bool
*sharedListener // embedded, so we also become a net.Listener
keepAliveConfig net.KeepAliveConfig
}
@ -131,7 +131,7 @@ type canSetKeepAliveConfig interface {
func (fcl *fakeCloseListener) Accept() (net.Conn, error) {
// if the listener is already "closed", return error
if atomic.LoadInt32(&fcl.closed) == 1 {
if fcl.closed.Load() {
return nil, fakeClosedErr(fcl)
}
@ -155,7 +155,7 @@ func (fcl *fakeCloseListener) Accept() (net.Conn, error) {
// that we set when Close() was called, and return a non-temporary and
// non-timeout error value to the caller, masking the "true" error, so
// that server loops / goroutines won't retry, linger, and leak
if atomic.LoadInt32(&fcl.closed) == 1 {
if fcl.closed.Load() {
// we dereference the sharedListener explicitly even though it's embedded
// so that it's clear in the code that side-effects are shared with other
// users of this listener, not just our own reference to it; we also don't
@ -175,7 +175,7 @@ func (fcl *fakeCloseListener) Accept() (net.Conn, error) {
// underlying listener. The underlying listener is only closed
// if the caller is the last known user of the socket.
func (fcl *fakeCloseListener) Close() error {
if atomic.CompareAndSwapInt32(&fcl.closed, 0, 1) {
if fcl.closed.CompareAndSwap(false, true) {
// There are two ways I know of to get an Accept()
// function to return to the server loop that called
// it: close the listener, or set a deadline in the
@ -238,13 +238,13 @@ func (sl *sharedListener) Destruct() error {
// fakeClosePacketConn is like fakeCloseListener, but for PacketConns,
// or more specifically, *net.UDPConn
type fakeClosePacketConn struct {
closed int32 // accessed atomically; belongs to this struct only
*sharedPacketConn // embedded, so we also become a net.PacketConn; its key is used in Close
closed atomic.Bool
*sharedPacketConn // embedded, so we also become a net.PacketConn; its key is used in Close
}
func (fcpc *fakeClosePacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
// if the listener is already "closed", return error
if atomic.LoadInt32(&fcpc.closed) == 1 {
if fcpc.closed.Load() {
return 0, nil, &net.OpError{
Op: "readfrom",
Net: fcpc.LocalAddr().Network(),
@ -258,7 +258,7 @@ func (fcpc *fakeClosePacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err e
if err != nil {
// this server was stopped, so clear the deadline and let
// any new server continue reading; but we will exit
if atomic.LoadInt32(&fcpc.closed) == 1 {
if fcpc.closed.Load() {
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
if err = fcpc.SetReadDeadline(time.Time{}); err != nil {
return n, addr, err
@ -273,7 +273,7 @@ func (fcpc *fakeClosePacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err e
// Close won't close the underlying socket unless there is no more reference, then listenerPool will close it.
func (fcpc *fakeClosePacketConn) Close() error {
if atomic.CompareAndSwapInt32(&fcpc.closed, 0, 1) {
if fcpc.closed.CompareAndSwap(false, true) {
_ = fcpc.SetReadDeadline(time.Now()) // unblock ReadFrom() calls to kick old servers out of their loops
_, _ = listenerPool.Delete(fcpc.sharedPacketConn.key)
}

View File

@ -63,7 +63,7 @@ func reuseUnixSocket(network, addr string) (any, error) {
if err != nil {
return nil, err
}
atomic.AddInt32(unixSocket.count, 1)
unixSocket.count.Add(1)
unixSockets[socketKey] = &unixListener{ln.(*net.UnixListener), socketKey, unixSocket.count}
case *unixConn:
@ -71,7 +71,7 @@ func reuseUnixSocket(network, addr string) (any, error) {
if err != nil {
return nil, err
}
atomic.AddInt32(unixSocket.count, 1)
unixSocket.count.Add(1)
unixSockets[socketKey] = &unixConn{pc.(*net.UnixConn), socketKey, unixSocket.count}
}
@ -165,8 +165,9 @@ func listenReusable(ctx context.Context, lnKey string, network, address string,
if !fd {
// TODO: Not 100% sure this is necessary, but we do this for net.UnixListener, so...
if unix, ok := ln.(*net.UnixConn); ok {
one := int32(1)
ln = &unixConn{unix, lnKey, &one}
cnt := new(atomic.Int32)
cnt.Store(1)
ln = &unixConn{unix, lnKey, cnt}
unixSockets[lnKey] = ln.(*unixConn)
}
}
@ -181,8 +182,9 @@ func listenReusable(ctx context.Context, lnKey string, network, address string,
// (we do our own "unlink on close" -- not required, but more tidy)
if unix, ok := ln.(*net.UnixListener); ok {
unix.SetUnlinkOnClose(false)
one := int32(1)
ln = &unixListener{unix, lnKey, &one}
cnt := new(atomic.Int32)
cnt.Store(1)
ln = &unixListener{unix, lnKey, cnt}
unixSockets[lnKey] = ln.(*unixListener)
}
}
@ -216,11 +218,11 @@ func reusePort(network, address string, conn syscall.RawConn) error {
type unixListener struct {
*net.UnixListener
mapKey string
count *int32 // accessed atomically
count *atomic.Int32
}
func (uln *unixListener) Close() error {
newCount := atomic.AddInt32(uln.count, -1)
newCount := uln.count.Add(-1)
if newCount == 0 {
file, err := uln.File()
var name string
@ -242,11 +244,11 @@ func (uln *unixListener) Close() error {
type unixConn struct {
*net.UnixConn
mapKey string
count *int32 // accessed atomically
count *atomic.Int32
}
func (uc *unixConn) Close() error {
newCount := atomic.AddInt32(uc.count, -1)
newCount := uc.count.Add(-1)
if newCount == 0 {
file, err := uc.File()
var name string

View File

@ -611,8 +611,8 @@ func fakeClosedErr(l interface{ Addr() net.Addr }) error {
var errFakeClosed = fmt.Errorf("QUIC listener 'closed' 😉")
type fakeCloseQuicListener struct {
closed int32 // accessed atomically; belongs to this struct only
*sharedQuicListener // embedded, so we also become a quic.EarlyListener
closed atomic.Int32
*sharedQuicListener // embedded, so we also become a quic.EarlyListener
context context.Context
contextCancel context.CancelCauseFunc
}
@ -629,16 +629,16 @@ func (fcql *fakeCloseQuicListener) Accept(_ context.Context) (*quic.Conn, error)
}
// if the listener is "closed", return a fake closed error instead
if atomic.LoadInt32(&fcql.closed) == 1 && errors.Is(err, context.Canceled) {
if fcql.closed.Load() == 1 && errors.Is(err, context.Canceled) {
return nil, fakeClosedErr(fcql)
}
return nil, err
}
func (fcql *fakeCloseQuicListener) Close() error {
if atomic.CompareAndSwapInt32(&fcql.closed, 0, 1) {
if fcql.closed.CompareAndSwap(0, 1) {
fcql.contextCancel(errFakeClosed)
} else if atomic.CompareAndSwapInt32(&fcql.closed, 1, 2) {
} else if fcql.closed.CompareAndSwap(1, 2) {
_, _ = listenerPool.Delete(fcql.sharedQuicListener.key)
}
return nil

View File

@ -79,14 +79,15 @@ func (up *UsagePool) LoadOrNew(key any, construct Constructor) (value any, loade
up.Lock()
upv, loaded = up.pool[key]
if loaded {
atomic.AddInt32(&upv.refs, 1)
upv.refs.Add(1)
up.Unlock()
upv.RLock()
value = upv.value
err = upv.err
upv.RUnlock()
} else {
upv = &usagePoolVal{refs: 1}
upv = &usagePoolVal{}
upv.refs.Store(1)
upv.Lock()
up.pool[key] = upv
up.Unlock()
@ -118,7 +119,7 @@ func (up *UsagePool) LoadOrStore(key, val any) (value any, loaded bool) {
up.Lock()
upv, loaded = up.pool[key]
if loaded {
atomic.AddInt32(&upv.refs, 1)
upv.refs.Add(1)
up.Unlock()
upv.Lock()
if upv.err == nil {
@ -129,7 +130,8 @@ func (up *UsagePool) LoadOrStore(key, val any) (value any, loaded bool) {
}
upv.Unlock()
} else {
upv = &usagePoolVal{refs: 1, value: val}
upv = &usagePoolVal{value: val}
upv.refs.Store(1)
up.pool[key] = upv
up.Unlock()
value = val
@ -173,7 +175,7 @@ func (up *UsagePool) Delete(key any) (deleted bool, err error) {
up.Unlock()
return false, nil
}
refs := atomic.AddInt32(&upv.refs, -1)
refs := upv.refs.Add(-1)
if refs == 0 {
delete(up.pool, key)
up.Unlock()
@ -188,7 +190,7 @@ func (up *UsagePool) Delete(key any) (deleted bool, err error) {
up.Unlock()
if refs < 0 {
panic(fmt.Sprintf("deleted more than stored: %#v (usage: %d)",
upv.value, upv.refs))
upv.value, upv.refs.Load()))
}
}
return deleted, err
@ -203,7 +205,7 @@ func (up *UsagePool) References(key any) (int, bool) {
if loaded {
// I wonder if it'd be safer to read this value during
// our lock on the UsagePool... guess we'll see...
refs := atomic.LoadInt32(&upv.refs)
refs := upv.refs.Load()
return int(refs), true
}
return 0, false
@ -220,7 +222,7 @@ type Destructor interface {
}
type usagePoolVal struct {
refs int32 // accessed atomically; must be 64-bit aligned for 32-bit systems
refs atomic.Int32
value any
err error
sync.RWMutex