2

Author: babyTsakhes

connection.go

@@ -22,12 +22,6 @@ import (
 
 const requestsMap = 128
 const ignoreStreamId = 0
-const (
-	connDisconnected = 0
-	connConnected    = 1
-	connShutdown     = 2
-	connClosed       = 3
-)
 
 const shutdownEventKey = "box.shutdown"
 
@@ -51,6 +45,32 @@ const (
 	Closed
 )
 
+// ConnectionState represents the state of a connection.
+type ConnectionState uint32
+
+const (
+	StateDisconnected ConnectionState = iota
+	StateConnected
+	StateShutdown
+	StateClosed
+)
+
+// String implements fmt.Stringer.
+func (s ConnectionState) String() string {
+	switch s {
+	case StateDisconnected:
+		return "disconnected"
+	case StateConnected:
+		return "connected"
+	case StateShutdown:
+		return "shutdown"
+	case StateClosed:
+		return "closed"
+	default:
+		return "unknown"
+	}
+}
+
 // ConnEvent is sent throw Notify channel specified in Opts.
 type ConnEvent struct {
 	Conn *Connection
@@ -355,30 +375,14 @@ func (conn *Connection) logEvent(event LogEvent) {
 	}
 }
 
-func (conn *Connection) stateToString() string {
-	state := atomic.LoadUint32(&conn.state)
-	switch state {
-	case connDisconnected:
-		return "disconnected"
-	case connConnected:
-		return "connected"
-	case connShutdown:
-		return "shutdown"
-	case connClosed:
-		return "closed"
-	default:
-		return "unknown"
-	}
-}
-
 // ConnectedNow reports if connection is established at the moment.
 func (conn *Connection) ConnectedNow() bool {
-	return atomic.LoadUint32(&conn.state) == connConnected
+	return ConnectionState(atomic.LoadUint32(&conn.state)) == StateConnected
 }
 
 // ClosedNow reports if connection is closed by user or after reconnect.
 func (conn *Connection) ClosedNow() bool {
-	return atomic.LoadUint32(&conn.state) == connClosed
+	return ConnectionState(atomic.LoadUint32(&conn.state)) == StateClosed
 }
 
 // Close closes Connection.
@@ -468,7 +472,7 @@ func (conn *Connection) dial(ctx context.Context) error {
 	// Only if connected and fully initialized.
 	conn.lockShards()
 	conn.c = c
-	atomic.StoreUint32(&conn.state, connConnected)
+	atomic.StoreUint32(&conn.state, uint32(StateConnected))
 	conn.cond.Broadcast()
 	conn.unlockShards()
 	go conn.writer(c, c)
@@ -542,7 +546,7 @@ func pack(h *smallWBuf, enc *msgpack.Encoder, reqid uint32,
 
 func (conn *Connection) connect(ctx context.Context) error {
 	var err error
-	if conn.c == nil && conn.state == connDisconnected {
+	if conn.c == nil && ConnectionState(atomic.LoadUint32(&conn.state)) == StateDisconnected {
 		if err = conn.dial(ctx); err == nil {
 
 			conn.logEvent(ConnectedEvent{
@@ -552,7 +556,7 @@ func (conn *Connection) connect(ctx context.Context) error {
 			return nil
 		}
 	}
-	if conn.state == connClosed {
+	if ConnectionState(atomic.LoadUint32(&conn.state)) == StateClosed {
 		err = ClientError{ErrConnectionClosed, "using closed connection"}
 	}
 	return err
@@ -563,9 +567,9 @@ func (conn *Connection) closeConnection(neterr error, forever bool) (err error)
 	defer conn.unlockShards()
 
 	if forever {
-		if conn.state != connClosed {
+		if ConnectionState(atomic.LoadUint32(&conn.state)) != StateClosed {
 			close(conn.control)
-			atomic.StoreUint32(&conn.state, connClosed)
+			atomic.StoreUint32(&conn.state, uint32(StateClosed))
 			conn.cond.Broadcast()
 			// Free the resources.
 			if conn.shutdownWatcher != nil {
@@ -578,7 +582,7 @@ func (conn *Connection) closeConnection(neterr error, forever bool) (err error)
 			conn.notify(Closed)
 		}
 	} else {
-		atomic.StoreUint32(&conn.state, connDisconnected)
+		atomic.StoreUint32(&conn.state, uint32(StateDisconnected))
 		conn.cond.Broadcast()
 		conn.logEvent(DisconnectedEvent{
 			baseEvent: newBaseEvent(conn.addr),
@@ -736,7 +740,7 @@ func (conn *Connection) notify(kind ConnEventKind) {
 func (conn *Connection) writer(w writeFlusher, c Conn) {
 	var shardn uint32
 	var packet smallWBuf
-	for atomic.LoadUint32(&conn.state) != connClosed {
+	for ConnectionState(atomic.LoadUint32(&conn.state)) != StateClosed {
 		select {
 		case shardn = <-conn.dirtyShard:
 		default:
@@ -829,7 +833,7 @@ func (conn *Connection) reader(r io.Reader, c Conn) {
 
 	go conn.eventer(events)
 
-	for atomic.LoadUint32(&conn.state) != connClosed {
+	for ConnectionState(atomic.LoadUint32(&conn.state)) != StateClosed {
 		respBytes, err := read(r, conn.lenbuf[:])
 		if err != nil {
 			err = ClientError{
@@ -932,24 +936,24 @@ func (conn *Connection) newFuture(req Request) (fut *future) {
 	shardn := fut.requestId & (conn.opts.Concurrency - 1)
 	shard := &conn.shard[shardn]
 	shard.rmut.Lock()
-	switch atomic.LoadUint32(&conn.state) {
-	case connClosed:
+	switch ConnectionState(atomic.LoadUint32(&conn.state)) {
+	case StateClosed:
 		fut.err = ClientError{
 			ErrConnectionClosed,
 			"using closed connection",
 		}
 		fut.finish()
 		shard.rmut.Unlock()
 		return
-	case connDisconnected:
+	case StateDisconnected:
 		fut.err = ClientError{
 			ErrConnectionNotReady,
 			"client connection is not ready",
 		}
 		fut.finish()
 		shard.rmut.Unlock()
 		return
-	case connShutdown:
+	case StateShutdown:
 		fut.err = ClientError{
 			ErrConnectionShutdown,
 			"server shutdown in progress",
@@ -1545,7 +1549,7 @@ func (conn *Connection) shutdown(forever bool) error {
 	conn.mutex.Lock()
 	defer conn.mutex.Unlock()
 
-	if !atomic.CompareAndSwapUint32(&conn.state, connConnected, connShutdown) {
+	if !atomic.CompareAndSwapUint32(&conn.state, uint32(StateConnected), uint32(StateShutdown)) {
 		if forever {
 			err := ClientError{ErrConnectionClosed, "connection closed by client"}
 			return conn.closeConnection(err, true)
@@ -1569,7 +1573,7 @@ func (conn *Connection) shutdown(forever bool) error {
 
 	c := conn.c
 	for {
-		if (atomic.LoadUint32(&conn.state) != connShutdown) || (c != conn.c) {
+		if (ConnectionState(atomic.LoadUint32(&conn.state)) != StateShutdown) || (c != conn.c) {
 			return nil
 		}
 		if atomic.LoadInt64(&conn.requestCnt) == 0 {

logger.go

@@ -3,6 +3,7 @@ package tarantool
 import (
 	"context"
 	"log/slog"
+	"sync/atomic"
 )
 
 type Logger interface {
@@ -46,7 +47,8 @@ func (l SlogLogger) Report(event LogEvent, conn *Connection) {
 		}
 
 		if !keys["connection_state"] {
-			attrs = append(attrs, slog.String("connection_state", conn.stateToString()))
+			state := ConnectionState(atomic.LoadUint32(&conn.state))
+			attrs = append(attrs, slog.String("connection_state", state.String()))
 		}
 
 		if conn.opts.MaxReconnects > 0 && !keys["max_reconnects"] {