Repository: open-telemetry/opamp-go
Pull request #240 (2024-01-13 12:02:53)

ReceiverLoop should stop on cancelling context
Here is test run that shows it doesn't stop without fix https://github.com/open-telemetry/opamp-go/actions/runs/7512286887/job/20452817400

Fixes https://github.com/open-telemetry/opamp-go/issues/239

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) Report
Attention: `131 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`efddaa2`)](https://app.codecov.io/gh/open-telemetry/opamp-go/commit/efddaa260895e7ebb2431deeefef3a3380d8d902?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) 76.11% compared to head [(`15eb486`)](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) 72.59%.
> Report is 43 commits behind head on main.

| [Files](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) | Patch % | Lines |
|---|---|---|
| [internal/certs.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-aW50ZXJuYWwvY2VydHMuZ28=) | 0.00% | [118 Missing :warning: ](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) |
| [client/internal/httpsender.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L2ludGVybmFsL2h0dHBzZW5kZXIuZ28=) | 88.00% | [2 Missing and 1 partial :warning: ](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) |
| [client/internal/mockserver.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L2ludGVybmFsL21vY2tzZXJ2ZXIuZ28=) | 85.71% | [2 Missing and 1 partial :warning: ](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) |
| [client/internal/wsreceiver.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L2ludGVybmFsL3dzcmVjZWl2ZXIuZ28=) | 92.59% | [2 Missing :warning: ](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) |
| [server/httpconnection.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-c2VydmVyL2h0dHBjb25uZWN0aW9uLmdv) | 0.00% | [2 Missing :warning: ](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) |
| [client/internal/clientcommon.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L2ludGVybmFsL2NsaWVudGNvbW1vbi5nbw==) | 90.90% | [1 Missing :warning: ](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) |
| [client/internal/packagessyncer.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L2ludGVybmFsL3BhY2thZ2Vzc3luY2VyLmdv) | 0.00% | [1 Missing :warning: ](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) |
| [server/serverimpl.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-c2VydmVyL3NlcnZlcmltcGwuZ28=) | 95.65% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #240      +/-   ##
==========================================
- Coverage   76.11%   72.59%   -3.53%     
==========================================
  Files          24       25       +1     
  Lines        1834     2036     +202     
==========================================
+ Hits         1396     1478      +82     
- Misses        326      448     +122     
+ Partials      112      110       -2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/240?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry).


Patch:
diff --git a/client/internal/wsreceiver.go b/client/internal/wsreceiver.go
index 30af697..b53b54d 100644
--- a/client/internal/wsreceiver.go
+++ b/client/internal/wsreceiver.go
@@ -43,22 +43,38 @@ func NewWSReceiver(
 
 // ReceiverLoop runs the receiver loop. To stop the receiver cancel the context.
 func (r *wsReceiver) ReceiverLoop(ctx context.Context) {
-	runContext, cancelFunc := context.WithCancel(ctx)
+	type receivedMessage struct {
+		message *protobufs.ServerToAgent
+		err     error
+	}
 
-out:
 	for {
-		var message protobufs.ServerToAgent
-		if err := r.receiveMessage(&message); err != nil {
-			if ctx.Err() == nil && !websocket.IsCloseError(err, websocket.CloseNormalClosure) {
-				r.logger.Errorf(ctx, "Unexpected error while receiving: %v", err)
+		select {
+		case <-ctx.Done():
+			return
+		default:
+			result := make(chan receivedMessage, 1)
+
+			go func() {
+				var message protobufs.ServerToAgent
+				err := r.receiveMessage(&message)
+				result <- receivedMessage{&message, err}
+			}()
+
+			select {
+			case <-ctx.Done():
+				return
+			case res := <-result:
+				if res.err != nil {
+					if !websocket.IsCloseError(res.err, websocket.CloseNormalClosure) {
+						r.logger.Errorf(ctx, "Unexpected error while receiving: %v", res.err)
+					}
+					return
+				}
+				r.processor.ProcessReceivedMessage(ctx, res.message)
 			}
-			break out
-		} else {
-			r.processor.ProcessReceivedMessage(runContext, &message)
 		}
 	}
-
-	cancelFunc()
 }
 
 func (r *wsReceiver) receiveMessage(msg *protobufs.ServerToAgent) error {


Test patch:
diff --git a/client/internal/wsreceiver_test.go b/client/internal/wsreceiver_test.go
index bd79c7a..b428879 100644
--- a/client/internal/wsreceiver_test.go
+++ b/client/internal/wsreceiver_test.go
@@ -3,8 +3,11 @@ package internal
 import (
 	"context"
 	"fmt"
+	"sync/atomic"
 	"testing"
+	"time"
 
+	"github.com/gorilla/websocket"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/proto"
@@ -181,3 +184,36 @@ func TestDecodeMessage(t *testing.T) {
 		}
 	}
 }
+
+func TestReceiverLoopStop(t *testing.T) {
+
+	srv := StartMockServer(t)
+
+	conn, _, err := websocket.DefaultDialer.DialContext(
+		context.Background(),
+		"ws://"+srv.Endpoint,
+		nil,
+	)
+	require.NoError(t, err)
+
+	var receiverLoopStopped atomic.Bool
+
+	callbacks := types.CallbacksStruct{}
+	clientSyncedState := ClientSyncedState{
+		remoteConfigStatus: &protobufs.RemoteConfigStatus{},
+	}
+	sender := WSSender{}
+	capabilities := protobufs.AgentCapabilities_AgentCapabilities_AcceptsRestartCommand
+	receiver := NewWSReceiver(TestLogger{t}, callbacks, conn, &sender, &clientSyncedState, nil, capabilities)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go func() {
+		receiver.ReceiverLoop(ctx)
+		receiverLoopStopped.Store(true)
+	}()
+	cancel()
+
+	assert.Eventually(t, func() bool {
+		return receiverLoopStopped.Load()
+	}, 2*time.Second, 100*time.Millisecond, "ReceiverLoop should stop when context is cancelled")
+}