Skip to content

Commit

Permalink
Adds attributes to startSpan (opensearch-project#9199)
Browse files Browse the repository at this point in the history
* Adds attributes to startSpan

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>

* Update Changelog

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>

* Adds attributes to startSpan

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>

* Refactor code

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>

* Add java doc

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>

* Refactor code

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>

* Refactor code

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>

* Refactor code

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>

* Removes dependency

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>

---------

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>
Co-authored-by: Gagan Juneja <gjjuneja@amazon.com>
  • Loading branch information
Gaganjuneja and Gagan Juneja authored Aug 11, 2023
1 parent 3ddff4c commit 00558e3
Show file tree
Hide file tree
Showing 19 changed files with 440 additions and 40 deletions.
3 changes: 2 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
- [Remove] Deprecated Fractional ByteSizeValue support #9005 ([#9005](https://github.com/opensearch-project/OpenSearch/pull/9005))
- Make MultiBucketConsumerService thread safe to use across slices during search ([#9047](https://github.com/opensearch-project/OpenSearch/pull/9047))
- Change shard_size and shard_min_doc_count evaluation to happen in shard level reduce phase ([#9085](https://github.com/opensearch-project/OpenSearch/pull/9085))
- Add attributes to startSpan methods ([#9199](https://github.com/opensearch-project/OpenSearch/pull/9199))

### Deprecated

Expand All @@ -138,4 +139,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
### Security

[Unreleased 3.0]: https://github.com/opensearch-project/OpenSearch/compare/2.x...HEAD
[Unreleased 2.x]: https://github.com/opensearch-project/OpenSearch/compare/2.10...2.x
[Unreleased 2.x]: https://github.com/opensearch-project/OpenSearch/compare/2.10...2.x
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

import java.io.Closeable;
import java.io.IOException;
import org.opensearch.telemetry.tracing.attributes.Attributes;

/**
*
Expand Down Expand Up @@ -37,16 +38,21 @@ public DefaultTracer(TracingTelemetry tracingTelemetry, TracerContextStorage<Str

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

@Override
public SpanScope startSpan(String spanName, SpanContext parentSpan) {
public SpanScope startSpan(String spanName, Attributes attributes) {
return startSpan(spanName, null, attributes);
}

@Override
public SpanScope startSpan(String spanName, SpanContext parentSpan, Attributes attributes) {
Span span = null;
if (parentSpan != null) {
span = createSpan(spanName, parentSpan.getSpan());
span = createSpan(spanName, parentSpan.getSpan(), attributes);
} else {
span = createSpan(spanName, getCurrentSpanInternal());
span = createSpan(spanName, getCurrentSpanInternal(), attributes);
}
setCurrentSpanInContext(span);
addDefaultAttributes(span);
Expand Down Expand Up @@ -74,8 +80,8 @@ private void endSpan(Span span) {
}
}

private Span createSpan(String spanName, Span parentSpan) {
return tracingTelemetry.createSpan(spanName, parentSpan);
private Span createSpan(String spanName, Span parentSpan, Attributes attributes) {
return tracingTelemetry.createSpan(spanName, parentSpan, attributes);
}

private void setCurrentSpanInContext(Span span) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
package org.opensearch.telemetry.tracing;

import java.io.Closeable;
import org.opensearch.telemetry.tracing.attributes.Attributes;

/**
* Tracer is the interface used to create a {@link Span}
Expand All @@ -27,12 +28,22 @@ public interface Tracer extends Closeable {
SpanScope startSpan(String spanName);

/**
* Started the {@link Span} with the given name and parent.
* Starts the {@link Span} with given name and attributes. This is required in cases when some attribute based
* decision needs to be made before starting the span. Very useful in the case of Sampling.
* @param spanName span name.
* @param attributes attributes to be added.
* @return scope of the span, must be closed with explicit close or with try-with-resource
*/
SpanScope startSpan(String spanName, Attributes attributes);

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

/**
* Returns the current span.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
package org.opensearch.telemetry.tracing;

import java.io.Closeable;
import org.opensearch.telemetry.tracing.attributes.Attributes;

/**
* Interface for tracing telemetry providers
Expand All @@ -21,9 +22,10 @@ public interface TracingTelemetry extends Closeable {
* Creates span with provided arguments
* @param spanName name of the span
* @param parentSpan span's parent span
* @param attributes attributes to be added.
* @return span instance
*/
Span createSpan(String spanName, Span parentSpan);
Span createSpan(String spanName, Span parentSpan, Attributes attributes);

/**
* provides tracing context propagator
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
/*
* 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.attributes;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

/**
* Class to create attributes for a span.
*/
public class Attributes {
private final Map<String, Object> attributesMap;
/**
* Empty value.
*/
public final static Attributes EMPTY = new Attributes(Collections.emptyMap());

/**
* Factory method.
* @return attributes.
*/
public static Attributes create() {
return new Attributes(new HashMap<>());
}

/**
* Constructor.
*/
private Attributes(Map<String, Object> attributesMap) {
this.attributesMap = attributesMap;
}

/**
* Add String attribute.
* @param key key
* @param value value
* @return Same instance.
*/
public Attributes addAttribute(String key, String value) {
Objects.requireNonNull(value, "value cannot be null");
attributesMap.put(key, value);
return this;
}

/**
* Add long attribute.
* @param key key
* @param value value
* @return Same instance.
*/
public Attributes addAttribute(String key, long value) {
attributesMap.put(key, value);
return this;
};

/**
* Add double attribute.
* @param key key
* @param value value
* @return Same instance.
*/
public Attributes addAttribute(String key, double value) {
attributesMap.put(key, value);
return this;
};

/**
* Add boolean attribute.
* @param key key
* @param value value
* @return Same instance.
*/
public Attributes addAttribute(String key, boolean value) {
attributesMap.put(key, value);
return this;
};

/**
* Returns the attribute map.
* @return attributes map
*/
public Map<String, ?> getAttributesMap() {
return Collections.unmodifiableMap(attributesMap);
}

}
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 No-op implementations
*/
package org.opensearch.telemetry.tracing.attributes;
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import org.opensearch.telemetry.tracing.SpanScope;
import org.opensearch.telemetry.tracing.Tracer;
import org.opensearch.telemetry.tracing.SpanContext;
import org.opensearch.telemetry.tracing.attributes.Attributes;

/**
* No-op implementation of Tracer
Expand All @@ -32,15 +33,20 @@ public SpanScope startSpan(String spanName) {
}

@Override
public SpanContext getCurrentSpan() {
return null;
public SpanScope startSpan(String spanName, Attributes attributes) {
return SpanScope.NO_OP;
}

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

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

@Override
public void close() {

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import org.opensearch.telemetry.tracing.SpanContext;
import org.opensearch.telemetry.tracing.SpanScope;
import org.opensearch.telemetry.tracing.Tracer;
import org.opensearch.telemetry.tracing.attributes.Attributes;

/**
* Wraps the runnable and add instrumentation to trace the {@link Runnable}
Expand All @@ -20,24 +21,27 @@ public class TraceableRunnable implements Runnable {
private final SpanContext parent;
private final Tracer tracer;
private final String spanName;
private final Attributes attributes;

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

@Override
public void run() {
try (SpanScope spanScope = tracer.startSpan(spanName, parent)) {
try (SpanScope spanScope = tracer.startSpan(spanName, parent, attributes)) {
runnable.run();
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,17 @@
import org.junit.Assert;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.util.concurrent.ThreadContext;
import org.opensearch.telemetry.tracing.attributes.Attributes;
import org.opensearch.test.OpenSearchTestCase;

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

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.verify;

Expand Down Expand Up @@ -46,6 +51,42 @@ public void testCreateSpan() {
Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
}

public void testCreateSpanWithAttributesWithMock() {
DefaultTracer defaultTracer = new DefaultTracer(mockTracingTelemetry, mockTracerContextStorage);
Attributes attributes = Attributes.create().addAttribute("name", "value");
when(mockTracingTelemetry.createSpan("span_name", mockParentSpan, attributes)).thenReturn(mockSpan);
defaultTracer.startSpan("span_name", attributes);
verify(mockTracingTelemetry).createSpan("span_name", mockParentSpan, attributes);
}

public void testCreateSpanWithAttributesWithParentMock() {
DefaultTracer defaultTracer = new DefaultTracer(mockTracingTelemetry, mockTracerContextStorage);
Attributes attributes = Attributes.create().addAttribute("name", "value");
when(mockTracingTelemetry.createSpan("span_name", mockParentSpan, attributes)).thenReturn(mockSpan);
defaultTracer.startSpan("span_name", new SpanContext(mockParentSpan), attributes);
verify(mockTracingTelemetry).createSpan("span_name", mockParentSpan, attributes);
verify(mockTracerContextStorage, never()).get(TracerContextStorage.CURRENT_SPAN);
}

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

defaultTracer.startSpan(
"span_name",
Attributes.create().addAttribute("key1", 1.0).addAttribute("key2", 2l).addAttribute("key3", true).addAttribute("key4", "key4")
);

Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
Assert.assertEquals(1.0, ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key1"));
Assert.assertEquals(2l, ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key2"));
Assert.assertEquals(true, ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key3"));
Assert.assertEquals("key4", ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key4"));
}

public void testCreateSpanWithParent() {
TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
DefaultTracer defaultTracer = new DefaultTracer(
Expand All @@ -57,7 +98,7 @@ public void testCreateSpanWithParent() {

SpanContext parentSpan = defaultTracer.getCurrentSpan();

defaultTracer.startSpan("span_name_1", parentSpan);
defaultTracer.startSpan("span_name_1", parentSpan, Attributes.EMPTY);

Assert.assertEquals("span_name_1", defaultTracer.getCurrentSpan().getSpan().getSpanName());
Assert.assertEquals(parentSpan.getSpan(), defaultTracer.getCurrentSpan().getSpan().getParentSpan());
Expand All @@ -70,7 +111,7 @@ public void testCreateSpanWithNullParent() {
new ThreadContextBasedTracerContextStorage(new ThreadContext(Settings.EMPTY), tracingTelemetry)
);

defaultTracer.startSpan("span_name", null);
defaultTracer.startSpan("span_name", null, Attributes.EMPTY);

Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
Assert.assertEquals(null, defaultTracer.getCurrentSpan().getSpan().getParentSpan());
Expand Down Expand Up @@ -105,6 +146,6 @@ private void setupMocks() {
when(mockParentSpan.getSpanId()).thenReturn("parent_span_id");
when(mockParentSpan.getTraceId()).thenReturn("trace_id");
when(mockTracerContextStorage.get(TracerContextStorage.CURRENT_SPAN)).thenReturn(mockParentSpan, mockSpan);
when(mockTracingTelemetry.createSpan("span_name", mockParentSpan)).thenReturn(mockSpan);
when(mockTracingTelemetry.createSpan(eq("span_name"), eq(mockParentSpan), any(Attributes.class))).thenReturn(mockSpan);
}
}
Loading

0 comments on commit 00558e3

Please sign in to comment.