add opamp command processor

Author: breedx-splk

opamp-client/src/main/java/io/opentelemetry/opamp/client/OpampClient.java

@@ -11,6 +11,7 @@
 import opamp.proto.AgentDescription;
 import opamp.proto.RemoteConfigStatus;
 import opamp.proto.ServerErrorResponse;
+import opamp.proto.ServerToAgentCommand;
 
 public interface OpampClient extends Closeable {
 
@@ -79,4 +80,14 @@ interface Callbacks {
      */
     void onMessage(OpampClient client, MessageData messageData);
   }
+
+  @FunctionalInterface
+  interface CommandProcessor {
+    /**
+     * Called when the Agent receives a command from the Server.
+     *
+     * @param command The command requested by the Server.
+     */
+    void onCommand(ServerToAgentCommand command);
+  }
 }

opamp-client/src/main/java/io/opentelemetry/opamp/client/OpampClientBuilder.java

@@ -37,6 +37,7 @@ public final class OpampClientBuilder {
       HttpRequestService.create(OkHttpSender.create("http://localhost:4320/v1/opamp"));
   @Nullable private byte[] instanceUid;
   @Nullable private State.EffectiveConfig effectiveConfigState;
+  @Nullable private OpampClient.CommandProcessor commandProcessor;
 
   OpampClientBuilder() {}
 
@@ -376,6 +377,20 @@ public OpampClientBuilder setEffectiveConfigState(State.EffectiveConfig effectiv
     return this;
   }
 
+  /**
+   * Sets the command processor. This command processor will be called for every ServerToAgent
+   * message that contains a non-null command.
+   *
+   * @param commandProcessor The command processor.
+   * @return this
+   */
+  @CanIgnoreReturnValue
+  public OpampClientBuilder setCommandProcessor(
+      @Nullable OpampClient.CommandProcessor commandProcessor) {
+    this.commandProcessor = commandProcessor;
+    return this;
+  }
+
   public OpampClient build(OpampClient.Callbacks callbacks) {
     List<KeyValue> protoIdentifyingAttributes = new ArrayList<>();
     List<KeyValue> protoNonIdentifyingAttributes = new ArrayList<>();
@@ -402,7 +417,7 @@ public OpampClient build(OpampClient.Callbacks callbacks) {
             new State.InstanceUid(instanceUid),
             new State.Flags(0L),
             effectiveConfigState);
-    return OpampClientImpl.create(service, state, callbacks);
+    return OpampClientImpl.create(service, state, callbacks, commandProcessor);
   }
 
   private static State.EffectiveConfig createEffectiveConfigNoop() {

opamp-client/src/main/java/io/opentelemetry/opamp/client/internal/impl/OpampClientImpl.java

@@ -31,6 +31,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import okio.ByteString;
 import opamp.proto.AgentDescription;
 import opamp.proto.AgentToServer;
@@ -45,12 +46,16 @@
  */
 public final class OpampClientImpl
     implements OpampClient, ObservableState.Listener, RequestService.Callback, Supplier<Request> {
+
+  private static final CommandProcessor NOOP_COMMAND_PROCESSOR = command -> {};
+
   private final RequestService requestService;
   private final AgentToServerAppenders appenders;
   private final OpampClientState state;
   private final RecipeManager recipeManager;
   private final AtomicBoolean hasStopped = new AtomicBoolean(false);
   private final Callbacks callbacks;
+  private final CommandProcessor commandProcessor;
 
   /** Fields that must always be sent. */
   private static final List<Field> REQUIRED_FIELDS;
@@ -83,6 +88,16 @@ public final class OpampClientImpl
 
   public static OpampClientImpl create(
       RequestService requestService, OpampClientState state, Callbacks callbacks) {
+    return create(requestService, state, callbacks, NOOP_COMMAND_PROCESSOR);
+  }
+
+  public static OpampClientImpl create(
+      RequestService requestService,
+      OpampClientState state,
+      Callbacks callbacks,
+      @Nullable CommandProcessor commandProcessor) {
+    CommandProcessor nonNullCommandProcessor =
+        commandProcessor == null ? NOOP_COMMAND_PROCESSOR : commandProcessor;
     AgentToServerAppenders appenders =
         new AgentToServerAppenders(
             AgentDescriptionAppender.create(state.agentDescription),
@@ -95,7 +110,12 @@ public static OpampClientImpl create(
             AgentDisconnectAppender.create());
     OpampClientImpl client =
         new OpampClientImpl(
-            requestService, appenders, state, RecipeManager.create(REQUIRED_FIELDS), callbacks);
+            requestService,
+            appenders,
+            state,
+            RecipeManager.create(REQUIRED_FIELDS),
+            callbacks,
+            nonNullCommandProcessor);
 
     // Start
     requestService.start(client, client);
@@ -111,12 +131,14 @@ private OpampClientImpl(
       AgentToServerAppenders appenders,
       OpampClientState state,
       RecipeManager recipeManager,
-      Callbacks callbacks) {
+      Callbacks callbacks,
+      CommandProcessor commandProcessor) {
     this.requestService = requestService;
     this.appenders = appenders;
     this.state = state;
     this.recipeManager = recipeManager;
     this.callbacks = callbacks;
+    this.commandProcessor = commandProcessor;
   }
 
   @Override
@@ -186,16 +208,15 @@ private void handleResponsePayload(ServerToAgent response) {
     }
     handleAgentIdentification(response);
 
-    boolean notifyOnMessage = false;
     MessageData.Builder messageBuilder = MessageData.builder();
 
     if (response.remote_config != null) {
-      notifyOnMessage = true;
       messageBuilder.setRemoteConfig(response.remote_config);
+      callbacks.onMessage(this, messageBuilder.build());
     }
 
-    if (notifyOnMessage) {
-      callbacks.onMessage(this, messageBuilder.build());
+    if (response.command != null) {
+      commandProcessor.onCommand(response.command);
     }
   }
 

opamp-client/src/test/java/io/opentelemetry/opamp/client/internal/impl/OpampClientImplTest.java

@@ -29,6 +29,7 @@
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import mockwebserver3.MockResponse;
@@ -51,6 +52,7 @@
 import opamp.proto.RemoteConfigStatuses;
 import opamp.proto.ServerErrorResponse;
 import opamp.proto.ServerToAgent;
+import opamp.proto.ServerToAgentCommand;
 import opamp.proto.ServerToAgentFlags;
 import org.jetbrains.annotations.Nullable;
 import org.junit.jupiter.api.AfterEach;
@@ -219,6 +221,36 @@ void onSuccess_withNoChangesToReport_doNotNotifyCallbackOnMessage() {
     verify(callbacks, never()).onMessage(eq(client), any());
   }
 
+  @Test
+  void onSuccess_withCommand_notifyCommandProcessor() {
+    TestCommandProcessor commandProcessor = new TestCommandProcessor();
+    initializeClient(new ServerToAgent.Builder().build(), commandProcessor);
+    ServerToAgentCommand command = new ServerToAgentCommand.Builder().build();
+    enqueueServerToAgentResponse(new ServerToAgent.Builder().command(command).build());
+
+    // Force request
+    requestService.sendRequest();
+
+    await().atMost(Duration.ofSeconds(5)).until(() -> commandProcessor.onCommandCalls.get() == 1);
+
+    assertThat(commandProcessor.latestCommand.get()).isEqualTo(command);
+    verify(callbacks, never()).onMessage(eq(client), any());
+  }
+
+  @Test
+  void onSuccess_withNoCommand_doNotNotifyCommandProcessor() {
+    TestCommandProcessor commandProcessor = new TestCommandProcessor();
+    initializeClient(new ServerToAgent.Builder().build(), commandProcessor);
+    enqueueServerToAgentResponse(new ServerToAgent.Builder().build());
+
+    // Force request
+    requestService.sendRequest();
+
+    await().atMost(Duration.ofSeconds(5)).until(() -> callbacks.onConnectCalls.get() == 2);
+
+    assertThat(commandProcessor.onCommandCalls.get()).isZero();
+  }
+
   @Test
   void verifyAgentDescriptionSetter() {
     initializeClient();
@@ -402,6 +434,17 @@ private RecordedRequest initializeClient(ServerToAgent initialResponse) {
     return takeRequest();
   }
 
+  private RecordedRequest initializeClient(
+      ServerToAgent initialResponse, OpampClient.CommandProcessor commandProcessor) {
+    // Prepare first request on start
+    enqueueServerToAgentResponse(initialResponse);
+
+    callbacks = spy(new TestCallbacks());
+    client = OpampClientImpl.create(requestService, state, callbacks, commandProcessor);
+
+    return takeRequest();
+  }
+
   private static class TestEffectiveConfig extends State.EffectiveConfig {
     private opamp.proto.EffectiveConfig config;
 
@@ -471,4 +514,15 @@ public void onMessage(OpampClient client, MessageData messageData) {
       onMessageCalls.incrementAndGet();
     }
   }
+
+  private static class TestCommandProcessor implements OpampClient.CommandProcessor {
+    private final AtomicInteger onCommandCalls = new AtomicInteger();
+    private final AtomicReference<ServerToAgentCommand> latestCommand = new AtomicReference<>();
+
+    @Override
+    public void onCommand(ServerToAgentCommand command) {
+      latestCommand.set(command);
+      onCommandCalls.incrementAndGet();
+    }
+  }
 }