Skip to content

Commit

Permalink
Tracing for deep search path
Browse files Browse the repository at this point in the history 
Signed-off-by: David Zane <[email protected]>
  • Loading branch information
@dzane17
dzane17 committed Jan 30, 2024
1 parent fb2c5f2 commit e6befb5
Show file tree
Hide file tree
Showing 16 changed files with 286 additions and 88 deletions.
15 changes: 10 additions & 5 deletions server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
import org.opensearch.search.internal.SearchContext;
import org.opensearch.search.internal.ShardSearchRequest;
import org.opensearch.search.pipeline.PipelinedRequest;
import org.opensearch.telemetry.tracing.SpanScope;
import org.opensearch.transport.Transport;

import java.util.ArrayDeque;
Expand Down Expand Up @@ -220,6 +221,7 @@ public final void start() {
null
)
);
searchRequestContext.getSearchRequestOperationsListener().onRequestEnd(searchRequestContext);
return;
}
executePhase(this);
Expand Down Expand Up @@ -442,18 +444,20 @@ private void onPhaseEnd(SearchRequestContext searchRequestContext) {
void onPhaseStart(SearchPhase phase) {
setCurrentPhase(phase);
if (SearchPhaseName.isValidName(phase.getName())) {
this.searchRequestContext.getSearchRequestOperationsListener().onPhaseStart(this);
this.searchRequestContext.getSearchRequestOperationsListener().onPhaseStart(this, searchRequestContext);
}
}

private void onRequestEnd(SearchRequestContext searchRequestContext) {
this.searchRequestContext.getSearchRequestOperationsListener().onRequestEnd(this, searchRequestContext);
this.searchRequestContext.getSearchRequestOperationsListener().onRequestEnd(searchRequestContext);
}

private void executePhase(SearchPhase phase) {
try {
onPhaseStart(phase);
phase.recordAndRun();
onPhaseStart(phase, searchRequestContext);
try (SpanScope spanScope = searchRequestContext.getTracer().withSpanInScope(searchRequestContext.getPhaseSpan())) {
phase.recordAndRun();
}
} catch (Exception e) {
if (logger.isDebugEnabled()) {
logger.debug(new ParameterizedMessage("Failed to execute [{}] while moving to [{}] phase", request, phase.getName()), e);
Expand Down Expand Up @@ -705,6 +709,7 @@ public void sendSearchResponse(InternalSearchResponse internalSearchResponse, At
searchContextId = null;
}
}
searchRequestContext.setSearchTask(getTask());
searchRequestContext.setTotalHits(internalSearchResponse.hits().getTotalHits());
searchRequestContext.setShardStats(results.getNumShards(), successfulOps.get(), skippedOps.get(), failures.length);
onPhaseEnd(searchRequestContext);
Expand All @@ -717,7 +722,7 @@ public void sendSearchResponse(InternalSearchResponse internalSearchResponse, At
@Override
public final void onPhaseFailure(SearchPhase phase, String msg, Throwable cause) {
if (SearchPhaseName.isValidName(phase.getName())) {
this.searchRequestContext.getSearchRequestOperationsListener().onPhaseFailure(this);
this.searchRequestContext.getSearchRequestOperationsListener().onPhaseFailure(this, searchRequestContext);
}
raisePhaseFailure(new SearchPhaseExecutionException(phase.getName(), msg, cause, buildShardFailures()));
}
Expand Down
75 changes: 70 additions & 5 deletions server/src/main/java/org/opensearch/action/search/SearchRequestContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,17 @@

package org.opensearch.action.search;

import org.apache.logging.log4j.LogManager;
import org.apache.lucene.search.TotalHits;
import org.opensearch.common.annotation.InternalApi;
import org.opensearch.telemetry.tracing.Span;
import org.opensearch.telemetry.tracing.Tracer;
import org.opensearch.telemetry.tracing.noop.NoopSpan;
import org.opensearch.telemetry.tracing.noop.NoopTracer;

import java.util.EnumMap;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;

Expand All @@ -23,20 +29,59 @@
*/
@InternalApi
class SearchRequestContext {
private final SearchRequest searchRequest;
private SearchTask searchTask;
private final SearchRequestOperationsListener searchRequestOperationsListener;
private long absoluteStartNanos;
private final Map<String, Long> phaseTookMap;
private TotalHits totalHits;
private final EnumMap<ShardStatsFieldNames, Integer> shardStats;
private Tracer tracer;
private Span requestSpan;
private Span phaseSpan;

private final SearchRequest searchRequest;
/**
* This constructor is for testing only
*/
SearchRequestContext() {
this(new SearchRequest());
}

/**
* This constructor is for testing only
*/
SearchRequestContext(SearchRequest searchRequest) {
this(new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()), searchRequest);
}

/**
* This constructor is for testing only
*/
SearchRequestContext(SearchRequestOperationsListener searchRequestOperationsListener, SearchRequest searchRequest) {
this(searchRequestOperationsListener, searchRequest, NoopTracer.INSTANCE);
}

SearchRequestContext(final SearchRequestOperationsListener searchRequestOperationsListener, final SearchRequest searchRequest) {
SearchRequestContext(SearchRequestOperationsListener searchRequestOperationsListener, SearchRequest searchRequest, Tracer tracer) {
this.searchRequestOperationsListener = searchRequestOperationsListener;
this.searchRequest = searchRequest;
this.tracer = tracer;
this.absoluteStartNanos = System.nanoTime();
this.phaseTookMap = new HashMap<>();
this.shardStats = new EnumMap<>(ShardStatsFieldNames.class);
this.searchRequest = searchRequest;
this.requestSpan = NoopSpan.INSTANCE;
this.phaseSpan = NoopSpan.INSTANCE;
}

public SearchRequest getSearchRequest() {
return searchRequest;
}

public void setSearchTask(SearchTask searchTask) {
this.searchTask = searchTask;
}

public SearchTask getSearchTask() {
return searchTask;
}

SearchRequestOperationsListener getSearchRequestOperationsListener() {
Expand Down Expand Up @@ -78,7 +123,7 @@ void setTotalHits(TotalHits totalHits) {
this.totalHits = totalHits;
}

TotalHits totalHits() {
public TotalHits totalHits() {
return totalHits;
}

Expand All @@ -89,7 +134,7 @@ void setShardStats(int total, int successful, int skipped, int failed) {
this.shardStats.put(ShardStatsFieldNames.SEARCH_REQUEST_SLOWLOG_SHARD_FAILED, failed);
}

String formattedShardStats() {
public String formattedShardStats() {
if (shardStats.isEmpty()) {
return "";
} else {
Expand All @@ -107,6 +152,26 @@ String formattedShardStats() {
);
}
}

public Tracer getTracer() {
return tracer;
}

public void setRequestSpan(Span requestSpan) {
this.requestSpan = requestSpan;
}

public Span getRequestSpan() {
return requestSpan;
}

public void setPhaseSpan(Span phaseSpan) {
this.phaseSpan = phaseSpan;
}

public Span getPhaseSpan() {
return phaseSpan;
}
}

enum ShardStatsFieldNames {
Expand Down
71 changes: 71 additions & 0 deletions server/src/main/java/org/opensearch/action/search/SearchRequestCoordinatorTrace.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
/*
* 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.action.search;

import org.opensearch.telemetry.tracing.AttributeNames;
import org.opensearch.telemetry.tracing.Span;
import org.opensearch.telemetry.tracing.SpanBuilder;
import org.opensearch.telemetry.tracing.SpanContext;
import org.opensearch.telemetry.tracing.Tracer;

import static org.opensearch.core.common.Strings.capitalize;

/**
* Listener for search request tracing on the coordinator node
*
* @opensearch.internal
*/
public final class SearchRequestCoordinatorTrace extends SearchRequestOperationsListener {
private final Tracer tracer;

public SearchRequestCoordinatorTrace(Tracer tracer) {
this.tracer = tracer;
}

@Override
void onPhaseStart(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
searchRequestContext.setPhaseSpan(
tracer.startSpan(
SpanBuilder.from(
"coord" + capitalize(context.getCurrentPhase().getName()),
new SpanContext(searchRequestContext.getRequestSpan())
)
)
);
}

@Override
void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
searchRequestContext.getPhaseSpan().endSpan();
}

@Override
void onPhaseFailure(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
searchRequestContext.getPhaseSpan().endSpan();
}

@Override
void onRequestEnd(SearchRequestContext searchRequestContext) {
Span requestSpan = searchRequestContext.getRequestSpan();

// add response-related attributes on request end
requestSpan.addAttribute(
AttributeNames.TOTAL_HITS,
searchRequestContext.totalHits() == null ? "0" : searchRequestContext.totalHits().toString()
);
requestSpan.addAttribute(
AttributeNames.SHARDS,
searchRequestContext.formattedShardStats().isEmpty() ? "null" : searchRequestContext.formattedShardStats()
);
requestSpan.addAttribute(
AttributeNames.SOURCE,
searchRequestContext.getSearchRequest().source() == null ? "null" : searchRequestContext.getSearchRequest().source().toString()
);
}
}
18 changes: 9 additions & 9 deletions server/src/main/java/org/opensearch/action/search/SearchRequestOperationsListener.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,15 +31,15 @@ protected SearchRequestOperationsListener(final boolean enabled) {
this.enabled = enabled;
}

abstract void onPhaseStart(SearchPhaseContext context);
abstract void onPhaseStart(SearchPhaseContext context, SearchRequestContext searchRequestContext);

abstract void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext);

abstract void onPhaseFailure(SearchPhaseContext context);
abstract void onPhaseFailure(SearchPhaseContext context, SearchRequestContext searchRequestContext);

void onRequestStart(SearchRequestContext searchRequestContext) {}

void onRequestEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {}
void onRequestEnd(SearchRequestContext searchRequestContext) {}

boolean isEnabled(SearchRequest searchRequest) {
return isEnabled();
Expand Down Expand Up @@ -69,10 +69,10 @@ static final class CompositeListener extends SearchRequestOperationsListener {
}

@Override
void onPhaseStart(SearchPhaseContext context) {
void onPhaseStart(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
for (SearchRequestOperationsListener listener : listeners) {
try {
listener.onPhaseStart(context);
listener.onPhaseStart(context, searchRequestContext);
} catch (Exception e) {
logger.warn(() -> new ParameterizedMessage("onPhaseStart listener [{}] failed", listener), e);
}
Expand All @@ -91,10 +91,10 @@ void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestCo
}

@Override
void onPhaseFailure(SearchPhaseContext context) {
void onPhaseFailure(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
for (SearchRequestOperationsListener listener : listeners) {
try {
listener.onPhaseFailure(context);
listener.onPhaseFailure(context, searchRequestContext);
} catch (Exception e) {
logger.warn(() -> new ParameterizedMessage("onPhaseFailure listener [{}] failed", listener), e);
}
Expand All @@ -113,10 +113,10 @@ void onRequestStart(SearchRequestContext searchRequestContext) {
}

@Override
public void onRequestEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
public void onRequestEnd(SearchRequestContext searchRequestContext) {
for (SearchRequestOperationsListener listener : listeners) {
try {
listener.onRequestEnd(context, searchRequestContext);
listener.onRequestEnd(searchRequestContext);
} catch (Exception e) {
logger.warn(() -> new ParameterizedMessage("onRequestEnd listener [{}] failed", listener), e);
}
Expand Down
Loading

0 comments on commit e6befb5

Please sign in to comment.