Repository: open-telemetry/opamp-go
Pull request #286 (2024-06-10 09:09:54)

Enable clients to set flags
Updates #198 

This PR adds a new method to the OpAMPClient that allows clients to set their own flags (such as `RequestInstanceUid`).

This new method deviates a little bit from other ones as it just sets a uint64 instead of a pointer struct field, and doesn't use the same errors and nil checks.

Thank you in advance for taking and look, and please let me know if I'm missing anything! 🙌 

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/286?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 75.82%. Comparing base [(`1a9603e`)](https://app.codecov.io/gh/open-telemetry/opamp-go/commit/1a9603e3d2ded1c7a63917e6d13c4d6e22a7af6d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) to head [(`b89e957`)](https://app.codecov.io/gh/open-telemetry/opamp-go/commit/b89e957e5a1e8fdc37a42784a210e74b53b97762?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry).
> Report is 1 commits behind head on main.


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


```diff
@@            Coverage Diff             @@
##             main     #286      +/-   ##
==========================================
- Coverage   75.96%   75.82%   -0.15%     
==========================================
  Files          25       25              
  Lines        1656     1671      +15     
==========================================
+ Hits         1258     1267       +9     
- Misses        291      294       +3     
- Partials      107      110       +3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/286?dropdown=coverage&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).

tpaschalis: cc @srikanthccv I followed your hint over at #198, let me know if I'm missing anything 🙌
tigrannajaryan: I think this is missing a couple additional things:

1. Use the Flags from ClientSyncedState in [PrepareFirstMessage](https://github.com/open-telemetry/opamp-go/blob/2fa19cf0e5eb7045c839f3024311d9bce700c165/client/internal/clientcommon.go#L210). This will allow calling SetFlags() before Start() too, which I imagine is useful so that we ask for a new instance id immediately, not after the first request is sent.

2. Also use the Flags from ClientSyncedState when the server asks `ServerToAgentFlags_ServerToAgentFlags_ReportFullState` [here](https://github.com/open-telemetry/opamp-go/blob/2fa19cf0e5eb7045c839f3024311d9bce700c165/client/internal/receivedprocessor.go#L187). This is important since the agent may reconnect to a different server instance after the first sending attempt is done but reply carrying the new instance id was not received.

3. Once 1 and 2 are added we need to reset the `RequestInstanceUid` flag after we receive `new_instance_uid` from the Server, otherwise we will keep asking for new ids every time. Perhaps do it in [rcvAgentIdentification](https://github.com/open-telemetry/opamp-go/blob/2fa19cf0e5eb7045c839f3024311d9bce700c165/client/internal/receivedprocessor.go#L234) after SetInstanceUid returns success.

I may be missing some more things. This needs a bit more thought and additional test cases that demonstrate the real use cases and behavior in presence of failures.
tpaschalis: Thanks for the feedback @tigrannajaryan

I focused only on the client-side implementation but I realize it's a more interconnected change. I'll try to set up an end-to-end test to have at hand and figure out the entire lifecycle of the IDs.
tpaschalis: Hey there @tigrannajaryan @srikanthccv 👋 Hope your week is going well!

I've pushed a couple of commits; bd001cf that handles the cases Tigran mentioned on your last message and ca66f1d that adds a test for how messages containing valid/invalid AgentIdentification fields interact with the state of the flag we're tracking.

Let me know how this looks, and if you have any concerns, or ideas. I'm still trying to understand all the moving pieces, so if I'm missing anything please let me know!

Patch:
diff --git a/client/client.go b/client/client.go
index da9a818..afc0831 100644
--- a/client/client.go
+++ b/client/client.go
@@ -107,6 +107,16 @@ type OpAMPClient interface {
 	// for more details.
 	SetCustomCapabilities(customCapabilities *protobufs.CustomCapabilities) error
 
+	// SetFlags modifies the set of flags supported by the client.
+	// May be called before or after Start(), including from OnMessage handler.
+	// The zero value of protobufs.AgentToServerFlags corresponds to FlagsUnspecified
+	// and is safe to use.
+	//
+	// See
+	// https://github.com/open-telemetry/opamp-spec/blob/main/specification.md#agenttoserverflags
+	// for more details.
+	SetFlags(flags protobufs.AgentToServerFlags)
+
 	// SendCustomMessage sends the custom message to the Server. May be called anytime after
 	// Start(), including from OnMessage handler.
 	//
diff --git a/client/httpclient.go b/client/httpclient.go
index 117c818..b27c249 100644
--- a/client/httpclient.go
+++ b/client/httpclient.go
@@ -105,11 +105,16 @@ func (c *httpClient) SetPackageStatuses(statuses *protobufs.PackageStatuses) err
 	return c.common.SetPackageStatuses(statuses)
 }
 
-// SendCustomMessage implements OpAMPClient.SetCustomCapabilities.
+// SendCustomCapabilities implements OpAMPClient.SetCustomCapabilities.
 func (c *httpClient) SetCustomCapabilities(customCapabilities *protobufs.CustomCapabilities) error {
 	return c.common.SetCustomCapabilities(customCapabilities)
 }
 
+// SetFlags implements OpAMPClient.SetFlags.
+func (c *httpClient) SetFlags(flags protobufs.AgentToServerFlags) {
+	c.common.SetFlags(flags)
+}
+
 // SendCustomMessage implements OpAMPClient.SendCustomMessage.
 func (c *httpClient) SendCustomMessage(message *protobufs.CustomMessage) (messageSendingChannel chan struct{}, err error) {
 	return c.common.SendCustomMessage(message)
diff --git a/client/internal/clientcommon.go b/client/internal/clientcommon.go
index 4521ff9..a8df364 100644
--- a/client/internal/clientcommon.go
+++ b/client/internal/clientcommon.go
@@ -214,6 +214,7 @@ func (c *ClientCommon) PrepareFirstMessage(ctx context.Context) error {
 			msg.PackageStatuses = c.ClientSyncedState.PackageStatuses()
 			msg.Capabilities = uint64(c.Capabilities)
 			msg.CustomCapabilities = c.ClientSyncedState.CustomCapabilities()
+			msg.Flags = c.ClientSyncedState.Flags()
 		},
 	)
 	return nil
@@ -385,6 +386,19 @@ func (c *ClientCommon) SetCustomCapabilities(customCapabilities *protobufs.Custo
 	return nil
 }
 
+func (c *ClientCommon) SetFlags(flags protobufs.AgentToServerFlags) {
+	// store the flags to send
+	c.ClientSyncedState.SetFlags(flags)
+
+	// send the new flags to the Server
+	c.sender.NextMessage().Update(
+		func(msg *protobufs.AgentToServer) {
+			msg.Flags = uint64(flags)
+		},
+	)
+	c.sender.ScheduleSend()
+}
+
 // SendCustomMessage sends the specified custom message to the server.
 func (c *ClientCommon) SendCustomMessage(message *protobufs.CustomMessage) (messageSendingChannel chan struct{}, err error) {
 	if message == nil {
diff --git a/client/internal/clientstate.go b/client/internal/clientstate.go
index 97eb4dc..93250c9 100644
--- a/client/internal/clientstate.go
+++ b/client/internal/clientstate.go
@@ -17,8 +17,8 @@ var (
 )
 
 // ClientSyncedState stores the state of the Agent messages that the OpAMP Client needs to
-// have access to synchronize to the Server. 5 messages can be stored in this store:
-// AgentDescription, ComponentHealth, RemoteConfigStatus, PackageStatuses and CustomCapabilities.
+// have access to synchronize to the Server. Six messages can be stored in this store:
+// AgentDescription, ComponentHealth, RemoteConfigStatus, PackageStatuses, CustomCapabilities and Flags.
 //
 // See OpAMP spec for more details on how status reporting works:
 // https://github.com/open-telemetry/opamp-spec/blob/main/specification.md#status-reporting
@@ -39,6 +39,7 @@ type ClientSyncedState struct {
 	remoteConfigStatus *protobufs.RemoteConfigStatus
 	packageStatuses    *protobufs.PackageStatuses
 	customCapabilities *protobufs.CustomCapabilities
+	flags              protobufs.AgentToServerFlags
 }
 
 func (s *ClientSyncedState) AgentDescription() *protobufs.AgentDescription {
@@ -71,6 +72,12 @@ func (s *ClientSyncedState) CustomCapabilities() *protobufs.CustomCapabilities {
 	return s.customCapabilities
 }
 
+func (s *ClientSyncedState) Flags() uint64 {
+	defer s.mutex.Unlock()
+	s.mutex.Lock()
+	return uint64(s.flags)
+}
+
 // SetAgentDescription sets the AgentDescription in the state.
 func (s *ClientSyncedState) SetAgentDescription(descr *protobufs.AgentDescription) error {
 	if descr == nil {
@@ -168,3 +175,11 @@ func (s *ClientSyncedState) HasCustomCapability(capability string) bool {
 
 	return false
 }
+
+// SetFlags sets the flags in the state.
+func (s *ClientSyncedState) SetFlags(flags protobufs.AgentToServerFlags) {
+	defer s.mutex.Unlock()
+	s.mutex.Lock()
+
+	s.flags = flags
+}
diff --git a/client/internal/receivedprocessor.go b/client/internal/receivedprocessor.go
index 736073d..bd18a14 100644
--- a/client/internal/receivedprocessor.go
+++ b/client/internal/receivedprocessor.go
@@ -190,8 +190,9 @@ func (r *receivedProcessor) rcvFlags(
 				msg.RemoteConfigStatus = r.clientSyncedState.RemoteConfigStatus()
 				msg.PackageStatuses = r.clientSyncedState.PackageStatuses()
 				msg.CustomCapabilities = r.clientSyncedState.CustomCapabilities()
+				msg.Flags = r.clientSyncedState.Flags()
 
-				// The logic for EffectiveConfig is similar to the previous 5 sub-messages however
+				// The logic for EffectiveConfig is similar to the previous 6 sub-messages however
 				// the EffectiveConfig is fetched using GetEffectiveConfig instead of
 				// from clientSyncedState. We do this to avoid keeping EffectiveConfig in-memory.
 				msg.EffectiveConfig = cfg
@@ -237,6 +238,9 @@ func (r *receivedProcessor) rcvAgentIdentification(ctx context.Context, agentId
 		return err
 	}
 
+	// If we set up a new instance ID, reset the RequestInstanceUid flag.
+	r.clientSyncedState.flags &^= protobufs.AgentToServerFlags_AgentToServerFlags_RequestInstanceUid
+
 	return nil
 }
 
diff --git a/client/wsclient.go b/client/wsclient.go
index b017a44..030972f 100644
--- a/client/wsclient.go
+++ b/client/wsclient.go
@@ -129,6 +129,10 @@ func (c *wsClient) SetCustomCapabilities(customCapabilities *protobufs.CustomCap
 	return c.common.SetCustomCapabilities(customCapabilities)
 }
 
+func (c *wsClient) SetFlags(flags protobufs.AgentToServerFlags) {
+	c.common.SetFlags(flags)
+}
+
 func (c *wsClient) SendCustomMessage(message *protobufs.CustomMessage) (messageSendingChannel chan struct{}, err error) {
 	return c.common.SendCustomMessage(message)
 }
diff --git a/internal/retryafter.go b/internal/retryafter.go
index 784473b..593cd2c 100644
--- a/internal/retryafter.go
+++ b/internal/retryafter.go
@@ -18,26 +18,32 @@ type OptionalDuration struct {
 }
 
 func parseDelaySeconds(s string) (time.Duration, error) {
+	// Verify duration parsed properly
 	n, err := strconv.Atoi(s)
+	if err != nil {
+		return 0, errCouldNotParseRetryAfterHeader
+	}
 
-	// Verify duration parsed properly and bigger than 0
-	if err == nil && n > 0 {
+	// If n > 0 return n seconds, otherwise return 0
+	if n > 0 {
 		duration := time.Duration(n) * time.Second
 		return duration, nil
 	}
-	return 0, errCouldNotParseRetryAfterHeader
+	return 0, nil
 }
 
 func parseHTTPDate(s string) (time.Duration, error) {
+	// Verify duration parsed properly
 	t, err := http.ParseTime(s)
+	if err != nil {
+		return 0, errCouldNotParseRetryAfterHeader
+	}
 
-	// Verify duration parsed properly and bigger than 0
-	if err == nil {
-		if duration := time.Until(t); duration > 0 {
-			return duration, nil
-		}
+	// If the date is in the future return that duration, otherwise return 0
+	if duration := time.Until(t); duration > 0 {
+		return duration, nil
 	}
-	return 0, errCouldNotParseRetryAfterHeader
+	return 0, nil
 }
 
 // ExtractRetryAfterHeader extracts Retry-After response header if the status


Test patch:
diff --git a/client/clientimpl_test.go b/client/clientimpl_test.go
index 46e8510..78ecb5c 100644
--- a/client/clientimpl_test.go
+++ b/client/clientimpl_test.go
@@ -93,7 +93,7 @@ func eventually(t *testing.T, f func() bool) {
 	assert.Eventually(t, f, 5*time.Second, 10*time.Millisecond)
 }
 
-func newInstanceUid(t *testing.T) types.InstanceUid {
+func genNewInstanceUid(t *testing.T) types.InstanceUid {
 	uid, err := uuid.NewV7()
 	require.NoError(t, err)
 	b, err := uid.MarshalBinary()
@@ -103,7 +103,7 @@ func newInstanceUid(t *testing.T) types.InstanceUid {
 
 func prepareSettings(t *testing.T, settings *types.StartSettings, c OpAMPClient) {
 	// Autogenerate instance id.
-	settings.InstanceUid = newInstanceUid(t)
+	settings.InstanceUid = genNewInstanceUid(t)
 
 	// Make sure correct URL scheme is used, based on the type of the OpAMP client.
 	u, err := url.Parse(settings.OpAMPServerURL)
@@ -630,27 +630,24 @@ func TestAgentIdentification(t *testing.T) {
 	testClients(t, func(t *testing.T, client OpAMPClient) {
 		// Start a server.
 		srv := internal.StartMockServer(t)
-		newInstanceUid := newInstanceUid(t)
+		newInstanceUid := genNewInstanceUid(t)
 		var rcvAgentInstanceUid atomic.Value
-		var sentInvalidId atomic.Bool
 		srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
-			rcvAgentInstanceUid.Store(msg.InstanceUid)
-			if sentInvalidId.Load() {
+			if msg.Flags&uint64(protobufs.AgentToServerFlags_AgentToServerFlags_RequestInstanceUid) == 1 {
+				newInstanceUid = genNewInstanceUid(t)
+				rcvAgentInstanceUid.Store(newInstanceUid[:])
 				return &protobufs.ServerToAgent{
 					InstanceUid: msg.InstanceUid,
 					AgentIdentification: &protobufs.AgentIdentification{
-						// If we sent the invalid one first, send a valid one now
+						// If the RequestInstanceUid flag was set, populate this field.
 						NewInstanceUid: newInstanceUid[:],
 					},
 				}
 			}
-			sentInvalidId.Store(true)
+			rcvAgentInstanceUid.Store(msg.InstanceUid)
+			// Start by sending just the old instance ID.
 			return &protobufs.ServerToAgent{
 				InstanceUid: msg.InstanceUid,
-				AgentIdentification: &protobufs.AgentIdentification{
-					// Start by sending an invalid id forcing an error.
-					NewInstanceUid: nil,
-				},
 			}
 		}
 
@@ -689,8 +686,8 @@ func TestAgentIdentification(t *testing.T) {
 			},
 		)
 
-		// Send a dummy message again to get the _new_ id
-		_ = client.SetAgentDescription(createAgentDescr())
+		// Set the flags to request a new ID.
+		client.SetFlags(protobufs.AgentToServerFlags_AgentToServerFlags_RequestInstanceUid)
 
 		// When it was sent, the new instance uid should have been used, which should
 		// have been observed by the Server
@@ -2122,3 +2119,104 @@ func TestSetCustomCapabilities(t *testing.T) {
 		assert.NoError(t, err)
 	})
 }
+
+// TestSetFlags tests the ability for the client to change the set of flags it sends.
+func TestSetFlags(t *testing.T) {
+	testClients(t, func(t *testing.T, client OpAMPClient) {
+
+		// Start a Server.
+		srv := internal.StartMockServer(t)
+		var rcvCustomFlags atomic.Value
+		var flags protobufs.AgentToServerFlags
+
+		srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
+			if msg.Flags != 0 {
+				rcvCustomFlags.Store(msg.Flags)
+			}
+			return nil
+		}
+
+		settings := types.StartSettings{}
+		settings.OpAMPServerURL = "ws://" + srv.Endpoint
+		prepareClient(t, &settings, client)
+
+		assert.NoError(t, client.Start(context.Background(), settings))
+
+		// The zero value of AgentToServerFlags is ready to use
+		client.SetFlags(flags)
+
+		// Update flags to send
+		flags |= protobufs.AgentToServerFlags_AgentToServerFlags_RequestInstanceUid
+		client.SetFlags(flags)
+
+		// Verify new flags were delivered to the server
+		eventually(
+			t,
+			func() bool {
+				msg, ok := rcvCustomFlags.Load().(uint64)
+				if !ok || msg == 0 {
+					return false
+				}
+				return uint64(flags) == msg
+			},
+		)
+
+		// Shutdown the Server.
+		srv.Close()
+
+		// Shutdown the client.
+		err := client.Stop(context.Background())
+		assert.NoError(t, err)
+	})
+}
+
+// TestSetFlags tests the ability for the client to set its flags before starting up.
+func TestSetFlagsBeforeStart(t *testing.T) {
+	testClients(t, func(t *testing.T, client OpAMPClient) {
+		// Start a Server.
+		flags := protobufs.AgentToServerFlags_AgentToServerFlags_RequestInstanceUid
+		srv := internal.StartMockServer(t)
+		var rcvCustomFlags atomic.Value
+		var isFirstMessage atomic.Bool
+		isFirstMessage.Store(true)
+
+		// Make sure we only record flags from the very first message.
+		srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
+			if isFirstMessage.Load() {
+				rcvCustomFlags.Store(msg.Flags)
+			}
+			isFirstMessage.Store(false)
+			return nil
+		}
+
+		settings := types.StartSettings{}
+		settings.OpAMPServerURL = "ws://" + srv.Endpoint
+		prepareClient(t, &settings, client)
+
+		// Set up the flags _before_ calling Start to verify that they're
+		// handled correctly in PrepareFirstMessage.
+		client.SetFlags(flags)
+
+		// Start the client.
+		assert.NoError(t, client.Start(context.Background(), settings))
+
+		// Verify the flags were delivered to the server during the first message.
+		eventually(
+			t,
+			func() bool {
+				msg, ok := rcvCustomFlags.Load().(uint64)
+				if !ok || msg == 0 {
+					return false
+				}
+				return uint64(flags) == msg
+			},
+		)
+
+		// Shutdown the Server.
+		srv.Close()
+
+		// Shutdown the client.
+		err := client.Stop(context.Background())
+		assert.NoError(t, err)
+	})
+}
diff --git a/client/internal/httpsender_test.go b/client/internal/httpsender_test.go
index 420bc9e..65df9c5 100644
--- a/client/internal/httpsender_test.go
+++ b/client/internal/httpsender_test.go
@@ -172,3 +172,39 @@ func TestHTTPSenderRetryForFailedRequests(t *testing.T) {
 	cancel()
 	srv.Close()
 }
+
+func TestRequestInstanceUidFlagReset(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	sender := NewHTTPSender(&sharedinternal.NopLogger{})
+	sender.callbacks = types.CallbacksStruct{}
+
+	// Set the RequestInstanceUid flag on the tracked state to request the server for a new ID to use.
+	clientSyncedState := &ClientSyncedState{}
+	clientSyncedState.SetFlags(protobufs.AgentToServerFlags_AgentToServerFlags_RequestInstanceUid)
+	capabilities := protobufs.AgentCapabilities_AgentCapabilities_Unspecified
+	sender.receiveProcessor = newReceivedProcessor(&sharedinternal.NopLogger{}, sender.callbacks, sender, clientSyncedState, nil, capabilities)
+
+	// If we process a message with a nil AgentIdentification, or an incorrect NewInstanceUid.
+	sender.receiveProcessor.ProcessReceivedMessage(ctx,
+		&protobufs.ServerToAgent{
+			AgentIdentification: nil,
+		})
+	sender.receiveProcessor.ProcessReceivedMessage(ctx,
+		&protobufs.ServerToAgent{
+			AgentIdentification: &protobufs.AgentIdentification{NewInstanceUid: []byte("foo")},
+		})
+
+	// Then the RequestInstanceUid flag stays intact.
+	assert.Equal(t, sender.receiveProcessor.clientSyncedState.flags, protobufs.AgentToServerFlags_AgentToServerFlags_RequestInstanceUid)
+
+	// If we process a message that contains a non-nil AgentIdentification that contains a NewInstanceUid.
+	sender.receiveProcessor.ProcessReceivedMessage(ctx,
+		&protobufs.ServerToAgent{
+			AgentIdentification: &protobufs.AgentIdentification{NewInstanceUid: []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}},
+		})
+
+	// Then the flag is reset so we don't request a new instance uid yet again.
+	assert.Equal(t, sender.receiveProcessor.clientSyncedState.flags, protobufs.AgentToServerFlags_AgentToServerFlags_Unspecified)
+	cancel()
+}
diff --git a/internal/retryafter_test.go b/internal/retryafter_test.go
index f2e36e9..123ada1 100644
--- a/internal/retryafter_test.go
+++ b/internal/retryafter_test.go
@@ -50,9 +50,10 @@ func TestExtractRetryAfterHeaderDelaySeconds(t *testing.T) {
 	resp.StatusCode = http.StatusBadGateway
 	assertUndefinedDuration(t, ExtractRetryAfterHeader(resp))
 
-	// Verify no duration is created for n < 0
+	// Verify a zero duration is created for n < 0
+	resp.StatusCode = http.StatusTooManyRequests
 	resp.Header.Set(retryAfterHTTPHeader, strconv.Itoa(-1))
-	assertUndefinedDuration(t, ExtractRetryAfterHeader(resp))
+	assertDuration(t, ExtractRetryAfterHeader(resp), 0)
 }
 
 func TestExtractRetryAfterHeaderHttpDate(t *testing.T) {
@@ -81,7 +82,11 @@ func TestExtractRetryAfterHeaderHttpDate(t *testing.T) {
 	resp.Header.Set(retryAfterHTTPHeader, retryAfter.Format(time.RFC1123))
 	assertUndefinedDuration(t, ExtractRetryAfterHeader(resp))
 
-	// Verify no duration is created for n < 0
+	// Verify a zero duration is created for n = 0
+	resp.Header.Set(retryAfterHTTPHeader, now.UTC().Format(http.TimeFormat))
+	assertDuration(t, ExtractRetryAfterHeader(resp), 0)
+
+	// Verify a zero duration is created for n < 0
 	resp.Header.Set(retryAfterHTTPHeader, now.Add(-1*time.Second).UTC().Format(http.TimeFormat))
-	assertUndefinedDuration(t, ExtractRetryAfterHeader(resp))
+	assertDuration(t, ExtractRetryAfterHeader(resp), 0)
 }