From 709f2a493f4e00d7ee7e3db11c5ed8e0e7b042f1 Mon Sep 17 00:00:00 2001 From: Ryan Caudy Date: Fri, 30 Jun 2023 15:02:37 -0400 Subject: [PATCH] Make AuthContext available to QPL, QOPL, and UPL (#4115) * Record AuthContext when creating QueryPerformanceNugget (for QPL and QOPL) and PerformanceEntry (for UPL). Add AuthContext (toString()ed) as a column in QPL, QOPL, and UPL (for the in-memory blink tables available in core). Enables registered loggers to also get AuthContext. * Add in the enclosing AuthContext in AbstractScriptSession.evaluateScript. --- .../table/impl/perf/PerformanceEntry.java | 13 +++++++++++ .../impl/perf/QueryPerformanceNugget.java | 23 +++++++++++++++---- .../UpdatePerformanceStreamPublisher.java | 4 +++- ...ryOperationPerformanceStreamPublisher.java | 4 +++- .../util/QueryPerformanceStreamPublisher.java | 6 ++++- .../engine/util/AbstractScriptSession.java | 6 +++-- 6 files changed, 46 insertions(+), 10 deletions(-) diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/PerformanceEntry.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/PerformanceEntry.java index c601599d7c1..5f6ce782001 100644 --- a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/PerformanceEntry.java +++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/PerformanceEntry.java @@ -3,7 +3,9 @@ */ package io.deephaven.engine.table.impl.perf; +import io.deephaven.auth.AuthContext; import io.deephaven.base.log.LogOutput; +import io.deephaven.engine.context.ExecutionContext; import io.deephaven.engine.rowset.RowSet; import io.deephaven.engine.rowset.RowSetShiftData; import io.deephaven.engine.table.TableListener; @@ -22,6 +24,8 @@ public class PerformanceEntry extends BasePerformanceEntry implements TableListe private final String description; private final String callerLine; + private final AuthContext authContext; + private long intervalInvocationCount; private long intervalAdded; @@ -44,6 +48,7 @@ public class PerformanceEntry extends BasePerformanceEntry implements TableListe this.operationNumber = operationNumber; this.description = description; this.callerLine = callerLine; + authContext = id == QueryConstants.NULL_INT ? null : ExecutionContext.getContext().getAuthContext(); startSample = new RuntimeMemory.Sample(); endSample = new RuntimeMemory.Sample(); maxTotalMemory = 0; @@ -114,6 +119,7 @@ public LogOutput append(final LogOutput logOutput) { .append(", operationNumber=").append(operationNumber) .append(", description='").append(description).append('\'') .append(", callerLine='").append(callerLine).append('\'') + .append(", authContext=").append(authContext) .append(", intervalUsageNanos=").append(getIntervalUsageNanos()) .append(", intervalCpuNanos=").append(getIntervalCpuNanos()) .append(", intervalUserCpuNanos=").append(getIntervalUserCpuNanos()) @@ -152,6 +158,13 @@ public String getCallerLine() { return callerLine; } + /** + * @return The {@link AuthContext} that was installed when this PerformanceEntry was constructed + */ + public AuthContext getAuthContext() { + return authContext; + } + public long getIntervalAdded() { return intervalAdded; } diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceNugget.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceNugget.java index 1f98fa15c68..76761e8b561 100644 --- a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceNugget.java +++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/QueryPerformanceNugget.java @@ -3,6 +3,8 @@ */ package io.deephaven.engine.table.impl.perf; +import io.deephaven.auth.AuthContext; +import io.deephaven.engine.context.ExecutionContext; import io.deephaven.time.DateTimeUtils; import io.deephaven.engine.table.impl.util.RuntimeMemory; import io.deephaven.util.QueryConstants; @@ -35,9 +37,11 @@ public class QueryPerformanceNugget implements Serializable, AutoCloseable { private final int depth; private final String description; private final boolean isUser; - private final String callerLine; private final long inputSize; + private final AuthContext authContext; + private final String callerLine; + private final long startClockTime; private final long startTimeNanos; @@ -98,14 +102,15 @@ public class QueryPerformanceNugget implements Serializable, AutoCloseable { this.isUser = isUser; this.inputSize = inputSize; + authContext = ExecutionContext.getContext().getAuthContext(); + callerLine = QueryPerformanceRecorder.getCallerLine(); + final RuntimeMemory runtimeMemory = RuntimeMemory.getInstance(); runtimeMemory.read(startMemorySample); startAllocatedBytes = ThreadProfiler.DEFAULT.getCurrentThreadAllocatedBytes(); startPoolAllocatedBytes = QueryPerformanceRecorder.getPoolAllocatedBytesForCurrentThread(); - callerLine = QueryPerformanceRecorder.getCallerLine(); - startClockTime = System.currentTimeMillis(); startTimeNanos = System.nanoTime(); @@ -128,11 +133,12 @@ private QueryPerformanceNugget() { isUser = false; inputSize = NULL_LONG; + authContext = null; + callerLine = null; + startAllocatedBytes = NULL_LONG; startPoolAllocatedBytes = NULL_LONG; - callerLine = null; - startClockTime = NULL_LONG; startTimeNanos = NULL_LONG; @@ -251,6 +257,13 @@ public long getInputSize() { return inputSize; } + /** + * @return The {@link AuthContext} that was installed when this QueryPerformanceNugget was constructed + */ + public AuthContext getAuthContext() { + return authContext; + } + public String getCallerLine() { return callerLine; } diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/UpdatePerformanceStreamPublisher.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/UpdatePerformanceStreamPublisher.java index c27380437cc..e4565185441 100644 --- a/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/UpdatePerformanceStreamPublisher.java +++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/perf/UpdatePerformanceStreamPublisher.java @@ -42,7 +42,8 @@ class UpdatePerformanceStreamPublisher implements StreamPublisher { ColumnDefinition.ofLong("Collections"), ColumnDefinition.ofLong("CollectionTimeNanos"), ColumnDefinition.ofLong("EntryIntervalAllocatedBytes"), - ColumnDefinition.ofLong("EntryIntervalPoolAllocatedBytes")); + ColumnDefinition.ofLong("EntryIntervalPoolAllocatedBytes"), + ColumnDefinition.ofString("AuthContext")); public static TableDefinition definition() { return DEFINITION; @@ -91,6 +92,7 @@ public synchronized void add(IntervalLevelDetails intervalLevelDetails, Performa chunks[20].asWritableLongChunk().add(performanceEntry.getCollectionTimeNanos()); chunks[21].asWritableLongChunk().add(performanceEntry.getIntervalAllocatedBytes()); chunks[22].asWritableLongChunk().add(performanceEntry.getIntervalPoolAllocatedBytes()); + chunks[23].asWritableObjectChunk().add(Objects.toString(performanceEntry.getAuthContext())); if (chunks[0].size() == CHUNK_SIZE) { flushInternal(); } diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/util/QueryOperationPerformanceStreamPublisher.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/util/QueryOperationPerformanceStreamPublisher.java index 8c0e970ad2e..bf32aec3e30 100644 --- a/engine/table/src/main/java/io/deephaven/engine/table/impl/util/QueryOperationPerformanceStreamPublisher.java +++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/util/QueryOperationPerformanceStreamPublisher.java @@ -42,7 +42,8 @@ class QueryOperationPerformanceStreamPublisher implements StreamPublisher { ColumnDefinition.ofLong("AllocatedBytes"), ColumnDefinition.ofLong("PoolAllocatedBytes"), ColumnDefinition.ofLong("InputSizeLong"), - ColumnDefinition.ofBoolean("WasInterrupted")); + ColumnDefinition.ofBoolean("WasInterrupted"), + ColumnDefinition.ofString("AuthContext")); private static final int CHUNK_SIZE = ArrayBackedColumnSource.BLOCK_SIZE; public static TableDefinition definition() { @@ -93,6 +94,7 @@ public synchronized void add( chunks[18].asWritableLongChunk().add(nugget.getAllocatedBytes()); chunks[19].asWritableLongChunk().add(nugget.getPoolAllocatedBytes()); chunks[20].asWritableByteChunk().add(BooleanUtils.booleanAsByte(nugget.wasInterrupted())); + chunks[21].asWritableObjectChunk().add(Objects.toString(nugget.getAuthContext())); if (chunks[0].size() == CHUNK_SIZE) { flushInternal(); } diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/util/QueryPerformanceStreamPublisher.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/util/QueryPerformanceStreamPublisher.java index 4ad1a4b67bd..03ad9d37541 100644 --- a/engine/table/src/main/java/io/deephaven/engine/table/impl/util/QueryPerformanceStreamPublisher.java +++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/util/QueryPerformanceStreamPublisher.java @@ -40,7 +40,8 @@ class QueryPerformanceStreamPublisher implements StreamPublisher { ColumnDefinition.ofLong("PoolAllocatedBytes"), ColumnDefinition.ofBoolean("WasInterrupted"), ColumnDefinition.ofBoolean("IsReplayer"), - ColumnDefinition.ofString("Exception")); + ColumnDefinition.ofString("Exception"), + ColumnDefinition.ofString("AuthContext")); private static final int CHUNK_SIZE = ArrayBackedColumnSource.BLOCK_SIZE; public static TableDefinition definition() { @@ -124,6 +125,9 @@ public synchronized void add( // ColumnDefinition.ofString("Exception") chunks[17].asWritableObjectChunk().add(queryProcessingResults.getException()); + // ColumnDefinition.ofString("AuthContext") + chunks[18].asWritableObjectChunk().add(Objects.toString(nugget.getAuthContext())); + if (chunks[0].size() == CHUNK_SIZE) { flushInternal(); } diff --git a/engine/table/src/main/java/io/deephaven/engine/util/AbstractScriptSession.java b/engine/table/src/main/java/io/deephaven/engine/util/AbstractScriptSession.java index 47ca7f47442..eaa6a3b8cad 100644 --- a/engine/table/src/main/java/io/deephaven/engine/util/AbstractScriptSession.java +++ b/engine/table/src/main/java/io/deephaven/engine/util/AbstractScriptSession.java @@ -146,8 +146,10 @@ public synchronized final Changes evaluateScript(final String script, final @Nul final SafeCloseable ignored = LivenessScopeStack.open(this, false)) { try { - // actually evaluate the script - executionContext.apply(() -> evaluate(script, scriptName)); + // Actually evaluate the script; use the enclosing auth context, since AbstractScriptSession's + // ExecutionContext never has a non-null AuthContext + executionContext.withAuthContext(ExecutionContext.getContext().getAuthContext()) + .apply(() -> evaluate(script, scriptName)); } catch (final RuntimeException err) { evaluateErr = err; }