fix: Add atomic.Bool to ensure wsConnection.Disconnect() is only called once

server/serverimpl.go

@@ -225,12 +225,12 @@ func (s *server) httpHandler(w http.ResponseWriter, req *http.Request) {
 }
 
 func (s *server) handleWSConnection(reqCtx context.Context, wsConn *websocket.Conn, connectionCallbacks *serverTypes.ConnectionCallbacks) {
-	agentConn := wsConnection{wsConn: wsConn, connMutex: &sync.Mutex{}}
+	agentConn := newWSConnection(wsConn)
 
 	defer func() {
 		// Close the connection when all is done.
 		defer func() {
-			err := wsConn.Close()
+			err := agentConn.Disconnect()
 			if err != nil {
 				s.logger.Errorf(context.Background(), "error closing the WebSocket connection: %v", err)
 			}

server/serverimpl_test.go

@@ -16,19 +16,26 @@ import (
 	"github.com/gorilla/websocket"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
 	"google.golang.org/protobuf/proto"
 
+	clienttypes "github.com/open-telemetry/opamp-go/client/types"
 	sharedinternal "github.com/open-telemetry/opamp-go/internal"
 	"github.com/open-telemetry/opamp-go/internal/testhelpers"
 	"github.com/open-telemetry/opamp-go/protobufs"
 	"github.com/open-telemetry/opamp-go/server/types"
 )
 
 func startServer(t *testing.T, settings *StartSettings) *server {
-	srv := New(&sharedinternal.NopLogger{})
+	return startServerWithLogger(t, settings, &sharedinternal.NopLogger{})
+
+}
+
+func startServerWithLogger(t *testing.T, settings *StartSettings, logger clienttypes.Logger) *server {
+	srv := New(logger)
 	require.NotNil(t, srv)
 	if settings.ListenEndpoint == "" {
-		// Find an avaiable port to listne on.
+		// Find an avaiable port to listen on.
 		settings.ListenEndpoint = testhelpers.GetAvailableLocalAddress()
 	}
 	if settings.ListenPath == "" {
@@ -239,7 +246,7 @@ func TestDisconnectHttpConnection(t *testing.T) {
 	assert.Equal(t, ErrInvalidHTTPConnection, err)
 }
 
-func TestDisconnectWSConnection(t *testing.T) {
+func TestDisconnectClientWSConnection(t *testing.T) {
 	connectionCloseCalled := int32(0)
 	callback := types.Callbacks{
 		OnConnecting: func(request *http.Request) types.ConnectionResponse {
@@ -263,10 +270,77 @@ func TestDisconnectWSConnection(t *testing.T) {
 	assert.NoError(t, err)
 	assert.True(t, atomic.LoadInt32(&connectionCloseCalled) == 0)
 
-	// Close connection from server side
-	srvConn := wsConnection{wsConn: conn}
-	err = srvConn.Disconnect()
+	// Close connection from client side
+	clientConn := newWSConnection(conn)
+	err = clientConn.Disconnect()
+	assert.NoError(t, err)
+
+	// Verify connection disconnected from server side
+	eventually(t, func() bool { return atomic.LoadInt32(&connectionCloseCalled) == 1 })
+	// Waiting for wsConnection to fail ReadMessage() over a Disconnected communication
+	eventually(t, func() bool {
+		_, _, err := conn.ReadMessage()
+		return err != nil
+	})
+}
+
+// testLogger is a struct that adapts a *zap.Logger to opamp-go's Logger interface.
+type testLogger struct {
+	logs []string
+}
+
+func newTestLogger() *testLogger {
+	return &testLogger{
+		logs: []string{},
+	}
+}
+
+func (o *testLogger) Debugf(_ context.Context, format string, v ...any) {
+	log := fmt.Sprintf(format, v...)
+	o.logs = append(o.logs, fmt.Sprintf("Debugf: %s\n", log))
+}
+
+func (o *testLogger) Errorf(_ context.Context, format string, v ...any) {
+	log := fmt.Sprintf(format, v...)
+	o.logs = append(o.logs, fmt.Sprintf("Errorf: %s\n", log))
+}
+
+func TestDisconnectServerWSConnection(t *testing.T) {
+	connectionCloseCalled := int32(0)
+	var serverConn types.Connection
+	connReady := make(chan struct{}) // Channel to signal when serverConn is assigned
+	callback := types.Callbacks{
+		OnConnecting: func(request *http.Request) types.ConnectionResponse {
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: types.ConnectionCallbacks{
+				OnConnected: func(ctx context.Context, conn types.Connection) {
+					serverConn = conn
+					close(connReady)
+				},
+				OnConnectionClose: func(conn types.Connection) {
+					atomic.StoreInt32(&connectionCloseCalled, 1)
+				},
+			}}
+		},
+	}
+
+	// Start a Server.
+	logger := newTestLogger()
+	settings := &StartSettings{Settings: Settings{Callbacks: callback}}
+	srv := startServerWithLogger(t, settings, logger)
+	defer srv.Stop(context.Background())
+
+	// Connect to the Server.
+	conn, _, err := dialClient(settings)
+
+	// Verify that the connection is successful.
 	assert.NoError(t, err)
+	assert.True(t, atomic.LoadInt32(&connectionCloseCalled) == 0)
+
+	// Wait for serverConn to be assigned
+	<-connReady
+
+	// Close connection from server side
+	serverConn.Disconnect()
 
 	// Verify connection disconnected from server side
 	eventually(t, func() bool { return atomic.LoadInt32(&connectionCloseCalled) == 1 })
@@ -275,6 +349,10 @@ func TestDisconnectWSConnection(t *testing.T) {
 		_, _, err := conn.ReadMessage()
 		return err != nil
 	})
+
+	require.Equal(t, 1, len(logger.logs))
+	require.Contains(t, logger.logs[0], "Errorf: Cannot read a message from WebSocket: read tcp")
+	require.Contains(t, logger.logs[0], "use of closed network connection")
 }
 
 var testInstanceUid = []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6}

server/wsconnection.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"net"
 	"sync"
+	"sync/atomic"
 
 	"github.com/gorilla/websocket"
 
@@ -17,23 +18,31 @@ type wsConnection struct {
 	// The websocket library does not allow multiple concurrent write operations,
 	// so ensure that we only have a single operation in progress at a time.
 	// For more: https://pkg.go.dev/github.com/gorilla/websocket#hdr-Concurrency
-	connMutex *sync.Mutex
+	connMutex sync.Mutex
 	wsConn    *websocket.Conn
+	closed    atomic.Bool
 }
 
 var _ types.Connection = (*wsConnection)(nil)
 
-func (c wsConnection) Connection() net.Conn {
+func newWSConnection(wsConn *websocket.Conn) types.Connection {
+	return &wsConnection{wsConn: wsConn, connMutex: sync.Mutex{}, closed: atomic.Bool{}}
+}
+
+func (c *wsConnection) Connection() net.Conn {
 	return c.wsConn.UnderlyingConn()
 }
 
-func (c wsConnection) Send(_ context.Context, message *protobufs.ServerToAgent) error {
+func (c *wsConnection) Send(_ context.Context, message *protobufs.ServerToAgent) error {
 	c.connMutex.Lock()
 	defer c.connMutex.Unlock()
 
 	return internal.WriteWSMessage(c.wsConn, message)
 }
 
-func (c wsConnection) Disconnect() error {
+func (c *wsConnection) Disconnect() error {
+	if !c.closed.CompareAndSwap(false, true) {
+		return nil
+	}
 	return c.wsConn.Close()
 }