From 0088a54dc0c74f7e9a556ee7581439049fa1968a Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sun, 2 Aug 2026 00:27:38 -0700 Subject: [PATCH] [FLINK-40294][table-planner] Instrument async scalar and table UDF calls with metrics Extend the FLIP-485 UDF metrics to async scalar and table user-defined functions. The sampling decision and the start timestamp are captured at dispatch on the task thread; the elapsed time and any exceptional completion are recorded at completion on the callback thread, in the per-invocation DelegatingAsyncResultFuture and DelegatingAsyncTableResultFuture. Both writes happen before the completion callback is registered, establishing a happens-before edge to the callback thread; the histogram is synchronized and the exception counter is thread-safe. udfProcessingTime for an async function therefore spans the full dispatch to completion, not just the synchronous hand-off. As on the sync path, the instrumentation is emitted at code generation only when table.exec.udf-metric-enabled is true. --- .../planner/codegen/AsyncCodeGenerator.java | 9 +- .../codegen/CodeGeneratorContext.scala | 15 ++ .../calls/BridgingFunctionGenUtil.scala | 98 ++++++++--- .../runtime/stream/sql/UdfMetricsITCase.java | 157 +++++++++++++++++- .../async/DelegatingAsyncResultFuture.java | 33 ++++ .../DelegatingAsyncTableResultFuture.java | 36 ++++ 6 files changed, 324 insertions(+), 24 deletions(-) diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/codegen/AsyncCodeGenerator.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/codegen/AsyncCodeGenerator.java index 5a4bbc1118d554..d027e14ca74ed0 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/codegen/AsyncCodeGenerator.java +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/codegen/AsyncCodeGenerator.java @@ -152,9 +152,16 @@ private static String generateProcessCode( index++; } + // The async scalar call generated above registers the shared UdfMetrics handle when metrics + // are enabled; pass it into the per-invocation delegating future, which does the timing and + // exception counting. Null (feature off) yields the original two-argument construction. + String udfMetricsTerm = ctx.getSingleUdfMetricsTerm(); + String metricsCtorArg = udfMetricsTerm == null ? "" : ", " + udfMetricsTerm; + Map values = new HashMap<>(); values.put("delegatingFutureTerm", delegatingFutureTerm); values.put("delegatingFutureType", DelegatingAsyncResultFuture.class.getCanonicalName()); + values.put("metricsCtorArg", metricsCtorArg); values.put("collectorTerm", collectorTerm); values.put("typeTerm", GenericRowData.class.getCanonicalName()); values.put("recordTerm", recordTerm); @@ -169,7 +176,7 @@ private static String generateProcessCode( "\n", new String[] { "final ${delegatingFutureType} ${delegatingFutureTerm} ", - " = new ${delegatingFutureType}(${collectorTerm}, ${fieldCount});", + " = new ${delegatingFutureType}(${collectorTerm}, ${fieldCount}${metricsCtorArg});", "final org.apache.flink.types.RowKind rowKind = ${inputTerm}.getRowKind();\n", "try {", // Ensure that metadata setup come first so that we know that they're diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/CodeGeneratorContext.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/CodeGeneratorContext.scala index 2edf0dbac83ee3..372713e8828f75 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/CodeGeneratorContext.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/CodeGeneratorContext.scala @@ -33,6 +33,7 @@ import org.apache.flink.table.types.logical._ import org.apache.flink.table.types.logical.LogicalTypeRoot._ import org.apache.flink.table.utils.{DateTimeUtils, EncodingUtils} import org.apache.flink.util.InstantiationUtil +import org.apache.flink.util.Preconditions import java.time.ZoneId import java.util.TimeZone @@ -931,6 +932,20 @@ class CodeGeneratorContext( ) } + /** + * Returns the sole [[UdfMetrics]] handle term registered in this context, or `null` if none is. + * An async fetcher hosts exactly one async UDF, so at most one handle is ever registered; this + * lets the async scalar generator pass the handle into the per-invocation delegating future + * without re-deriving the UDF name. + */ + def getSingleUdfMetricsTerm: String = { + val errorMessage: Any = + s"An async fetcher hosts exactly one async UDF, but ${reusableUdfMetricsTerms.size} UDF " + + "metrics handles were registered in one context." + Preconditions.checkState(reusableUdfMetricsTerms.size <= 1, errorMessage) + if (reusableUdfMetricsTerms.size == 1) reusableUdfMetricsTerms.values.head else null + } + /** * Adds a reusable [[DataStructureConverter]] to the member area of the generated class. * diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/BridgingFunctionGenUtil.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/BridgingFunctionGenUtil.scala index 770934d4ed95ff..defbe8ad37c029 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/BridgingFunctionGenUtil.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/codegen/calls/BridgingFunctionGenUtil.scala @@ -278,18 +278,21 @@ object BridgingFunctionGenUtil { ) } else if (udf.getKind == FunctionKind.ASYNC_TABLE) { generateAsyncTableFunctionCall( + ctx, functionTerm, externalOperands, returnType, outputDataType, - skipIfArgsNull) + skipIfArgsNull, + udfMetricName) } else if (udf.getKind == FunctionKind.ASYNC_SCALAR) { generateAsyncScalarFunctionCall( ctx, functionTerm, externalOperands, returnType, - outputDataType) + outputDataType, + udfMetricName) } else { generateScalarFunctionCall(ctx, functionTerm, externalOperands, outputDataType, udfMetricName) } @@ -409,11 +412,13 @@ object BridgingFunctionGenUtil { } private def generateAsyncTableFunctionCall( + ctx: CodeGeneratorContext, functionTerm: String, externalOperands: Seq[GeneratedExpression], returnType: LogicalType, outputDataType: DataType, - skipIfArgsNull: Boolean): GeneratedExpression = { + skipIfArgsNull: Boolean, + udfMetricName: Option[String]): GeneratedExpression = { val DELEGATE_ASYNC_TABLE = className[DelegatingAsyncTableResultFuture] val outputType = outputDataType.getLogicalType @@ -428,6 +433,17 @@ object BridgingFunctionGenUtil { ) ++ externalOperands.map(_.resultTerm) val anyNull = externalOperands.map(_.nullTerm) ++ Seq("false") + // When metrics are enabled the handle is passed into the delegating future, which takes the + // sample decision at construction (dispatch, task thread) and records the completion span in + // its callback. The extra ctor argument is omitted when off, keeping the code byte-identical. + val metricsTerm = udfMetricsTermIfEnabled(ctx, udfMetricName) + val metricsCtorArg = metricsTerm.map(t => s", $t").getOrElse("") + val constructDelegate = + s"""$DELEGATE_ASYNC_TABLE delegates = new $DELEGATE_ASYNC_TABLE($DEFAULT_COLLECTOR_TERM, + | $needsWrapping, $isInternal$metricsCtorArg);""".stripMargin + val instrumentedEval = + instrumentAsyncDispatch(ctx, metricsTerm, s"$functionTerm.eval(${arguments.mkString(", ")});") + val functionCallCode = { if (skipIfArgsNull) { s""" @@ -435,17 +451,15 @@ object BridgingFunctionGenUtil { |if (${anyNull.mkString(" || ")}) { | $DEFAULT_COLLECTOR_TERM.complete(java.util.Collections.emptyList()); |} else { - | $DELEGATE_ASYNC_TABLE delegates = new $DELEGATE_ASYNC_TABLE($DEFAULT_COLLECTOR_TERM, - | $needsWrapping, $isInternal); - | $functionTerm.eval(${arguments.mkString(", ")}); + | $constructDelegate + | $instrumentedEval |} |""".stripMargin } else { s""" |${externalOperands.map(_.code).mkString("\n")} - |$DELEGATE_ASYNC_TABLE delegates = new $DELEGATE_ASYNC_TABLE($DEFAULT_COLLECTOR_TERM, - | $needsWrapping, $isInternal); - | $functionTerm.eval(${arguments.mkString(", ")}); + |$constructDelegate + | $instrumentedEval |""".stripMargin } } @@ -459,17 +473,25 @@ object BridgingFunctionGenUtil { functionTerm: String, externalOperands: Seq[GeneratedExpression], outputType: LogicalType, - outputDataType: DataType): GeneratedExpression = { + outputDataType: DataType, + udfMetricName: Option[String]): GeneratedExpression = { val converterTerm = ctx.addReusableConverter(outputDataType) + // Registering the handle here lets the async scalar fetcher pass it into the delegating future + // (see AsyncCodeGenerator); the future then takes the sample decision in createAsyncFuture on + // the task thread and records the completion span in its callback. + val metricsTerm = udfMetricsTermIfEnabled(ctx, udfMetricName) + val evalStatement = + s"""$functionTerm.eval( + | $DEFAULT_DELEGATING_FUTURE_TERM.createAsyncFuture($converterTerm), + | ${externalOperands.map(_.resultTerm).mkString(", ")});""".stripMargin + val instrumentedEval = instrumentAsyncDispatch(ctx, metricsTerm, evalStatement) val functionCallCode = s""" |${externalOperands.map(_.code).mkString("\n")} |if (${externalOperands.map(_.nullTerm).mkString(" || ")}) { | $DEFAULT_DELEGATING_FUTURE_TERM.createAsyncFuture($converterTerm).complete(null); |} else { - | $functionTerm.eval( - | $DEFAULT_DELEGATING_FUTURE_TERM.createAsyncFuture($converterTerm), - | ${externalOperands.map(_.resultTerm).mkString(", ")}); + | $instrumentedEval |} |""".stripMargin @@ -535,14 +557,8 @@ object BridgingFunctionGenUtil { ctx: CodeGeneratorContext, udfMetricName: Option[String], evalStatement: String): String = { - udfMetricName match { - case Some(name) - if ctx.tableConfig.get(ExecutionConfigOptions.TABLE_EXEC_UDF_METRIC_ENABLED) => - val sampleInterval = - ctx.tableConfig - .get(ExecutionConfigOptions.TABLE_EXEC_UDF_METRIC_SAMPLE_INTERVAL) - .intValue() - val metricsTerm = ctx.addReusableUdfMetrics(name, sampleInterval) + udfMetricsTermIfEnabled(ctx, udfMetricName) match { + case Some(metricsTerm) => val sampleTerm = ctx.addReusableLocalVariable("boolean", "udfSample") val startNanosTerm = ctx.addReusableLocalVariable("long", "udfStartNanos") val exceptionTerm = newName(ctx, "udfException") @@ -557,10 +573,48 @@ object BridgingFunctionGenUtil { |if ($sampleTerm) { | $metricsTerm.update(System.nanoTime() - $startNanosTerm); |}""".stripMargin - case _ => evalStatement + case None => evalStatement } } + /** + * Returns the shared [[UdfMetrics]] handle term when metrics are enabled for this call, else + * [[None]]. Acquiring the term registers the handle member and its `open()` registration exactly + * once per UDF name in the current context. + */ + private def udfMetricsTermIfEnabled( + ctx: CodeGeneratorContext, + udfMetricName: Option[String]): Option[String] = udfMetricName match { + case Some(name) if ctx.tableConfig.get(ExecutionConfigOptions.TABLE_EXEC_UDF_METRIC_ENABLED) => + val sampleInterval = + ctx.tableConfig + .get(ExecutionConfigOptions.TABLE_EXEC_UDF_METRIC_SAMPLE_INTERVAL) + .intValue() + Some(ctx.addReusableUdfMetrics(name, sampleInterval)) + case _ => None + } + + /** + * Brackets an async UDF dispatch `eval` so a synchronous throw (before the future is handed to + * the framework) still increments the exception counter, mirroring the sync path. Exceptional + * completions are counted separately in the delegating future's completion callback. Returns the + * statement unchanged when metrics are disabled, keeping the generated code byte-identical. + */ + private def instrumentAsyncDispatch( + ctx: CodeGeneratorContext, + metricsTerm: Option[String], + evalStatement: String): String = metricsTerm match { + case Some(term) => + val exceptionTerm = newName(ctx, "udfException") + s"""try { + | $evalStatement + |} catch (Throwable $exceptionTerm) { + | $term.markException(); + | throw $exceptionTerm; + |}""".stripMargin + case None => evalStatement + } + private def generateScalarFunctionCall( ctx: CodeGeneratorContext, functionTerm: String, diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/runtime/stream/sql/UdfMetricsITCase.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/runtime/stream/sql/UdfMetricsITCase.java index 0d76140c316c28..dadcfcddc6041e 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/runtime/stream/sql/UdfMetricsITCase.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/runtime/stream/sql/UdfMetricsITCase.java @@ -31,6 +31,9 @@ import org.apache.flink.table.api.TableResult; import org.apache.flink.table.api.bridge.java.StreamTableEnvironment; import org.apache.flink.table.api.config.ExecutionConfigOptions; +import org.apache.flink.table.functions.AsyncScalarFunction; +import org.apache.flink.table.functions.AsyncTableFunction; +import org.apache.flink.table.functions.FunctionContext; import org.apache.flink.table.functions.ScalarFunction; import org.apache.flink.table.functions.TableFunction; import org.apache.flink.table.planner.factories.TestValuesTableFactory; @@ -41,16 +44,22 @@ import org.junit.jupiter.api.extension.RegisterExtension; import java.util.Arrays; +import java.util.Collection; import java.util.List; import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; /** * End-to-end tests for the opt-in per-operator UDF metrics (FLIP-485) registered under {@code - * .udf.} for sync scalar and table user-defined functions. + * .udf.} for synchronous and asynchronous scalar and table user-defined + * functions. */ class UdfMetricsITCase { @@ -123,6 +132,88 @@ public void eval(Integer i) { } } + /** Doubles an int off the task thread; the metered async scalar path. */ + public static class AsyncIntDoubler extends AsyncScalarFunction { + private transient ScheduledExecutorService executor; + + @Override + public void open(FunctionContext context) { + executor = Executors.newSingleThreadScheduledExecutor(); + } + + @Override + public void close() { + if (executor != null) { + executor.shutdownNow(); + } + } + + public void eval(CompletableFuture future, Integer i) { + executor.schedule( + () -> future.complete(i == null ? null : i * 2), 5, TimeUnit.MILLISECONDS); + } + } + + /** + * Completes exceptionally the first {@code numFailures} invocations, then succeeds. Exercises + * the async completion-exception counter while the async operator's retry keeps the job alive. + */ + public static class AsyncFlaky extends AsyncScalarFunction { + private final int numFailures; + private final AtomicInteger failures = new AtomicInteger(); + private transient ScheduledExecutorService executor; + + public AsyncFlaky(int numFailures) { + this.numFailures = numFailures; + } + + @Override + public void open(FunctionContext context) { + executor = Executors.newSingleThreadScheduledExecutor(); + } + + @Override + public void close() { + if (executor != null) { + executor.shutdownNow(); + } + } + + public void eval(CompletableFuture future, Integer i) { + executor.schedule( + () -> { + if (failures.getAndIncrement() < numFailures) { + future.completeExceptionally(new RuntimeException("boom")); + } else { + future.complete(i); + } + }, + 5, + TimeUnit.MILLISECONDS); + } + } + + /** Emits each input twice off the task thread; the metered async table path. */ + public static class AsyncDuplicateRows extends AsyncTableFunction { + private transient ScheduledExecutorService executor; + + @Override + public void open(FunctionContext context) { + executor = Executors.newSingleThreadScheduledExecutor(); + } + + @Override + public void close() { + if (executor != null) { + executor.shutdownNow(); + } + } + + public void eval(CompletableFuture> future, Integer i) { + executor.schedule(() -> future.complete(Arrays.asList(i, i)), 5, TimeUnit.MILLISECONDS); + } + } + @Test void testSyncScalarMetricsRecorded() throws Exception { StreamTableEnvironment tEnv = createTableEnv(true); @@ -278,6 +369,70 @@ void testDistinctFunctionsGetSeparateHandles() throws Exception { .isEqualTo(SOURCE_ROWS.size()); } + @Test + void testAsyncScalarMetricsRecorded() throws Exception { + StreamTableEnvironment tEnv = createTableEnv(true); + tEnv.createTemporarySystemFunction("asyncudf", AsyncIntDoubler.class); + createSource(tEnv, "src", "id INT"); + createBlackHoleSink(tEnv, "sink", "v INT"); + + JobID jobId = execute(tEnv, "INSERT INTO sink SELECT asyncudf(id) FROM src"); + + // The processing time spans dispatch to off-thread completion; one sample per input row. + assertThat(histogram(jobId, processingTimePattern("asyncudf")).getCount()) + .isEqualTo(SOURCE_ROWS.size()); + assertThat(counter(jobId, exceptionCountPattern("asyncudf")).getCount()).isZero(); + } + + @Test + void testAsyncCompletionExceptionSurvivesJob() throws Exception { + StreamTableEnvironment tEnv = createTableEnv(true); + // Serialize invocations so the shared failure counter drives a deterministic retry. + tEnv.getConfig() + .set(ExecutionConfigOptions.TABLE_EXEC_ASYNC_SCALAR_MAX_CONCURRENT_OPERATIONS, 1); + tEnv.createTemporarySystemFunction("flakyudf", new AsyncFlaky(2)); + createSource(tEnv, "src", "id INT"); + createBlackHoleSink(tEnv, "sink", "v INT"); + + // Two exceptional completions are counted, then the async retry succeeds and the job + // finishes normally: an exceptional completion is a soft error, not a job failure. + JobID jobId = execute(tEnv, "INSERT INTO sink SELECT flakyudf(id) FROM src"); + + assertThat(counter(jobId, exceptionCountPattern("flakyudf")).getCount()) + .isGreaterThanOrEqualTo(1); + } + + @Test + void testAsyncTableMetricsRecorded() throws Exception { + StreamTableEnvironment tEnv = createTableEnv(true); + tEnv.createTemporarySystemFunction("asynctableudf", AsyncDuplicateRows.class); + createSource(tEnv, "src", "id INT"); + createBlackHoleSink(tEnv, "sink", "v INT"); + + JobID jobId = + execute( + tEnv, + "INSERT INTO sink SELECT x FROM src, " + + "LATERAL TABLE(asynctableudf(id)) AS T(x)"); + + // eval is called once per input row (it completes two rows); one sample each. + assertThat(histogram(jobId, processingTimePattern("asynctableudf")).getCount()) + .isEqualTo(SOURCE_ROWS.size()); + } + + @Test + void testAsyncDisabledRegistersNoUdfMetrics() throws Exception { + StreamTableEnvironment tEnv = createTableEnv(false); + tEnv.createTemporarySystemFunction("asyncoffudf", AsyncIntDoubler.class); + createSource(tEnv, "src", "id INT"); + createBlackHoleSink(tEnv, "sink", "v INT"); + + JobID jobId = execute(tEnv, "INSERT INTO sink SELECT asyncoffudf(id) FROM src"); + + assertThat(reporter.findMetrics(jobId, ANY_PROCESSING_TIME_PATTERN)).isEmpty(); + assertThat(reporter.findMetrics(jobId, ANY_EXCEPTION_COUNT_PATTERN)).isEmpty(); + } + private static StreamTableEnvironment createTableEnv(boolean udfMetricEnabled) { Configuration conf = new Configuration(); conf.set(RestartStrategyOptions.RESTART_STRATEGY, "disable"); diff --git a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/calc/async/DelegatingAsyncResultFuture.java b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/calc/async/DelegatingAsyncResultFuture.java index 1421be97ce9f92..b9396f34953fdd 100644 --- a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/calc/async/DelegatingAsyncResultFuture.java +++ b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/calc/async/DelegatingAsyncResultFuture.java @@ -22,9 +22,12 @@ import org.apache.flink.table.data.GenericRowData; import org.apache.flink.table.data.RowData; import org.apache.flink.table.data.conversion.DataStructureConverter; +import org.apache.flink.table.runtime.operators.metrics.UdfMetrics; import org.apache.flink.types.RowKind; import org.apache.flink.util.Preconditions; +import javax.annotation.Nullable; + import java.util.Collections; import java.util.HashMap; import java.util.Map; @@ -47,10 +50,26 @@ public class DelegatingAsyncResultFuture implements BiConsumer delegatedResultFuture, int totalResultSize) { + this(delegatedResultFuture, totalResultSize, null); + } + + public DelegatingAsyncResultFuture( + ResultFuture delegatedResultFuture, + int totalResultSize, + @Nullable UdfMetrics udfMetrics) { this.delegatedResultFuture = delegatedResultFuture; this.totalResultSize = totalResultSize; + this.udfMetrics = udfMetrics; } public synchronized void setRowKind(RowKind rowKind) { @@ -73,12 +92,26 @@ public CompletableFuture createAsyncFuture( Preconditions.checkState(this.asyncIndex >= 0); future = new CompletableFuture<>(); this.converter = converter; + // Sample decision taken on the task thread; the sampler counter is never touched + // off-thread. + if (udfMetrics != null) { + sample = udfMetrics.shouldSample(); + startNanos = sample ? System.nanoTime() : 0L; + } future.whenComplete(this); return future; } @Override public void accept(Object o, Throwable throwable) { + if (udfMetrics != null) { + if (throwable != null) { + udfMetrics.markException(); + } + if (sample) { + udfMetrics.update(System.nanoTime() - startNanos); + } + } if (throwable != null) { delegatedResultFuture.completeExceptionally(throwable); } else { diff --git a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/correlate/async/DelegatingAsyncTableResultFuture.java b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/correlate/async/DelegatingAsyncTableResultFuture.java index e745eca01a80cf..0350933619e4c7 100644 --- a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/correlate/async/DelegatingAsyncTableResultFuture.java +++ b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/correlate/async/DelegatingAsyncTableResultFuture.java @@ -20,8 +20,11 @@ import org.apache.flink.streaming.api.functions.async.ResultFuture; import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.runtime.operators.metrics.UdfMetrics; import org.apache.flink.types.Row; +import javax.annotation.Nullable; + import java.util.ArrayList; import java.util.Collection; import java.util.List; @@ -41,21 +44,54 @@ public class DelegatingAsyncTableResultFuture implements BiConsumer> completableFuture; + // Null unless UDF metrics are enabled. The sample decision and start-time are taken on the task + // thread in the constructor (invoked at dispatch, before eval); the histogram/counter are + // updated at completion (accept, callback thread). The two updated metrics are internally + // thread-safe; sample/startNanos are published to the callback via the future's completion. + @Nullable private final UdfMetrics udfMetrics; + private boolean sample; + private long startNanos; + public DelegatingAsyncTableResultFuture( ResultFuture delegatedResultFuture, boolean needsWrapping, boolean isInternalResultType) { + this(delegatedResultFuture, needsWrapping, isInternalResultType, null); + } + + public DelegatingAsyncTableResultFuture( + ResultFuture delegatedResultFuture, + boolean needsWrapping, + boolean isInternalResultType, + @Nullable UdfMetrics udfMetrics) { this.delegatedResultFuture = delegatedResultFuture; this.wrapFunction = needsWrapping ? (isInternalResultType ? this::wrapInternal : this::wrapExternal) : outs -> outs; this.completableFuture = new CompletableFuture<>(); + this.udfMetrics = udfMetrics; + // Sample decision taken on the task thread; the sampler counter is never touched + // off-thread. These writes must precede whenComplete below: the callback registration + // performs the volatile completion-stack push that establishes the happens-before edge + // carrying sample/startNanos to the completing thread's accept(). + if (udfMetrics != null) { + sample = udfMetrics.shouldSample(); + startNanos = sample ? System.nanoTime() : 0L; + } this.completableFuture.whenComplete(this); } @Override public void accept(Collection outs, Throwable throwable) { + if (udfMetrics != null) { + if (throwable != null) { + udfMetrics.markException(); + } + if (sample) { + udfMetrics.update(System.nanoTime() - startNanos); + } + } if (throwable != null) { delegatedResultFuture.completeExceptionally(throwable); return;