Skip to content

Commit

Permalink
Adds support for tracing runnable scenarios (#9187)
Browse files Browse the repository at this point in the history
Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>
  • Loading branch information
Gaganjuneja authored Aug 9, 2023
1 parent ab8e306 commit c3a0784
Show file tree
Hide file tree
Showing 11 changed files with 234 additions and 6 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
- [Refactor] MediaTypeParser to MediaTypeParserRegistry ([#8636](https://github.com/opensearch-project/OpenSearch/pull/8636))
- Make MultiBucketConsumerService thread safe to use across slices during search ([#9047](https://github.com/opensearch-project/OpenSearch/pull/9047))
- Removed blocking wait in TransportGetSnapshotsAction which was exhausting generic threadpool ([#8377](https://github.com/opensearch-project/OpenSearch/pull/8377))
- Adds support for tracing runnable scenarios ([#8831](https://github.com/opensearch-project/OpenSearch/pull/8831))

### Deprecated

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,17 @@ public DefaultTracer(TracingTelemetry tracingTelemetry, TracerContextStorage<Str

@Override
public SpanScope startSpan(String spanName) {
Span span = createSpan(spanName, getCurrentSpan());
return startSpan(spanName, null);
}

@Override
public SpanScope startSpan(String spanName, SpanContext parentSpan) {
Span span = null;
if (parentSpan != null) {
span = createSpan(spanName, parentSpan.getSpan());
} else {
span = createSpan(spanName, getCurrentSpanInternal());
}
setCurrentSpanInContext(span);
addDefaultAttributes(span);
return new DefaultSpanScope(span, (scopeSpan) -> endSpan(scopeSpan));
Expand All @@ -48,11 +58,15 @@ public void close() throws IOException {
((Closeable) tracingTelemetry).close();
}

// Visible for testing
Span getCurrentSpan() {
private Span getCurrentSpanInternal() {
return tracerContextStorage.get(TracerContextStorage.CURRENT_SPAN);
}

public SpanContext getCurrentSpan() {
final Span currentSpan = tracerContextStorage.get(TracerContextStorage.CURRENT_SPAN);
return (currentSpan == null) ? null : new SpanContext(currentSpan);
}

private void endSpan(Span span) {
if (span != null) {
span.endSpan();
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*/

package org.opensearch.telemetry.tracing;

/**
* Wrapped Span will be exposed to the code outside of tracing package for sharing the {@link Span} without having access to
* its properties.
*/
public final class SpanContext {
private final Span span;

/**
* Constructor.
* @param span span to be wrapped.
*/
public SpanContext(Span span) {
this.span = span;
}

Span getSpan() {
return span;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -26,4 +26,17 @@ public interface Tracer extends Closeable {
*/
SpanScope startSpan(String spanName);

/**
* Started the {@link Span} with the given name and parent.
* @param spanName span name.
* @param parentSpan parent span.
* @return scope of the span, must be closed with explicit close or with try-with-resource
*/
SpanScope startSpan(String spanName, SpanContext parentSpan);

/**
* Returns the current span.
* @return current wrapped span.
*/
SpanContext getCurrentSpan();
}
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

import org.opensearch.telemetry.tracing.SpanScope;
import org.opensearch.telemetry.tracing.Tracer;
import org.opensearch.telemetry.tracing.SpanContext;

/**
* No-op implementation of Tracer
Expand All @@ -30,6 +31,16 @@ public SpanScope startSpan(String spanName) {
return SpanScope.NO_OP;
}

@Override
public SpanContext getCurrentSpan() {
return null;
}

@Override
public SpanScope startSpan(String spanName, SpanContext parentSpan) {
return SpanScope.NO_OP;
}

@Override
public void close() {

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*/

package org.opensearch.telemetry.tracing.runnable;

import org.opensearch.telemetry.tracing.SpanContext;
import org.opensearch.telemetry.tracing.SpanScope;
import org.opensearch.telemetry.tracing.Tracer;

/**
* Wraps the runnable and add instrumentation to trace the {@link Runnable}
*/
public class TraceableRunnable implements Runnable {
private final Runnable runnable;
private final SpanContext parent;
private final Tracer tracer;
private final String spanName;

/**
* Constructor.
* @param tracer tracer
* @param spanName spanName
* @param parent parent Span.
* @param runnable runnable.
*/
public TraceableRunnable(Tracer tracer, String spanName, SpanContext parent, Runnable runnable) {
this.tracer = tracer;
this.spanName = spanName;
this.parent = parent;
this.runnable = runnable;
}

@Override
public void run() {
try (SpanScope spanScope = tracer.startSpan(spanName, parent)) {
runnable.run();
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*/

/**
* Contains tracing related classes
*/
package org.opensearch.telemetry.tracing.runnable;
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,12 @@
package org.opensearch.telemetry.tracing;

import org.junit.Assert;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.util.concurrent.ThreadContext;
import org.opensearch.test.OpenSearchTestCase;

import java.io.IOException;
import org.opensearch.test.telemetry.tracing.MockTracingTelemetry;

import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
Expand Down Expand Up @@ -40,7 +43,37 @@ public void testCreateSpan() {

defaultTracer.startSpan("span_name");

Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpanName());
Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
}

public void testCreateSpanWithParent() {
TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
DefaultTracer defaultTracer = new DefaultTracer(
tracingTelemetry,
new ThreadContextBasedTracerContextStorage(new ThreadContext(Settings.EMPTY), tracingTelemetry)
);

defaultTracer.startSpan("span_name", null);

SpanContext parentSpan = defaultTracer.getCurrentSpan();

defaultTracer.startSpan("span_name_1", parentSpan);

Assert.assertEquals("span_name_1", defaultTracer.getCurrentSpan().getSpan().getSpanName());
Assert.assertEquals(parentSpan.getSpan(), defaultTracer.getCurrentSpan().getSpan().getParentSpan());
}

public void testCreateSpanWithNullParent() {
TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
DefaultTracer defaultTracer = new DefaultTracer(
tracingTelemetry,
new ThreadContextBasedTracerContextStorage(new ThreadContext(Settings.EMPTY), tracingTelemetry)
);

defaultTracer.startSpan("span_name", null);

Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
Assert.assertEquals(null, defaultTracer.getCurrentSpan().getSpan().getParentSpan());
}

public void testEndSpanByClosingScope() {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*/

package org.opensearch.telemetry.tracing;

import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.util.concurrent.ThreadContext;
import org.opensearch.telemetry.tracing.runnable.TraceableRunnable;
import org.opensearch.test.OpenSearchTestCase;
import org.opensearch.test.telemetry.tracing.MockTracingTelemetry;

public class TraceableRunnableTests extends OpenSearchTestCase {

private final ThreadContextBasedTracerContextStorage contextStorage = new ThreadContextBasedTracerContextStorage(
new ThreadContext(Settings.EMPTY),
new MockTracingTelemetry()
);

public void testRunnableWithNullParent() throws Exception {
String spanName = "testRunnable";
DefaultTracer defaultTracer = new DefaultTracer(new MockTracingTelemetry(), contextStorage);
final AtomicBoolean isRunnableCompleted = new AtomicBoolean(false);
TraceableRunnable traceableRunnable = new TraceableRunnable(
defaultTracer,
spanName,
null,
() -> { isRunnableCompleted.set(true); }
);
traceableRunnable.run();
assertTrue(isRunnableCompleted.get());
assertEquals(spanName, defaultTracer.getCurrentSpan().getSpan().getSpanName());
assertEquals(null, defaultTracer.getCurrentSpan().getSpan().getParentSpan());
}

public void testRunnableWithParent() throws Exception {
String spanName = "testRunnable";
String parentSpanName = "parentSpan";
DefaultTracer defaultTracer = new DefaultTracer(new MockTracingTelemetry(), contextStorage);
defaultTracer.startSpan(parentSpanName);
SpanContext parentSpan = defaultTracer.getCurrentSpan();
AtomicReference<SpanContext> currrntSpan = new AtomicReference<>(new SpanContext(null));
final AtomicBoolean isRunnableCompleted = new AtomicBoolean(false);
TraceableRunnable traceableRunnable = new TraceableRunnable(defaultTracer, spanName, parentSpan, () -> {
isRunnableCompleted.set(true);
currrntSpan.set(defaultTracer.getCurrentSpan());
});
traceableRunnable.run();
assertTrue(isRunnableCompleted.get());
assertEquals(spanName, currrntSpan.get().getSpan().getSpanName());
assertEquals(parentSpan.getSpan(), currrntSpan.get().getSpan().getParentSpan());
assertEquals(parentSpan.getSpan(), defaultTracer.getCurrentSpan().getSpan());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,19 @@ public WrappedTracer(TelemetrySettings telemetrySettings, Tracer defaultTracer)

@Override
public SpanScope startSpan(String spanName) {
return startSpan(spanName, null);
}

@Override
public SpanContext getCurrentSpan() {
Tracer delegateTracer = getDelegateTracer();
return delegateTracer.getCurrentSpan();
}

@Override
public SpanScope startSpan(String spanName, SpanContext parentSpan) {
Tracer delegateTracer = getDelegateTracer();
return delegateTracer.startSpan(spanName);
return delegateTracer.startSpan(spanName, parentSpan);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ public void testStartSpanWithTracingEnabledInvokesDefaultTracer() throws Excepti
wrappedTracer.startSpan("foo");

assertTrue(wrappedTracer.getDelegateTracer() instanceof DefaultTracer);
verify(mockDefaultTracer).startSpan("foo");
verify(mockDefaultTracer).startSpan("foo", null);
}
}

Expand Down

0 comments on commit c3a0784

Please sign in to comment.