diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ActiveSpanTracker.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ActiveSpanTracker.java index 01ae49601..7f8c199bf 100644 --- a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ActiveSpanTracker.java +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ActiveSpanTracker.java @@ -22,16 +22,18 @@ import io.opentelemetry.context.ContextStorage; import io.opentelemetry.context.Scope; import io.opentelemetry.instrumentation.api.internal.cache.Cache; +import java.util.Objects; import java.util.Optional; +import java.util.function.Supplier; import javax.annotation.Nullable; class ActiveSpanTracker implements ContextStorage, SpanTracker { - private final Cache cache = Cache.weak(); + private final Cache cache = Cache.weak(); private final ContextStorage delegate; - private final TraceRegistry registry; + private final Supplier registry; - ActiveSpanTracker(ContextStorage delegate, TraceRegistry registry) { + ActiveSpanTracker(ContextStorage delegate, Supplier registry) { this.delegate = delegate; this.registry = registry; } @@ -39,14 +41,15 @@ class ActiveSpanTracker implements ContextStorage, SpanTracker { @Override public Scope attach(Context toAttach) { Scope scope = delegate.attach(toAttach); - SpanContext newSpanContext = Span.fromContext(toAttach).getSpanContext(); - if (doNotTrack(newSpanContext)) { + SpanContext spanContext = Span.fromContext(toAttach).getSpanContext(); + if (doNotTrack(spanContext)) { return scope; } Thread thread = Thread.currentThread(); - SpanContext oldSpanContext = cache.get(thread); - if (oldSpanContext == newSpanContext) { + ProfilingSpanContext oldSpanContext = cache.get(thread); + ProfilingSpanContext newSpanContext = ProfilingSpanContext.from(spanContext); + if (Objects.equals(oldSpanContext, newSpanContext)) { return scope; } @@ -62,7 +65,7 @@ public Scope attach(Context toAttach) { } private boolean doNotTrack(SpanContext spanContext) { - return !spanContext.isSampled() || !registry.isRegistered(spanContext); + return !spanContext.isSampled() || !registry.get().isRegistered(spanContext); } @Nullable @@ -71,7 +74,8 @@ public Context current() { return delegate.current(); } - public Optional getActiveSpan(Thread thread) { + @Override + public Optional getActiveSpan(Thread thread) { return Optional.ofNullable(cache.get(thread)); } } diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/InterceptingContextStorageSpanTrackingActivator.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/InterceptingContextStorageSpanTrackingActivator.java index 282454785..198d961c5 100644 --- a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/InterceptingContextStorageSpanTrackingActivator.java +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/InterceptingContextStorageSpanTrackingActivator.java @@ -19,6 +19,7 @@ import com.google.common.annotations.VisibleForTesting; import io.opentelemetry.context.ContextStorage; import java.util.function.Consumer; +import java.util.function.Supplier; import java.util.function.UnaryOperator; class InterceptingContextStorageSpanTrackingActivator implements SpanTrackingActivator { @@ -35,7 +36,7 @@ class InterceptingContextStorageSpanTrackingActivator implements SpanTrackingAct } @Override - public void activate(TraceRegistry registry) { + public void activate(Supplier registry) { contextStorageWrappingFunction.accept( contextStorage -> { ActiveSpanTracker tracker = new ActiveSpanTracker(contextStorage, registry); diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/OrphanedTraceDetectingTraceRegistry.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/OrphanedTraceDetectingTraceRegistry.java new file mode 100644 index 000000000..90ce61fe7 --- /dev/null +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/OrphanedTraceDetectingTraceRegistry.java @@ -0,0 +1,161 @@ +/* + * 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.snapshot; + +import com.google.common.annotations.VisibleForTesting; +import io.opentelemetry.api.trace.SpanContext; +import java.lang.ref.ReferenceQueue; +import java.lang.ref.WeakReference; +import java.util.Collections; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Supplier; + +class OrphanedTraceDetectingTraceRegistry implements TraceRegistry { + private final Set traces = Collections.newSetFromMap(new ConcurrentHashMap<>()); + private final ReferenceQueue referenceQueue = new ReferenceQueue<>(); + + private final TraceRegistry delegate; + private final Supplier sampler; + private final Thread thread; + + OrphanedTraceDetectingTraceRegistry(TraceRegistry delegate, Supplier sampler) { + this.delegate = delegate; + this.sampler = sampler; + + thread = new Thread(this::unregisterOrphanedTraces); + thread.setName("orphaned-trace-detector"); + thread.setDaemon(true); + thread.start(); + } + + @Override + public void register(SpanContext spanContext) { + delegate.register(spanContext); + traces.add(new WeakSpanContext(spanContext, referenceQueue)); + } + + @Override + public boolean isRegistered(SpanContext spanContext) { + return delegate.isRegistered(spanContext); + } + + @Override + public void unregister(SpanContext spanContext) { + delegate.unregister(spanContext); + traces.remove(new LookupKey(spanContext)); + } + + /** Test-specific method to inspect the weak references. DO NOT USE IN PRODUCTION CODE! */ + @VisibleForTesting + boolean isRegisteredInternal(SpanContext spanContext) { + return traces.contains(new LookupKey(spanContext)); + } + + public void unregisterOrphanedTraces() { + while (!Thread.interrupted()) { + try { + Object reference = referenceQueue.remove(); + if (reference != null) { + Key key = (Key) reference; + traces.remove(key); + delegate.unregister(key.getSpanContext()); + sampler.get().stop(key.getSpanContext()); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + } + + @Override + public void close() { + delegate.close(); + thread.interrupt(); + } + + public interface Key { + SpanContext getSpanContext(); + } + + public static class LookupKey implements Key { + private final SpanContext spanContext; + + public LookupKey(SpanContext spanContext) { + this.spanContext = spanContext; + } + + @Override + public SpanContext getSpanContext() { + return spanContext; + } + + @Override + public int hashCode() { + return spanContext.hashCode(); + } + + @Override + public boolean equals(Object o) { + if (o == this) { + return true; + } + if (!(o instanceof Key)) { + return false; + } + Key other = (Key) o; + return Objects.equals(spanContext, other.getSpanContext()); + } + } + + public static class WeakSpanContext extends WeakReference implements Key { + private final SpanContext spanContext; + + public WeakSpanContext(SpanContext referent, ReferenceQueue q) { + super(referent, q); + this.spanContext = + SpanContext.create( + referent.getTraceId(), + referent.getSpanId(), + referent.getTraceFlags(), + referent.getTraceState()); + } + + @Override + public SpanContext getSpanContext() { + return spanContext; + } + + @Override + public int hashCode() { + return spanContext.hashCode(); + } + + @Override + public boolean equals(Object o) { + if (o == this) { + return true; + } + if (!(o instanceof Key)) { + return false; + } + Key other = (Key) o; + return Objects.equals(spanContext, other.getSpanContext()); + } + } +} diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ProfilingSpanContext.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ProfilingSpanContext.java new file mode 100644 index 000000000..a63412504 --- /dev/null +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ProfilingSpanContext.java @@ -0,0 +1,58 @@ +/* + * 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.snapshot; + +import io.opentelemetry.api.trace.SpanContext; +import java.util.Objects; + +class ProfilingSpanContext { + static final ProfilingSpanContext INVALID = from(SpanContext.getInvalid()); + + static ProfilingSpanContext from(SpanContext spanContext) { + return new ProfilingSpanContext(spanContext.getTraceId(), spanContext.getSpanId()); + } + + private final String traceId; + private final String spanId; + + private ProfilingSpanContext(String traceId, String spanId) { + this.traceId = traceId; + this.spanId = spanId; + } + + public String getTraceId() { + return traceId; + } + + public String getSpanId() { + return spanId; + } + + @Override + public int hashCode() { + return Objects.hash(traceId, spanId); + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) { + return false; + } + ProfilingSpanContext that = (ProfilingSpanContext) o; + return Objects.equals(traceId, that.traceId) && Objects.equals(spanId, that.spanId); + } +} diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ScheduledExecutorStackTraceSampler.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ScheduledExecutorStackTraceSampler.java index 626ac4e6c..e8b8071be 100644 --- a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ScheduledExecutorStackTraceSampler.java +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ScheduledExecutorStackTraceSampler.java @@ -59,7 +59,8 @@ public void start(SpanContext spanContext) { } samplers.computeIfAbsent( - spanContext.getTraceId(), id -> new ThreadSampler(spanContext, samplingPeriod)); + spanContext.getTraceId(), + id -> new ThreadSampler(ProfilingSpanContext.from(spanContext), samplingPeriod)); } @Override @@ -67,7 +68,8 @@ public void stop(SpanContext spanContext) { samplers.computeIfPresent( spanContext.getTraceId(), (traceId, sampler) -> { - if (spanContext.equals(sampler.getSpanContext())) { + ProfilingSpanContext context = ProfilingSpanContext.from(spanContext); + if (context.equals(sampler.context)) { sampler.shutdown(); waitForShutdown(sampler); return null; @@ -98,17 +100,15 @@ private void waitForShutdown(ThreadSampler sampler) { private class ThreadSampler { private final ScheduledExecutorService scheduler; - private final SpanContext spanContext; + private final ProfilingSpanContext context; private final StackTraceGatherer gatherer; - ThreadSampler(SpanContext spanContext, Duration samplingPeriod) { - this.spanContext = spanContext; - gatherer = - new StackTraceGatherer( - spanContext.getTraceId(), Thread.currentThread(), System.nanoTime()); + ThreadSampler(ProfilingSpanContext context, Duration samplingPeriod) { + this.context = context; + gatherer = new StackTraceGatherer(context, Thread.currentThread(), System.nanoTime()); scheduler = HelpfulExecutors.newSingleThreadedScheduledExecutor( - "stack-trace-sampler-" + spanContext.getTraceId()); + "stack-trace-sampler-" + context.getSpanId()); scheduler.scheduleAtFixedRate( gatherer, SCHEDULER_INITIAL_DELAY, samplingPeriod.toMillis(), TimeUnit.MILLISECONDS); } @@ -125,19 +125,15 @@ void shutdownNow() { boolean awaitTermination(Duration timeout) throws InterruptedException { return scheduler.awaitTermination(timeout.toMillis(), TimeUnit.MILLISECONDS); } - - SpanContext getSpanContext() { - return spanContext; - } } private class StackTraceGatherer implements Runnable { - private final String traceId; + private final ProfilingSpanContext context; private final Thread thread; private volatile long timestampNanos; - StackTraceGatherer(String traceId, Thread thread, long timestampNanos) { - this.traceId = traceId; + StackTraceGatherer(ProfilingSpanContext context, Thread thread, long timestampNanos) { + this.context = context; this.thread = thread; this.timestampNanos = timestampNanos; } @@ -151,23 +147,24 @@ public void run() { Duration samplingPeriod = Duration.ofNanos(currentSampleTimestamp - previousTimestampNanos); String spanId = retrieveActiveSpan(thread).getSpanId(); StackTrace stackTrace = - StackTrace.from(Instant.now(), samplingPeriod, threadInfo, traceId, spanId); + StackTrace.from( + Instant.now(), samplingPeriod, threadInfo, context.getTraceId(), spanId); stagingArea.get().stage(stackTrace); } catch (Exception e) { - logger.log(Level.SEVERE, e, samplerErrorMessage(traceId, thread.getId())); + logger.log(Level.SEVERE, e, samplerErrorMessage(context, thread.getId())); } finally { timestampNanos = currentSampleTimestamp; } } - private SpanContext retrieveActiveSpan(Thread thread) { - return spanTracker.get().getActiveSpan(thread).orElse(SpanContext.getInvalid()); + private ProfilingSpanContext retrieveActiveSpan(Thread thread) { + return spanTracker.get().getActiveSpan(thread).orElse(ProfilingSpanContext.INVALID); } - private Supplier samplerErrorMessage(String traceId, long threadId) { + private Supplier samplerErrorMessage(ProfilingSpanContext context, long threadId) { return () -> "Exception thrown attempting to stage callstacks for trace ID ' " - + traceId + + context.getTraceId() + "' on profiled thread " + threadId; } diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SdkShutdownHook.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SdkShutdownHook.java index 2b05b0646..058702728 100644 --- a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SdkShutdownHook.java +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SdkShutdownHook.java @@ -24,14 +24,32 @@ import java.io.Closeable; import java.util.ArrayList; import java.util.List; +import java.util.function.Supplier; class SdkShutdownHook implements SpanProcessor { + private final Supplier registry; + private final Supplier sampler; + private final Supplier stagingArea; + private final Supplier exporter; + + SdkShutdownHook( + Supplier registry, + Supplier sampler, + Supplier stagingArea, + Supplier exporter) { + this.registry = registry; + this.sampler = sampler; + this.stagingArea = stagingArea; + this.exporter = exporter; + } + @Override public CompletableResultCode shutdown() { List results = new ArrayList<>(); - results.add(close(StackTraceSampler.SUPPLIER.get())); - results.add(close(StagingArea.SUPPLIER.get())); - results.add(close(StackTraceExporter.SUPPLIER.get())); + results.add(close(registry.get())); + results.add(close(sampler.get())); + results.add(close(stagingArea.get())); + results.add(close(exporter.get())); return CompletableResultCode.ofAll(results); } diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SimpleTraceRegistry.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SimpleTraceRegistry.java new file mode 100644 index 000000000..f054b5012 --- /dev/null +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SimpleTraceRegistry.java @@ -0,0 +1,51 @@ +/* + * 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.snapshot; + +import io.opentelemetry.api.trace.SpanContext; +import java.util.Collections; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; + +class SimpleTraceRegistry implements TraceRegistry { + private final Set traceIds = Collections.newSetFromMap(new ConcurrentHashMap<>()); + private volatile boolean closed = false; + + @Override + public void register(SpanContext spanContext) { + if (closed) { + return; + } + traceIds.add(spanContext.getTraceId()); + } + + @Override + public boolean isRegistered(SpanContext spanContext) { + return traceIds.contains(spanContext.getTraceId()); + } + + @Override + public void unregister(SpanContext spanContext) { + traceIds.remove(spanContext.getTraceId()); + } + + @Override + public void close() { + closed = true; + traceIds.clear(); + } +} diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSdkCustomizer.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSdkCustomizer.java index bc7f51e2c..8922c24a5 100644 --- a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSdkCustomizer.java +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSdkCustomizer.java @@ -30,16 +30,17 @@ import java.util.Set; import java.util.function.BiFunction; import java.util.function.Function; +import java.util.function.Supplier; @AutoService(AutoConfigurationCustomizerProvider.class) public class SnapshotProfilingSdkCustomizer implements AutoConfigurationCustomizerProvider { - private final TraceRegistry registry; + private final ConfigurableSupplier registry; private final Function samplerProvider; private final SpanTrackingActivator spanTrackingActivator; public SnapshotProfilingSdkCustomizer() { this( - new TraceRegistry(), + new ConfigurableSupplier<>(new SimpleTraceRegistry()), stackTraceSamplerProvider(), new InterceptingContextStorageSpanTrackingActivator()); } @@ -61,12 +62,14 @@ private static StagingArea createStagingArea(ConfigProperties properties) { @VisibleForTesting SnapshotProfilingSdkCustomizer( - TraceRegistry registry, StackTraceSampler sampler, SpanTrackingActivator activator) { + ConfigurableSupplier registry, + StackTraceSampler sampler, + SpanTrackingActivator activator) { this(registry, properties -> sampler, activator); } private SnapshotProfilingSdkCustomizer( - TraceRegistry registry, + ConfigurableSupplier registry, Function samplerProvider, SpanTrackingActivator spanTrackingActivator) { this.registry = registry; @@ -78,23 +81,42 @@ private SnapshotProfilingSdkCustomizer( public void customize(AutoConfigurationCustomizer autoConfigurationCustomizer) { autoConfigurationCustomizer .addPropertiesCustomizer(autoConfigureSnapshotVolumePropagator()) + .addPropertiesCustomizer(configureTraceRegistry(registry)) .addTracerProviderCustomizer(snapshotProfilingSpanProcessor(registry)) .addPropertiesCustomizer(startTrackingActiveSpans(registry)) - .addTracerProviderCustomizer(addShutdownHook()); + .addTracerProviderCustomizer(addShutdownHook(registry)); + } + + private Function> configureTraceRegistry( + ConfigurableSupplier registry) { + return properties -> { + if (snapshotProfilingEnabled(properties)) { + TraceRegistry current = registry.get(); + TraceRegistry orphanedTraceDetector = + new OrphanedTraceDetectingTraceRegistry(current, StackTraceSampler.SUPPLIER); + registry.configure(orphanedTraceDetector); + } + return Collections.emptyMap(); + }; } private BiFunction - addShutdownHook() { + addShutdownHook(Supplier registry) { return (builder, properties) -> { if (snapshotProfilingEnabled(properties)) { - builder.addSpanProcessor(new SdkShutdownHook()); + builder.addSpanProcessor( + new SdkShutdownHook( + registry, + StackTraceSampler.SUPPLIER, + StagingArea.SUPPLIER, + StackTraceExporter.SUPPLIER)); } return builder; }; } private BiFunction - snapshotProfilingSpanProcessor(TraceRegistry registry) { + snapshotProfilingSpanProcessor(Supplier registry) { return (builder, properties) -> { if (snapshotProfilingEnabled(properties)) { StackTraceSampler sampler = samplerProvider.apply(properties); @@ -141,7 +163,7 @@ private boolean includeTraceContextPropagator(Set configuredPropagators) } private Function> startTrackingActiveSpans( - TraceRegistry registry) { + Supplier registry) { return properties -> { if (snapshotProfilingEnabled(properties)) { spanTrackingActivator.activate(registry); diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSpanProcessor.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSpanProcessor.java index 3905df64b..0026f2d2d 100644 --- a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSpanProcessor.java +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSpanProcessor.java @@ -36,10 +36,11 @@ * be profiled at a time. */ public class SnapshotProfilingSpanProcessor implements SpanProcessor { - private final TraceRegistry registry; + private final Supplier registry; private final Supplier sampler; - SnapshotProfilingSpanProcessor(TraceRegistry registry, Supplier sampler) { + SnapshotProfilingSpanProcessor( + Supplier registry, Supplier sampler) { this.registry = registry; this.sampler = sampler; } @@ -49,11 +50,11 @@ public void onStart(Context context, ReadWriteSpan span) { if (isEntry(span)) { Volume volume = Volume.from(context); if (volume == Volume.HIGHEST) { - registry.register(span.getSpanContext()); + registry.get().register(span.getSpanContext()); } } - if (isEntry(span) && registry.isRegistered(span.getSpanContext())) { + if (isEntry(span) && registry.get().isRegistered(span.getSpanContext())) { sampler.get().start(span.getSpanContext()); span.setAttribute(SNAPSHOT_PROFILING, true); } @@ -67,14 +68,17 @@ public boolean isStartRequired() { /** * Relying solely on the OpenTelemetry instrumentation to correctly notify this SpanProcessor when * a span has ended opens up the possibility of a memory leak in the event a bug is encountered - * within the instrumentation layer that prevents a span from being ended. + * within the instrumentation layer that prevents a span from being ended. To protect against this + * {@link OrphanedTraceDetectingTraceRegistry}, a specialized {@link TraceRegistry}, is installed + * to search for garbage collected {@link io.opentelemetry.api.trace.SpanContext} instances and + * automatically unregister traces and stop sampling. * - *

Will follow up with a more robust solution to this potential problem. + * @see OrphanedTraceDetectingTraceRegistry */ @Override public void onEnd(ReadableSpan span) { if (isEntry(span)) { - registry.unregister(span.getSpanContext()); + registry.get().unregister(span.getSpanContext()); sampler.get().stop(span.getSpanContext()); } } diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SpanTracker.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SpanTracker.java index 6c5f370d5..2dcab12f1 100644 --- a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SpanTracker.java +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SpanTracker.java @@ -16,12 +16,11 @@ package com.splunk.opentelemetry.profiler.snapshot; -import io.opentelemetry.api.trace.SpanContext; import java.util.Optional; interface SpanTracker { SpanTracker NOOP = thread -> Optional.empty(); ConfigurableSupplier SUPPLIER = new ConfigurableSupplier<>(SpanTracker.NOOP); - Optional getActiveSpan(Thread thread); + Optional getActiveSpan(Thread thread); } diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SpanTrackingActivator.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SpanTrackingActivator.java index 87d098d76..3dec8d672 100644 --- a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SpanTrackingActivator.java +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SpanTrackingActivator.java @@ -16,6 +16,8 @@ package com.splunk.opentelemetry.profiler.snapshot; +import java.util.function.Supplier; + interface SpanTrackingActivator { - void activate(TraceRegistry registry); + void activate(Supplier registry); } diff --git a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/TraceRegistry.java b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/TraceRegistry.java index 610f721f1..989bac523 100644 --- a/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/TraceRegistry.java +++ b/profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/TraceRegistry.java @@ -17,22 +17,15 @@ package com.splunk.opentelemetry.profiler.snapshot; import io.opentelemetry.api.trace.SpanContext; -import java.util.Collections; -import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; +import java.io.Closeable; -class TraceRegistry { - private final Set traceIds = Collections.newSetFromMap(new ConcurrentHashMap<>()); +interface TraceRegistry extends Closeable { + void register(SpanContext spanContext); - public void register(SpanContext spanContext) { - traceIds.add(spanContext.getTraceId()); - } + boolean isRegistered(SpanContext spanContext); - public boolean isRegistered(SpanContext spanContext) { - return traceIds.contains(spanContext.getTraceId()); - } + void unregister(SpanContext spanContext); - public void unregister(SpanContext spanContext) { - traceIds.remove(spanContext.getTraceId()); - } + @Override + default void close() {} } diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/ActiveSpanTrackerTest.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/ActiveSpanTrackerTest.java index dc498418f..ffe0bc72b 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/ActiveSpanTrackerTest.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/ActiveSpanTrackerTest.java @@ -40,7 +40,7 @@ class ActiveSpanTrackerTest { private final ContextStorage storage = ContextStorage.get(); private final TogglableTraceRegistry registry = new TogglableTraceRegistry(); - private final ActiveSpanTracker spanTracker = new ActiveSpanTracker(storage, registry); + private final ActiveSpanTracker spanTracker = new ActiveSpanTracker(storage, () -> registry); @Test void currentContextComesFromOpenTelemetryContextStorage() { @@ -63,7 +63,9 @@ void trackActiveSpanWhenNewContextAttached() { registry.register(spanContext); try (var ignored = spanTracker.attach(context)) { - assertEquals(Optional.of(spanContext), spanTracker.getActiveSpan(Thread.currentThread())); + var profilingSpanContext = ProfilingSpanContext.from(spanContext); + assertEquals( + Optional.of(profilingSpanContext), spanTracker.getActiveSpan(Thread.currentThread())); } } @@ -91,7 +93,9 @@ void trackActiveSpanAcrossMultipleContextChanges() { var spanContext = span.getSpanContext(); context = context.with(span); try (var ignoredScope2 = spanTracker.attach(context)) { - assertEquals(Optional.of(spanContext), spanTracker.getActiveSpan(Thread.currentThread())); + var profilingSpanContext = ProfilingSpanContext.from(spanContext); + assertEquals( + Optional.of(profilingSpanContext), spanTracker.getActiveSpan(Thread.currentThread())); } } } @@ -110,7 +114,9 @@ void restoreActiveSpanToPreviousSpanAfterScopeClosing() { scope.close(); var rootSpanContext = root.getSpanContext(); - assertEquals(Optional.of(rootSpanContext), spanTracker.getActiveSpan(Thread.currentThread())); + var profilingSpanContext = ProfilingSpanContext.from(rootSpanContext); + assertEquals( + Optional.of(profilingSpanContext), spanTracker.getActiveSpan(Thread.currentThread())); } } @@ -130,8 +136,11 @@ void trackActiveSpanForMultipleTraces() throws Exception { try (var scope1 = f1.get(); var scope2 = f2.get()) { - assertEquals(Optional.of(span1.getSpanContext()), spanTracker.getActiveSpan(scope1.thread)); - assertEquals(Optional.of(span2.getSpanContext()), spanTracker.getActiveSpan(scope2.thread)); + var profilingSpanContext1 = ProfilingSpanContext.from(span1.getSpanContext()); + assertEquals(Optional.of(profilingSpanContext1), spanTracker.getActiveSpan(scope1.thread)); + + var profilingSpanContext2 = ProfilingSpanContext.from(span2.getSpanContext()); + assertEquals(Optional.of(profilingSpanContext2), spanTracker.getActiveSpan(scope2.thread)); } finally { executor.shutdown(); } @@ -148,8 +157,11 @@ void trackMultipleActiveSpansForSameTraceFromDifferentThreads() throws Exception var executor = Executors.newFixedThreadPool(2); try (var scope1 = executor.submit(attach(span1)).get(); var scope2 = executor.submit(attach(span2)).get()) { - assertEquals(Optional.of(span1.getSpanContext()), spanTracker.getActiveSpan(scope1.thread)); - assertEquals(Optional.of(span2.getSpanContext()), spanTracker.getActiveSpan(scope2.thread)); + var profilingSpanContext1 = ProfilingSpanContext.from(span1.getSpanContext()); + assertEquals(Optional.of(profilingSpanContext1), spanTracker.getActiveSpan(scope1.thread)); + + var profilingSpanContext2 = ProfilingSpanContext.from(span2.getSpanContext()); + assertEquals(Optional.of(profilingSpanContext2), spanTracker.getActiveSpan(scope2.thread)); } finally { executor.shutdown(); } @@ -194,8 +206,12 @@ void activeSpanForThreadIsUnchangedWhenTraceStartsSpanInAnotherThread() throws E var executor = Executors.newSingleThreadExecutor(); try (var scope1 = attach(root).call(); var scope2 = executor.submit(attach(child)).get()) { - assertEquals(Optional.of(root.getSpanContext()), spanTracker.getActiveSpan(scope1.thread)); - assertEquals(Optional.of(child.getSpanContext()), spanTracker.getActiveSpan(scope2.thread)); + var rootProfilingSpanContext = ProfilingSpanContext.from(root.getSpanContext()); + assertEquals(Optional.of(rootProfilingSpanContext), spanTracker.getActiveSpan(scope1.thread)); + + var childProfilingSpanContext = ProfilingSpanContext.from(child.getSpanContext()); + assertEquals( + Optional.of(childProfilingSpanContext), spanTracker.getActiveSpan(scope2.thread)); } finally { executor.shutdown(); } diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/InMemorySpanTracker.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/InMemorySpanTracker.java index be1a78a95..8c7978251 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/InMemorySpanTracker.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/InMemorySpanTracker.java @@ -22,14 +22,15 @@ import java.util.Optional; class InMemorySpanTracker implements SpanTracker { - private final Map activeSpans = new HashMap<>(); + private final Map activeSpans = new HashMap<>(); void store(long threadId, SpanContext spanContext) { - activeSpans.put(threadId, spanContext); + ProfilingSpanContext context = ProfilingSpanContext.from(spanContext); + activeSpans.put(threadId, context); } @Override - public Optional getActiveSpan(Thread thread) { + public Optional getActiveSpan(Thread thread) { return Optional.ofNullable(activeSpans.get(thread.getId())); } } diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/InterceptingContextStorageSpanTrackingActivatorTest.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/InterceptingContextStorageSpanTrackingActivatorTest.java index 4c14498a0..5a1b44954 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/InterceptingContextStorageSpanTrackingActivatorTest.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/InterceptingContextStorageSpanTrackingActivatorTest.java @@ -30,13 +30,13 @@ class InterceptingContextStorageSpanTrackingActivatorTest { @Test void interceptContextStorage() { - activator.activate(new TraceRegistry()); + activator.activate(SimpleTraceRegistry::new); assertInstanceOf(ActiveSpanTracker.class, delegate.storage); } @Test void activateSpanTracker() { - activator.activate(new TraceRegistry()); + activator.activate(SimpleTraceRegistry::new); assertInstanceOf(ActiveSpanTracker.class, SpanTracker.SUPPLIER.get()); } diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/OrphanedTraceDetectingTraceRegistryTest.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/OrphanedTraceDetectingTraceRegistryTest.java new file mode 100644 index 000000000..0f63e42b9 --- /dev/null +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/OrphanedTraceDetectingTraceRegistryTest.java @@ -0,0 +1,129 @@ +/* + * 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.snapshot; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; + +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.instrumentation.test.utils.GcUtils; +import java.lang.ref.WeakReference; +import java.time.Duration; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +class OrphanedTraceDetectingTraceRegistryTest { + private final TraceRegistry delegate = new SimpleTraceRegistry(); + private final ObservableStackTraceSampler sampler = new ObservableStackTraceSampler(); + private final OrphanedTraceDetectingTraceRegistry registry = + new OrphanedTraceDetectingTraceRegistry(delegate, () -> sampler); + + @AfterEach + void teardown() { + registry.close(); + } + + @Test + void trackRegisteredSpanContext() { + var spanContext = Snapshotting.spanContext().build(); + registry.register(spanContext); + assertThat(registry.isRegisteredInternal(spanContext)).isTrue(); + } + + @Test + void delegateTraceRegistration() { + var spanContext = Snapshotting.spanContext().build(); + registry.register(spanContext); + assertThat(delegate.isRegistered(spanContext)).isTrue(); + } + + @Test + void stopTrackingUnregisteredSpanContext() { + var spanContext = Snapshotting.spanContext().build(); + + registry.register(spanContext); + registry.unregister(spanContext); + + assertThat(registry.isRegisteredInternal(spanContext)).isFalse(); + } + + @Test + void delegateTraceDeregistration() { + var spanContext = Snapshotting.spanContext().build(); + + registry.register(spanContext); + registry.unregister(spanContext); + + assertThat(delegate.isRegistered(spanContext)).isFalse(); + } + + @Test + void delegateTraceRegistrationLookup() { + var spanContext = Snapshotting.spanContext().build(); + delegate.register(spanContext); + assertThat(registry.isRegistered(spanContext)).isTrue(); + } + + @Test + void closeDelegateTraceRegistryWhenClosed() { + var spanContext = Snapshotting.spanContext().build(); + + registry.register(spanContext); + registry.close(); + + assertThat(delegate.isRegistered(spanContext)).isFalse(); + } + + @Test + void unregisterOrphanedTraces() throws Exception { + var spanContext = Snapshotting.spanContext().build(); + registry.register(spanContext); + var spanContextCopy = + SpanContext.create( + spanContext.getTraceId(), + spanContext.getSpanId(), + spanContext.getTraceFlags(), + spanContext.getTraceState()); + + var spanContextReference = new WeakReference<>(spanContext); + spanContext = null; + GcUtils.awaitGc(spanContextReference, Duration.ofSeconds(10)); + + await().untilAsserted(() -> assertThat(delegate.isRegistered(spanContextCopy)).isFalse()); + await().untilAsserted(() -> assertThat(registry.isRegistered(spanContextCopy)).isFalse()); + } + + @Test + void stopSamplingForOrphanedTraces() throws Exception { + var spanContext = Snapshotting.spanContext().build(); + registry.register(spanContext); + sampler.start(spanContext); + var spanContextCopy = + SpanContext.create( + spanContext.getTraceId(), + spanContext.getSpanId(), + spanContext.getTraceFlags(), + spanContext.getTraceState()); + await().untilAsserted(() -> assertThat(sampler.isBeingSampled(spanContextCopy)).isTrue()); + + var spanContextReference = new WeakReference<>(spanContext); + spanContext = null; + GcUtils.awaitGc(spanContextReference, Duration.ofSeconds(10)); + + await().untilAsserted(() -> assertThat(sampler.isBeingSampled(spanContextCopy)).isFalse()); + } +} diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/ProfilingSpanContextTest.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/ProfilingSpanContextTest.java new file mode 100644 index 000000000..f9fce7acb --- /dev/null +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/ProfilingSpanContextTest.java @@ -0,0 +1,83 @@ +/* + * 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.snapshot; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; + +import io.opentelemetry.api.trace.SpanId; +import io.opentelemetry.api.trace.TraceId; +import org.junit.jupiter.api.Test; + +class ProfilingSpanContextTest { + @Test + void createFromOpenTelemetrySpanContext() { + var spanContext = Snapshotting.spanContext().build(); + var profilingSpanContext = ProfilingSpanContext.from(spanContext); + + assertEquals(spanContext.getTraceId(), profilingSpanContext.getTraceId()); + assertEquals(spanContext.getSpanId(), profilingSpanContext.getSpanId()); + } + + @Test + void equals() { + var spanContext = Snapshotting.spanContext().build(); + + var one = ProfilingSpanContext.from(spanContext); + var two = ProfilingSpanContext.from(spanContext); + + assertThat(one.equals(two)).isTrue(); + assertThat(two.equals(one)).isTrue(); + } + + @Test + void notEquals() { + var one = ProfilingSpanContext.from(Snapshotting.spanContext().build()); + var two = ProfilingSpanContext.from(Snapshotting.spanContext().build()); + + assertThat(one.equals(two)).isFalse(); + assertThat(one.equals(new Object())).isFalse(); + assertThat(two.equals(one)).isFalse(); + assertThat(two.equals(new Object())).isFalse(); + } + + @Test + void hashCodeEquals() { + var spanContext = Snapshotting.spanContext().build(); + + var one = ProfilingSpanContext.from(spanContext); + var two = ProfilingSpanContext.from(spanContext); + + assertEquals(one.hashCode(), two.hashCode()); + } + + @Test + void hasCodeNotEquals() { + var one = ProfilingSpanContext.from(Snapshotting.spanContext().build()); + var two = ProfilingSpanContext.from(Snapshotting.spanContext().build()); + + assertNotEquals(one.hashCode(), two.hashCode()); + } + + @Test + void invalidProfilingSpanContext() { + var invalid = ProfilingSpanContext.INVALID; + assertEquals(TraceId.getInvalid(), invalid.getTraceId()); + assertEquals(SpanId.getInvalid(), invalid.getSpanId()); + } +} diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/RecordingTraceRegistry.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/RecordingTraceRegistry.java index 8a6ff9ae5..5d00de036 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/RecordingTraceRegistry.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/RecordingTraceRegistry.java @@ -18,33 +18,33 @@ import io.opentelemetry.api.trace.SpanContext; import java.util.Collections; -import java.util.HashSet; +import java.util.Map; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; /** * Test only version of {@link TraceRegistry} that keeps a record of every trace ID registered over * the lifetime of the instance. */ -class RecordingTraceRegistry extends TraceRegistry { - private final Set registeredTraceIds = new HashSet<>(); +class RecordingTraceRegistry implements TraceRegistry { + private final Map traceIds = new ConcurrentHashMap<>(); @Override public void register(SpanContext spanContext) { - registeredTraceIds.add(spanContext.getTraceId()); - super.register(spanContext); + traceIds.put(spanContext.getTraceId(), true); } @Override public boolean isRegistered(SpanContext spanContext) { - return super.isRegistered(spanContext); + return traceIds.getOrDefault(spanContext.getTraceId(), false); } @Override public void unregister(SpanContext spanContext) { - super.unregister(spanContext); + traceIds.put(spanContext.getTraceId(), false); } Set registeredTraceIds() { - return Collections.unmodifiableSet(registeredTraceIds); + return Collections.unmodifiableSet(traceIds.keySet()); } } diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SdkShutdownHookTest.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SdkShutdownHookTest.java index 552c85500..cb711eb81 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SdkShutdownHookTest.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SdkShutdownHookTest.java @@ -24,7 +24,18 @@ class SdkShutdownHookTest { private final ClosingObserver observer = new ClosingObserver(); - private final SdkShutdownHook shutdownHook = new SdkShutdownHook(); + private final ConfigurableSupplier registry = + new ConfigurableSupplier<>(new SimpleTraceRegistry()); + private final SdkShutdownHook shutdownHook = + new SdkShutdownHook( + registry, StackTraceSampler.SUPPLIER, StagingArea.SUPPLIER, StackTraceExporter.SUPPLIER); + + @Test + void shutdownTraceRegistrySampling() { + registry.configure(observer); + shutdownHook.shutdown(); + assertThat(observer.isClosed).isTrue(); + } @Test void shutdownStackTraceSampling() { @@ -60,7 +71,7 @@ void shutdownStackTraceExporting() { } private static class ClosingObserver - implements StackTraceSampler, StagingArea, StackTraceExporter { + implements TraceRegistry, StackTraceSampler, StagingArea, StackTraceExporter { private boolean isClosed = false; @Override @@ -68,6 +79,17 @@ public void close() { isClosed = true; } + @Override + public void register(SpanContext spanContext) {} + + @Override + public boolean isRegistered(SpanContext spanContext) { + return false; + } + + @Override + public void unregister(SpanContext spanContext) {} + @Override public void start(SpanContext spanContext) {} diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/TraceRegistryTest.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SimpleTraceRegistryTest.java similarity index 63% rename from profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/TraceRegistryTest.java rename to profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SimpleTraceRegistryTest.java index bcd6d5350..798bb8bdd 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/TraceRegistryTest.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SimpleTraceRegistryTest.java @@ -16,13 +16,14 @@ package com.splunk.opentelemetry.profiler.snapshot; +import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; -class TraceRegistryTest { - private final TraceRegistry registry = new TraceRegistry(); +class SimpleTraceRegistryTest { + private final SimpleTraceRegistry registry = new SimpleTraceRegistry(); @Test void registerTrace() { @@ -46,4 +47,28 @@ void unregisterTraceForProfiling() { assertFalse(registry.isRegistered(spanContext)); } + + @Test + void unregisterTracesWhenClosed() { + var spanContext1 = Snapshotting.spanContext().build(); + var spanContext2 = Snapshotting.spanContext().build(); + + registry.register(spanContext1); + registry.register(spanContext2); + + registry.close(); + + assertThat(registry.isRegistered(spanContext1)).isFalse(); + assertThat(registry.isRegistered(spanContext2)).isFalse(); + } + + @Test + void doNotRegisterNewTracesWhenClosed() { + var spanContext = Snapshotting.spanContext().build(); + + registry.close(); + registry.register(spanContext); + + assertThat(registry.isRegistered(spanContext)).isFalse(); + } } diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingFeatureFlagTest.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingFeatureFlagTest.java index 2f961e950..0b6f202a8 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingFeatureFlagTest.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingFeatureFlagTest.java @@ -26,7 +26,7 @@ import org.junit.jupiter.api.extension.RegisterExtension; class SnapshotProfilingFeatureFlagTest { - private final TraceRegistry registry = new TraceRegistry(); + private final TraceRegistry registry = new SimpleTraceRegistry(); private final SnapshotProfilingSdkCustomizer customizer = Snapshotting.customizer().with(registry).build(); diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingLogExportingTest.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingLogExportingTest.java index 8f5480a8b..2b696c10f 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingLogExportingTest.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingLogExportingTest.java @@ -37,6 +37,7 @@ import java.util.Set; import java.util.function.Function; import java.util.function.Predicate; +import java.util.function.Supplier; import java.util.stream.Collectors; import java.util.stream.Stream; import org.junit.jupiter.api.AfterEach; @@ -106,7 +107,7 @@ private Predicate> label(AttributeKey key) { private static class ResetContextStorage implements SpanTrackingActivator, AfterEachCallback { @Override - public void activate(TraceRegistry registry) { + public void activate(Supplier registry) { ActiveSpanTracker spanTracker = new ActiveSpanTracker(ContextStorage.defaultStorage(), registry); SpanTracker.SUPPLIER.configure(spanTracker); diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSdkCustomizerBuilder.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSdkCustomizerBuilder.java index a9479a96a..64382375d 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSdkCustomizerBuilder.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSdkCustomizerBuilder.java @@ -19,7 +19,7 @@ import java.time.Duration; class SnapshotProfilingSdkCustomizerBuilder { - private TraceRegistry registry = new TraceRegistry(); + private TraceRegistry registry = new SimpleTraceRegistry(); private StackTraceSampler sampler = new ObservableStackTraceSampler(); private SpanTrackingActivator spanTrackingActivator = registry -> {}; @@ -54,6 +54,7 @@ SnapshotProfilingSdkCustomizerBuilder with(SpanTrackingActivator spanTrackingAct } SnapshotProfilingSdkCustomizer build() { - return new SnapshotProfilingSdkCustomizer(registry, sampler, spanTrackingActivator); + return new SnapshotProfilingSdkCustomizer( + new ConfigurableSupplier<>(registry), sampler, spanTrackingActivator); } } diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SpanSamplingTest.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SpanSamplingTest.java index 2b85fb4e7..23bb43a0c 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SpanSamplingTest.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/SpanSamplingTest.java @@ -27,7 +27,7 @@ import org.junit.jupiter.api.extension.RegisterExtension; class SpanSamplingTest { - private final TraceRegistry registry = new TraceRegistry(); + private final TraceRegistry registry = new SimpleTraceRegistry(); private final SnapshotProfilingSdkCustomizer customizer = Snapshotting.customizer().with(registry).build(); diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/TogglableTraceRegistry.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/TogglableTraceRegistry.java index 594579416..9244b4f23 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/TogglableTraceRegistry.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/TogglableTraceRegistry.java @@ -17,8 +17,12 @@ package com.splunk.opentelemetry.profiler.snapshot; import io.opentelemetry.api.trace.SpanContext; +import java.util.HashSet; +import java.util.Set; + +class TogglableTraceRegistry implements TraceRegistry { + private final Set traceIds = new HashSet<>(); -class TogglableTraceRegistry extends TraceRegistry { enum State { ON, OFF @@ -29,7 +33,7 @@ enum State { @Override public void register(SpanContext spanContext) { if (state == State.ON) { - super.register(spanContext); + traceIds.add(spanContext.getTraceId()); } } @@ -39,11 +43,11 @@ public void toggle(State state) { @Override public boolean isRegistered(SpanContext spanContext) { - return super.isRegistered(spanContext); + return traceIds.contains(spanContext.getTraceId()); } @Override public void unregister(SpanContext spanContext) { - super.unregister(spanContext); + traceIds.remove(spanContext.getTraceId()); } } diff --git a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/TraceRegistrationTest.java b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/TraceRegistrationTest.java index 56a15b9a8..9eb60b8d9 100644 --- a/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/TraceRegistrationTest.java +++ b/profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/TraceRegistrationTest.java @@ -26,7 +26,7 @@ import org.junit.jupiter.api.extension.RegisterExtension; class TraceRegistrationTest { - private final TraceRegistry registry = new TraceRegistry(); + private final TraceRegistry registry = new SimpleTraceRegistry(); private final SnapshotProfilingSdkCustomizer customizer = Snapshotting.customizer().with(registry).build();