Wenjun Ruan
10 months ago
committed by
GitHub
49 changed files with 1453 additions and 918 deletions
@ -0,0 +1,62 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.dao.repository; |
||||
|
||||
import org.apache.dolphinscheduler.dao.entity.TaskGroup; |
||||
|
||||
import java.util.List; |
||||
|
||||
public interface TaskGroupDao extends IDao<TaskGroup> { |
||||
|
||||
/** |
||||
* Query all TaskGroups |
||||
* |
||||
* @return all TaskGroups |
||||
*/ |
||||
List<TaskGroup> queryAllTaskGroups(); |
||||
|
||||
/** |
||||
* Query all TaskGroups which useSize > 0 |
||||
* |
||||
* @return the TaskGroups which useSize > 0 |
||||
*/ |
||||
List<TaskGroup> queryUsedTaskGroups(); |
||||
|
||||
/** |
||||
* Query all TaskGroups which useSize < groupSize |
||||
* |
||||
* @return the TaskGroups which useSize < groupSize |
||||
*/ |
||||
List<TaskGroup> queryAvailableTaskGroups(); |
||||
|
||||
/** |
||||
* Acquire a slot for the TaskGroup which useSize should < groupSize, set the useSize = useSize + 1. |
||||
* |
||||
* @param taskGroupId taskGroupId which shouldn't be null |
||||
* @return true if acquire successfully, false otherwise. |
||||
*/ |
||||
boolean acquireTaskGroupSlot(Integer taskGroupId); |
||||
|
||||
/** |
||||
* Release a slot for the TaskGroup which useSize should > 0, set the useSize = useSize - 1. |
||||
* |
||||
* @param taskGroupId taskGroupId which shouldn't be null |
||||
* @return true if release successfully, false otherwise. |
||||
*/ |
||||
boolean releaseTaskGroupSlot(Integer taskGroupId); |
||||
} |
||||
@ -0,0 +1,64 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.dao.repository; |
||||
|
||||
import org.apache.dolphinscheduler.dao.entity.ProcessInstance; |
||||
import org.apache.dolphinscheduler.dao.entity.TaskGroupQueue; |
||||
|
||||
import java.util.List; |
||||
|
||||
public interface TaskGroupQueueDao extends IDao<TaskGroupQueue> { |
||||
|
||||
/** |
||||
* Delete {@link TaskGroupQueue} by {@link ProcessInstance#getId()} |
||||
* |
||||
* @param workflowInstanceIds workflowInstanceIds |
||||
*/ |
||||
void deleteByWorkflowInstanceIds(List<Integer> workflowInstanceIds); |
||||
|
||||
/** |
||||
* Query all {@link TaskGroupQueue} which in_queue is {@link org.apache.dolphinscheduler.common.enums.Flag#YES} |
||||
* |
||||
* @return TaskGroupQueue ordered by priority desc |
||||
*/ |
||||
List<TaskGroupQueue> queryAllInQueueTaskGroupQueue(); |
||||
|
||||
/** |
||||
* Query all {@link TaskGroupQueue} which in_queue is {@link org.apache.dolphinscheduler.common.enums.Flag#YES} and taskGroupId is taskGroupId |
||||
* |
||||
* @param taskGroupId taskGroupId |
||||
* @return TaskGroupQueue ordered by priority desc |
||||
*/ |
||||
List<TaskGroupQueue> queryAllInQueueTaskGroupQueueByGroupId(Integer taskGroupId); |
||||
|
||||
/** |
||||
* Query all {@link TaskGroupQueue} which taskId is taskInstanceId |
||||
* |
||||
* @param taskInstanceId taskInstanceId |
||||
* @return TaskGroupQueue ordered by priority desc |
||||
*/ |
||||
List<TaskGroupQueue> queryByTaskInstanceId(Integer taskInstanceId); |
||||
|
||||
/** |
||||
* Query all {@link TaskGroupQueue} which status is TaskGroupQueueStatus.ACQUIRE_SUCCESS and forceStart is {@link org.apache.dolphinscheduler.common.enums.Flag#NO}. |
||||
* |
||||
* @param taskGroupId taskGroupId |
||||
* @return TaskGroupQueue |
||||
*/ |
||||
List<TaskGroupQueue> queryAcquiredTaskGroupQueueByGroupId(Integer taskGroupId); |
||||
} |
||||
@ -0,0 +1,70 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.dao.repository.impl; |
||||
|
||||
import org.apache.dolphinscheduler.dao.entity.TaskGroup; |
||||
import org.apache.dolphinscheduler.dao.mapper.TaskGroupMapper; |
||||
import org.apache.dolphinscheduler.dao.repository.BaseDao; |
||||
import org.apache.dolphinscheduler.dao.repository.TaskGroupDao; |
||||
|
||||
import java.util.List; |
||||
|
||||
import lombok.NonNull; |
||||
import lombok.extern.slf4j.Slf4j; |
||||
|
||||
import org.springframework.stereotype.Repository; |
||||
|
||||
@Slf4j |
||||
@Repository |
||||
public class TaskGroupDaoImpl extends BaseDao<TaskGroup, TaskGroupMapper> implements TaskGroupDao { |
||||
|
||||
public TaskGroupDaoImpl(@NonNull TaskGroupMapper taskGroupMapper) { |
||||
super(taskGroupMapper); |
||||
} |
||||
|
||||
@Override |
||||
public List<TaskGroup> queryAllTaskGroups() { |
||||
return mybatisMapper.selectList(null); |
||||
} |
||||
|
||||
@Override |
||||
public List<TaskGroup> queryUsedTaskGroups() { |
||||
return mybatisMapper.queryUsedTaskGroups(); |
||||
} |
||||
|
||||
@Override |
||||
public List<TaskGroup> queryAvailableTaskGroups() { |
||||
return mybatisMapper.queryAvailableTaskGroups(); |
||||
} |
||||
|
||||
@Override |
||||
public boolean acquireTaskGroupSlot(Integer taskGroupId) { |
||||
if (taskGroupId == null) { |
||||
throw new IllegalArgumentException("taskGroupId cannot be null"); |
||||
} |
||||
return mybatisMapper.acquireTaskGroupSlot(taskGroupId) > 0; |
||||
} |
||||
|
||||
@Override |
||||
public boolean releaseTaskGroupSlot(Integer taskGroupId) { |
||||
if (taskGroupId == null) { |
||||
throw new IllegalArgumentException("taskGroupId cannot be null"); |
||||
} |
||||
return mybatisMapper.releaseTaskGroupSlot(taskGroupId) > 0; |
||||
} |
||||
} |
||||
@ -0,0 +1,73 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.dao.repository.impl; |
||||
|
||||
import org.apache.dolphinscheduler.common.enums.Flag; |
||||
import org.apache.dolphinscheduler.common.enums.TaskGroupQueueStatus; |
||||
import org.apache.dolphinscheduler.dao.entity.TaskGroupQueue; |
||||
import org.apache.dolphinscheduler.dao.mapper.TaskGroupQueueMapper; |
||||
import org.apache.dolphinscheduler.dao.repository.BaseDao; |
||||
import org.apache.dolphinscheduler.dao.repository.TaskGroupQueueDao; |
||||
|
||||
import org.apache.commons.collections4.CollectionUtils; |
||||
|
||||
import java.util.List; |
||||
|
||||
import lombok.NonNull; |
||||
|
||||
import org.springframework.stereotype.Repository; |
||||
|
||||
@Repository |
||||
public class TaskGroupQueueDaoImpl extends BaseDao<TaskGroupQueue, TaskGroupQueueMapper> implements TaskGroupQueueDao { |
||||
|
||||
public TaskGroupQueueDaoImpl(@NonNull TaskGroupQueueMapper taskGroupQueueMapper) { |
||||
super(taskGroupQueueMapper); |
||||
} |
||||
|
||||
@Override |
||||
public void deleteByWorkflowInstanceIds(List<Integer> workflowInstanceIds) { |
||||
if (CollectionUtils.isEmpty(workflowInstanceIds)) { |
||||
return; |
||||
} |
||||
mybatisMapper.deleteByWorkflowInstanceIds(workflowInstanceIds); |
||||
} |
||||
|
||||
@Override |
||||
public List<TaskGroupQueue> queryAllInQueueTaskGroupQueue() { |
||||
return mybatisMapper.queryAllTaskGroupQueueByInQueue(Flag.YES.getCode()); |
||||
} |
||||
|
||||
@Override |
||||
public List<TaskGroupQueue> queryAllInQueueTaskGroupQueueByGroupId(Integer taskGroupId) { |
||||
return mybatisMapper.queryAllInQueueTaskGroupQueueByGroupId(taskGroupId, Flag.YES.getCode()); |
||||
} |
||||
|
||||
@Override |
||||
public List<TaskGroupQueue> queryByTaskInstanceId(Integer taskInstanceId) { |
||||
return mybatisMapper.queryByTaskInstanceId(taskInstanceId); |
||||
} |
||||
|
||||
@Override |
||||
public List<TaskGroupQueue> queryAcquiredTaskGroupQueueByGroupId(Integer taskGroupId) { |
||||
return mybatisMapper.queryUsingTaskGroupQueueByGroupId( |
||||
taskGroupId, |
||||
TaskGroupQueueStatus.ACQUIRE_SUCCESS.getCode(), |
||||
Flag.YES.getCode(), |
||||
Flag.NO.getCode()); |
||||
} |
||||
} |
||||
@ -0,0 +1,118 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.dao.repository.impl; |
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals; |
||||
import static org.junit.jupiter.api.Assertions.assertFalse; |
||||
import static org.junit.jupiter.api.Assertions.assertTrue; |
||||
|
||||
import org.apache.dolphinscheduler.common.enums.Flag; |
||||
import org.apache.dolphinscheduler.dao.BaseDaoTest; |
||||
import org.apache.dolphinscheduler.dao.entity.TaskGroup; |
||||
import org.apache.dolphinscheduler.dao.repository.TaskGroupDao; |
||||
|
||||
import java.util.Date; |
||||
import java.util.List; |
||||
|
||||
import org.junit.jupiter.api.Test; |
||||
import org.springframework.beans.factory.annotation.Autowired; |
||||
|
||||
class TaskGroupDaoImplTest extends BaseDaoTest { |
||||
|
||||
@Autowired |
||||
private TaskGroupDao taskGroupDao; |
||||
|
||||
@Test |
||||
void queryAllTaskGroups() { |
||||
TaskGroup taskGroup = createTaskGroup("test", 0, 1); |
||||
taskGroupDao.insert(taskGroup); |
||||
List<TaskGroup> taskGroups = taskGroupDao.queryAllTaskGroups(); |
||||
assertEquals(1, taskGroups.size()); |
||||
} |
||||
|
||||
@Test |
||||
void queryUsedTaskGroups() { |
||||
// Insert a unused task group
|
||||
TaskGroup taskGroup = createTaskGroup("testUnused", 0, 1); |
||||
taskGroupDao.insert(taskGroup); |
||||
assertEquals(0, taskGroupDao.queryUsedTaskGroups().size()); |
||||
|
||||
// Insert a used task group
|
||||
taskGroup = createTaskGroup("testUsed", 1, 1); |
||||
taskGroupDao.insert(taskGroup); |
||||
assertEquals(1, taskGroupDao.queryUsedTaskGroups().size()); |
||||
} |
||||
|
||||
@Test |
||||
void queryAvailableTaskGroups() { |
||||
// Insert a full task group
|
||||
TaskGroup taskGroup = createTaskGroup("testFull", 1, 1); |
||||
taskGroupDao.insert(taskGroup); |
||||
assertEquals(0, taskGroupDao.queryAvailableTaskGroups().size()); |
||||
|
||||
// Insert a used task group
|
||||
taskGroup = createTaskGroup("testNotFull", 0, 1); |
||||
taskGroupDao.insert(taskGroup); |
||||
assertEquals(1, taskGroupDao.queryAvailableTaskGroups().size()); |
||||
} |
||||
|
||||
@Test |
||||
void acquireTaskGroupSlot() { |
||||
// Insert a full task group will acquire failed
|
||||
TaskGroup taskGroup = createTaskGroup("testFull", 1, 1); |
||||
taskGroupDao.insert(taskGroup); |
||||
assertFalse(taskGroupDao.acquireTaskGroupSlot(taskGroup.getId())); |
||||
|
||||
taskGroup.setUseSize(0); |
||||
taskGroupDao.updateById(taskGroup); |
||||
assertTrue(taskGroupDao.acquireTaskGroupSlot(taskGroup.getId())); |
||||
|
||||
taskGroup = taskGroupDao.queryById(taskGroup.getId()); |
||||
assertEquals(1, taskGroup.getUseSize()); |
||||
} |
||||
|
||||
@Test |
||||
void releaseTaskGroupSlot() { |
||||
// Insert an empty task group will release failed
|
||||
TaskGroup taskGroup = createTaskGroup("testEmpty", 0, 1); |
||||
taskGroupDao.insert(taskGroup); |
||||
assertFalse(taskGroupDao.releaseTaskGroupSlot(taskGroup.getId())); |
||||
|
||||
taskGroup.setUseSize(1); |
||||
taskGroupDao.updateById(taskGroup); |
||||
assertTrue(taskGroupDao.releaseTaskGroupSlot(taskGroup.getId())); |
||||
|
||||
taskGroup = taskGroupDao.queryById(taskGroup.getId()); |
||||
assertEquals(0, taskGroup.getUseSize()); |
||||
} |
||||
|
||||
private TaskGroup createTaskGroup(String name, int useSize, int groupSize) { |
||||
return TaskGroup.builder() |
||||
.name(name) |
||||
.description("test") |
||||
.groupSize(groupSize) |
||||
.useSize(useSize) |
||||
.userId(1) |
||||
.status(Flag.YES) |
||||
.createTime(new Date()) |
||||
.updateTime(new Date()) |
||||
.projectCode(1) |
||||
.build(); |
||||
} |
||||
|
||||
} |
||||
@ -0,0 +1,108 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.dao.repository.impl; |
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals; |
||||
import static org.junit.jupiter.api.Assertions.assertNotNull; |
||||
import static org.junit.jupiter.api.Assertions.assertNull; |
||||
|
||||
import org.apache.dolphinscheduler.common.enums.Flag; |
||||
import org.apache.dolphinscheduler.common.enums.TaskGroupQueueStatus; |
||||
import org.apache.dolphinscheduler.dao.BaseDaoTest; |
||||
import org.apache.dolphinscheduler.dao.entity.TaskGroupQueue; |
||||
import org.apache.dolphinscheduler.dao.repository.TaskGroupQueueDao; |
||||
|
||||
import java.util.Date; |
||||
|
||||
import org.assertj.core.util.Lists; |
||||
import org.junit.jupiter.api.Test; |
||||
import org.springframework.beans.factory.annotation.Autowired; |
||||
|
||||
class TaskGroupQueueDaoImplTest extends BaseDaoTest { |
||||
|
||||
@Autowired |
||||
private TaskGroupQueueDao taskGroupQueueDao; |
||||
|
||||
@Test |
||||
void deleteByWorkflowInstanceIds() { |
||||
TaskGroupQueue taskGroupQueue = createTaskGroupQueue(Flag.NO, TaskGroupQueueStatus.ACQUIRE_SUCCESS); |
||||
taskGroupQueueDao.insert(taskGroupQueue); |
||||
assertNotNull(taskGroupQueueDao.queryById(taskGroupQueue.getId())); |
||||
|
||||
taskGroupQueueDao.deleteByWorkflowInstanceIds(Lists.newArrayList(1)); |
||||
assertNull(taskGroupQueueDao.queryById(taskGroupQueue.getId())); |
||||
} |
||||
|
||||
@Test |
||||
void queryAllInQueueTaskGroupQueue() { |
||||
TaskGroupQueue taskGroupQueue = createTaskGroupQueue(Flag.NO, TaskGroupQueueStatus.ACQUIRE_SUCCESS); |
||||
taskGroupQueueDao.insert(taskGroupQueue); |
||||
assertEquals(1, taskGroupQueueDao.queryAllInQueueTaskGroupQueue().size()); |
||||
} |
||||
|
||||
@Test |
||||
void queryAllInQueueTaskGroupQueueByGroupId() { |
||||
TaskGroupQueue taskGroupQueue = createTaskGroupQueue(Flag.NO, TaskGroupQueueStatus.ACQUIRE_SUCCESS); |
||||
taskGroupQueueDao.insert(taskGroupQueue); |
||||
assertEquals(1, taskGroupQueueDao.queryAllInQueueTaskGroupQueueByGroupId(1).size()); |
||||
} |
||||
|
||||
@Test |
||||
void updateById() { |
||||
TaskGroupQueue taskGroupQueue = createTaskGroupQueue(Flag.NO, TaskGroupQueueStatus.WAIT_QUEUE); |
||||
taskGroupQueueDao.insert(taskGroupQueue); |
||||
|
||||
taskGroupQueue.setStatus(TaskGroupQueueStatus.ACQUIRE_SUCCESS); |
||||
taskGroupQueueDao.updateById(taskGroupQueue); |
||||
assertEquals(TaskGroupQueueStatus.ACQUIRE_SUCCESS, |
||||
taskGroupQueueDao.queryById(taskGroupQueue.getId()).getStatus()); |
||||
} |
||||
|
||||
@Test |
||||
void queryByTaskInstanceId() { |
||||
TaskGroupQueue taskGroupQueue = createTaskGroupQueue(Flag.NO, TaskGroupQueueStatus.ACQUIRE_SUCCESS); |
||||
taskGroupQueueDao.insert(taskGroupQueue); |
||||
assertEquals(1, taskGroupQueueDao.queryByTaskInstanceId(1).size()); |
||||
} |
||||
|
||||
@Test |
||||
void queryUsingTaskGroupQueueByGroupId() { |
||||
TaskGroupQueue taskGroupQueue = createTaskGroupQueue(Flag.NO, TaskGroupQueueStatus.ACQUIRE_SUCCESS); |
||||
taskGroupQueueDao.insert(taskGroupQueue); |
||||
assertEquals(1, taskGroupQueueDao.queryAcquiredTaskGroupQueueByGroupId(1).size()); |
||||
|
||||
taskGroupQueue = createTaskGroupQueue(Flag.YES, TaskGroupQueueStatus.WAIT_QUEUE); |
||||
taskGroupQueueDao.insert(taskGroupQueue); |
||||
assertEquals(1, taskGroupQueueDao.queryAcquiredTaskGroupQueueByGroupId(1).size()); |
||||
} |
||||
|
||||
private TaskGroupQueue createTaskGroupQueue(Flag forceStart, TaskGroupQueueStatus taskGroupQueueStatus) { |
||||
return TaskGroupQueue.builder() |
||||
.taskId(1) |
||||
.taskName("test") |
||||
.groupId(1) |
||||
.processId(1) |
||||
.priority(0) |
||||
.forceStart(forceStart.getCode()) |
||||
.inQueue(Flag.YES.getCode()) |
||||
.status(taskGroupQueueStatus) |
||||
.createTime(new Date()) |
||||
.updateTime(new Date()) |
||||
.build(); |
||||
} |
||||
} |
||||
@ -1,42 +0,0 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.extract.master.transportor; |
||||
|
||||
import lombok.Data; |
||||
import lombok.NoArgsConstructor; |
||||
|
||||
@Data |
||||
@NoArgsConstructor |
||||
public class TaskInstanceForceStartRequest { |
||||
|
||||
private String key; |
||||
|
||||
private int processInstanceId; |
||||
|
||||
private int taskInstanceId; |
||||
|
||||
public TaskInstanceForceStartRequest( |
||||
int processInstanceId, |
||||
int taskInstanceId) { |
||||
this.key = String.format("%d-%d", processInstanceId, taskInstanceId); |
||||
|
||||
this.processInstanceId = processInstanceId; |
||||
this.taskInstanceId = taskInstanceId; |
||||
} |
||||
|
||||
} |
||||
@ -1,41 +0,0 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.extract.master.transportor; |
||||
|
||||
import lombok.AllArgsConstructor; |
||||
import lombok.Data; |
||||
import lombok.NoArgsConstructor; |
||||
|
||||
@Data |
||||
@NoArgsConstructor |
||||
@AllArgsConstructor |
||||
public class TaskInstanceForceStartResponse { |
||||
|
||||
private boolean success; |
||||
|
||||
private String message; |
||||
|
||||
public static TaskInstanceForceStartResponse success() { |
||||
return new TaskInstanceForceStartResponse(true, "dispatch success"); |
||||
} |
||||
|
||||
public static TaskInstanceForceStartResponse failed(String message) { |
||||
return new TaskInstanceForceStartResponse(false, message); |
||||
} |
||||
|
||||
} |
||||
@ -1,47 +0,0 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.server.master.event; |
||||
|
||||
import org.apache.dolphinscheduler.common.enums.StateEventType; |
||||
import org.apache.dolphinscheduler.server.master.runner.WorkflowExecuteRunnable; |
||||
|
||||
import lombok.extern.slf4j.Slf4j; |
||||
|
||||
import com.google.auto.service.AutoService; |
||||
|
||||
@AutoService(StateEventHandler.class) |
||||
@Slf4j |
||||
public class TaskWaitTaskGroupStateHandler implements StateEventHandler { |
||||
|
||||
@Override |
||||
public boolean handleStateEvent(WorkflowExecuteRunnable workflowExecuteRunnable, |
||||
StateEvent stateEvent) { |
||||
log.info("Handle task instance wait task group event, taskInstanceId: {}", stateEvent.getTaskInstanceId()); |
||||
if (workflowExecuteRunnable.checkForceStartAndWakeUp(stateEvent)) { |
||||
log.info("Success wake up task instance, taskInstanceId: {}", stateEvent.getTaskInstanceId()); |
||||
} else { |
||||
log.info("Failed to wake up task instance, taskInstanceId: {}", stateEvent.getTaskInstanceId()); |
||||
} |
||||
return true; |
||||
} |
||||
|
||||
@Override |
||||
public StateEventType getEventType() { |
||||
return StateEventType.WAKE_UP_TASK_GROUP; |
||||
} |
||||
} |
||||
@ -1,58 +0,0 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.server.master.rpc; |
||||
|
||||
import org.apache.dolphinscheduler.common.enums.StateEventType; |
||||
import org.apache.dolphinscheduler.extract.master.transportor.TaskInstanceForceStartRequest; |
||||
import org.apache.dolphinscheduler.extract.master.transportor.TaskInstanceForceStartResponse; |
||||
import org.apache.dolphinscheduler.plugin.task.api.utils.LogUtils; |
||||
import org.apache.dolphinscheduler.server.master.event.TaskStateEvent; |
||||
import org.apache.dolphinscheduler.server.master.processor.queue.StateEventResponseService; |
||||
|
||||
import lombok.extern.slf4j.Slf4j; |
||||
|
||||
import org.springframework.beans.factory.annotation.Autowired; |
||||
import org.springframework.stereotype.Component; |
||||
|
||||
@Slf4j |
||||
@Component |
||||
public class TaskInstanceForceStartOperationFunction |
||||
implements |
||||
ITaskInstanceOperationFunction<TaskInstanceForceStartRequest, TaskInstanceForceStartResponse> { |
||||
|
||||
@Autowired |
||||
private StateEventResponseService stateEventResponseService; |
||||
|
||||
@Override |
||||
public TaskInstanceForceStartResponse operate(TaskInstanceForceStartRequest taskInstanceForceStartRequest) { |
||||
TaskStateEvent stateEvent = TaskStateEvent.builder() |
||||
.processInstanceId(taskInstanceForceStartRequest.getProcessInstanceId()) |
||||
.taskInstanceId(taskInstanceForceStartRequest.getTaskInstanceId()) |
||||
.key(taskInstanceForceStartRequest.getKey()) |
||||
.type(StateEventType.WAKE_UP_TASK_GROUP) |
||||
.build(); |
||||
try { |
||||
LogUtils.setWorkflowAndTaskInstanceIDMDC(stateEvent.getProcessInstanceId(), stateEvent.getTaskInstanceId()); |
||||
log.info("Received forceStartTaskInstance, event: {}", stateEvent); |
||||
stateEventResponseService.addEvent2WorkflowExecute(stateEvent); |
||||
return TaskInstanceForceStartResponse.success(); |
||||
} finally { |
||||
LogUtils.removeWorkflowAndTaskInstanceIdMDC(); |
||||
} |
||||
} |
||||
} |
||||
@ -0,0 +1,451 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.server.master.runner.taskgroup; |
||||
|
||||
import org.apache.dolphinscheduler.common.constants.Constants; |
||||
import org.apache.dolphinscheduler.common.enums.Flag; |
||||
import org.apache.dolphinscheduler.common.enums.TaskGroupQueueStatus; |
||||
import org.apache.dolphinscheduler.common.enums.WorkflowExecutionStatus; |
||||
import org.apache.dolphinscheduler.common.lifecycle.ServerLifeCycleManager; |
||||
import org.apache.dolphinscheduler.common.thread.BaseDaemonThread; |
||||
import org.apache.dolphinscheduler.common.thread.ThreadUtils; |
||||
import org.apache.dolphinscheduler.dao.entity.ProcessInstance; |
||||
import org.apache.dolphinscheduler.dao.entity.TaskGroup; |
||||
import org.apache.dolphinscheduler.dao.entity.TaskGroupQueue; |
||||
import org.apache.dolphinscheduler.dao.entity.TaskInstance; |
||||
import org.apache.dolphinscheduler.dao.repository.ProcessInstanceDao; |
||||
import org.apache.dolphinscheduler.dao.repository.TaskGroupDao; |
||||
import org.apache.dolphinscheduler.dao.repository.TaskGroupQueueDao; |
||||
import org.apache.dolphinscheduler.dao.repository.TaskInstanceDao; |
||||
import org.apache.dolphinscheduler.extract.base.client.SingletonJdkDynamicRpcClientProxyFactory; |
||||
import org.apache.dolphinscheduler.extract.master.IWorkflowInstanceService; |
||||
import org.apache.dolphinscheduler.extract.master.transportor.TaskInstanceWakeupRequest; |
||||
import org.apache.dolphinscheduler.extract.master.transportor.TaskInstanceWakeupResponse; |
||||
import org.apache.dolphinscheduler.plugin.task.api.enums.TaskExecutionStatus; |
||||
import org.apache.dolphinscheduler.plugin.task.api.utils.LogUtils; |
||||
import org.apache.dolphinscheduler.registry.api.RegistryClient; |
||||
import org.apache.dolphinscheduler.registry.api.enums.RegistryNodeType; |
||||
|
||||
import org.apache.commons.collections4.CollectionUtils; |
||||
import org.apache.commons.lang3.time.StopWatch; |
||||
|
||||
import java.util.Date; |
||||
import java.util.List; |
||||
import java.util.Map; |
||||
import java.util.function.Function; |
||||
import java.util.stream.Collectors; |
||||
|
||||
import lombok.extern.slf4j.Slf4j; |
||||
|
||||
import org.springframework.beans.factory.annotation.Autowired; |
||||
import org.springframework.stereotype.Component; |
||||
|
||||
/** |
||||
* The TaskGroupCoordinator use to manage the task group slot. The task group slot is used to limit the number of {@link TaskInstance} that can be run at the same time. |
||||
* <p> |
||||
* The {@link TaskGroupQueue} is used to represent the task group slot. When a {@link TaskGroupQueue} which inQueue is YES means the {@link TaskGroupQueue} is using by a {@link TaskInstance}. |
||||
* <p> |
||||
* When the {@link TaskInstance} need to use task group, we should use @{@link TaskGroupCoordinator#acquireTaskGroupSlot(TaskInstance)} to acquire the task group slot, |
||||
* this method doesn't block should always acquire successfully, and you should directly stop dispatch the task instance. |
||||
* When the task group slot is available, the TaskGroupCoordinator will wake up the waiting {@link TaskInstance} to dispatch. |
||||
* <pre> |
||||
* if(needAcquireTaskGroupSlot(taskInstance)) { |
||||
* taskGroupCoordinator.acquireTaskGroupSlot(taskInstance); |
||||
* return; |
||||
* } |
||||
* </pre> |
||||
* <p> |
||||
* When the {@link TaskInstance} is finished, we should use @{@link TaskGroupCoordinator#releaseTaskGroupSlot(TaskInstance)} to release the task group slot. |
||||
* <pre> |
||||
* if(needToReleaseTaskGroupSlot(taskInstance)) { |
||||
* taskGroupCoordinator.releaseTaskGroupSlot(taskInstance); |
||||
* } |
||||
* </pre> |
||||
*/ |
||||
@Slf4j |
||||
@Component |
||||
public class TaskGroupCoordinator extends BaseDaemonThread { |
||||
|
||||
@Autowired |
||||
private RegistryClient registryClient; |
||||
|
||||
@Autowired |
||||
private TaskGroupDao taskGroupDao; |
||||
|
||||
@Autowired |
||||
private TaskGroupQueueDao taskGroupQueueDao; |
||||
|
||||
@Autowired |
||||
private TaskInstanceDao taskInstanceDao; |
||||
|
||||
@Autowired |
||||
private ProcessInstanceDao processInstanceDao; |
||||
|
||||
public TaskGroupCoordinator() { |
||||
super("TaskGroupCoordinator"); |
||||
} |
||||
|
||||
@Override |
||||
public synchronized void start() { |
||||
log.info("TaskGroupCoordinator starting..."); |
||||
super.start(); |
||||
log.info("TaskGroupCoordinator started..."); |
||||
} |
||||
|
||||
@Override |
||||
public void run() { |
||||
while (!ServerLifeCycleManager.isStopped()) { |
||||
try { |
||||
if (!ServerLifeCycleManager.isRunning()) { |
||||
continue; |
||||
} |
||||
try { |
||||
registryClient.getLock(RegistryNodeType.MASTER_TASK_GROUP_COORDINATOR_LOCK.getRegistryPath()); |
||||
StopWatch taskGroupCoordinatorRoundTimeCost = StopWatch.createStarted(); |
||||
|
||||
amendTaskGroupUseSize(); |
||||
amendTaskGroupQueueStatus(); |
||||
dealWithForceStartTaskGroupQueue(); |
||||
dealWithWaitingTaskGroupQueue(); |
||||
|
||||
taskGroupCoordinatorRoundTimeCost.stop(); |
||||
log.info("TaskGroupCoordinator round time cost: {}/ms", |
||||
taskGroupCoordinatorRoundTimeCost.getTime()); |
||||
} finally { |
||||
registryClient.releaseLock(RegistryNodeType.MASTER_TASK_GROUP_COORDINATOR_LOCK.getRegistryPath()); |
||||
} |
||||
} catch (Throwable e) { |
||||
log.error("TaskGroupCoordinator error", e); |
||||
} finally { |
||||
// sleep 5s
|
||||
ThreadUtils.sleep(Constants.SLEEP_TIME_MILLIS * 5); |
||||
} |
||||
} |
||||
} |
||||
|
||||
/** |
||||
* Make sure the TaskGroup useSize is equal to the TaskGroupQueue which status is {@link TaskGroupQueueStatus#ACQUIRE_SUCCESS} and forceStart is {@link org.apache.dolphinscheduler.common.enums.Flag#NO}. |
||||
*/ |
||||
private void amendTaskGroupUseSize() { |
||||
// The TaskGroup useSize should equal to the TaskGroupQueue which inQueue is YES and forceStart is NO
|
||||
List<TaskGroup> taskGroups = taskGroupDao.queryAllTaskGroups(); |
||||
if (CollectionUtils.isEmpty(taskGroups)) { |
||||
return; |
||||
} |
||||
for (TaskGroup taskGroup : taskGroups) { |
||||
List<TaskGroupQueue> taskGroupQueues = |
||||
taskGroupQueueDao.queryAcquiredTaskGroupQueueByGroupId(taskGroup.getId()); |
||||
int actualUseSize = taskGroupQueues.size(); |
||||
if (taskGroup.getUseSize() == actualUseSize) { |
||||
continue; |
||||
} |
||||
log.warn("The TaskGroup: {} useSize is {}, but the actual use size is {}, will amend it", |
||||
taskGroup.getName(), |
||||
taskGroup.getUseSize(), actualUseSize); |
||||
taskGroup.setUseSize(actualUseSize); |
||||
taskGroupDao.updateById(taskGroup); |
||||
} |
||||
} |
||||
|
||||
/** |
||||
* Make sure the TaskGroupQueue status is {@link TaskGroupQueueStatus#RELEASE} when the related {@link TaskInstance} is not exist or status is finished. |
||||
*/ |
||||
private void amendTaskGroupQueueStatus() { |
||||
List<TaskGroupQueue> taskGroupQueues = taskGroupQueueDao.queryAllInQueueTaskGroupQueue(); |
||||
List<Integer> taskInstanceIds = taskGroupQueues.stream() |
||||
.map(TaskGroupQueue::getTaskId) |
||||
.collect(Collectors.toList()); |
||||
Map<Integer, TaskInstance> taskInstanceMap = taskInstanceDao.queryByIds(taskInstanceIds) |
||||
.stream() |
||||
.collect(Collectors.toMap(TaskInstance::getId, Function.identity())); |
||||
|
||||
for (TaskGroupQueue taskGroupQueue : taskGroupQueues) { |
||||
int taskId = taskGroupQueue.getTaskId(); |
||||
TaskInstance taskInstance = taskInstanceMap.get(taskId); |
||||
|
||||
if (taskInstance == null) { |
||||
log.warn("The TaskInstance: {} is not exist, will release the TaskGroupQueue: {}", taskId, |
||||
taskGroupQueue); |
||||
releaseTaskGroupQueueSlot(taskGroupQueue); |
||||
continue; |
||||
} |
||||
|
||||
if (taskInstance.getState().isFinished()) { |
||||
log.warn("The TaskInstance: {} state: {} finished, will release the TaskGroupQueue: {}", |
||||
taskInstance.getName(), taskInstance.getState(), taskGroupQueue); |
||||
releaseTaskGroupQueueSlot(taskGroupQueue); |
||||
continue; |
||||
} |
||||
} |
||||
} |
||||
|
||||
private void dealWithForceStartTaskGroupQueue() { |
||||
// Find the force start task group queue(Which is inQueue and forceStart is YES)
|
||||
// Notify the related waiting task instance
|
||||
// Set the taskGroupQueue status to RELEASE and remove it from queue
|
||||
List<TaskGroupQueue> taskGroupQueues = taskGroupQueueDao.queryAllInQueueTaskGroupQueue() |
||||
.stream() |
||||
.filter(taskGroupQueue -> Flag.YES.getCode() == taskGroupQueue.getForceStart()) |
||||
.collect(Collectors.toList()); |
||||
for (TaskGroupQueue taskGroupQueue : taskGroupQueues) { |
||||
try { |
||||
LogUtils.setTaskInstanceIdMDC(taskGroupQueue.getTaskId()); |
||||
// notify the waiting task instance
|
||||
// We notify first, it notify failed, the taskGroupQueue will be in queue, and then we will retry it
|
||||
// next time.
|
||||
notifyWaitingTaskInstance(taskGroupQueue); |
||||
log.info("Notify the ForceStart waiting TaskInstance: {} for taskGroupQueue: {} success", |
||||
taskGroupQueue.getTaskName(), |
||||
taskGroupQueue.getId()); |
||||
|
||||
taskGroupQueue.setInQueue(Flag.NO.getCode()); |
||||
taskGroupQueue.setStatus(TaskGroupQueueStatus.RELEASE); |
||||
taskGroupQueue.setUpdateTime(new Date()); |
||||
taskGroupQueueDao.updateById(taskGroupQueue); |
||||
log.info("Release the force start TaskGroupQueue {}", taskGroupQueue); |
||||
} catch (UnsupportedOperationException unsupportedOperationException) { |
||||
releaseTaskGroupQueueSlot(taskGroupQueue); |
||||
log.info( |
||||
"Notify the ForceStart TaskInstance: {} for taskGroupQueue: {} failed, will release the taskGroupQueue", |
||||
taskGroupQueue.getTaskName(), taskGroupQueue.getId(), unsupportedOperationException); |
||||
} catch (Throwable throwable) { |
||||
log.info("Notify the force start TaskGroupQueue {} failed", taskGroupQueue, throwable); |
||||
} finally { |
||||
LogUtils.removeTaskInstanceIdMDC(); |
||||
} |
||||
} |
||||
} |
||||
|
||||
private void dealWithWaitingTaskGroupQueue() { |
||||
// Find the TaskGroup which usage < maxSize.
|
||||
// Find the highest priority inQueue task group queue(Which is inQueue and status is Waiting and force start is
|
||||
// NO) belong to the
|
||||
// task group.
|
||||
List<TaskGroup> taskGroups = taskGroupDao.queryAvailableTaskGroups(); |
||||
if (CollectionUtils.isEmpty(taskGroups)) { |
||||
log.debug("There is no available task group"); |
||||
return; |
||||
} |
||||
for (TaskGroup taskGroup : taskGroups) { |
||||
int availableSize = taskGroup.getGroupSize() - taskGroup.getUseSize(); |
||||
if (availableSize <= 0) { |
||||
log.info("TaskGroup {} is full, available size is {}", taskGroup, availableSize); |
||||
continue; |
||||
} |
||||
List<TaskGroupQueue> taskGroupQueues = |
||||
taskGroupQueueDao.queryAllInQueueTaskGroupQueueByGroupId(taskGroup.getId()) |
||||
.stream() |
||||
.filter(taskGroupQueue -> Flag.NO.getCode() == taskGroupQueue.getForceStart()) |
||||
.filter(taskGroupQueue -> TaskGroupQueueStatus.WAIT_QUEUE == taskGroupQueue.getStatus()) |
||||
.limit(availableSize) |
||||
.collect(Collectors.toList()); |
||||
if (CollectionUtils.isEmpty(taskGroupQueues)) { |
||||
log.debug("There is no waiting task group queue for task group {}", taskGroup.getName()); |
||||
continue; |
||||
} |
||||
for (TaskGroupQueue taskGroupQueue : taskGroupQueues) { |
||||
try { |
||||
LogUtils.setTaskInstanceIdMDC(taskGroupQueue.getTaskId()); |
||||
// Reduce the taskGroupSize
|
||||
boolean acquireResult = taskGroupDao.acquireTaskGroupSlot(taskGroup.getId()); |
||||
if (!acquireResult) { |
||||
log.error("Failed to acquire task group slot for task group {}", taskGroup); |
||||
continue; |
||||
} |
||||
// Notify the waiting task instance
|
||||
// We notify first, it notify failed, the taskGroupQueue will be in queue, and then we will retry it
|
||||
// next time.
|
||||
notifyWaitingTaskInstance(taskGroupQueue); |
||||
|
||||
// Set the taskGroupQueue status to RUNNING and remove from queue
|
||||
taskGroupQueue.setInQueue(Flag.YES.getCode()); |
||||
taskGroupQueue.setStatus(TaskGroupQueueStatus.ACQUIRE_SUCCESS); |
||||
taskGroupQueue.setUpdateTime(new Date()); |
||||
taskGroupQueueDao.updateById(taskGroupQueue); |
||||
} catch (UnsupportedOperationException unsupportedOperationException) { |
||||
releaseTaskGroupQueueSlot(taskGroupQueue); |
||||
log.info( |
||||
"Notify the Waiting TaskInstance: {} for taskGroupQueue: {} failed, will release the taskGroupQueue", |
||||
taskGroupQueue.getTaskName(), taskGroupQueue.getId(), unsupportedOperationException); |
||||
} catch (Throwable throwable) { |
||||
log.error("Notify Waiting TaskGroupQueue: {} failed", taskGroupQueue, throwable); |
||||
} finally { |
||||
LogUtils.removeTaskInstanceIdMDC(); |
||||
} |
||||
} |
||||
} |
||||
} |
||||
|
||||
/** |
||||
* If the {@link TaskInstance#getTaskGroupId()} > 0, and the TaskGroup flag is {@link Flag#YES} then the task instance need to use task group. |
||||
* |
||||
* @param taskInstance task instance |
||||
* @return true if the TaskInstance need to acquireTaskGroupSlot |
||||
*/ |
||||
public boolean needAcquireTaskGroupSlot(TaskInstance taskInstance) { |
||||
if (taskInstance == null) { |
||||
throw new IllegalArgumentException("The TaskInstance is null"); |
||||
} |
||||
if (taskInstance.getTaskGroupId() <= 0) { |
||||
log.debug("The current TaskInstance doesn't use TaskGroup, no need to acquire TaskGroupSlot"); |
||||
return false; |
||||
} |
||||
TaskGroup taskGroup = taskGroupDao.queryById(taskInstance.getTaskGroupId()); |
||||
if (taskGroup == null) { |
||||
log.warn("The current TaskGroup: {} does not exist, will not acquire TaskGroupSlot", |
||||
taskInstance.getTaskGroupId()); |
||||
return false; |
||||
} |
||||
return Flag.YES.equals(taskGroup.getStatus()); |
||||
} |
||||
|
||||
/** |
||||
* Acquire the task group slot for the given {@link TaskInstance}. |
||||
* <p> |
||||
* When taskInstance want to acquire a TaskGroup slot, should call this method. If acquire successfully, will create a TaskGroupQueue in db which is in queue and status is {@link TaskGroupQueueStatus#WAIT_QUEUE}. |
||||
* The TaskInstance shouldn't dispatch until there exist available slot, the taskGroupCoordinator notify it. |
||||
* |
||||
* @param taskInstance the task instance which want to acquire task group slot. |
||||
* @throws IllegalArgumentException if the taskInstance is null or the used taskGroup doesn't exist. |
||||
*/ |
||||
public void acquireTaskGroupSlot(TaskInstance taskInstance) { |
||||
if (taskInstance == null || taskInstance.getTaskGroupId() <= 0) { |
||||
throw new IllegalArgumentException("The current TaskInstance does not use task group"); |
||||
} |
||||
TaskGroup taskGroup = taskGroupDao.queryById(taskInstance.getTaskGroupId()); |
||||
if (taskGroup == null) { |
||||
throw new IllegalArgumentException( |
||||
"The current TaskGroup: " + taskInstance.getTaskGroupId() + " does not exist"); |
||||
} |
||||
// Write TaskGroupQueue in db, and then return wait TaskGroupCoordinator to notify it
|
||||
// Set the taskGroupQueue status to WAIT_QUEUE and add to queue
|
||||
// The queue only contains the taskGroupQueue which status is WAIT_QUEUE or ACQUIRE_SUCCESS
|
||||
Date now = new Date(); |
||||
TaskGroupQueue taskGroupQueue = TaskGroupQueue |
||||
.builder() |
||||
.taskId(taskInstance.getId()) |
||||
.taskName(taskInstance.getName()) |
||||
.groupId(taskInstance.getTaskGroupId()) |
||||
.processId(taskInstance.getProcessInstanceId()) |
||||
.priority(taskInstance.getTaskGroupPriority()) |
||||
.inQueue(Flag.YES.getCode()) |
||||
.forceStart(Flag.NO.getCode()) |
||||
.status(TaskGroupQueueStatus.WAIT_QUEUE) |
||||
.createTime(now) |
||||
.updateTime(now) |
||||
.build(); |
||||
log.info("Success insert TaskGroupQueue: {} for TaskInstance: {}", taskGroupQueue, taskInstance.getName()); |
||||
taskGroupQueueDao.insert(taskGroupQueue); |
||||
} |
||||
|
||||
/** |
||||
* If the TaskInstance is using TaskGroup then it need to release TaskGroupSlot. |
||||
* |
||||
* @param taskInstance taskInsatnce |
||||
* @return true if the TaskInstance need to release TaskGroupSlot |
||||
*/ |
||||
public boolean needToReleaseTaskGroupSlot(TaskInstance taskInstance) { |
||||
if (taskInstance == null) { |
||||
throw new IllegalArgumentException("The TaskInstance is null"); |
||||
} |
||||
if (taskInstance.getTaskGroupId() <= 0) { |
||||
log.debug("The current TaskInstance doesn't use TaskGroup, no need to release TaskGroupSlot"); |
||||
return false; |
||||
} |
||||
return true; |
||||
} |
||||
|
||||
/** |
||||
* Release the task group slot for the given {@link TaskInstance}. |
||||
* <p> |
||||
* When taskInstance want to release a TaskGroup slot, should call this method. The release method will move the TaskGroupQueue out queue and set status to {@link TaskGroupQueueStatus#RELEASE}. |
||||
* This method is idempotent, this means that if the task group slot is already released, this method will do nothing. |
||||
* |
||||
* @param taskInstance the task instance which want to release task group slot. |
||||
* @throws IllegalArgumentException If the taskInstance is null or the task doesn't use task group. |
||||
*/ |
||||
public void releaseTaskGroupSlot(TaskInstance taskInstance) { |
||||
if (taskInstance == null || taskInstance.getTaskGroupId() <= 0) { |
||||
throw new IllegalArgumentException("The current TaskInstance does not use task group"); |
||||
} |
||||
List<TaskGroupQueue> taskGroupQueues = taskGroupQueueDao.queryByTaskInstanceId(taskInstance.getId()); |
||||
for (TaskGroupQueue taskGroupQueue : taskGroupQueues) { |
||||
releaseTaskGroupQueueSlot(taskGroupQueue); |
||||
} |
||||
} |
||||
|
||||
private void notifyWaitingTaskInstance(TaskGroupQueue taskGroupQueue) { |
||||
// Find the related waiting task instance
|
||||
// send RPC to notify the waiting task instance
|
||||
TaskInstance taskInstance = taskInstanceDao.queryById(taskGroupQueue.getTaskId()); |
||||
if (taskInstance == null) { |
||||
throw new UnsupportedOperationException( |
||||
"The TaskInstance: " + taskGroupQueue.getTaskId() + " is not exist, no need to notify"); |
||||
} |
||||
// todo: We may need to add a new status to represent the task instance is waiting for task group slot
|
||||
if (taskInstance.getState() != TaskExecutionStatus.SUBMITTED_SUCCESS) { |
||||
throw new UnsupportedOperationException( |
||||
"The TaskInstance: " + taskInstance.getId() + " state is " + taskInstance.getState() |
||||
+ ", no need to notify"); |
||||
} |
||||
ProcessInstance processInstance = processInstanceDao.queryById(taskInstance.getProcessInstanceId()); |
||||
if (processInstance == null) { |
||||
throw new UnsupportedOperationException( |
||||
"The WorkflowInstance: " + taskInstance.getProcessInstanceId() |
||||
+ " is not exist, no need to notify"); |
||||
} |
||||
if (processInstance.getState() != WorkflowExecutionStatus.RUNNING_EXECUTION) { |
||||
throw new UnsupportedOperationException( |
||||
"The WorkflowInstance: " + processInstance.getId() + " state is " + processInstance.getState() |
||||
+ ", no need to notify"); |
||||
} |
||||
if (processInstance.getHost() == null || Constants.NULL.equals(processInstance.getHost())) { |
||||
throw new UnsupportedOperationException( |
||||
"WorkflowInstance host is null, maybe it is in failover: " + processInstance); |
||||
} |
||||
|
||||
TaskInstanceWakeupRequest taskInstanceWakeupRequest = TaskInstanceWakeupRequest.builder() |
||||
.processInstanceId(processInstance.getId()) |
||||
.taskInstanceId(taskInstance.getId()) |
||||
.build(); |
||||
|
||||
IWorkflowInstanceService iWorkflowInstanceService = SingletonJdkDynamicRpcClientProxyFactory |
||||
.getProxyClient(processInstance.getHost(), IWorkflowInstanceService.class); |
||||
TaskInstanceWakeupResponse taskInstanceWakeupResponse = |
||||
iWorkflowInstanceService.wakeupTaskInstance(taskInstanceWakeupRequest); |
||||
if (!taskInstanceWakeupResponse.isSuccess()) { |
||||
throw new UnsupportedOperationException( |
||||
"Notify TaskInstance: " + taskInstance.getId() + " failed: " + taskInstanceWakeupResponse); |
||||
} |
||||
log.info("Wake up TaskInstance: {} success", taskInstance.getName()); |
||||
} |
||||
|
||||
private void releaseTaskGroupQueueSlot(TaskGroupQueue taskGroupQueue) { |
||||
if (TaskGroupQueueStatus.RELEASE.equals(taskGroupQueue.getStatus()) |
||||
&& Flag.NO.getCode() == taskGroupQueue.getInQueue()) { |
||||
log.info("The TaskGroupQueue: {} is already released", taskGroupQueue); |
||||
return; |
||||
} |
||||
taskGroupQueue.setInQueue(Flag.NO.getCode()); |
||||
taskGroupQueue.setStatus(TaskGroupQueueStatus.RELEASE); |
||||
taskGroupQueue.setUpdateTime(new Date()); |
||||
taskGroupQueueDao.updateById(taskGroupQueue); |
||||
log.info("Success release TaskGroupQueue: {}", taskGroupQueue); |
||||
} |
||||
|
||||
} |
||||
@ -0,0 +1,182 @@
|
||||
/* |
||||
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||
* contributor license agreements. See the NOTICE file distributed with |
||||
* this work for additional information regarding copyright ownership. |
||||
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||
* (the "License"); you may not use this file except in compliance with |
||||
* the License. You may obtain a copy of the License at |
||||
* |
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* |
||||
* Unless required by applicable law or agreed to in writing, software |
||||
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
* See the License for the specific language governing permissions and |
||||
* limitations under the License. |
||||
*/ |
||||
|
||||
package org.apache.dolphinscheduler.server.master.runner.taskgroup; |
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals; |
||||
import static org.junit.jupiter.api.Assertions.assertFalse; |
||||
import static org.junit.jupiter.api.Assertions.assertThrows; |
||||
import static org.junit.jupiter.api.Assertions.assertTrue; |
||||
import static org.mockito.Mockito.verify; |
||||
import static org.mockito.Mockito.when; |
||||
|
||||
import org.apache.dolphinscheduler.common.enums.Flag; |
||||
import org.apache.dolphinscheduler.common.enums.TaskGroupQueueStatus; |
||||
import org.apache.dolphinscheduler.dao.entity.TaskGroup; |
||||
import org.apache.dolphinscheduler.dao.entity.TaskGroupQueue; |
||||
import org.apache.dolphinscheduler.dao.entity.TaskInstance; |
||||
import org.apache.dolphinscheduler.dao.repository.ProcessInstanceDao; |
||||
import org.apache.dolphinscheduler.dao.repository.TaskGroupDao; |
||||
import org.apache.dolphinscheduler.dao.repository.TaskGroupQueueDao; |
||||
import org.apache.dolphinscheduler.dao.repository.TaskInstanceDao; |
||||
import org.apache.dolphinscheduler.registry.api.RegistryClient; |
||||
import org.apache.dolphinscheduler.registry.api.enums.RegistryNodeType; |
||||
|
||||
import java.util.List; |
||||
|
||||
import org.junit.jupiter.api.Assertions; |
||||
import org.junit.jupiter.api.Test; |
||||
import org.junit.jupiter.api.extension.ExtendWith; |
||||
import org.mockito.InjectMocks; |
||||
import org.mockito.Mock; |
||||
import org.mockito.Mockito; |
||||
import org.mockito.junit.jupiter.MockitoExtension; |
||||
import org.mockito.junit.jupiter.MockitoSettings; |
||||
import org.mockito.quality.Strictness; |
||||
|
||||
import com.google.common.collect.Lists; |
||||
|
||||
@ExtendWith(MockitoExtension.class) |
||||
@MockitoSettings(strictness = Strictness.LENIENT) |
||||
class TaskGroupCoordinatorTest { |
||||
|
||||
@InjectMocks |
||||
private TaskGroupCoordinator taskGroupCoordinator; |
||||
|
||||
@Mock |
||||
private RegistryClient registryClient; |
||||
|
||||
@Mock |
||||
private TaskGroupDao taskGroupDao; |
||||
|
||||
@Mock |
||||
private TaskGroupQueueDao taskGroupQueueDao; |
||||
|
||||
@Mock |
||||
private TaskInstanceDao taskInstanceDao; |
||||
|
||||
@Mock |
||||
private ProcessInstanceDao processInstanceDao; |
||||
|
||||
@Test |
||||
void start() throws InterruptedException { |
||||
// Get the Lock from Registry
|
||||
taskGroupCoordinator.start(); |
||||
Thread.sleep(1_000); |
||||
verify(registryClient, Mockito.times(1)) |
||||
.getLock(RegistryNodeType.MASTER_TASK_GROUP_COORDINATOR_LOCK.getRegistryPath()); |
||||
verify(registryClient, Mockito.times(1)) |
||||
.releaseLock(RegistryNodeType.MASTER_TASK_GROUP_COORDINATOR_LOCK.getRegistryPath()); |
||||
|
||||
} |
||||
|
||||
@Test |
||||
void needAcquireTaskGroupSlot() { |
||||
// TaskInstance is null
|
||||
IllegalArgumentException illegalArgumentException = |
||||
assertThrows(IllegalArgumentException.class, () -> taskGroupCoordinator.needAcquireTaskGroupSlot(null)); |
||||
assertEquals("The TaskInstance is null", illegalArgumentException.getMessage()); |
||||
|
||||
// TaskGroupId < 0
|
||||
TaskInstance taskInstance = new TaskInstance(); |
||||
assertFalse(taskGroupCoordinator.needAcquireTaskGroupSlot(taskInstance)); |
||||
|
||||
// TaskGroup not exist
|
||||
taskInstance.setTaskGroupId(1); |
||||
when(taskGroupDao.queryById(taskInstance.getTaskGroupId())).thenReturn(null); |
||||
assertFalse(taskGroupCoordinator.needAcquireTaskGroupSlot(taskInstance)); |
||||
|
||||
// TaskGroup is closed
|
||||
TaskGroup taskGroup = new TaskGroup(); |
||||
taskGroup.setStatus(Flag.NO); |
||||
when(taskGroupDao.queryById(taskInstance.getTaskGroupId())).thenReturn(taskGroup); |
||||
assertFalse(taskGroupCoordinator.needAcquireTaskGroupSlot(taskInstance)); |
||||
|
||||
// TaskGroup is open
|
||||
taskGroup.setStatus(Flag.YES); |
||||
when(taskGroupDao.queryById(taskInstance.getTaskGroupId())).thenReturn(taskGroup); |
||||
assertTrue(taskGroupCoordinator.needToReleaseTaskGroupSlot(taskInstance)); |
||||
|
||||
} |
||||
|
||||
@Test |
||||
void acquireTaskGroupSlot() { |
||||
// TaskInstance is NULL
|
||||
IllegalArgumentException illegalArgumentException = |
||||
assertThrows(IllegalArgumentException.class, () -> taskGroupCoordinator.acquireTaskGroupSlot(null)); |
||||
assertEquals("The current TaskInstance does not use task group", illegalArgumentException.getMessage()); |
||||
|
||||
// TaskGroupId is NULL
|
||||
TaskInstance taskInstance = new TaskInstance(); |
||||
illegalArgumentException = assertThrows(IllegalArgumentException.class, |
||||
() -> taskGroupCoordinator.acquireTaskGroupSlot(taskInstance)); |
||||
assertEquals("The current TaskInstance does not use task group", illegalArgumentException.getMessage()); |
||||
|
||||
// TaskGroup not exist
|
||||
taskInstance.setTaskGroupId(1); |
||||
taskInstance.setId(1); |
||||
when(taskGroupDao.queryById(taskInstance.getTaskGroupId())).thenReturn(null); |
||||
illegalArgumentException = assertThrows(IllegalArgumentException.class, |
||||
() -> taskGroupCoordinator.acquireTaskGroupSlot(taskInstance)); |
||||
assertEquals("The current TaskGroup: 1 does not exist", illegalArgumentException.getMessage()); |
||||
|
||||
// TaskGroup exist
|
||||
when(taskGroupDao.queryById(taskInstance.getTaskGroupId())).thenReturn(new TaskGroup()); |
||||
Assertions.assertDoesNotThrow(() -> taskGroupCoordinator.acquireTaskGroupSlot(taskInstance)); |
||||
|
||||
} |
||||
|
||||
@Test |
||||
void needToReleaseTaskGroupSlot() { |
||||
IllegalArgumentException illegalArgumentException = assertThrows(IllegalArgumentException.class, |
||||
() -> taskGroupCoordinator.needToReleaseTaskGroupSlot(null)); |
||||
assertEquals("The TaskInstance is null", illegalArgumentException.getMessage()); |
||||
|
||||
TaskInstance taskInstance = new TaskInstance(); |
||||
assertFalse(taskGroupCoordinator.needToReleaseTaskGroupSlot(taskInstance)); |
||||
|
||||
taskInstance.setTaskGroupId(1); |
||||
assertTrue(taskGroupCoordinator.needToReleaseTaskGroupSlot(taskInstance)); |
||||
} |
||||
|
||||
@Test |
||||
void releaseTaskGroupSlot() { |
||||
// TaskInstance is NULL
|
||||
IllegalArgumentException illegalArgumentException = |
||||
assertThrows(IllegalArgumentException.class, () -> taskGroupCoordinator.releaseTaskGroupSlot(null)); |
||||
assertEquals("The current TaskInstance does not use task group", illegalArgumentException.getMessage()); |
||||
|
||||
// TaskGroupId is NULL
|
||||
TaskInstance taskInstance = new TaskInstance(); |
||||
illegalArgumentException = assertThrows(IllegalArgumentException.class, |
||||
() -> taskGroupCoordinator.releaseTaskGroupSlot(taskInstance)); |
||||
assertEquals("The current TaskInstance does not use task group", illegalArgumentException.getMessage()); |
||||
|
||||
// Release TaskGroupQueue
|
||||
taskInstance.setId(1); |
||||
taskInstance.setTaskGroupId(1); |
||||
TaskGroupQueue taskGroupQueue = new TaskGroupQueue(); |
||||
List<TaskGroupQueue> taskGroupQueues = Lists.newArrayList(taskGroupQueue); |
||||
when(taskGroupQueueDao.queryByTaskInstanceId(taskInstance.getId())).thenReturn(taskGroupQueues); |
||||
taskGroupCoordinator.releaseTaskGroupSlot(taskInstance); |
||||
|
||||
assertEquals(Flag.NO.getCode(), taskGroupQueue.getInQueue()); |
||||
assertEquals(TaskGroupQueueStatus.RELEASE, taskGroupQueue.getStatus()); |
||||
verify(taskGroupQueueDao, Mockito.times(1)).updateById(taskGroupQueue); |
||||
|
||||
} |
||||
} |
||||
Loading…
Reference in new issue