ReceiverLoop should stop on cancelling context (#240)

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 #239

client/internal/wsreceiver.go

@@ -58,23 +58,38 @@ func (r *wsReceiver) IsStopped() <-chan struct{} {
 // ReceiverLoop runs the receiver loop.
 // To stop the receiver cancel the context and close the websocket connection
 func (r *wsReceiver) ReceiverLoop(ctx context.Context) {
-	processorCtx, stopProcessor := 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(processorCtx, &message)
 		}
 	}
-
-	stopProcessor()
-	close(r.stopped)
 }
 
 func (r *wsReceiver) receiveMessage(msg *protobufs.ServerToAgent) error {

client/internal/wsreceiver_test.go

@@ -212,7 +212,6 @@ func TestReceiverLoopStop(t *testing.T) {
 		receiverLoopStopped.Store(true)
 	}()
 	cancel()
-	conn.Close()
 
 	assert.Eventually(t, func() bool {
 		return receiverLoopStopped.Load()