refactor: replace atomic operations with atomic types

Author: alexandear

db_test.go

@@ -717,9 +717,9 @@ var _ = Describe("CopyFrom/CopyTo", func() {
 		Expect(res.RowsAffected()).To(Equal(n))
 
 		st := db.PoolStats()
-		Expect(st.Hits).To(Equal(uint32(4)))
-		Expect(st.Misses).To(Equal(uint32(1)))
-		Expect(st.Timeouts).To(Equal(uint32(0)))
+		Expect(st.Hits.Load()).To(Equal(uint32(4)))
+		Expect(st.Misses.Load()).To(Equal(uint32(1)))
+		Expect(st.Timeouts.Load()).To(Equal(uint32(0)))
 		Expect(st.TotalConns).To(Equal(uint32(1)))
 		Expect(st.IdleConns).To(Equal(uint32(1)))
 
@@ -736,9 +736,9 @@ var _ = Describe("CopyFrom/CopyTo", func() {
 		Expect(res).To(BeNil())
 
 		st := db.Pool().Stats()
-		Expect(st.Hits).To(Equal(uint32(3)))
-		Expect(st.Misses).To(Equal(uint32(1)))
-		Expect(st.Timeouts).To(Equal(uint32(0)))
+		Expect(st.Hits.Load()).To(Equal(uint32(3)))
+		Expect(st.Misses.Load()).To(Equal(uint32(1)))
+		Expect(st.Timeouts.Load()).To(Equal(uint32(0)))
 		Expect(st.TotalConns).To(Equal(uint32(1)))
 		Expect(st.IdleConns).To(Equal(uint32(1)))
 

internal/pool/conn.go

@@ -20,7 +20,7 @@ type Conn struct {
 	lastID    int64
 
 	createdAt time.Time
-	usedAt    uint32 // atomic
+	usedAt    atomic.Uint32
 	pooled    bool
 	Inited    bool
 }
@@ -36,12 +36,12 @@ func NewConn(netConn net.Conn, pool *ConnPool) *Conn {
 }
 
 func (cn *Conn) UsedAt() time.Time {
-	unix := atomic.LoadUint32(&cn.usedAt)
+	unix := cn.usedAt.Load()
 	return time.Unix(int64(unix), 0)
 }
 
 func (cn *Conn) SetUsedAt(tm time.Time) {
-	atomic.StoreUint32(&cn.usedAt, uint32(tm.Unix()))
+	cn.usedAt.Store(uint32(tm.Unix()))
 }
 
 func (cn *Conn) RemoteAddr() net.Addr {

internal/pool/pool.go

@@ -26,13 +26,13 @@ var timers = sync.Pool{
 
 // Stats contains pool state information and accumulated stats.
 type Stats struct {
-	Hits     uint32 // number of times free connection was found in the pool
-	Misses   uint32 // number of times free connection was NOT found in the pool
-	Timeouts uint32 // number of times a wait timeout occurred
+	Hits     atomic.Uint32 // number of times free connection was found in the pool
+	Misses   atomic.Uint32 // number of times free connection was NOT found in the pool
+	Timeouts atomic.Uint32 // number of times a wait timeout occurred
 
-	TotalConns uint32 // number of total connections in the pool
-	IdleConns  uint32 // number of idle connections in the pool
-	StaleConns uint32 // number of stale connections removed from the pool
+	TotalConns uint32        // number of total connections in the pool
+	IdleConns  uint32        // number of idle connections in the pool
+	StaleConns atomic.Uint32 // number of stale connections removed from the pool
 }
 
 type Pooler interface {
@@ -71,9 +71,9 @@ type Options struct {
 type ConnPool struct {
 	opt *Options
 
-	dialErrorsNum uint32 // atomic
+	dialErrorsNum atomic.Uint32
 
-	_closed uint32 // atomic
+	_closed atomic.Bool
 
 	lastDialErrorMu sync.RWMutex
 	lastDialError   error
@@ -188,14 +188,14 @@ func (p *ConnPool) dialConn(c context.Context, pooled bool) (*Conn, error) {
 		return nil, ErrClosed
 	}
 
-	if atomic.LoadUint32(&p.dialErrorsNum) >= uint32(p.opt.PoolSize) {
+	if p.dialErrorsNum.Load() >= uint32(p.opt.PoolSize) {
 		return nil, p.getLastDialError()
 	}
 
 	netConn, err := p.opt.Dialer(c)
 	if err != nil {
 		p.setLastDialError(err)
-		if atomic.AddUint32(&p.dialErrorsNum, 1) == uint32(p.opt.PoolSize) {
+		if p.dialErrorsNum.Add(1) == uint32(p.opt.PoolSize) {
 			go p.tryDial()
 		}
 		return nil, err
@@ -219,7 +219,7 @@ func (p *ConnPool) tryDial() {
 			continue
 		}
 
-		atomic.StoreUint32(&p.dialErrorsNum, 0)
+		p.dialErrorsNum.Store(0)
 		_ = conn.Close()
 		return
 	}
@@ -263,11 +263,11 @@ func (p *ConnPool) Get(ctx context.Context) (*Conn, error) {
 			continue
 		}
 
-		atomic.AddUint32(&p.stats.Hits, 1)
+		p.stats.Hits.Add(1)
 		return cn, nil
 	}
 
-	atomic.AddUint32(&p.stats.Misses, 1)
+	p.stats.Misses.Add(1)
 
 	newcn, err := p.newConn(ctx, true)
 	if err != nil {
@@ -313,7 +313,7 @@ func (p *ConnPool) waitTurn(c context.Context) error {
 		return nil
 	case <-timer.C:
 		timers.Put(timer)
-		atomic.AddUint32(&p.stats.Timeouts, 1)
+		p.stats.Timeouts.Add(1)
 		return ErrPoolTimeout
 	}
 }
@@ -402,20 +402,19 @@ func (p *ConnPool) IdleLen() int {
 }
 
 func (p *ConnPool) Stats() *Stats {
-	idleLen := p.IdleLen()
-	return &Stats{
-		Hits:     atomic.LoadUint32(&p.stats.Hits),
-		Misses:   atomic.LoadUint32(&p.stats.Misses),
-		Timeouts: atomic.LoadUint32(&p.stats.Timeouts),
-
+	stats := &Stats{
 		TotalConns: uint32(p.Len()),
-		IdleConns:  uint32(idleLen),
-		StaleConns: atomic.LoadUint32(&p.stats.StaleConns),
+		IdleConns:  uint32(p.IdleLen()),
 	}
+	stats.Hits.Store(p.stats.Hits.Load())
+	stats.Misses.Store(p.stats.Misses.Load())
+	stats.Timeouts.Store(p.stats.Timeouts.Load())
+	stats.StaleConns.Store(p.stats.StaleConns.Load())
+	return stats
 }
 
 func (p *ConnPool) closed() bool {
-	return atomic.LoadUint32(&p._closed) == 1
+	return p._closed.Load()
 }
 
 func (p *ConnPool) Filter(fn func(*Conn) bool) error {
@@ -433,7 +432,7 @@ func (p *ConnPool) Filter(fn func(*Conn) bool) error {
 }
 
 func (p *ConnPool) Close() error {
-	if !atomic.CompareAndSwapUint32(&p._closed, 0, 1) {
+	if !p._closed.CompareAndSwap(false, true) {
 		return ErrClosed
 	}
 
@@ -466,7 +465,7 @@ func (p *ConnPool) reaper(frequency time.Duration) {
 			internal.Logger.Printf(context.TODO(), "ReapStaleConns failed: %s", err)
 			continue
 		}
-		atomic.AddUint32(&p.stats.StaleConns, uint32(n))
+		p.stats.StaleConns.Add(uint32(n))
 	}
 }
 

internal/pool/pool_sticky.go

@@ -35,9 +35,9 @@ func (e BadConnError) Unwrap() error {
 
 type StickyConnPool struct {
 	pool   Pooler
-	shared int32 // atomic
+	shared atomic.Int32
 
-	state uint32 // atomic
+	state atomic.Uint32
 	ch    chan *Conn
 
 	_badConnError atomic.Value
@@ -53,7 +53,7 @@ func NewStickyConnPool(pool Pooler) *StickyConnPool {
 			ch:   make(chan *Conn, 1),
 		}
 	}
-	atomic.AddInt32(&p.shared, 1)
+	p.shared.Add(1)
 	return p
 }
 
@@ -68,13 +68,13 @@ func (p *StickyConnPool) CloseConn(cn *Conn) error {
 func (p *StickyConnPool) Get(ctx context.Context) (*Conn, error) {
 	// In worst case this races with Close which is not a very common operation.
 	for i := 0; i < 1000; i++ {
-		switch atomic.LoadUint32(&p.state) {
+		switch p.state.Load() {
 		case stateDefault:
 			cn, err := p.pool.Get(ctx)
 			if err != nil {
 				return nil, err
 			}
-			if atomic.CompareAndSwapUint32(&p.state, stateDefault, stateInited) {
+			if p.state.CompareAndSwap(stateDefault, stateInited) {
 				return cn, nil
 			}
 			p.pool.Remove(ctx, cn, ErrClosed)
@@ -124,16 +124,16 @@ func (p *StickyConnPool) Remove(ctx context.Context, cn *Conn, reason error) {
 }
 
 func (p *StickyConnPool) Close() error {
-	if shared := atomic.AddInt32(&p.shared, -1); shared > 0 {
+	if shared := p.shared.Add(-1); shared > 0 {
 		return nil
 	}
 
 	for i := 0; i < 1000; i++ {
-		state := atomic.LoadUint32(&p.state)
+		state := p.state.Load()
 		if state == stateClosed {
 			return ErrClosed
 		}
-		if atomic.CompareAndSwapUint32(&p.state, state, stateClosed) {
+		if p.state.CompareAndSwap(state, stateClosed) {
 			close(p.ch)
 			cn, ok := <-p.ch
 			if ok {
@@ -162,8 +162,8 @@ func (p *StickyConnPool) Reset(ctx context.Context) error {
 		return errors.New("pg: StickyConnPool does not have a Conn")
 	}
 
-	if !atomic.CompareAndSwapUint32(&p.state, stateInited, stateDefault) {
-		state := atomic.LoadUint32(&p.state)
+	if !p.state.CompareAndSwap(stateInited, stateDefault) {
+		state := p.state.Load()
 		return fmt.Errorf("pg: invalid StickyConnPool state: %d", state)
 	}
 
@@ -181,7 +181,7 @@ func (p *StickyConnPool) badConnError() error {
 }
 
 func (p *StickyConnPool) Len() int {
-	switch atomic.LoadUint32(&p.state) {
+	switch p.state.Load() {
 	case stateDefault:
 		return 0
 	case stateInited:

listener_test.go

@@ -43,9 +43,9 @@ var _ = Context("Listener", func() {
 		}
 
 		st := db.PoolStats()
-		Expect(st.Hits).To(Equal(uint32(0)))
-		Expect(st.Misses).To(Equal(uint32(0)))
-		Expect(st.Timeouts).To(Equal(uint32(0)))
+		Expect(st.Hits.Load()).To(Equal(uint32(0)))
+		Expect(st.Misses.Load()).To(Equal(uint32(0)))
+		Expect(st.Timeouts.Load()).To(Equal(uint32(0)))
 		Expect(st.TotalConns).To(Equal(uint32(1)))
 		Expect(st.IdleConns).To(Equal(uint32(0)))
 	})

tx.go

@@ -33,7 +33,7 @@ type Tx struct {
 	stmtsMu sync.Mutex
 	stmts   []*Stmt
 
-	_closed int32
+	_closed atomic.Bool
 }
 
 var _ orm.DB = (*Tx)(nil)
@@ -368,7 +368,7 @@ func (tx *Tx) CloseContext(ctx context.Context) error {
 }
 
 func (tx *Tx) close() {
-	if !atomic.CompareAndSwapInt32(&tx._closed, 0, 1) {
+	if !tx._closed.CompareAndSwap(false, true) {
 		return
 	}
 
@@ -384,5 +384,5 @@ func (tx *Tx) close() {
 }
 
 func (tx *Tx) closed() bool {
-	return atomic.LoadInt32(&tx._closed) == 1
+	return tx._closed.Load()
 }