Skip to content

Commit 89ab3a0

Browse files
AnandShivanshAnandShivansh
authored
chore | Add SpanJoiner based on list of spanId (#177)
* Add SpanJoiner based on list of spanId * Add SpanJoiner based on list of spanId * Revert "Add SpanJoiner based on list of spanId" This reverts commit 14ba29c. * Revert "Add SpanJoiner based on list of spanId" This reverts commit 8e9dfae. * Add SpanJoiner based on list of spanId * Resolve PR reviews * Resolve PR reviews
1 parent d6b4b57 commit 89ab3a0

File tree

4 files changed

+176
-63
lines changed

4 files changed

+176
-63
lines changed

‎hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/DefaultSpanJoinerBuilder.java

+72-42
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,11 @@
33
import static com.google.common.collect.ImmutableList.copyOf;
44
import static com.google.common.collect.Iterables.concat;
55
import static org.hypertrace.core.graphql.atttributes.scopes.HypertraceCoreAttributeScopeString.SPAN;
6+
import static org.hypertrace.core.graphql.span.joiner.MultipleSpanJoin.SPANS_KEY;
67
import static org.hypertrace.core.graphql.span.joiner.SpanJoin.SPAN_KEY;
78

9+
import com.google.common.collect.ImmutableListMultimap;
10+
import com.google.common.collect.ListMultimap;
811
import graphql.schema.DataFetchingFieldSelectionSet;
912
import graphql.schema.SelectedField;
1013
import io.reactivex.rxjava3.core.Observable;
@@ -46,9 +49,9 @@
4649
public class DefaultSpanJoinerBuilder implements SpanJoinerBuilder {
4750

4851
private static final int ZERO_OFFSET = 0;
49-
5052
private final SpanDao spanDao;
5153
private final GraphQlSelectionFinder selectionFinder;
54+
5255
private final ResultSetRequestBuilder resultSetRequestBuilder;
5356
private final FilterRequestBuilder filterRequestBuilder;
5457

@@ -70,61 +73,100 @@ public Single<SpanJoiner> build(
7073
TimeRangeArgument timeRange,
7174
DataFetchingFieldSelectionSet selectionSet,
7275
List<String> pathToSpanJoin) {
73-
return Single.just(
74-
new DefaultSpanJoiner(
75-
context, timeRange, this.getSelections(selectionSet, pathToSpanJoin)));
76-
}
77-
78-
private List<SelectedField> getSelections(
79-
DataFetchingFieldSelectionSet selectionSet, List<String> pathToSpanJoin) {
80-
List<String> fullPath = copyOf(concat(pathToSpanJoin, List.of(SPAN_KEY)));
81-
return selectionFinder
82-
.findSelections(selectionSet, SelectionQuery.builder().selectionPath(fullPath).build())
83-
.collect(Collectors.toUnmodifiableList());
76+
return Single.just(new DefaultSpanJoiner(context, timeRange, selectionSet, pathToSpanJoin));
8477
}
8578

8679
@AllArgsConstructor
8780
private class DefaultSpanJoiner implements SpanJoiner {
8881

8982
private final GraphQlRequestContext context;
9083
private final TimeRangeArgument timeRange;
91-
private final List<SelectedField> selectedFields;
84+
private final DataFetchingFieldSelectionSet selectionSet;
85+
private final List<String> pathToJoin;
9286

9387
@Override
94-
public <T> Single<Map<T, Span>> joinSpans(
88+
public <T> Single<Map<T, Span>> joinSpan(
9589
Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
96-
return this.buildSourceToIdMap(joinSources, spanIdGetter).flatMap(this::joinSpans);
90+
Function<T, Single<List<String>>> idsGetter =
91+
source -> spanIdGetter.getSpanId(source).map(List::of);
92+
return this.joinSpans(joinSources, idsGetter, SPAN_KEY).map(this::reduceMap);
9793
}
9894

99-
private <T> Single<Map<T, Span>> joinSpans(Map<T, String> sourceToSpanIdMap) {
100-
return this.buildSpanRequest(sourceToSpanIdMap)
101-
.flatMap(spanDao::getSpans)
102-
.map(this::buildSpanIdToSpanMap)
103-
.map(spanIdToSpanMap -> buildSourceToSpanMap(sourceToSpanIdMap, spanIdToSpanMap));
95+
@Override
96+
public <T> Single<ListMultimap<T, Span>> joinSpans(
97+
Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
98+
return this.joinSpans(joinSources, multipleSpanIdGetter::getSpanIds, SPANS_KEY);
10499
}
105100

106-
private <T> Map<T, Span> buildSourceToSpanMap(
107-
Map<T, String> sourceToSpanIdMap, Map<String, Span> spanIdToSpanMap) {
108-
return sourceToSpanIdMap.entrySet().stream()
109-
.filter(entry -> spanIdToSpanMap.containsKey(entry.getValue()))
101+
private <T> Single<ListMultimap<T, Span>> joinSpans(
102+
Collection<T> joinSources,
103+
Function<T, Single<List<String>>> idsGetter,
104+
String joinSpanKey) {
105+
return this.buildSourceToIdsMap(joinSources, idsGetter)
106+
.flatMap(
107+
sourceToSpanIdsMap ->
108+
this.buildSpanRequest(sourceToSpanIdsMap, joinSpanKey)
109+
.flatMap(spanDao::getSpans)
110+
.map(this::buildSpanIdToSpanMap)
111+
.map(
112+
spanIdToSpanMap ->
113+
this.buildSourceToSpanListMultiMap(
114+
sourceToSpanIdsMap, spanIdToSpanMap)));
115+
}
116+
117+
private <T> Map<T, Span> reduceMap(ListMultimap<T, Span> listMultimap) {
118+
return listMultimap.entries().stream()
110119
.collect(
111120
Collectors.toUnmodifiableMap(
121+
Entry::getKey, Entry::getValue, (first, second) -> first));
122+
}
123+
124+
private <T> Single<ImmutableListMultimap<T, String>> buildSourceToIdsMap(
125+
Collection<T> joinSources, Function<T, Single<List<String>>> idsGetter) {
126+
return Observable.fromIterable(joinSources)
127+
.flatMapSingle(source -> idsGetter.apply(source).map(ids -> Map.entry(source, ids)))
128+
.collect(
129+
ImmutableListMultimap.flatteningToImmutableListMultimap(
130+
Entry::getKey, entry -> entry.getValue().stream()));
131+
}
132+
133+
private <T> ImmutableListMultimap<T, Span> buildSourceToSpanListMultiMap(
134+
ListMultimap<T, String> sourceToSpanIdsMultimap, Map<String, Span> spanIdToSpanMap) {
135+
return sourceToSpanIdsMultimap.entries().stream()
136+
.filter(entry -> spanIdToSpanMap.containsKey(entry.getValue()))
137+
.collect(
138+
ImmutableListMultimap.toImmutableListMultimap(
112139
Entry::getKey, entry -> spanIdToSpanMap.get(entry.getValue())));
113140
}
114141

142+
private List<SelectedField> getSelections(String joinSpanKey) {
143+
List<String> fullPath = copyOf(concat(pathToJoin, List.of(joinSpanKey)));
144+
return selectionFinder
145+
.findSelections(selectionSet, SelectionQuery.builder().selectionPath(fullPath).build())
146+
.collect(Collectors.toUnmodifiableList());
147+
}
148+
115149
private Map<String, Span> buildSpanIdToSpanMap(SpanResultSet resultSet) {
116150
return resultSet.results().stream()
117151
.collect(Collectors.toUnmodifiableMap(Identifiable::id, Function.identity()));
118152
}
119153

120-
private <T> Single<SpanRequest> buildSpanRequest(Map<T, String> sourceToSpanIdMap) {
121-
Collection<String> spanIds = sourceToSpanIdMap.values();
122-
return buildSpanIdsFilter(spanIds)
123-
.flatMap(filterArguments -> buildSpanRequest(spanIds.size(), filterArguments));
154+
private <T> Single<SpanRequest> buildSpanRequest(
155+
ListMultimap<T, String> sourceToSpanIdsMultimap, String joinSpanKey) {
156+
Collection<String> spanIds =
157+
sourceToSpanIdsMultimap.values().stream()
158+
.distinct()
159+
.collect(Collectors.toUnmodifiableList());
160+
List<SelectedField> selectedFields = getSelections(joinSpanKey);
161+
return buildSpanIdsFilter(context, spanIds)
162+
.flatMap(
163+
filterArguments -> buildSpanRequest(spanIds.size(), filterArguments, selectedFields));
124164
}
125165

126166
private Single<SpanRequest> buildSpanRequest(
127-
int size, List<AttributeAssociation<FilterArgument>> filterArguments) {
167+
int size,
168+
List<AttributeAssociation<FilterArgument>> filterArguments,
169+
List<SelectedField> selectedFields) {
128170
return resultSetRequestBuilder
129171
.build(
130172
context,
@@ -140,21 +182,9 @@ private Single<SpanRequest> buildSpanRequest(
140182
}
141183

142184
private Single<List<AttributeAssociation<FilterArgument>>> buildSpanIdsFilter(
143-
Collection<String> spanIds) {
185+
GraphQlRequestContext context, Collection<String> spanIds) {
144186
return filterRequestBuilder.build(context, SPAN, Set.of(new SpanIdFilter(spanIds)));
145187
}
146-
147-
private <T> Single<Map<T, String>> buildSourceToIdMap(
148-
Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
149-
return Observable.fromIterable(joinSources)
150-
.flatMapSingle(source -> this.maybeBuildMapEntry(source, spanIdGetter))
151-
.collect(Collectors.toMap(Entry::getKey, Entry::getValue));
152-
}
153-
154-
private <T> Single<Entry<T, String>> maybeBuildMapEntry(
155-
T source, SpanIdGetter<T> spanIdGetter) {
156-
return spanIdGetter.getSpanId(source).map(id -> Map.entry(source, id));
157-
}
158188
}
159189

160190
@Value

‎hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/MultipleSpanJoin.java

+14
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
package org.hypertrace.core.graphql.span.joiner;
2+
3+
import graphql.annotations.annotationTypes.GraphQLField;
4+
import graphql.annotations.annotationTypes.GraphQLName;
5+
import java.util.List;
6+
import org.hypertrace.core.graphql.span.schema.Span;
7+
8+
public interface MultipleSpanJoin {
9+
String SPANS_KEY = "spans";
10+
11+
@GraphQLField
12+
@GraphQLName(SPANS_KEY)
13+
List<Span> spans();
14+
}

‎hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoiner.java

+19-2
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
package org.hypertrace.core.graphql.span.joiner;
22

3+
import com.google.common.collect.ArrayListMultimap;
4+
import com.google.common.collect.ListMultimap;
35
import io.reactivex.rxjava3.core.Single;
46
import java.util.Collection;
57
import java.util.Collections;
8+
import java.util.List;
69
import java.util.Map;
710
import org.hypertrace.core.graphql.span.schema.Span;
811

@@ -12,16 +15,30 @@ public interface SpanJoiner {
1215
SpanJoiner NO_OP_JOINER =
1316
new SpanJoiner() {
1417
@Override
15-
public <T> Single<Map<T, Span>> joinSpans(
18+
public <T> Single<Map<T, Span>> joinSpan(
1619
Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
1720
return Single.just(Collections.emptyMap());
1821
}
22+
23+
@Override
24+
public <T> Single<ListMultimap<T, Span>> joinSpans(
25+
Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
26+
return Single.just(ArrayListMultimap.create());
27+
}
1928
};
2029

21-
<T> Single<Map<T, Span>> joinSpans(Collection<T> joinSources, SpanIdGetter<T> spanIdGetter);
30+
<T> Single<Map<T, Span>> joinSpan(Collection<T> joinSources, SpanIdGetter<T> spanIdGetter);
31+
32+
<T> Single<ListMultimap<T, Span>> joinSpans(
33+
Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter);
2234

2335
@FunctionalInterface
2436
interface SpanIdGetter<T> {
2537
Single<String> getSpanId(T source);
2638
}
39+
40+
@FunctionalInterface
41+
interface MultipleSpanIdGetter<T> {
42+
Single<List<String>> getSpanIds(T source);
43+
}
2744
}

‎hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilderTest.java

+71-19
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
import static org.mockito.Mockito.mock;
1111
import static org.mockito.Mockito.when;
1212

13+
import com.google.common.collect.ArrayListMultimap;
14+
import com.google.common.collect.ListMultimap;
1315
import graphql.schema.DataFetchingFieldSelectionSet;
1416
import graphql.schema.SelectedField;
1517
import io.reactivex.rxjava3.core.Single;
@@ -30,6 +32,7 @@
3032
import org.hypertrace.core.graphql.context.GraphQlRequestContext;
3133
import org.hypertrace.core.graphql.log.event.schema.LogEventResultSet;
3234
import org.hypertrace.core.graphql.span.dao.SpanDao;
35+
import org.hypertrace.core.graphql.span.joiner.SpanJoiner.MultipleSpanIdGetter;
3336
import org.hypertrace.core.graphql.span.joiner.SpanJoiner.SpanIdGetter;
3437
import org.hypertrace.core.graphql.span.request.SpanRequest;
3538
import org.hypertrace.core.graphql.span.schema.Span;
@@ -45,8 +48,10 @@
4548
@ExtendWith(MockitoExtension.class)
4649
public class SpanJoinerBuilderTest {
4750

48-
private static final String FIRST_SPAN_ID = "spanId1";
49-
private static final String SECOND_SPAN_ID = "spanId2";
51+
private static final String SPAN_ID1 = "spanId1";
52+
private static final String SPAN_ID2 = "spanId2";
53+
private static final String SPAN_ID3 = "spanId3";
54+
private static final String SPAN_ID4 = "spanId4";
5055

5156
@Mock private SpanDao mockSpanDao;
5257
@Mock private GraphQlSelectionFinder mockSelectionFinder;
@@ -72,16 +77,31 @@ void setup() {
7277

7378
@Test
7479
void fetchSpans() {
75-
Span span1 = new TestSpan(FIRST_SPAN_ID);
76-
Span span2 = new TestSpan(SECOND_SPAN_ID);
77-
TestJoinSource joinSource1 = new TestJoinSource(FIRST_SPAN_ID);
78-
TestJoinSource joinSource2 = new TestJoinSource(SECOND_SPAN_ID);
80+
Span span1 = new TestSpan(SPAN_ID1);
81+
Span span2 = new TestSpan(SPAN_ID2);
82+
TestJoinSource joinSource1 = new TestJoinSource(SPAN_ID1);
83+
TestJoinSource joinSource2 = new TestJoinSource(SPAN_ID2);
7984
Map<TestJoinSource, Span> expected =
8085
Map.ofEntries(entry(joinSource1, span1), entry(joinSource2, span2));
8186
List<TestJoinSource> joinSources = List.of(joinSource1, joinSource2);
82-
mockRequestedSelectionFields(
83-
List.of(mock(SelectedField.class), mock(SelectedField.class)), "pathToSpan");
84-
mockRequestBuilding();
87+
when(mockSelectionFinder.findSelections(
88+
mockSelectionSet,
89+
SelectionQuery.builder().selectionPath(List.of("pathToSpan", "span")).build()))
90+
.thenReturn(Stream.of(mock(SelectedField.class), mock(SelectedField.class)));
91+
when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anySet()))
92+
.thenReturn(Single.just(List.of(mockFilter)));
93+
94+
when(mockResultSetRequestBuilder.build(
95+
eq(mockRequestContext),
96+
eq(SPAN),
97+
eq(2),
98+
eq(0),
99+
eq(mockTimeRangeArgument),
100+
eq(emptyList()),
101+
eq(List.of(mockFilter)),
102+
any(Stream.class),
103+
eq(Optional.empty())))
104+
.thenReturn(Single.just(mockResultSetRequest));
85105
mockResult(List.of(span1, span2));
86106
SpanJoiner joiner =
87107
this.spanJoinerBuilder
@@ -92,36 +112,55 @@ void fetchSpans() {
92112
List.of("pathToSpan"))
93113
.blockingGet();
94114
assertEquals(
95-
expected, joiner.joinSpans(joinSources, new TestJoinSourceIdGetter()).blockingGet());
115+
expected, joiner.joinSpan(joinSources, new TestJoinSourceIdGetter()).blockingGet());
96116
}
97117

98-
private void mockRequestBuilding() {
118+
@Test
119+
void fetchMultipleSpans() {
120+
Span span1 = new TestSpan(SPAN_ID1);
121+
Span span2 = new TestSpan(SPAN_ID2);
122+
TestMultipleJoinSource joinSource1 = new TestMultipleJoinSource(List.of(SPAN_ID1, SPAN_ID2));
123+
TestMultipleJoinSource joinSource2 = new TestMultipleJoinSource(List.of(SPAN_ID3, SPAN_ID4));
124+
ListMultimap<TestMultipleJoinSource, Span> expected = ArrayListMultimap.create();
125+
expected.put(joinSource1, span1);
126+
expected.put(joinSource1, span2);
127+
List<TestMultipleJoinSource> joinSources = List.of(joinSource1, joinSource2);
128+
when(mockSelectionFinder.findSelections(
129+
mockSelectionSet,
130+
SelectionQuery.builder().selectionPath(List.of("pathToSpans", "spans")).build()))
131+
.thenReturn(Stream.of(mock(SelectedField.class), mock(SelectedField.class)));
99132
when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anySet()))
100133
.thenReturn(Single.just(List.of(mockFilter)));
101134

102135
when(mockResultSetRequestBuilder.build(
103136
eq(mockRequestContext),
104137
eq(SPAN),
105-
eq(2),
138+
eq(4),
106139
eq(0),
107140
eq(mockTimeRangeArgument),
108141
eq(emptyList()),
109142
eq(List.of(mockFilter)),
110143
any(Stream.class),
111144
eq(Optional.empty())))
112145
.thenReturn(Single.just(mockResultSetRequest));
113-
}
114146

115-
private void mockRequestedSelectionFields(List<SelectedField> selectedFields, String location) {
116-
when(mockSelectionFinder.findSelections(
117-
mockSelectionSet,
118-
SelectionQuery.builder().selectionPath(List.of(location, "span")).build()))
119-
.thenReturn(selectedFields.stream());
147+
mockResult(List.of(span1, span2));
148+
SpanJoiner joiner =
149+
this.spanJoinerBuilder
150+
.build(
151+
this.mockRequestContext,
152+
this.mockTimeRangeArgument,
153+
this.mockSelectionSet,
154+
List.of("pathToSpans"))
155+
.blockingGet();
156+
assertEquals(
157+
expected,
158+
joiner.joinSpans(joinSources, new TestMultipleJoinSourceIdGetter()).blockingGet());
120159
}
121160

122161
private void mockResult(List<Span> spans) {
123162
when(mockSpanDao.getSpans(any(SpanRequest.class)))
124-
.thenAnswer(invocation -> Single.just(new TestSpanResultSet(spans)));
163+
.thenReturn(Single.just(new TestSpanResultSet(spans)));
125164
}
126165

127166
@Value
@@ -162,4 +201,17 @@ public LogEventResultSet logEvents() {
162201
return null;
163202
}
164203
}
204+
205+
@Value
206+
private static class TestMultipleJoinSource {
207+
List<String> spanIds;
208+
}
209+
210+
private static class TestMultipleJoinSourceIdGetter
211+
implements MultipleSpanIdGetter<TestMultipleJoinSource> {
212+
@Override
213+
public Single<List<String>> getSpanIds(TestMultipleJoinSource source) {
214+
return Single.just(source.getSpanIds());
215+
}
216+
}
165217
}

0 commit comments

Comments
 (0)