Align atomically-accessed integer fields correctly (#723)

Uses Int64 instead of Int32 for atomically accessed integer fields and places those fields at the beginning of their encasing structs. This avoids alignment-based panics on 32-bit architectures.

Author: xoac

x/mongo/driver/topology/connection.go

@@ -39,6 +39,11 @@ var (
 func nextConnectionID() uint64 { return atomic.AddUint64(&globalConnectionID, 1) }
 
 type connection struct {
+	// connected must be accessed using the atomic package and should be at the beginning of the struct.
+	// - atomic bug: https://pkg.go.dev/sync/atomic#pkg-note-BUG
+	// - suggested layout: https://go101.org/article/memory-layout.html
+	connected int64
+
 	id                   string
 	nc                   net.Conn // When nil, the connection is closed.
 	addr                 address.Address
@@ -52,7 +57,6 @@ type connection struct {
 	compressor           wiremessage.CompressorID
 	zliblevel            int
 	zstdLevel            int
-	connected            int32 // must be accessed using the sync/atomic package
 	connectDone          chan struct{}
 	connectErr           error
 	config               *connectionConfig
@@ -97,13 +101,13 @@ func newConnection(addr address.Address, opts ...ConnectionOption) (*connection,
 	if !c.config.loadBalanced {
 		c.setGenerationNumber()
 	}
-	atomic.StoreInt32(&c.connected, initialized)
+	atomic.StoreInt64(&c.connected, initialized)
 
 	return c, nil
 }
 
 func (c *connection) processInitializationError(opCtx context.Context, err error) {
-	atomic.StoreInt32(&c.connected, disconnected)
+	atomic.StoreInt64(&c.connected, disconnected)
 	if c.nc != nil {
 		_ = c.nc.Close()
 	}
@@ -138,7 +142,7 @@ func (c *connection) hasGenerationNumber() bool {
 // connect handles the I/O for a connection. It will dial, configure TLS, and perform
 // initialization handshakes.
 func (c *connection) connect(ctx context.Context) {
-	if !atomic.CompareAndSwapInt32(&c.connected, initialized, connected) {
+	if !atomic.CompareAndSwapInt64(&c.connected, initialized, connected) {
 		return
 	}
 	defer close(c.connectDone)
@@ -345,7 +349,7 @@ func (c *connection) cancellationListenerCallback() {
 
 func (c *connection) writeWireMessage(ctx context.Context, wm []byte) error {
 	var err error
-	if atomic.LoadInt32(&c.connected) != connected {
+	if atomic.LoadInt64(&c.connected) != connected {
 		return ConnectionError{ConnectionID: c.id, message: "connection is closed"}
 	}
 	select {
@@ -402,7 +406,7 @@ func (c *connection) write(ctx context.Context, wm []byte) (err error) {
 
 // readWireMessage reads a wiremessage from the connection. The dst parameter will be overwritten.
 func (c *connection) readWireMessage(ctx context.Context, dst []byte) ([]byte, error) {
-	if atomic.LoadInt32(&c.connected) != connected {
+	if atomic.LoadInt64(&c.connected) != connected {
 		return dst, ConnectionError{ConnectionID: c.id, message: "connection is closed"}
 	}
 
@@ -505,7 +509,7 @@ func (c *connection) read(ctx context.Context, dst []byte) (bytesRead []byte, er
 
 func (c *connection) close() error {
 	// Overwrite the connection state as the first step so only the first close call will execute.
-	if !atomic.CompareAndSwapInt32(&c.connected, connected, disconnected) {
+	if !atomic.CompareAndSwapInt64(&c.connected, connected, disconnected) {
 		return nil
 	}
 
@@ -518,7 +522,7 @@ func (c *connection) close() error {
 }
 
 func (c *connection) closed() bool {
-	return atomic.LoadInt32(&c.connected) == disconnected
+	return atomic.LoadInt64(&c.connected) == disconnected
 }
 
 func (c *connection) idleTimeoutExpired() bool {

x/mongo/driver/topology/connection_test.go

@@ -81,7 +81,7 @@ func TestConnection(t *testing.T) {
 				if !cmp.Equal(got, want, cmp.Comparer(compareErrors)) {
 					t.Errorf("errors do not match. got %v; want %v", got, want)
 				}
-				connState := atomic.LoadInt32(&conn.connected)
+				connState := atomic.LoadInt64(&conn.connected)
 				assert.Equal(t, disconnected, connState, "expected connection state %v, got %v", disconnected, connState)
 			})
 			t.Run("handshaker error", func(t *testing.T) {
@@ -109,7 +109,7 @@ func TestConnection(t *testing.T) {
 				if !cmp.Equal(got, want, cmp.Comparer(compareErrors)) {
 					t.Errorf("errors do not match. got %v; want %v", got, want)
 				}
-				connState := atomic.LoadInt32(&conn.connected)
+				connState := atomic.LoadInt64(&conn.connected)
 				assert.Equal(t, disconnected, connState, "expected connection state %v, got %v", disconnected, connState)
 			})
 			t.Run("calls error callback", func(t *testing.T) {
@@ -753,7 +753,7 @@ func TestConnection(t *testing.T) {
 				conn.connect(context.Background())
 				err = conn.wait()
 				assert.NotNil(t, err, "expected handshake error from wait, got nil")
-				connState := atomic.LoadInt32(&conn.connected)
+				connState := atomic.LoadInt64(&conn.connected)
 				assert.Equal(t, disconnected, connState, "expected connection state %v, got %v", disconnected, connState)
 
 				err = conn.close()

x/mongo/driver/topology/pool.go

@@ -58,21 +58,23 @@ type checkOutResult struct {
 
 // pool is a wrapper of resource pool that follows the CMAP spec for connection pools
 type pool struct {
-	address    address.Address
-	opts       []ConnectionOption
-	conns      *resourcePool // pool for non-checked out connections
-	generation *poolGenerationMap
-	monitor    *event.PoolMonitor
-
-	// Must be accessed using the atomic package.
-	connected                    int32
+	// These fields must be accessed using the atomic package and should be at the beginning of the struct.
+	// - atomic bug: https://pkg.go.dev/sync/atomic#pkg-note-BUG
+	// - suggested layout: https://go101.org/article/memory-layout.html
+	connected                    int64
 	pinnedCursorConnections      uint64
 	pinnedTransactionConnections uint64
 
 	nextid uint64
 	opened map[uint64]*connection // opened holds all of the currently open connections.
 	sem    *semaphore.Weighted
 	sync.Mutex
+
+	address    address.Address
+	opts       []ConnectionOption
+	conns      *resourcePool // pool for non-checked out connections
+	generation *poolGenerationMap
+	monitor    *event.PoolMonitor
 }
 
 // connectionExpiredFunc checks if a given connection is stale and should be removed from the resource pool
@@ -87,7 +89,7 @@ func connectionExpiredFunc(v interface{}) bool {
 	}
 
 	switch {
-	case atomic.LoadInt32(&c.pool.connected) != connected:
+	case atomic.LoadInt64(&c.pool.connected) != connected:
 		c.expireReason = event.ReasonPoolClosed
 	case c.closed():
 		// A connection would only be closed if it encountered a network error during an operation and closed itself.
@@ -208,7 +210,7 @@ func (p *pool) stale(c *connection) bool {
 
 // connect puts the pool into the connected state, allowing it to be used and will allow items to begin being processed from the wait queue
 func (p *pool) connect() error {
-	if !atomic.CompareAndSwapInt32(&p.connected, disconnected, connected) {
+	if !atomic.CompareAndSwapInt64(&p.connected, disconnected, connected) {
 		return ErrPoolConnected
 	}
 	p.generation.connect()
@@ -218,7 +220,7 @@ func (p *pool) connect() error {
 
 // disconnect disconnects the pool and closes all connections including those both in and out of the pool
 func (p *pool) disconnect(ctx context.Context) error {
-	if !atomic.CompareAndSwapInt32(&p.connected, connected, disconnecting) {
+	if !atomic.CompareAndSwapInt64(&p.connected, connected, disconnecting) {
 		return ErrPoolDisconnected
 	}
 
@@ -267,7 +269,7 @@ func (p *pool) disconnect(ctx context.Context) error {
 		_ = p.removeConnection(pc, event.ReasonPoolClosed)
 		_ = p.closeConnection(pc) // We don't care about errors while closing the connection.
 	}
-	atomic.StoreInt32(&p.connected, disconnected)
+	atomic.StoreInt64(&p.connected, disconnected)
 	p.conns.clearTotal()
 
 	if p.monitor != nil {
@@ -301,7 +303,7 @@ func (p *pool) makeNewConnection() (*connection, string, error) {
 		})
 	}
 
-	if atomic.LoadInt32(&p.connected) != connected {
+	if atomic.LoadInt64(&p.connected) != connected {
 		// Manually publish a ConnectionClosed event here because the connection reference hasn't been stored and we
 		// need to ensure each ConnectionCreated event has a corresponding ConnectionClosed event.
 		if p.monitor != nil {
@@ -348,7 +350,7 @@ func (p *pool) get(ctx context.Context) (*connection, error) {
 		ctx = context.Background()
 	}
 
-	if atomic.LoadInt32(&p.connected) != connected {
+	if atomic.LoadInt64(&p.connected) != connected {
 		if p.monitor != nil {
 			p.monitor.Event(&event.PoolEvent{
 				Type:    event.GetFailed,
@@ -380,7 +382,7 @@ func (p *pool) get(ctx context.Context) (*connection, error) {
 	// This loop is so that we don't end up with more than maxPoolSize connections if p.conns.Maintain runs between
 	// calling p.conns.Get() and making the new connection
 	for {
-		if atomic.LoadInt32(&p.connected) != connected {
+		if atomic.LoadInt64(&p.connected) != connected {
 			if p.monitor != nil {
 				p.monitor.Event(&event.PoolEvent{
 					Type:    event.GetFailed,
@@ -395,7 +397,7 @@ func (p *pool) get(ctx context.Context) (*connection, error) {
 		connVal := p.conns.Get()
 		if c, ok := connVal.(*connection); ok && connVal != nil {
 			// call connect if not connected
-			if atomic.LoadInt32(&c.connected) == initialized {
+			if atomic.LoadInt64(&c.connected) == initialized {
 				c.connect(ctx)
 			}
 
@@ -499,12 +501,12 @@ func (p *pool) closeConnection(c *connection) error {
 		return ErrWrongPool
 	}
 
-	if atomic.LoadInt32(&c.connected) == connected {
+	if atomic.LoadInt64(&c.connected) == connected {
 		c.closeConnectContext()
 		_ = c.wait() // Make sure that the connection has finished connecting
 	}
 
-	if !atomic.CompareAndSwapInt32(&c.connected, connected, disconnected) {
+	if !atomic.CompareAndSwapInt64(&c.connected, connected, disconnected) {
 		return nil // We're closing an already closed connection
 	}
 

x/mongo/driver/topology/pool_generation_counter.go

@@ -24,8 +24,10 @@ type generationStats struct {
 // load balancer, there is only one service ID: primitive.NilObjectID. For load-balanced deployments, each server behind
 // the load balancer will have a unique service ID.
 type poolGenerationMap struct {
-	// state must be accessed using the atomic package.
-	state         int32
+	// state must be accessed using the atomic package and should be at the beginning of the struct.
+	// - atomic bug: https://pkg.go.dev/sync/atomic#pkg-note-BUG
+	// - suggested layout: https://go101.org/article/memory-layout.html
+	state         int64
 	generationMap map[primitive.ObjectID]*generationStats
 
 	sync.Mutex
@@ -40,11 +42,11 @@ func newPoolGenerationMap() *poolGenerationMap {
 }
 
 func (p *poolGenerationMap) connect() {
-	atomic.StoreInt32(&p.state, connected)
+	atomic.StoreInt64(&p.state, connected)
 }
 
 func (p *poolGenerationMap) disconnect() {
-	atomic.StoreInt32(&p.state, disconnected)
+	atomic.StoreInt64(&p.state, disconnected)
 }
 
 // addConnection increments the connection count for the generation associated with the given service ID and returns the
@@ -100,7 +102,7 @@ func (p *poolGenerationMap) clear(serviceIDPtr *primitive.ObjectID) {
 
 func (p *poolGenerationMap) stale(serviceIDPtr *primitive.ObjectID, knownGeneration uint64) bool {
 	// If the map has been disconnected, all connections should be considered stale to ensure that they're closed.
-	if atomic.LoadInt32(&p.state) == disconnected {
+	if atomic.LoadInt64(&p.state) == disconnected {
 		return true
 	}
 

x/mongo/driver/topology/pool_test.go

@@ -193,7 +193,7 @@ func TestPool(t *testing.T) {
 				t.Errorf("Should have closed 1 connections, but didn't. got %d; want %d", d.lenclosed(), 1)
 			}
 			close(cleanup)
-			state := atomic.LoadInt32(&p.connected)
+			state := atomic.LoadInt64(&p.connected)
 			if state != disconnected {
 				t.Errorf("Should have set the connection state on return. got %d; want %d", state, disconnected)
 			}
@@ -250,10 +250,10 @@ func TestPool(t *testing.T) {
 	})
 	t.Run("connect", func(t *testing.T) {
 		t.Run("can reconnect a disconnected pool", func(t *testing.T) {
-			assertGenerationMapState := func(t *testing.T, p *pool, expectedState int32) {
+			assertGenerationMapState := func(t *testing.T, p *pool, expectedState int64) {
 				t.Helper()
 
-				actualState := atomic.LoadInt32(&p.generation.state)
+				actualState := atomic.LoadInt64(&p.generation.state)
 				assert.Equal(t, expectedState, actualState, "expected generation map state %d, got %d", expectedState, actualState)
 			}
 
@@ -296,7 +296,7 @@ func TestPool(t *testing.T) {
 				t.Errorf("Pool should have 0 total connections. got %d; want %d", p.conns.totalSize, 0)
 			}
 			close(cleanup)
-			state := atomic.LoadInt32(&p.connected)
+			state := atomic.LoadInt64(&p.connected)
 			if state != disconnected {
 				t.Errorf("Should have set the connection state on return. got %d; want %d", state, disconnected)
 			}

x/mongo/driver/topology/server.go

@@ -56,14 +56,14 @@ func (ss *SelectedServer) Description() description.SelectedServer {
 
 // These constants represent the connection states of a server.
 const (
-	disconnected int32 = iota
+	disconnected int64 = iota
 	disconnecting
 	connected
 	connecting
 	initialized
 )
 
-func connectionStateString(state int32) string {
+func connectionStateString(state int64) string {
 	switch state {
 	case 0:
 		return "Disconnected"
@@ -84,7 +84,7 @@ func connectionStateString(state int32) string {
 type Server struct {
 	cfg             *serverConfig
 	address         address.Address
-	connectionstate int32
+	connectionstate int64
 
 	// connection related fields
 	pool *pool
@@ -195,7 +195,7 @@ func NewServer(addr address.Address, topologyID primitive.ObjectID, opts ...Serv
 // Connect initializes the Server by starting background monitoring goroutines.
 // This method must be called before a Server can be used.
 func (s *Server) Connect(updateCallback updateTopologyCallback) error {
-	if !atomic.CompareAndSwapInt32(&s.connectionstate, disconnected, connected) {
+	if !atomic.CompareAndSwapInt64(&s.connectionstate, disconnected, connected) {
 		return ErrServerConnected
 	}
 
@@ -225,7 +225,7 @@ func (s *Server) Connect(updateCallback updateTopologyCallback) error {
 // any in flight read or write operations. If this method returns with no
 // errors, all connections associated with this Server have been closed.
 func (s *Server) Disconnect(ctx context.Context) error {
-	if !atomic.CompareAndSwapInt32(&s.connectionstate, connected, disconnecting) {
+	if !atomic.CompareAndSwapInt64(&s.connectionstate, connected, disconnecting) {
 		return ErrServerClosed
 	}
 
@@ -246,7 +246,7 @@ func (s *Server) Disconnect(ctx context.Context) error {
 	}
 
 	s.closewg.Wait()
-	atomic.StoreInt32(&s.connectionstate, disconnected)
+	atomic.StoreInt64(&s.connectionstate, disconnected)
 
 	return nil
 }
@@ -261,7 +261,7 @@ func (s *Server) Connection(ctx context.Context) (driver.Connection, error) {
 		})
 	}
 
-	if atomic.LoadInt32(&s.connectionstate) != connected {
+	if atomic.LoadInt64(&s.connectionstate) != connected {
 		return nil, ErrServerClosed
 	}
 
@@ -379,7 +379,7 @@ func (s *Server) SelectedDescription() description.SelectedServer {
 // updated server descriptions will be sent. The channel will have a buffer
 // size of one, and will be pre-populated with the current description.
 func (s *Server) Subscribe() (*ServerSubscription, error) {
-	if atomic.LoadInt32(&s.connectionstate) != connected {
+	if atomic.LoadInt64(&s.connectionstate) != connected {
 		return nil, ErrSubscribeAfterClosed
 	}
 	ch := make(chan description.Server, 1)
@@ -577,7 +577,7 @@ func (s *Server) update() {
 		// Perform the next check.
 		desc, err := s.check()
 		if err == errCheckCancelled {
-			if atomic.LoadInt32(&s.connectionstate) != connected {
+			if atomic.LoadInt64(&s.connectionstate) != connected {
 				continue
 			}
 
@@ -843,7 +843,7 @@ func extractTopologyVersion(err error) *description.TopologyVersion {
 // String implements the Stringer interface.
 func (s *Server) String() string {
 	desc := s.Description()
-	connState := atomic.LoadInt32(&s.connectionstate)
+	connState := atomic.LoadInt64(&s.connectionstate)
 	str := fmt.Sprintf("Addr: %s, Type: %s, State: %s",
 		s.address, desc.Kind, connectionStateString(connState))
 	if len(desc.Tags) != 0 {