Skip to content

[FLINK-40294][table-planner] Instrument async scalar and table UDF calls with metrics (fork reference) #10

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
weiqingy wants to merge 1 commit into
base: flink-38071-pr2-sync
Choose a base branch
from
Open

[FLINK-40294][table-planner] Instrument async scalar and table UDF calls with metrics (fork reference) #10

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
9 changes: 8 additions & 1 deletion ...able-planner/src/main/java/org/apache/flink/table/planner/codegen/AsyncCodeGenerator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, String> 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);
Expand All @@ -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
Expand Down
15 changes: 15 additions & 0 deletions ...-planner/src/main/scala/org/apache/flink/table/planner/codegen/CodeGeneratorContext.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
*
Expand Down
98 changes: 76 additions & 22 deletions ...src/main/scala/org/apache/flink/table/planner/codegen/calls/BridgingFunctionGenUtil.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down Expand Up @@ -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
Expand All @@ -428,24 +433,33 @@ 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"""
|${externalOperands.map(_.code).mkString("\n")}
|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
}
}
Expand All @@ -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

Expand Down Expand Up @@ -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")
Expand All @@ -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,
Expand Down
Loading
Loading