[OpAMP] Enable/disable profiler via remote config (#2881)

* work in progress

* simplify by making a test config class impl

* factor out the creational stuff into a builder, add tests

* default

* rework tests

* rename

* import logger

* work on tests

* rename sequencer -> flusher.

* spotless

* POC

* POC code refactored into production

* Test fixed

* Added possibility to turn off profiler from remote config

* fix

* code review followup

* Additional JFR check added

* Small refactoring of JfrContextStorage initialization.
Initial health report sent

* Code review followup

---------

Co-authored-by: Jason Plumb <jplumb@splunk.com>

Author: robsunday

opamp/src/main/java/com/splunk/opentelemetry/opamp/OpampActivator.java

@@ -22,6 +22,7 @@
 import static java.util.logging.Level.WARNING;
 
 import com.google.auto.service.AutoService;
+import com.splunk.opentelemetry.profiler.ProfilingSupervisor;
 import io.opentelemetry.javaagent.extension.AgentListener;
 import io.opentelemetry.opamp.client.OpampClient;
 import io.opentelemetry.opamp.client.OpampClientBuilder;
@@ -35,7 +36,9 @@
 import io.opentelemetry.sdk.resources.Resource;
 import java.io.IOException;
 import java.time.Duration;
+import java.time.Instant;
 import java.util.logging.Logger;
+import opamp.proto.ComponentHealth;
 import opamp.proto.ServerErrorResponse;
 import org.jetbrains.annotations.Nullable;
 
@@ -56,7 +59,8 @@ public void afterAgent(AutoConfiguredOpenTelemetrySdk autoConfiguredOpenTelemetr
         createEffectiveConfigFactory(autoConfiguredOpenTelemetrySdk);
     State.EffectiveConfig effectiveConfig = buildEffectiveConfig(effectiveConfigFactory);
 
-    ServerToAgentMessageHandler serverToAgentMessageHandler = new ServerToAgentMessageHandler();
+    ServerToAgentMessageHandler serverToAgentMessageHandler =
+        new ServerToAgentMessageHandler(ProfilingSupervisor.SUPPLIER.get());
 
     OpampClient client =
         startOpampClient(
@@ -99,6 +103,11 @@ public void onMessage(OpampClient opampClient, MessageData messageData) {
                 }));
   }
 
+  @Override
+  public int order() {
+    return Integer.MAX_VALUE;
+  }
+
   private EffectiveConfigFactory createEffectiveConfigFactory(AutoConfiguredOpenTelemetrySdk sdk) {
     if (AutoConfigureUtil.isDeclarativeConfig(sdk)) {
       return new DeclarativeEffectiveConfigFileFactory();
@@ -132,6 +141,8 @@ static OpampClient startOpampClient(
     agentAttributes.addNonIdentifyingAttributes(builder);
 
     builder.setEffectiveConfigState(effectiveConfig);
+    builder.enableHealthReporting(createInitialHealthReport());
+
     return builder.build(callbacks);
   }
 
@@ -143,4 +154,15 @@ public opamp.proto.EffectiveConfig get() {
       }
     };
   }
+
+  private static ComponentHealth createInitialHealthReport() {
+    Instant now = Instant.now();
+    long nowNanos = now.getEpochSecond() * 1_000_000_000L + now.getNano();
+    return new ComponentHealth.Builder()
+        .healthy(true)
+        .status("OK")
+        .start_time_unix_nano(nowNanos)
+        .status_time_unix_nano(nowNanos)
+        .build();
+  }
 }

opamp/src/main/java/com/splunk/opentelemetry/opamp/RemoteConfigProcessor.java

@@ -16,7 +16,15 @@
 
 package com.splunk.opentelemetry.opamp;
 
+import static io.opentelemetry.api.incubator.config.DeclarativeConfigProperties.empty;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.splunk.opentelemetry.profiler.ProfilerDeclarativeConfiguration;
+import com.splunk.opentelemetry.profiler.ProfilingSupervisor;
+import io.opentelemetry.api.incubator.config.DeclarativeConfigProperties;
 import io.opentelemetry.opamp.client.OpampClient;
+import io.opentelemetry.sdk.autoconfigure.declarativeconfig.DeclarativeConfiguration;
+import java.io.ByteArrayInputStream;
 import java.util.Map;
 import java.util.Objects;
 import java.util.logging.Logger;
