go: Make socketConn.Close thread-safe

Client: go

We used to rely on setting the connection inside TSocket/TSSLSocket as
nil after Close is called to mark the connection as closed, but that is
not thread safe and causing TSocket.Close/TSSLSocket.Close cannot be
called concurrently. Use an atomic int to mark closure instead.

lib/go/thrift/socket.go

@@ -194,15 +194,7 @@ func (p *TSocket) IsOpen() bool {
 
 // Closes the socket.
 func (p *TSocket) Close() error {
-	// Close the socket
-	if p.conn != nil {
-		err := p.conn.Close()
-		if err != nil {
-			return err
-		}
-		p.conn = nil
-	}
-	return nil
+	return p.conn.Close()
 }
 
 //Returns the remote address of the socket.

lib/go/thrift/socket_conn.go

@@ -21,13 +21,15 @@ package thrift
 
 import (
 	"net"
+	"sync/atomic"
 )
 
 // socketConn is a wrapped net.Conn that tries to do connectivity check.
 type socketConn struct {
 	net.Conn
 
 	buffer [1]byte
+	closed int32
 }
 
 var _ net.Conn = (*socketConn)(nil)
@@ -64,7 +66,7 @@ func wrapSocketConn(conn net.Conn) *socketConn {
 // It's the same as the previous implementation of TSocket.IsOpen and
 // TSSLSocket.IsOpen before we added connectivity check.
 func (sc *socketConn) isValid() bool {
-	return sc != nil && sc.Conn != nil
+	return sc != nil && sc.Conn != nil && atomic.LoadInt32(&sc.closed) == 0
 }
 
 // IsOpen checks whether the connection is open.
@@ -100,3 +102,12 @@ func (sc *socketConn) Read(p []byte) (n int, err error) {
 
 	return sc.Conn.Read(p)
 }
+
+func (sc *socketConn) Close() error {
+	if !sc.isValid() {
+		// Already closed
+		return net.ErrClosed
+	}
+	atomic.StoreInt32(&sc.closed, 1)
+	return sc.Conn.Close()
+}

lib/go/thrift/ssl_socket.go

@@ -220,15 +220,7 @@ func (p *TSSLSocket) IsOpen() bool {
 
 // Closes the socket.
 func (p *TSSLSocket) Close() error {
-	// Close the socket
-	if p.conn != nil {
-		err := p.conn.Close()
-		if err != nil {
-			return err
-		}
-		p.conn = nil
-	}
-	return nil
+	return p.conn.Close()
 }
 
 func (p *TSSLSocket) Read(buf []byte) (int, error) {