Skip to content

Commit 6085ce7

Browse files
aiceflowerv-kkhuangcasionone
authored
#AI commit# feat: 实现starrocks任务重试机制 (#5386)
- 引入ComputationEngineUtils工具类，提供向entrance发送消息的功能 - 添加配置项支持上下文语句和JDBC SET语句的错误索引调整 - 实现任务重试逻辑，支持JDBC和aiSQL类型的断点续跑 - 优化结果集别名数量管理，确保重试时状态正确恢复 - 新增TaskRetryInterceptor拦截器，动态为任务添加重试开关 - 增强错误处理机制，支持多种错误场景的正则匹配重试 - 修复任务重试时结果目录重复创建的问题 - 扩展ResponseTaskStatusWithExecuteCodeIndex协议，增加aliasNum字段 - 优化JDBC任务重试时的SQL索引调整逻辑，处理SET语句场景 - 增加StarRocks引擎类型支持和相关配置选项 Co-authored-by: v-kkhuang <420895376@qq.com> Co-authored-by: Casion <casionone@gmail.com>
1 parent 0ec73b1 commit 6085ce7

16 files changed

Lines changed: 365 additions & 61 deletions

File tree

‎linkis-commons/linkis-common/src/main/scala/org/apache/linkis/common/conf/Configuration.scala‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,12 @@ object Configuration extends Logging {
104104
val METRICS_INCREMENTAL_UPDATE_ENABLE =
105105
CommonVars[Boolean]("linkis.jobhistory.metrics.incremental.update.enable", false)
106106

107+
val EXECUTE_ERROR_CODE_INDEX =
108+
CommonVars("execute.error.code.index", "-1")
109+
110+
val EXECUTE_RESULTSET_ALIAS_NUM =
111+
CommonVars("execute.resultset.alias.num", "0")
112+
107113
val GLOBAL_CONF_CHN_NAME = "全局设置"
108114

109115
val GLOBAL_CONF_CHN_OLDNAME = "通用设置"

‎linkis-computation-governance/linkis-computation-governance-common/src/main/scala/org/apache/linkis/governance/common/protocol/task/ResponseTaskExecute.scala‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,10 +52,14 @@ case class ResponseTaskStatus(execId: String, status: ExecutionNodeStatus)
5252
class ResponseTaskStatusWithExecuteCodeIndex(
5353
execId: String,
5454
status: ExecutionNodeStatus,
55-
private var _errorIndex: Int = -1
55+
private var _errorIndex: Int = -1,
56+
private var _aliasNum: Int = 0 // 新增：aliasNum字段
5657
) extends ResponseTaskStatus(execId, status) {
5758
def errorIndex: Int = _errorIndex
5859
def errorIndex_=(value: Int): Unit = _errorIndex = value
60+
// 新增：aliasNum的getter和setter
61+
def aliasNum: Int = _aliasNum
62+
def aliasNum_=(value: Int): Unit = _aliasNum = value
5963
}
6064

6165
case class ResponseTaskResultSet(execId: String, output: String, alias: String)

‎linkis-computation-governance/linkis-engineconn/linkis-computation-engineconn/src/main/scala/org/apache/linkis/engineconn/computation/executor/conf/ComputationExecutorConf.scala‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,4 +155,16 @@ object ComputationExecutorConf {
155155
val SUPPORT_PARTIAL_RETRY_FOR_FAILED_TASKS_ENABLED: Boolean =
156156
CommonVars[Boolean]("linkis.partial.retry.for.failed.task.enabled", false).getValue
157157

158+
val CONTEXT_STATEMENT_PREFIXES = CommonVars(
159+
"linkis.engineconn.context.statement.prefixes",
160+
"USE ,SET ,ALTER SESSION ,SET ROLE ,SET SCHEMA ",
161+
"SQL context statement prefixes for partial retry"
162+
)
163+
164+
val JDBC_SET_STATEMENT_PREFIXES = CommonVars(
165+
"linkis.engineconn.jdbc.set.statement.prefixes",
166+
"SET QUERY_TIMEOUT,SET QUERY_QUEUE_PENDING_TIMEOUT,SET NEW_PLANNER",
167+
"JDBC SET statement prefixes for error index adjustment"
168+
)
169+
158170
}

‎linkis-computation-governance/linkis-engineconn/linkis-computation-engineconn/src/main/scala/org/apache/linkis/engineconn/computation/executor/execute/ComputationExecutor.scala‎

Lines changed: 125 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
package org.apache.linkis.engineconn.computation.executor.execute
1919

2020
import org.apache.linkis.DataWorkCloudApplication
21+
import org.apache.linkis.common.conf.Configuration
2122
import org.apache.linkis.common.log.LogUtils
2223
import org.apache.linkis.common.utils.{Logging, Utils}
2324
import org.apache.linkis.engineconn.acessible.executor.entity.AccessibleExecutor
@@ -35,16 +36,22 @@ import org.apache.linkis.engineconn.computation.executor.exception.HookExecuteEx
3536
import org.apache.linkis.engineconn.computation.executor.hook.ComputationExecutorHook
3637
import org.apache.linkis.engineconn.computation.executor.metrics.ComputationEngineConnMetrics
3738
import org.apache.linkis.engineconn.computation.executor.upstream.event.TaskStatusChangedForUpstreamMonitorEvent
39+
import org.apache.linkis.engineconn.computation.executor.utlis.ComputationEngineUtils
3840
import org.apache.linkis.engineconn.core.EngineConnObject
3941
import org.apache.linkis.engineconn.core.executor.ExecutorManager
4042
import org.apache.linkis.engineconn.executor.entity.{LabelExecutor, ResourceExecutor}
4143
import org.apache.linkis.engineconn.executor.listener.ExecutorListenerBusContext
4244
import org.apache.linkis.governance.common.entity.ExecutionNodeStatus
4345
import org.apache.linkis.governance.common.paser.CodeParser
44-
import org.apache.linkis.governance.common.protocol.task.{EngineConcurrentInfo, RequestTask}
46+
import org.apache.linkis.governance.common.protocol.task.{
47+
EngineConcurrentInfo,
48+
RequestTask,
49+
ResponseTaskError,
50+
ResponseTaskStatusWithExecuteCodeIndex
51+
}
4552
import org.apache.linkis.governance.common.utils.{JobUtils, LoggerUtils}
4653
import org.apache.linkis.manager.common.entity.enumeration.NodeStatus
47-
import org.apache.linkis.manager.label.entity.engine.{EngineType, UserCreatorLabel}
54+
import org.apache.linkis.manager.label.entity.engine.{EngineType, EngineTypeLabel, UserCreatorLabel}
4855
import org.apache.linkis.manager.label.utils.LabelUtil
4956
import org.apache.linkis.protocol.engine.JobProgressInfo
5057
import org.apache.linkis.scheduler.executer._
@@ -260,17 +267,40 @@ abstract class ComputationExecutor(val outputPrintLimit: Int = 1000)
260267
}
261268
var executeFlag = true
262269
val errorIndex: Int = Integer.valueOf(
263-
engineConnTask.getProperties.getOrDefault("execute.error.code.index", "-1").toString
270+
engineConnTask.getProperties
271+
.getOrDefault(Configuration.EXECUTE_ERROR_CODE_INDEX.key, "-1")
272+
.toString
264273
)
265-
engineExecutionContext.getProperties.put("execute.error.code.index", errorIndex.toString)
274+
engineExecutionContext.getProperties
275+
.put(Configuration.EXECUTE_ERROR_CODE_INDEX.key, errorIndex.toString)
276+
// jdbc执行任务重试，如果sql有被set进sql，会导致sql的index错位，这里会将日志打印的index进行减一，保证用户看的index是正常的，然后重试的errorIndex需要加一，保证重试的位置是正确的
277+
var newIndex = index
278+
var newErrorIndex = errorIndex
279+
if (adjustErrorIndexForSetScenarios(engineConnTask)) {
280+
newIndex = index - 1
281+
newErrorIndex = errorIndex + 1
282+
}
266283
// 重试的时候如果执行过则跳过执行
267-
if (retryEnable && errorIndex > 0 && index < errorIndex) {
268-
engineExecutionContext.appendStdout(
269-
LogUtils.generateInfo(
270-
s"aisql retry with errorIndex: ${errorIndex}, current sql index: ${index} will skip."
284+
if (retryEnable && errorIndex > 0 && index < newErrorIndex) {
285+
val code = codes(index).trim.toUpperCase()
286+
val shouldSkip = !isContextStatement(code)
287+
288+
if (shouldSkip) {
289+
engineExecutionContext.appendStdout(
290+
LogUtils.generateInfo(
291+
s"task retry with errorIndex: ${errorIndex}, current sql index: ${newIndex} will skip."
292+
)
271293
)
272-
)
273-
executeFlag = false
294+
executeFlag = false
295+
} else {
296+
if (newIndex >= 0) {
297+
engineExecutionContext.appendStdout(
298+
LogUtils.generateInfo(
299+
s"task retry with errorIndex: ${errorIndex}, current sql index: ${newIndex} is a context statement, will execute."
300+
)
301+
)
302+
}
303+
}
274304
}
275305
if (executeFlag) {
276306
val code = codes(index)
@@ -285,21 +315,23 @@ abstract class ComputationExecutor(val outputPrintLimit: Int = 1000)
285315
response match {
286316
case e: ErrorExecuteResponse =>
287317
val props: util.Map[String, String] = engineCreationContext.getOptions
288-
val aiSqlEnable: String = props.getOrDefault("linkis.ai.sql.enable", "false").toString
318+
val taskRetry: String =
319+
props.getOrDefault("linkis.task.retry.switch", "false").toString
289320
val retryNum: Int =
290-
Integer.valueOf(props.getOrDefault("linkis.ai.retry.num", "0").toString)
321+
Integer.valueOf(props.getOrDefault("linkis.task.retry.num", "0").toString)
291322

292-
if (retryEnable && !props.isEmpty && "true".equals(aiSqlEnable) && retryNum > 0) {
323+
if (retryEnable && !props.isEmpty && "true".equals(taskRetry) && retryNum > 0) {
293324
logger.info(
294-
s"aisql execute failed, with index: ${index} retryNum: ${retryNum}, and will retry",
325+
s"task execute failed, with index: ${index} retryNum: ${retryNum}, and will retry",
295326
e.t
296327
)
297328
engineExecutionContext.appendStdout(
298329
LogUtils.generateInfo(
299-
s"aisql execute failed, with index: ${index} retryNum: ${retryNum}, and will retry"
330+
s"task execute failed, with index: ${index} retryNum: ${retryNum}, and will retry"
300331
)
301332
)
302-
engineConnTask.getProperties.put("execute.error.code.index", index.toString)
333+
engineConnTask.getProperties
334+
.put(Configuration.EXECUTE_ERROR_CODE_INDEX.key, index.toString)
303335
return ErrorRetryExecuteResponse(e.message, index, e.t)
304336
} else {
305337
failedTasks.increase()
@@ -316,7 +348,14 @@ abstract class ComputationExecutor(val outputPrintLimit: Int = 1000)
316348
e.getOutput.substring(0, outputPrintLimit)
317349
} else e.getOutput
318350
engineExecutionContext.appendStdout(output)
319-
if (StringUtils.isNotBlank(e.getOutput)) engineExecutionContext.sendResultSet(e)
351+
if (StringUtils.isNotBlank(e.getOutput)) {
352+
engineConnTask.getProperties
353+
.put(
354+
Configuration.EXECUTE_RESULTSET_ALIAS_NUM.key,
355+
engineExecutionContext.getAliasNum.toString
356+
)
357+
engineExecutionContext.sendResultSet(e)
358+
}
320359
case _: IncompleteExecuteResponse =>
321360
incomplete ++= incompleteSplitter
322361
}
@@ -365,6 +404,31 @@ abstract class ComputationExecutor(val outputPrintLimit: Int = 1000)
365404
executeResponse match {
366405
case successExecuteResponse: SuccessExecuteResponse =>
367406
transformTaskStatus(engineConnTask, ExecutionNodeStatus.Succeed)
407+
case ErrorRetryExecuteResponse(message, index, throwable) =>
408+
ComputationEngineUtils.sendToEntrance(
409+
engineConnTask,
410+
ResponseTaskError(engineConnTask.getTaskId, message)
411+
)
412+
logger.warn(
413+
s"The task begins executing retries,jobId:${engineConnTask.getTaskId},index:${index} ,message:${message}",
414+
throwable
415+
)
416+
417+
val currentAliasNum = Integer.valueOf(
418+
engineConnTask.getProperties
419+
.getOrDefault(Configuration.EXECUTE_RESULTSET_ALIAS_NUM.key, "0")
420+
.toString
421+
)
422+
423+
ComputationEngineUtils.sendToEntrance(
424+
engineConnTask,
425+
new ResponseTaskStatusWithExecuteCodeIndex(
426+
engineConnTask.getTaskId,
427+
ExecutionNodeStatus.Failed,
428+
index,
429+
currentAliasNum
430+
)
431+
)
368432
case errorExecuteResponse: ErrorExecuteResponse =>
369433
listenerBusContext.getEngineConnSyncListenerBus.postToAll(
370434
TaskResponseErrorEvent(engineConnTask.getTaskId, errorExecuteResponse.message)
@@ -400,6 +464,22 @@ abstract class ComputationExecutor(val outputPrintLimit: Int = 1000)
400464

401465
def getProgressInfo(taskID: String): Array[JobProgressInfo]
402466

467+
/**
468+
* 调整错误索引：直接匹配三种SET语句场景 因为SET语句会被解析器视为第一条SQL
469+
*/
470+
protected def adjustErrorIndexForSetScenarios(engineConnTask: EngineConnTask): Boolean = {
471+
val executionCode = engineConnTask.getCode
472+
val engineTypeLabel = engineConnTask.getLables.find(_.isInstanceOf[EngineTypeLabel]).get
473+
val engineType = engineTypeLabel.asInstanceOf[EngineTypeLabel].getEngineType
474+
var result = false
475+
if (executionCode != null && engineType.equals(EngineType.JDBC.toString)) {
476+
val upperCode = executionCode.toUpperCase().trim
477+
val jdbcSetPrefixes = ComputationExecutorConf.JDBC_SET_STATEMENT_PREFIXES.getValue.split(",")
478+
result = jdbcSetPrefixes.exists(upperCode.startsWith)
479+
}
480+
result
481+
}
482+
403483
protected def createEngineExecutionContext(
404484
engineConnTask: EngineConnTask
405485
): EngineExecutionContext = {
@@ -423,6 +503,22 @@ abstract class ComputationExecutor(val outputPrintLimit: Int = 1000)
423503
engineExecutionContext.setJobId(engineConnTask.getTaskId)
424504
engineExecutionContext.getProperties.putAll(engineConnTask.getProperties)
425505
engineExecutionContext.setLabels(engineConnTask.getLables)
506+
507+
val errorIndex: Int = Integer.valueOf(
508+
engineConnTask.getProperties
509+
.getOrDefault(Configuration.EXECUTE_ERROR_CODE_INDEX.key, "-1")
510+
.toString
511+
)
512+
if (errorIndex > 0) {
513+
val savedAliasNum = Integer.valueOf(
514+
engineConnTask.getProperties
515+
.getOrDefault(Configuration.EXECUTE_RESULTSET_ALIAS_NUM.key, "0")
516+
.toString
517+
)
518+
engineExecutionContext.setResultSetNum(savedAliasNum)
519+
logger.info(s"Restore aliasNum to $savedAliasNum for retry task")
520+
}
521+
426522
engineExecutionContext
427523
}
428524

@@ -435,6 +531,18 @@ abstract class ComputationExecutor(val outputPrintLimit: Int = 1000)
435531
}
436532
}
437533

534+
/**
535+
* 判断是否为上下文语句，重试时需要保留执行
536+
*
537+
* @param code
538+
* SQL代码（已转换为大写并去除首尾空格）
539+
* @return
540+
* true表示是上下文语句，false表示不是
541+
*/
542+
private def isContextStatement(code: String): Boolean = {
543+
ComputationExecutorConf.CONTEXT_STATEMENT_PREFIXES.getValue.split(",").exists(code.startsWith)
544+
}
545+
438546
/**
439547
* job task log print task params info
440548
*

‎linkis-computation-governance/linkis-engineconn/linkis-computation-engineconn/src/main/scala/org/apache/linkis/engineconn/computation/executor/service/TaskExecutionServiceImpl.scala‎

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -249,17 +249,6 @@ class TaskExecutionServiceImpl
249249
sendToEntrance(task, ResponseTaskError(task.getTaskId, message))
250250
logger.error(message, throwable)
251251
sendToEntrance(task, ResponseTaskStatus(task.getTaskId, ExecutionNodeStatus.Failed))
252-
case ErrorRetryExecuteResponse(message, index, throwable) =>
253-
sendToEntrance(task, ResponseTaskError(task.getTaskId, message))
254-
logger.error(message, throwable)
255-
sendToEntrance(
256-
task,
257-
new ResponseTaskStatusWithExecuteCodeIndex(
258-
task.getTaskId,
259-
ExecutionNodeStatus.Failed,
260-
index
261-
)
262-
)
263252
case _ =>
264253
}
265254
LoggerUtils.removeJobIdMDC()

‎linkis-computation-governance/linkis-engineconn/linkis-computation-engineconn/src/main/scala/org/apache/linkis/engineconn/computation/executor/utlis/ComputationEngineUtils.scala‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,19 @@
1717

1818
package org.apache.linkis.engineconn.computation.executor.utlis
1919

20+
import org.apache.linkis.common.utils.{Logging, Utils}
21+
import org.apache.linkis.engineconn.computation.executor.entity.EngineConnTask
22+
import org.apache.linkis.governance.common.exception.engineconn.{
23+
EngineConnExecutorErrorCode,
24+
EngineConnExecutorErrorException
25+
}
26+
import org.apache.linkis.protocol.message.RequestProtocol
27+
import org.apache.linkis.rpc.Sender
2028
import org.apache.linkis.server.BDPJettyServerHelper
2129

2230
import com.google.gson.{Gson, GsonBuilder}
2331

24-
object ComputationEngineUtils {
32+
object ComputationEngineUtils extends Logging {
2533

2634
def GSON: Gson = BDPJettyServerHelper.gson
2735

@@ -30,4 +38,24 @@ object ComputationEngineUtils {
3038
private val WORK_DIR_STR = "user.dir"
3139
def getCurrentWorkDir: String = System.getProperty(WORK_DIR_STR)
3240

41+
def sendToEntrance(task: EngineConnTask, msg: RequestProtocol): Unit = {
42+
Utils.tryCatch {
43+
var sender: Sender = null
44+
if (null != task && null != task.getCallbackServiceInstance() && null != msg) {
45+
sender = Sender.getSender(task.getCallbackServiceInstance())
46+
sender.send(msg)
47+
} else {
48+
// todo
49+
logger.debug("SendtoEntrance error, cannot find entrance instance.")
50+
}
51+
} { t =>
52+
val errorMsg = s"SendToEntrance error. $msg" + t.getCause
53+
logger.error(errorMsg, t)
54+
throw new EngineConnExecutorErrorException(
55+
EngineConnExecutorErrorCode.SEND_TO_ENTRANCE_ERROR,
56+
errorMsg
57+
)
58+
}
59+
}
60+
3361
}

‎linkis-computation-governance/linkis-engineconn/linkis-engineconn-executor/executor-core/src/main/scala/org/apache/linkis/engineconn/executor/ExecutorExecutionContext.scala‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,4 +127,6 @@ trait ExecutorExecutionContext {
127127

128128
def setResultSetNum(num: Int): Unit = aliasNum.set(num)
129129

130+
def getAliasNum: Int = aliasNum.get()
131+
130132
}

‎linkis-computation-governance/linkis-entrance/src/main/java/org/apache/linkis/entrance/conf/EntranceSpringConfiguration.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,8 @@ public EntranceInterceptor[] entranceInterceptors() {
151151
new SQLLimitEntranceInterceptor(),
152152
new CommentInterceptor(),
153153
new SetTenantLabelInterceptor(),
154-
new UserCreatorIPCheckInterceptor()
154+
new UserCreatorIPCheckInterceptor(),
155+
new TaskRetryInterceptor()
155156
};
156157
}
157158

‎linkis-computation-governance/linkis-entrance/src/main/java/org/apache/linkis/entrance/job/EntranceExecutionJob.java‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -157,9 +157,15 @@ public ExecuteRequest jobToExecuteRequest() throws EntranceErrorException {
157157
String resultSetPathRoot = GovernanceCommonConf.RESULT_SET_STORE_PATH().getValue(runtimeMapTmp);
158158

159159
if (!runtimeMapTmp.containsKey(GovernanceCommonConf.RESULT_SET_STORE_PATH().key())) {
160-
String resultParentPath = CommonLogPathUtils.getResultParentPath(jobRequest);
161-
CommonLogPathUtils.buildCommonPath(resultParentPath, true);
162-
resultSetPathRoot = CommonLogPathUtils.getResultPath(jobRequest);
160+
// 修复：任务重试背景下，10：59分提交任务执行，重试时时间变成11：00，重试任务会重新生成结果目录，导致查询结果集时，重试之前执行的结果集丢失
161+
// 新增判断：生成结果目录之前，判断任务之前是否生成结果集，生成过就复用
162+
if (org.apache.commons.lang3.StringUtils.isNotEmpty(jobRequest.getResultLocation())) {
163+
resultSetPathRoot = jobRequest.getResultLocation();
164+
} else {
165+
String resultParentPath = CommonLogPathUtils.getResultParentPath(jobRequest);
166+
CommonLogPathUtils.buildCommonPath(resultParentPath, true);
167+
resultSetPathRoot = CommonLogPathUtils.getResultPath(jobRequest);
168+
}
163169
}
164170

165171
Map<String, Object> jobMap = new HashMap<String, Object>();

0 commit comments

Comments
 (0)