Handle profiler registration messages (#204)

Author: JonasKunz

universal-profiling-integration/src/main/java/co/elastic/otel/SpanProfilingSamplesCorrelator.java

@@ -53,17 +53,18 @@ public class SpanProfilingSamplesCorrelator {
   private final RingBuffer<DelayedSpan> delayedSpans;
   private final PeekingPoller<DelayedSpan> delayedSpansPoller;
 
-  private volatile long spanDelayNanos;
+  private volatile long spanBufferDurationNanos;
+  private volatile boolean shuttingDown = false;
 
   private final WriterReaderPhaser shutdownPhaser = new WriterReaderPhaser();
 
   public SpanProfilingSamplesCorrelator(
       int bufferCapacity,
       LongSupplier nanoClock,
-      long spanDelayNanos,
+      long initialSpanDelayNanos,
       Consumer<ReadableSpan> sendSpan) {
     this.nanoClock = nanoClock;
-    this.spanDelayNanos = spanDelayNanos;
+    this.spanBufferDurationNanos = initialSpanDelayNanos;
     this.sendSpan = sendSpan;
 
     bufferCapacity = nextPowerOf2(bufferCapacity);
@@ -78,6 +79,13 @@ public SpanProfilingSamplesCorrelator(
     delayedSpansPoller = new PeekingPoller<>(nonPeekingPoller, DelayedSpan::new);
   }
 
+  public void setSpanBufferDurationNanos(long nanos) {
+    if (nanos < 0) {
+      throw new IllegalArgumentException("nanos must be positive but was " + nanos);
+    }
+    spanBufferDurationNanos = nanos;
+  }
+
   public void onSpanStart(ReadableSpan span, Context parentCtx) {
     boolean sampled = span.getSpanContext().getTraceFlags().isSampled();
     boolean isLocalRoot = LocalRootSpan.getFor(span) == span;
@@ -95,7 +103,7 @@ public void sendOrBufferSpan(ReadableSpan span) {
 
     long criticalPhaseVal = shutdownPhaser.writerCriticalSectionEnter();
     try {
-      if (spanDelayNanos == 0) {
+      if (spanBufferDurationNanos == 0 || shuttingDown) {
         correlateAndSendSpan(span);
         return;
       }
@@ -133,14 +141,14 @@ public void correlate(
     }
   }
 
-  public synchronized void flushPendingDelayedSpans() {
+  public synchronized void flushPendingBufferedSpans() {
     try {
       delayedSpansPoller.poll(
-          delayedSpan -> {
-            long elapsed = nanoClock.getAsLong() - delayedSpan.endNanoTimestamp;
-            if (elapsed >= spanDelayNanos) {
-              correlateAndSendSpan(delayedSpan.span);
-              delayedSpan.clear();
+          bufferedSpan -> {
+            long elapsed = nanoClock.getAsLong() - bufferedSpan.endNanoTimestamp;
+            if (elapsed >= spanBufferDurationNanos || shuttingDown) {
+              correlateAndSendSpan(bufferedSpan.span);
+              bufferedSpan.clear();
               return true;
             }
             return false; // span is not yet ready to be sent
@@ -154,16 +162,16 @@ public synchronized void flushPendingDelayedSpans() {
   public synchronized void shutdownAndFlushAll() {
     shutdownPhaser.readerLock();
     try {
-      spanDelayNanos = 0L; // This will cause new ended spans to not be buffered anymore
+      shuttingDown = true; // This will cause new ended spans to not be buffered anymore
 
       // avoid race condition: we wait until we are
       // sure that no more spans will be added to the ringbuffer
       shutdownPhaser.flipPhase();
     } finally {
       shutdownPhaser.readerUnlock();
     }
-    // every span is now pending because the desired delay is zero
-    flushPendingDelayedSpans();
+    // This will flush all pending spans because shuttingDown=true
+    flushPendingBufferedSpans();
   }
 
   private void correlateAndSendSpan(ReadableSpan span) {

universal-profiling-integration/src/main/java/co/elastic/otel/UniversalProfilingProcessor.java

@@ -22,8 +22,10 @@
 import co.elastic.otel.common.LocalRootSpan;
 import co.elastic.otel.common.util.ExecutorUtils;
 import co.elastic.otel.common.util.HexUtils;
+import co.elastic.otel.hostid.ProfilerProvidedHostId;
 import co.elastic.otel.profiler.DecodeException;
 import co.elastic.otel.profiler.ProfilerMessage;
+import co.elastic.otel.profiler.ProfilerRegistrationMessage;
 import co.elastic.otel.profiler.TraceCorrelationMessage;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
@@ -66,20 +68,24 @@ public class UniversalProfilingProcessor extends AbstractChainingSpanProcessor {
 
   private static final Logger log = Logger.getLogger(UniversalProfilingProcessor.class.getName());
 
+  private static final long INITIAL_SPAN_DELAY_NANOS = Duration.ofSeconds(1).toNanos();
+
   /**
    * The frequency at which the processor polls the unix domain socket for new messages from the
    * profiler.
    */
-  private static final long POLL_FREQUENCY_MS = 20;
+  static final long POLL_FREQUENCY_MS = 20;
 
   private static boolean anyInstanceActive = false;
 
   private final SpanProfilingSamplesCorrelator correlator;
   private final ScheduledExecutorService messagePollAndSpanFlushExecutor;
 
-  // Visibile for testing
+  // Visible for testing
   String socketPath;
 
+  private volatile boolean tlsPropagationActive = false;
+
   public static UniversalProfilingProcessorBuilder builder(SpanProcessor next, Resource resource) {
     return new UniversalProfilingProcessorBuilder(next, resource);
   }
@@ -88,7 +94,7 @@ public static UniversalProfilingProcessorBuilder builder(SpanProcessor next, Res
       SpanProcessor next,
       Resource serviceResource,
       int bufferSize,
-      Duration spanDelay,
+      boolean activeOnlyAfterProfilerRegistration,
       String socketDir,
       LongSupplier nanoClock) {
     super(next);
@@ -99,9 +105,18 @@ public static UniversalProfilingProcessorBuilder builder(SpanProcessor next, Res
                 + " There must be at most one processor of this type active at a time!");
       }
 
+      long initialSpanDelay;
+      if (activeOnlyAfterProfilerRegistration) {
+        initialSpanDelay = 0; // do not buffer spans until we know that a profiler is running
+        tlsPropagationActive = false;
+      } else {
+        initialSpanDelay = INITIAL_SPAN_DELAY_NANOS; // delay conservatively to not miss any data
+        tlsPropagationActive = true;
+      }
+
       correlator =
           new SpanProfilingSamplesCorrelator(
-              bufferSize, nanoClock, spanDelay.toNanos(), this.next::onEnd);
+              bufferSize, nanoClock, initialSpanDelay, this.next::onEnd);
 
       socketPath = openProfilerSocket(socketDir);
       try {
@@ -199,6 +214,9 @@ public boolean isEndRequired() {
 
   @Nullable
   private void onContextChange(@Nullable Context previous, @Nullable Context next) {
+    if (!tlsPropagationActive) {
+      return;
+    }
     try {
       Span oldSpan = safeSpanFromContext(previous);
       Span newSpan = safeSpanFromContext(next);
@@ -222,7 +240,7 @@ synchronized void pollMessagesAndFlushPendingSpans() {
     // Order is important: we only want to flush spans after we have consumed all pending messages
     // otherwise the data for the spans to be flushed might be incomplete
     consumeProfilerMessages();
-    correlator.flushPendingDelayedSpans();
+    correlator.flushPendingBufferedSpans();
   }
 
   private void consumeProfilerMessages() {
@@ -236,6 +254,8 @@ private void consumeProfilerMessages() {
             break;
           } else if (message instanceof TraceCorrelationMessage) {
             handleMessage((TraceCorrelationMessage) message, tempBuffer);
+          } else if (message instanceof ProfilerRegistrationMessage) {
+            handleMessage((ProfilerRegistrationMessage) message);
           } else {
             log.log(Level.FINE, "Received unknown message type from profiler: {0}", message);
           }
@@ -249,6 +269,20 @@ private void consumeProfilerMessages() {
     }
   }
 
+  private void handleMessage(ProfilerRegistrationMessage message) {
+    log.log(
+        Level.FINE,
+        "Received profiler registration message! host.id is {0} and the span delay is {1} ms",
+        new Object[] {message.getHostId(), message.getSamplesDelayMillis()});
+
+    tlsPropagationActive = true;
+    long spanDelayNanos =
+        Duration.ofMillis(message.getSamplesDelayMillis() + POLL_FREQUENCY_MS).toNanos();
+    correlator.setSpanBufferDurationNanos(spanDelayNanos);
+
+    ProfilerProvidedHostId.set(message.getHostId());
+  }
+
   private void handleMessage(TraceCorrelationMessage message, StringBuilder tempBuffer) {
     tempBuffer.setLength(0);
     HexUtils.appendAsHex(message.getTraceId(), tempBuffer);

universal-profiling-integration/src/main/java/co/elastic/otel/UniversalProfilingProcessorBuilder.java

@@ -20,14 +20,13 @@
 
 import io.opentelemetry.sdk.resources.Resource;
 import io.opentelemetry.sdk.trace.SpanProcessor;
-import java.time.Duration;
 import java.util.function.LongSupplier;
 
 public class UniversalProfilingProcessorBuilder {
 
   private final Resource resource;
   private final SpanProcessor nextProcessor;
-  private Duration spanDelay = Duration.ofSeconds(10);
+  private boolean delayActivationAfterProfilerRegistration = true;
 
   private LongSupplier nanoClock = System::nanoTime;
 
@@ -42,16 +41,32 @@ public class UniversalProfilingProcessorBuilder {
 
   public UniversalProfilingProcessor build() {
     return new UniversalProfilingProcessor(
-        nextProcessor, resource, bufferSize, spanDelay, socketDir, nanoClock);
+        nextProcessor,
+        resource,
+        bufferSize,
+        delayActivationAfterProfilerRegistration,
+        socketDir,
+        nanoClock);
   }
 
   UniversalProfilingProcessorBuilder clock(LongSupplier nanoClock) {
     this.nanoClock = nanoClock;
     return this;
   }
 
-  public UniversalProfilingProcessorBuilder spanDelay(Duration delay) {
-    this.spanDelay = delay;
+  /**
+   * If enabled, the profiling integration will remain inactive until the presence of a profiler is
+   * actually detected. This safes a bit of overhead in the case no profiler is there.
+   *
+   * <p>The downside is if the application starts a span immediately after startup, the profiler
+   * might not be detected in time and therefore this first span might not be correlated correctly.
+   * This can be avoided by setting this option to {@code false}. In this case the {@link
+   * UniversalProfilingProcessor} will assume a profiler will be eventually running and start the
+   * correlation eagerly.
+   */
+  public UniversalProfilingProcessorBuilder delayActivationAfterProfilerRegistration(
+      boolean value) {
+    this.delayActivationAfterProfilerRegistration = value;
     return this;
   }
 

universal-profiling-integration/src/main/java/co/elastic/otel/hostid/ProfilerHostIdApplyingSpanExporter.java

@@ -0,0 +1,79 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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 co.elastic.otel.hostid;
+
+import io.opentelemetry.sdk.common.CompletableResultCode;
+import io.opentelemetry.sdk.resources.Resource;
+import io.opentelemetry.sdk.trace.data.DelegatingSpanData;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import io.opentelemetry.sdk.trace.export.SpanExporter;
+import java.util.Collection;
+import java.util.stream.Collectors;
+
+public class ProfilerHostIdApplyingSpanExporter implements SpanExporter {
+
+  private final SpanExporter delegate;
+
+  public ProfilerHostIdApplyingSpanExporter(SpanExporter delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public CompletableResultCode export(Collection<SpanData> collection) {
+    // do a pre-check to avoid unnecessary allocations
+    boolean anyUpdatesRequired =
+        collection.stream()
+            .anyMatch(
+                data ->
+                    data.getResource()
+                        != ProfilerHostIdResourceUpdater.applyHostId(data.getResource()));
+    Collection<SpanData> updatedSpanData = collection;
+    if (anyUpdatesRequired) {
+      updatedSpanData =
+          collection.stream()
+              .map(
+                  data -> {
+                    Resource original = data.getResource();
+                    Resource updated = ProfilerHostIdResourceUpdater.applyHostId(original);
+                    if (original != updated) {
+                      return new DelegatingSpanData(data) {
+                        @Override
+                        public Resource getResource() {
+                          return updated;
+                        }
+                      };
+                    } else {
+                      return data;
+                    }
+                  })
+              .collect(Collectors.toList());
+    }
+    return delegate.export(updatedSpanData);
+  }
+
+  @Override
+  public CompletableResultCode flush() {
+    return delegate.flush();
+  }
+
+  @Override
+  public CompletableResultCode shutdown() {
+    return delegate.flush();
+  }
+}