Skip to content

Correctly associate async queries with the HTTP request for Datadog #7996

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
labkey-jeckels wants to merge 3 commits into
base: release26.7-SNAPSHOT
Choose a base branch
from
+182 −43
Open

Correctly associate async queries with the HTTP request for Datadog #7996

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
fc85042
Correctly associate async queries with the HTTP request for Datadog
labkey-jeckels Aug 28, 2026
5850067
Code review
labkey-jeckels Aug 28, 2026
0eae3d3
Better materialization tracing
labkey-jeckels Aug 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 45 additions & 10 deletions api/src/org/labkey/api/data/AsyncQueryRequest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,16 @@
package org.labkey.api.data;

import datadog.trace.api.CorrelationIdentifier;
import datadog.trace.api.DDTags;
import io.opentracing.Scope;
import io.opentracing.Span;
import io.opentracing.Tracer;
import io.opentracing.log.Fields;
import io.opentracing.tag.Tags;
import io.opentracing.util.GlobalTracer;
import org.apache.logging.log4j.Logger;
import org.apache.logging.log4j.ThreadContext;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.labkey.api.miniprofiler.MiniProfiler;
import org.labkey.api.miniprofiler.RequestInfo;
Expand All @@ -39,6 +43,7 @@
import java.io.IOException;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Map;
import java.util.concurrent.Callable;

public class AsyncQueryRequest<T>
Expand All @@ -52,16 +57,23 @@ private static class CancelledException extends RuntimeException
@Nullable
private final StackTraceElement[] _creationStackTrace;
private final HttpServletResponse _rootResponse;
/** Names the APM span; without it every async query aggregates under the bare operation name */
private final @Nullable String _resourceName;
/** Extra tags for APM search and grouping, beyond the resource name */
private final @NotNull Map<String, String> _spanTags;

boolean _cancelled;
@Nullable Statement _statement;
@Nullable Span _span;

T _result;
Throwable _exception;

public AsyncQueryRequest(HttpServletResponse response)
public AsyncQueryRequest(HttpServletResponse response, @Nullable String resourceName, @NotNull Map<String, String> spanTags)
{
_creationStackTrace = MiniProfiler.getTroubleshootingStackTrace();
_resourceName = resourceName;
_spanTags = spanTags;

_rootResponse = getRootResponse(response);
}
Expand Down Expand Up @@ -103,29 +115,38 @@ synchronized public T waitForResult(final Callable<T> callable) throws SQLExcept
final Object state = qs.cloneEnvironment();
final RequestInfo current = MemTracker.getInstance().current();

String traceId = CorrelationIdentifier.getTraceId();

// Create the span, so we can connect the async query with its owning thread
final Tracer tracer = GlobalTracer.get();
final Span span = tracer.buildSpan("AsyncRequest").start();
String spanId = CorrelationIdentifier.getSpanId();
if (_resourceName != null)
span.setTag(DDTags.RESOURCE_NAME, _resourceName);
_spanTags.forEach(span::setTag);
_span = span;

Runnable runnable = () -> {
if (current != null)
MemTracker.get().startProfiler("async query");

assert ThreadContext.isEmpty(); // Prevent/detect leaks
// Connect log messages with the active trace and span
ThreadContext.put(CorrelationIdentifier.getTraceIdKey(), traceId);
ThreadContext.put(CorrelationIdentifier.getSpanIdKey(), spanId);

qs.copyEnvironment(state);
try
// Activate on this thread, not the caller's, or the JDBC integration starts a new trace instead of parenting here
try (Scope ignored = tracer.activateSpan(span))
{
// Connect log messages with the active trace and span
ThreadContext.put(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId());
ThreadContext.put(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId());

setResult(callable.call());
}
catch (Throwable t)
{
// Cancellation arrives here as a CancelledException or the driver's "statement cancelled" error; tagging that as an APM error makes every client disconnect look like a failure
if (!isCancelled())
{
Tags.ERROR.set(span, true);
span.log(Map.of(Fields.ERROR_OBJECT, t));
}
setException(t);
}
finally
Expand All @@ -150,10 +171,11 @@ synchronized public T waitForResult(final Callable<T> callable) throws SQLExcept
Thread thread = new Thread(runnable, threadName);
// We want the async thread to use the same database connection, in case we have a transaction open, and
// so that when the original thread finishes processing the results it ends up closing the right connection
try (DbScope.ConnectionSharingCloseable ignored = DbScope.shareConnections(Thread.currentThread(), thread);
Scope ignore = tracer.activateSpan(span))
boolean started = false;
try (DbScope.ConnectionSharingCloseable ignored = DbScope.shareConnections(Thread.currentThread(), thread))
{
thread.start();
started = true;

// Stash the original disconnect exception if the client has dropped off
IOException clientDisconnectException = null;
Expand Down Expand Up @@ -209,6 +231,12 @@ synchronized public T waitForResult(final Callable<T> callable) throws SQLExcept
}
}
}
finally
{
// The runnable finishes the span; without a thread to run it the span stays open and its parent trace never completes
if (!started)
span.finish();
}
}