@@ -29,6 +37,13 @@ public class RemoteConfigProcessor {
   private static final Logger logger = Logger.getLogger(RemoteConfigProcessor.class.getName());
 
   private static final String REMOTE_CONFIG_FILE_NAME = "splunk.remote.config";
+  private static final String PROFILING_NODE_NAME = "profiling";
+
+  public ProfilingSupervisor profilingSupervisor;
+
+  public RemoteConfigProcessor(ProfilingSupervisor profilingSupervisor) {
+    this.profilingSupervisor = Objects.requireNonNull(profilingSupervisor);
+  }
 
   public void applyConfig(AgentRemoteConfig remoteConfig, OpampClient opampClient) {
     Objects.requireNonNull(opampClient);
@@ -44,7 +59,27 @@ public void applyConfig(AgentRemoteConfig remoteConfig, OpampClient opampClient)
       return;
     }
 
-    // TODO: provide implementation
+    DeclarativeConfigProperties remoteConfigProperties = toDeclarativeConfigProperties(configFile);
+    DeclarativeConfigProperties splunkDistributionConfigProperties =
+        remoteConfigProperties
+            .getStructured("distribution", empty())
+            .getStructured("splunk", empty());
+
+    // Update profiler configuration only when profiling node exists
+    if (splunkDistributionConfigProperties.getPropertyKeys().contains(PROFILING_NODE_NAME)) {
+      DeclarativeConfigProperties profilingConfigProperties =
+          splunkDistributionConfigProperties.getStructured(PROFILING_NODE_NAME, empty());
+      ProfilerDeclarativeConfiguration profilingConfig =
+          new ProfilerDeclarativeConfiguration(profilingConfigProperties);
+      // TODO: should be merged with current profiling config. Probably we will need profiler
+      //       configuration refactoring and some listeners implemented for profiler configuration
+      //       changes. For POC use this temporary solution
+      if (profilingConfig.isEnabled()) {
+        profilingSupervisor.requestStartProfiling();
+      } else {
+        profilingSupervisor.requestStopProfiling();
+      }
+    }
 
     // Confirm to the OpAMP Server that remote config has been applied.
     opampClient.setRemoteConfigStatus(
@@ -53,4 +88,10 @@ public void applyConfig(AgentRemoteConfig remoteConfig, OpampClient opampClient)
             .status(RemoteConfigStatuses.RemoteConfigStatuses_APPLIED)
             .build());
   }
+
+  @VisibleForTesting
+  static DeclarativeConfigProperties toDeclarativeConfigProperties(AgentConfigFile configFile) {
+    return DeclarativeConfiguration.toConfigProperties(
+        new ByteArrayInputStream(configFile.body.toByteArray()));
+  }
 }

opamp/src/main/java/com/splunk/opentelemetry/opamp/ServerToAgentMessageHandler.java

