Disable scheduleSend when onMessage callback is running

client/internal/httpsender_test.go

@@ -13,6 +13,39 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestDelaySchedule(t *testing.T) {
+	sender := NewHTTPSender(&sharedinternal.NopLogger{})
+	pendingMessageChan := sender.hasPendingMessage
+	scheduleSendDelayChan := sender.registerScheduleSend
+	sender.DisableScheduleSend()
+
+	// Verify ScheduleSend is not writing to message channel when disabled
+	sender.ScheduleSend()
+	assert.Equal(t, 0, len(pendingMessageChan))
+	assert.Equal(t, 1, len(scheduleSendDelayChan))
+
+	// Repeat process to verify non-blocking and no change in channel length
+	sender.ScheduleSend()
+	assert.Equal(t, 0, len(pendingMessageChan))
+	assert.Equal(t, 1, len(scheduleSendDelayChan))
+
+	// Verify ScheduleSend is writing to message channel when enabled
+	sender.EnableScheduleSend()
+	assert.Equal(t, 1, len(pendingMessageChan))
+	assert.Equal(t, 0, len(scheduleSendDelayChan))
+
+	// Repeat process to verify non-blocking and no change in channel length
+	sender.EnableScheduleSend()
+	assert.Equal(t, 1, len(pendingMessageChan))
+	assert.Equal(t, 0, len(scheduleSendDelayChan))
+
+	// ScheduleSend sanity check after enabling
+	sender.ScheduleSend()
+	assert.Equal(t, 1, len(pendingMessageChan))
+	assert.Equal(t, 0, len(scheduleSendDelayChan))
+
+}
+
 func TestHTTPSenderRetryForStatusTooManyRequests(t *testing.T) {
 
 	var connectionAttempts int64

client/internal/receivedprocessor.go

@@ -50,6 +50,9 @@ func newReceivedProcessor(
 // the received message and performs any processing necessary based on what fields are set.
 // This function will call any relevant callbacks.
 func (r *receivedProcessor) ProcessReceivedMessage(ctx context.Context, msg *protobufs.ServerToAgent) {
+	r.sender.DisableScheduleSend()
+	defer r.sender.EnableScheduleSend()
+
 	if r.callbacks != nil {
 		if msg.Command != nil {
 			r.rcvCommand(msg.Command)
@@ -127,9 +130,9 @@ func (r *receivedProcessor) ProcessReceivedMessage(ctx context.Context, msg *pro
 				msgData.AgentIdentification = msg.AgentIdentification
 			}
 		}
-
 		r.callbacks.OnMessage(ctx, msgData)
 
+		r.sender.EnableScheduleSend()
 		r.rcvOpampConnectionSettings(ctx, msg.ConnectionSettings)
 
 		if scheduled {

client/internal/sender.go

@@ -2,6 +2,7 @@ package internal
 
 import (
 	"errors"
+	"sync/atomic"
 
 	"github.com/oklog/ulid/v2"
 	"github.com/open-telemetry/opamp-go/protobufs"
@@ -20,6 +21,12 @@ type Sender interface {
 	// "pending" flag is reset) then no message will be sent.
 	ScheduleSend()
 
+	// DisableScheduleSend temporary preventing ScheduleSend from writing to channel
+	DisableScheduleSend()
+
+	// EnableScheduleSend re-enables ScheduleSend and checks if it was called during onMessage callback
+	EnableScheduleSend()
+
 	// SetInstanceUid sets a new instanceUid to be used for all subsequent messages to be sent.
 	SetInstanceUid(instanceUid string) error
 }
@@ -31,6 +38,12 @@ type SenderCommon struct {
 	// Indicates that there is a pending message to send.
 	hasPendingMessage chan struct{}
 
+	// Indicates onMessage callback is running
+	isSendingDisabled int32
+
+	// Indicates ScheduleSend() was called during onMessage callback run
+	registerScheduleSend chan struct{}
+
 	// The next message to send.
 	nextMessage NextMessage
 }
@@ -39,15 +52,27 @@ type SenderCommon struct {
 // the WebSocket and HTTP Sender implementations.
 func NewSenderCommon() SenderCommon {
 	return SenderCommon{
-		hasPendingMessage: make(chan struct{}, 1),
-		nextMessage:       NewNextMessage(),
+		hasPendingMessage:    make(chan struct{}, 1),
+		registerScheduleSend: make(chan struct{}, 1),
+		nextMessage:          NewNextMessage(),
+		isSendingDisabled:    0,
 	}
 }
 
 // ScheduleSend signals to HTTPSender that the message in NextMessage struct
 // is now ready to be sent. If there is no pending message (e.g. the NextMessage was
 // already sent and "pending" flag is reset) then no message will be sent.
 func (h *SenderCommon) ScheduleSend() {
+	if h.IsSendingDisabled() {
+		// onMessage callback is running, ScheduleSend() will rerun after it is done
+		select {
+		case h.registerScheduleSend <- struct{}{}:
+		default:
+			break
+		}
+		return
+	}
+
 	// Set pending flag. Don't block on writing to channel.
 	select {
 	case h.hasPendingMessage <- struct{}{}:
@@ -62,6 +87,28 @@ func (h *SenderCommon) NextMessage() *NextMessage {
 	return &h.nextMessage
 }
 
+// IsSendingDisabled returns true when onMessage callback is running
+func (h *SenderCommon) IsSendingDisabled() bool {
+	return atomic.LoadInt32(&h.isSendingDisabled) != 0
+}
+
+// DisableScheduleSend temporary preventing ScheduleSend from writing to channel
+func (h *SenderCommon) DisableScheduleSend() {
+
+	atomic.StoreInt32(&h.isSendingDisabled, 1)
+}
+
+// EnableScheduleSend re-enables ScheduleSend and checks if it was called during onMessage callback
+func (h *SenderCommon) EnableScheduleSend() {
+	atomic.StoreInt32(&h.isSendingDisabled, 0)
+	select {
+	case <-h.registerScheduleSend:
+		h.ScheduleSend()
+	default:
+		break
+	}
+}
+
 // SetInstanceUid sets a new instanceUid to be used for all subsequent messages to be sent.
 // Can be called concurrently, normally is called when a message is received from the
 // Server that instructs us to change our instance UID.

client/internal/wsreceiver_test.go

@@ -97,7 +97,7 @@ func TestServerToAgentCommandExclusive(t *testing.T) {
 		},
 	}
 	clientSyncedState := ClientSyncedState{}
-	receiver := NewWSReceiver(TestLogger{t}, callbacks, nil, nil, &clientSyncedState, nil, 0)
+	receiver := NewWSReceiver(TestLogger{t}, callbacks, nil, &WSSender{}, &clientSyncedState, nil, 0)
 	receiver.processor.ProcessReceivedMessage(context.Background(), &protobufs.ServerToAgent{
 		Command: &protobufs.ServerToAgentCommand{
 			Type: protobufs.CommandType_CommandType_Restart,

client/wsclient.go

@@ -193,9 +193,10 @@ func (c *wsClient) ensureConnected(ctx context.Context) error {
 }
 
 // runOneCycle performs the following actions:
-//   1. connect (try until succeeds).
-//   2. send first status report.
-//   3. receive and process messages until error happens.
+//  1. connect (try until succeeds).
+//  2. send first status report.
+//  3. receive and process messages until error happens.
+//
 // If it encounters an error it closes the connection and returns.
 // Will stop and return if Stop() is called (ctx is cancelled, isStopping is set).
 func (c *wsClient) runOneCycle(ctx context.Context) {