synchronized public void setResult(T result)
Expand All @@ -234,9 +262,16 @@ synchronized public void setException(Throwable exception)
notify();
}

synchronized private boolean isCancelled()
{
return _cancelled;
}

synchronized private void cancel()
{
_cancelled = true;
if (_span != null)
_span.setTag("labkey.async_query.cancelled", true);
if (_statement != null)
{
try
Expand Down
122 changes: 97 additions & 25 deletions api/src/org/labkey/api/data/MaterializedQueryHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,18 @@
*/
package org.labkey.api.data;

import datadog.trace.api.CorrelationIdentifier;
import datadog.trace.api.DDTags;
import io.opentracing.Scope;
import io.opentracing.Span;
import io.opentracing.Tracer;
import io.opentracing.log.Fields;
import io.opentracing.tag.Tags;
import io.opentracing.util.GlobalTracer;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.Strings;
import org.apache.logging.log4j.Logger;
import org.apache.logging.log4j.ThreadContext;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.junit.After;
Expand Down Expand Up @@ -153,31 +162,33 @@ boolean load(SQLFragment selectQuery, boolean isSelectInto)
DbSchema temp = DbSchema.getTemp();
TempTableTracker.track(_tableName, this);

SQLFragment selectInto;
if (isSelectInto)
{
String sql = selectQuery.getSQL().replace("${NAME}", _tableName);
List<Object> params = selectQuery.getParams();
selectInto = new SQLFragment(sql,params);
}
else
{
// UNLOGGED skips WAL when populating and indexing the table; only supported in PostgreSQL.
selectInto = new SQLFragment("SELECT * INTO ")
.append(_mqh._unlogged && _mqh._scope.getSqlDialect().isPostgreSQL() ? "UNLOGGED " : "")
.appendIdentifier(temp.getName()).append(".").appendIdentifier(_tableName).append("\nFROM (\n");
selectInto.append(selectQuery);
selectInto.append("\n) _sql_");
}
new SqlExecutor(_mqh._scope).execute(selectInto);
traced("full", _mqh.getMaterializationName(), () -> {
SQLFragment selectInto;
if (isSelectInto)
{
String sql = selectQuery.getSQL().replace("${NAME}", _tableName);
List<Object> params = selectQuery.getParams();
selectInto = new SQLFragment(sql,params);
}
else
{
// UNLOGGED skips WAL when populating and indexing the table; only supported in PostgreSQL.
selectInto = new SQLFragment("SELECT * INTO ")
.append(_mqh._unlogged && _mqh._scope.getSqlDialect().isPostgreSQL() ? "UNLOGGED " : "")
.appendIdentifier(temp.getName()).append(".").appendIdentifier(_tableName).append("\nFROM (\n");
selectInto.append(selectQuery);
selectInto.append("\n) _sql_");
}
new SqlExecutor(_mqh._scope).execute(selectInto);

try (var ignored = SpringActionController.ignoreSqlUpdates())
{
for (String index : _mqh._indexes)
try (var ignored = SpringActionController.ignoreSqlUpdates())
{
new SqlExecutor(_mqh._scope).execute(StringUtils.replace(index, "${NAME}", _tableName));
for (String index : _mqh._indexes)
{
new SqlExecutor(_mqh._scope).execute(StringUtils.replace(index, "${NAME}", _tableName));
}
}
}
});

_loadingState.set(LoadingState.LOADED);
return true;
Expand Down Expand Up @@ -422,17 +433,43 @@ public void materializeAsync()
{
if (_backgroundTaskRunning.compareAndSet(false, true))
{
String triggeringTraceId = CorrelationIdentifier.getTraceId();
String triggeringSpanId = CorrelationIdentifier.getSpanId();

_materializationRunner.execute(() -> {
try
// Materialization outlives the request that triggers it and then serves every later request, so it gets its own trace; the trigger is recorded as tags rather than as a parent
Tracer tracer = GlobalTracer.get();
Span span = tracer.buildSpan("MaterializeAsync").ignoreActiveSpan().start();
span.setTag(DDTags.RESOURCE_NAME, StringUtils.defaultIfEmpty(_prefix, getClass().getSimpleName()));
span.setTag("labkey.materialized_view", getMaterializationName());
if (!"0".equals(triggeringTraceId))
{
span.setTag("labkey.triggering_trace_id", triggeringTraceId);
span.setTag("labkey.triggering_span_id", triggeringSpanId);
}

try (Scope ignored = tracer.activateSpan(span))
{
// Connect log messages with the active trace and span
ThreadContext.put(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId());
ThreadContext.put(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId());

getFromSql("_bg_");
}
catch (Exception e)
catch (Throwable t)
{
LOG.warn("Background materialization failed.", e);
Tags.ERROR.set(span, true);
span.log(Map.of(Fields.ERROR_OBJECT, t));
LOG.warn("Background materialization failed.", t);
// Broad enough to tag an Error on the span, but only Exceptions are swallowed
if (t instanceof Error e)
throw e;
}
finally
{
span.finish();
ThreadContext.remove(CorrelationIdentifier.getTraceIdKey());
ThreadContext.remove(CorrelationIdentifier.getSpanIdKey());
_backgroundTaskRunning.set(false);
}
});
Expand Down Expand Up @@ -552,6 +589,41 @@ protected void incrementalUpdateBeforeSelect(Materialized m)
{
}

/**
* Runs DB work inside its own Datadog APM span so each kind of materialization is a separate resource. Nests under
* MaterializeAsync on the background thread and under the HTTP request when a stale view is rebuilt inline.
*/
protected static void traced(String resource, String viewName, Runnable work)
{
Tracer tracer = GlobalTracer.get();
Span span = tracer.buildSpan("labkey.materialize").start();
span.setTag(DDTags.RESOURCE_NAME, resource);
// Never a service-entry span, so Datadog computes no hits/duration/error metrics for it without this
span.setTag(DDTags.MEASURED, true);
span.setTag("labkey.materialized_view", viewName);

try (Scope ignored = tracer.activateSpan(span))
{
work.run();
}
catch (Throwable t)
{
Tags.ERROR.set(span, true);
span.log(Map.of(Fields.ERROR_OBJECT, t));
throw t;
}
finally
{
span.finish();
}
}

/** Identifies the view in APM. Carried as a tag, never a resource name, to keep per-view cardinality out of trace metrics. */
protected String getMaterializationName()
{
return StringUtils.defaultIfEmpty(_prefix, getClass().getSimpleName());
}

/**
* A Materialized represents a particular instance of materialized view (stored in a temp table).
* We want to avoid two threads materializing the same view. This is why we synchronize first creating the
Expand Down
30 changes: 28 additions & 2 deletions api/src/org/labkey/api/data/TableSelector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
Expand Down Expand Up @@ -379,10 +380,34 @@ public Results getResults(boolean cache, boolean scrollable)
return new ResultsImpl(rs, tableSqlFactory.getSelectedColumns());
}

/** @return "schema.query", using the public (Query) names when the table has them, otherwise the DB schema and table */
private String getAsyncQueryName()
{
String schema = _table.getPublicSchemaName();
String name = _table.getPublicName();
if (null == schema || null == name)
{
schema = null != _table.getSchema() ? _table.getSchema().getName() : null;
name = _table.getName();
}
return (null != schema ? schema + "." : "") + name;
}

/** Extra APM span tags. resource.name concatenates the operation, so these are what let one query be grouped across getResults and getAggregates. */
private Map<String, String> getAsyncSpanTags(String queryName)
{
Map<String, String> tags = new HashMap<>();
tags.put("labkey.query", queryName);
if (null != _table.getSchema())
tags.put("labkey.db_schema", _table.getSchema().getName());
return tags;
}

public Results getResultsAsync(final boolean cache, final boolean scrollable, HttpServletResponse response) throws SQLException
{
setLogger(ConnectionWrapper.getConnectionLogger());
AsyncQueryRequest<Results> asyncRequest = new AsyncQueryRequest<>(response);
String queryName = getAsyncQueryName();
AsyncQueryRequest<Results> asyncRequest = new AsyncQueryRequest<>(response, "getResults " + queryName, getAsyncSpanTags(queryName));
setAsyncRequest(asyncRequest);

try
Expand Down Expand Up @@ -554,7 +579,8 @@ public Map<String, List<Result>> getAggregates(final List<Aggregate> aggregates)
public Map<String, List<Result>> getAggregatesAsync(final List<Aggregate> aggregates, HttpServletResponse response)
{
setLogger(ConnectionWrapper.getConnectionLogger());
AsyncQueryRequest<Map<String, List<Result>>> asyncRequest = new AsyncQueryRequest<>(response);
String queryName = getAsyncQueryName();
AsyncQueryRequest<Map<String, List<Result>>> asyncRequest = new AsyncQueryRequest<>(response, "getAggregates " + queryName, getAsyncSpanTags(queryName));
setAsyncRequest(asyncRequest);

try
Expand Down
18 changes: 12 additions & 6 deletions experiment/src/org/labkey/experiment/api/ExpMaterialTableImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1493,10 +1493,10 @@ protected void incrementalUpdateBeforeSelect(Materialized m)
if (Materialized.LoadingState.ERROR == materialized._loadingState.get())
throw materialized._loadException;

runIncremental(materialized.incrementalDeleteCheck, this::executeIncrementalDelete);
runIncremental(materialized.incrementalUpdateCheck, this::executeIncrementalUpdate);
runIncremental(materialized.incrementalRollupCheck, this::executeIncrementalRollup);
runIncremental(materialized.incrementalInsertCheck, this::executeIncrementalInsert);
runIncremental("delete", materialized.incrementalDeleteCheck, this::executeIncrementalDelete);
runIncremental("update", materialized.incrementalUpdateCheck, this::executeIncrementalUpdate);
runIncremental("rollup", materialized.incrementalRollupCheck, this::executeIncrementalRollup);
runIncremental("insert", materialized.incrementalInsertCheck, this::executeIncrementalInsert);
}
catch (RuntimeException|InterruptedException ex)
{
Expand Down Expand Up @@ -1524,15 +1524,21 @@ protected void incrementalUpdateBeforeSelect(Materialized m)
* reflects the change, instead of briefly serving a stale materialized view. Must be called while holding the
* materialized's loading lock so only one updater runs at a time.
*/
private static void runIncremental(MaterializedQueryHelper.SupplierInvalidator check, Runnable work)
private void runIncremental(String kind, MaterializedQueryHelper.SupplierInvalidator check, Runnable work)
{
if (check.peekValid())
return;
String token = check.current();
work.run();
traced("incremental." + kind, getMaterializationName(), work);
check.markValidAs(token);
}

@Override
protected String getMaterializationName()
{
return StringUtils.defaultIfEmpty(Lsid.parse(_lsid).getObjectId(), _lsid);
}

void upsertWithRetry(SQLFragment sql)
{
// not actually read-only, but we don't want to start an explicit transaction
Expand Down