@@ -17,14 +17,15 @@
 package com.splunk.opentelemetry.opamp;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.splunk.opentelemetry.profiler.ProfilingSupervisor;
 import io.opentelemetry.opamp.client.OpampClient;
 import io.opentelemetry.opamp.client.internal.response.MessageData;
 
 public class ServerToAgentMessageHandler {
   private final RemoteConfigProcessor remoteConfigProcessor;
 
-  public ServerToAgentMessageHandler() {
-    this(new RemoteConfigProcessor());
+  public ServerToAgentMessageHandler(ProfilingSupervisor profilingSupervisor) {
+    this(new RemoteConfigProcessor(profilingSupervisor));
   }
 
   @VisibleForTesting

opamp/src/test/java/com/splunk/opentelemetry/opamp/RemoteConfigProcessorTest.java

@@ -17,6 +17,7 @@
 package com.splunk.opentelemetry.opamp;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
@@ -32,19 +33,22 @@
 import org.mockito.ArgumentCaptor;
 
 class RemoteConfigProcessorTest {
-  private final RemoteConfigProcessor handler = new RemoteConfigProcessor();
-  private final OpampClient opampClient = org.mockito.Mockito.mock(OpampClient.class);
+  private final RemoteConfigProcessor handler = new RemoteConfigProcessor(mock());
+  private final OpampClient opampClient = mock(OpampClient.class);
 
   @Test
   void shouldMarkRemoteConfigAsApplied() {
     // given
+    String remoteConfigYaml = "test-config:";
     ByteString configHash = ByteString.encodeUtf8("test-config-hash");
     AgentRemoteConfig remoteConfig =
         createRemoteConfig(
             configHash,
             Map.of(
                 "splunk.remote.config",
-                new AgentConfigFile.Builder().body(ByteString.encodeUtf8("test-config")).build()));
+                new AgentConfigFile.Builder()
+                    .body(ByteString.encodeUtf8(remoteConfigYaml))
+                    .build()));
 
     // when
     handler.applyConfig(remoteConfig, opampClient);

profiler/src/main/java/com/splunk/opentelemetry/profiler/JfrAgentListener.java

@@ -0,0 +1,91 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.opentelemetry.profiler;
+
+import com.google.auto.service.AutoService;
+import com.google.common.annotations.VisibleForTesting;
+import io.opentelemetry.javaagent.extension.AgentListener;
+import io.opentelemetry.sdk.autoconfigure.AutoConfigureUtil;
+import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
+import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
+import java.util.logging.Logger;
+
+@AutoService(AgentListener.class)
+public class JfrAgentListener implements AgentListener {
+  private static final Logger logger = Logger.getLogger(JfrAgentListener.class.getName());
+  private final JFR jfr;
+
+  public JfrAgentListener() {
+    this(JFR.getInstance());
+  }
+
+  @VisibleForTesting
+  JfrAgentListener(JFR jfr) {
+    this.jfr = jfr;
+  }
+
+  @Override
+  public void afterAgent(AutoConfiguredOpenTelemetrySdk sdk) {
+    ProfilingSupervisor.setupJfrContextStorage();
+
+    ProfilerConfiguration config = getProfilerConfiguration(sdk);
+    // Always start the supervisor, so it can start profiling later elsewhere.
+    ProfilingSupervisor supervisor = makeProfilingSupervisor(sdk, config);
+
+    if (notClearForTakeoff(config)) {
+      return;
+    }
+
+    supervisor.requestStartProfiling();
+  }
+
+  @Override
+  public int order() {
+    // Run it a bit earlier than listeners with default priority
+    return -1;
+  }
+
+  // Exists for testing
+  ProfilingSupervisor makeProfilingSupervisor(
+      AutoConfiguredOpenTelemetrySdk sdk, ProfilerConfiguration config) {
+    return ProfilingSupervisor.createAndStart(sdk, config);
+  }
+
+  private static ProfilerConfiguration getProfilerConfiguration(
+      AutoConfiguredOpenTelemetrySdk sdk) {
+    if (ProfilerDeclarativeConfiguration.SUPPLIER.isConfigured()) {
+      return ProfilerDeclarativeConfiguration.SUPPLIER.get();
+    } else {
+      ConfigProperties configProperties = AutoConfigureUtil.getConfig(sdk);
+      return new ProfilerEnvVarsConfiguration(configProperties);
+    }
+  }
+
+  private boolean notClearForTakeoff(ProfilerConfiguration config) {
+    if (!config.isEnabled()) {
+      logger.fine("Profiler is not enabled.");
+      return true;
+    }
+    if (!jfr.isAvailable()) {
+      logger.warning(
+          "JDK Flight Recorder (JFR) is not available in this JVM. Profiling is disabled.");
+      return true;
+    }
+
+    return false;
+  }
+}

profiler/src/main/java/com/splunk/opentelemetry/profiler/JfrContextStorage.java

@@ -28,11 +28,12 @@
 import javax.annotation.Nullable;
 
 class JfrContextStorage implements ContextStorage {
-
   private final ContextStorage delegate;
   private final Function<SpanContext, ContextAttached> newEvent;
   private final ThreadLocal<Span> activeSpan = ThreadLocal.withInitial(Span::getInvalid);
 
+  private volatile boolean enabled = false;
+
   JfrContextStorage(ContextStorage delegate) {
     this(delegate, JfrContextStorage::newEvent);
   }
@@ -43,6 +44,14 @@ class JfrContextStorage implements ContextStorage {
     this.newEvent = newEvent;
   }
 
+  public void setEnabled(boolean enabled) {
+    this.enabled = enabled;
+  }
+
+  public boolean isEnabled() {
+    return enabled;
+  }
+
   static ContextAttached newEvent(SpanContext spanContext) {
     if (spanContext.isValid()) {
       return new ContextAttached(
@@ -54,6 +63,9 @@ static ContextAttached newEvent(SpanContext spanContext) {
   @Override
   public Scope attach(Context toAttach) {
     Scope delegatedScope = delegate.attach(toAttach);
+    if (!isEnabled()) {
+      return delegatedScope;
+    }
     Span span = Span.fromContext(toAttach);
     Span current = activeSpan.get();
     // do nothing when active span didn't change

profiler/src/main/java/com/splunk/opentelemetry/profiler/JfrRecorder.java

@@ -60,6 +60,9 @@ class JfrRecorder {
   }
 
   public void start() {
+    if (isStarted()) {
+      throw new IllegalStateException("Already started");
+    }
     logger.fine("Profiler is starting a JFR recording");
     recording = newRecording();
     recording.setSettings(settings);
@@ -70,6 +73,13 @@ public void start() {
     recording.start();
   }
 
+  public void stop() {
+    if (isStarted()) {
+      recording.stop();
+    }
+    recording = null;
+  }
+
   @VisibleForTesting
   Recording newRecording() {
     return new Recording();
@@ -119,11 +129,6 @@ public boolean isStarted() {
     return (recording != null) && RecordingState.RUNNING.equals(recording.getState());
   }
 
-  public void stop() {
-    recording.stop();
-    recording = null;
-  }
-
   public static Builder builder() {
     return new Builder();
   }