Repository: open-telemetry/opamp-go
Pull request #200 (2023-09-15 15:54:06)

Embed mutex in the server's websocket client
The Gorilla websocket package states that [only one concurrent write](https://pkg.go.dev/github.com/gorilla/websocket#hdr-Concurrency) can be made to a connection at a time. The library will try to detect concurrent writes and will panic if it detects that one is happening.

We might see concurrent writes in two places:
1. If the server application doesn't guard writes when writing to an agent connection.
2. If the OpAMP server library is [responding to a message](https://github.com/open-telemetry/opamp-go/blob/146b79221c76637368e7937b50a82689733f46a2/server/serverimpl.go#L233) from the agent while the server application is also sending a message using its reference to the agent connection obtained from `OnConnectedFunc`.

In my opinion it would make sense to prevent concurrent writes within the library to relieve the application from having to worry about this detail. This does necessitate a mutex per connection, but I don't see any significant performance impact from the mutexes.

Without the mutex in `wsConnection`:
```
go test -bench=. -run=^# -count=5
goos: linux
goarch: amd64
pkg: github.com/open-telemetry/opamp-go/server
cpu: 12th Gen Intel(R) Core(TM) i9-12900
BenchmarkSendToClient-24           14610            118877 ns/op
BenchmarkSendToClient-24           17510            326918 ns/op
BenchmarkSendToClient-24           16641            267099 ns/op
BenchmarkSendToClient-24           16287            239085 ns/op
BenchmarkSendToClient-24           16243            238157 ns/op
PASS
ok      github.com/open-telemetry/opamp-go/server       23.877s
```

With the mutex:
```
go test -bench=. -run=^# -count=5
goos: linux
goarch: amd64
pkg: github.com/open-telemetry/opamp-go/server
cpu: 12th Gen Intel(R) Core(TM) i9-12900
BenchmarkSendToClient-24           15170            175727 ns/op
BenchmarkSendToClient-24           16916            307997 ns/op
BenchmarkSendToClient-24           16446            249592 ns/op
BenchmarkSendToClient-24           16563            260181 ns/op
BenchmarkSendToClient-24           17136            315825 ns/op
PASS
ok      github.com/open-telemetry/opamp-go/server       25.690s
```

This is a fairly rudimentary solution, but it does solve the issue I was seeing and is straightforward. I'm open to more sophisticated ways of managing the write concurrency if there's a deficiency with this approach.


Patch:
diff --git a/go.mod b/go.mod
index 134b96a..93ae25c 100644
--- a/go.mod
+++ b/go.mod
@@ -1,6 +1,6 @@
 module github.com/open-telemetry/opamp-go
 
-go 1.17
+go 1.20
 
 require (
 	github.com/cenkalti/backoff/v4 v4.1.2
diff --git a/internal/examples/go.mod b/internal/examples/go.mod
index f79cc48..d3cfb4f 100644
--- a/internal/examples/go.mod
+++ b/internal/examples/go.mod
@@ -1,12 +1,14 @@
 module github.com/open-telemetry/opamp-go/internal/examples
 
-go 1.17
+go 1.20
 
 require (
+	github.com/cenkalti/backoff/v4 v4.1.2
 	github.com/knadh/koanf v1.3.3
 	github.com/oklog/ulid/v2 v2.0.2
 	github.com/open-telemetry/opamp-go v0.1.0
 	github.com/shirou/gopsutil v3.21.11+incompatible
+	github.com/stretchr/testify v1.7.0
 	go.opentelemetry.io/otel v1.3.0
 	go.opentelemetry.io/otel/exporters/otlp/otlpmetric v0.26.0
 	go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp v0.26.0
@@ -17,7 +19,6 @@ require (
 )
 
 require (
-	github.com/cenkalti/backoff/v4 v4.1.2 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fsnotify/fsnotify v1.4.9 // indirect
 	github.com/go-logr/logr v1.2.1 // indirect
@@ -30,7 +31,6 @@ require (
 	github.com/mitchellh/mapstructure v1.4.1 // indirect
 	github.com/mitchellh/reflectwalk v1.0.2 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/stretchr/testify v1.7.0 // indirect
 	github.com/tklauser/go-sysconf v0.3.9 // indirect
 	github.com/tklauser/numcpus v0.3.0 // indirect
 	github.com/yusufpapurcu/wmi v1.2.2 // indirect
diff --git a/internal/examples/server/data/agent.go b/internal/examples/server/data/agent.go
index e920f03..b562bd3 100644
--- a/internal/examples/server/data/agent.go
+++ b/internal/examples/server/data/agent.go
@@ -25,8 +25,6 @@ type Agent struct {
 
 	// Connection to the Agent.
 	conn types.Connection
-	// Mutex to protect Send() operation.
-	connMutex sync.Mutex
 
 	// mutex for the fields that follow it.
 	mux sync.RWMutex
@@ -421,9 +419,6 @@ func (agent *Agent) calcConnectionSettings(response *protobufs.ServerToAgent) {
 }
 
 func (agent *Agent) SendToAgent(msg *protobufs.ServerToAgent) {
-	agent.connMutex.Lock()
-	defer agent.connMutex.Unlock()
-
 	agent.conn.Send(context.Background(), msg)
 }
 
diff --git a/internal/tools/go.mod b/internal/tools/go.mod
index 0384b31..2554bd8 100644
--- a/internal/tools/go.mod
+++ b/internal/tools/go.mod
@@ -1,6 +1,6 @@
 module github.com/open-telemetry/opamp-go/internal/tools
 
-go 1.17
+go 1.20
 
 require github.com/ory/go-acc v0.2.8
 
diff --git a/server/serverimpl.go b/server/serverimpl.go
index 5fc7cac..cd19d68 100644
--- a/server/serverimpl.go
+++ b/server/serverimpl.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"sync"
 
 	"github.com/gorilla/websocket"
 	"google.golang.org/protobuf/proto"
@@ -179,7 +180,7 @@ func (s *server) httpHandler(w http.ResponseWriter, req *http.Request) {
 }
 
 func (s *server) handleWSConnection(wsConn *websocket.Conn, connectionCallbacks serverTypes.ConnectionCallbacks) {
-	agentConn := wsConnection{wsConn: wsConn}
+	agentConn := wsConnection{wsConn: wsConn, connMutex: &sync.Mutex{}}
 
 	defer func() {
 		// Close the connection when all is done.
diff --git a/server/wsconnection.go b/server/wsconnection.go
index 9ce53a4..b2b9921 100644
--- a/server/wsconnection.go
+++ b/server/wsconnection.go
@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"net"
+	"sync"
 
 	"github.com/gorilla/websocket"
 
@@ -13,7 +14,11 @@ import (
 
 // wsConnection represents a persistent OpAMP connection over a WebSocket.
 type wsConnection struct {
-	wsConn *websocket.Conn
+	// 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
+	wsConn    *websocket.Conn
 }
 
 var _ types.Connection = (*wsConnection)(nil)
@@ -22,10 +27,10 @@ func (c wsConnection) Connection() net.Conn {
 	return c.wsConn.UnderlyingConn()
 }
 
-// Message header is currently uint64 zero value.
-const wsMsgHeader = uint64(0)
-
 func (c wsConnection) Send(_ context.Context, message *protobufs.ServerToAgent) error {
+	c.connMutex.Lock()
+	defer c.connMutex.Unlock()
+
 	return internal.WriteWSMessage(c.wsConn, message)
 }
 


Test patch:
diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
index 546e056..c45ba20 100644
--- a/.github/workflows/build-and-test.yml
+++ b/.github/workflows/build-and-test.yml
@@ -15,7 +15,7 @@ jobs:
       - name: Setup Go
         uses: actions/setup-go@v2.1.4
         with:
-          go-version: 1.17
+          go-version: "1.20"
       - name: Setup Go Environment
         run: |
           echo "GOPATH=$(go env GOPATH)" >> $GITHUB_ENV
@@ -48,7 +48,7 @@ jobs:
       - name: Setup Go
         uses: actions/setup-go@v2.1.4
         with:
-          go-version: 1.17
+          go-version: "1.20"
       - name: Setup Go Environment
         run: |
           echo "GOPATH=$(go env GOPATH)" >> $GITHUB_ENV
@@ -76,7 +76,7 @@ jobs:
       - name: Setup Go
         uses: actions/setup-go@v3
         with:
-          go-version: 1.17
+          go-version: "1.20"
       - name: Setup Go Environment
         run: |
           echo "GOPATH=$(go env GOPATH)" >> $GITHUB_ENV
diff --git a/server/serverimpl_test.go b/server/serverimpl_test.go
index 7ebaaa3..04827d3 100644
--- a/server/serverimpl_test.go
+++ b/server/serverimpl_test.go
@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -723,3 +724,124 @@ func TestDecodeMessage(t *testing.T) {
 		}
 	}
 }
+
+func TestConnectionAllowsConcurrentWrites(t *testing.T) {
+	srvConnVal := atomic.Value{}
+	callbacks := CallbacksStruct{
+		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnConnectedFunc: func(conn types.Connection) {
+					srvConnVal.Store(conn)
+				},
+			}}
+		},
+	}
+
+	// Start a Server.
+	settings := &StartSettings{Settings: Settings{Callbacks: callbacks}}
+	srv := startServer(t, settings)
+	defer srv.Stop(context.Background())
+
+	// Connect to the Server.
+	conn, _, err := dialClient(settings)
+
+	// Verify that the connection is successful.
+	assert.NoError(t, err)
+	assert.NotNil(t, conn)
+
+	defer conn.Close()
+
+	timeout, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+
+	select {
+	case <-timeout.Done():
+		t.Error("Client failed to connect before timeout")
+	default:
+		if _, ok := srvConnVal.Load().(types.Connection); ok == true {
+			break
+		}
+	}
+
+	cancel()
+
+	srvConn := srvConnVal.Load().(types.Connection)
+	for i := 0; i < 20; i++ {
+		go func() {
+			defer func() {
+				if recover() != nil {
+					require.Fail(t, "Sending to client panicked")
+				}
+			}()
+
+			srvConn.Send(context.Background(), &protobufs.ServerToAgent{})
+		}()
+	}
+}
+
+func BenchmarkSendToClient(b *testing.B) {
+	clientConnections := []*websocket.Conn{}
+	serverConnections := []types.Connection{}
+	srvConnectionsMutex := sync.Mutex{}
+	callbacks := CallbacksStruct{
+		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnConnectedFunc: func(conn types.Connection) {
+					srvConnectionsMutex.Lock()
+					serverConnections = append(serverConnections, conn)
+					srvConnectionsMutex.Unlock()
+				},
+			}}
+		},
+	}
+
+	// Start a Server.
+	settings := &StartSettings{
+		Settings:       Settings{Callbacks: callbacks},
+		ListenEndpoint: testhelpers.GetAvailableLocalAddress(),
+		ListenPath:     "/",
+	}
+	srv := New(&sharedinternal.NopLogger{})
+	err := srv.Start(*settings)
+
+	if err != nil {
+		b.Error(err)
+	}
+
+	defer srv.Stop(context.Background())
+
+	for i := 0; i < b.N; i++ {
+		conn, resp, err := dialClient(settings)
+
+		if err != nil || resp == nil || conn == nil {
+			b.Error("Could not establish connection:", err)
+		}
+
+		clientConnections = append(clientConnections, conn)
+	}
+
+	timeout, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+
+	select {
+	case <-timeout.Done():
+		b.Error("Connections failed to establish in time")
+	default:
+		if len(serverConnections) == b.N {
+			break
+		}
+	}
+
+	cancel()
+
+	for _, conn := range serverConnections {
+		err := conn.Send(context.Background(), &protobufs.ServerToAgent{})
+
+		if err != nil {
+			b.Error(err)
+		}
+	}
+
+	for _, conn := range clientConnections {
+		conn.Close()
+	}
+
+}