Skip to content

Move some lifecycle management from doTask -> shutdown for the mm-less task runner #14895

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

Merged
dclim merged 9 commits into from
Aug 25, 2023
Merged

Move some lifecycle management from doTask -> shutdown for the mm-less task runner #14895

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
a8e921f
save work
Aug 9, 2023
633f391
Merge branch 'master' of github.com:georgew5656/druid into saveTaskLo…
Aug 21, 2023
ba4d1e2
Add syncronized
Aug 21, 2023
137f740
Don't shutdown in run
Aug 22, 2023
3cb910b
Adding unit tests
Aug 22, 2023
7f5caef
Cleanup lifecycle
Aug 22, 2023
7d0a3ea
Fix tests
Aug 22, 2023
dccca17
Fix merge conflicts
Aug 25, 2023
f72a3ca
remove newline
Aug 25, 2023
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: 3 additions & 6 deletions ...rlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycle.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,6 @@ protected synchronized TaskStatus run(Job job, long launchTimeout, long timeout)
}
catch (Exception e) {
log.info("Failed to run task: %s", taskId.getOriginalTaskId());
shutdown();
throw e;
}
finally {
Expand Down Expand Up @@ -168,10 +167,9 @@ protected synchronized TaskStatus join(long timeout) throws IllegalStateExceptio
finally {
try {
saveLogs();
shutdown();
}
catch (Exception e) {
log.warn(e, "Task [%s] cleanup failed", taskId);
log.warn(e, "Log processing failed for task [%s]", taskId);
}

stopTask();
Expand All @@ -188,7 +186,7 @@ protected synchronized TaskStatus join(long timeout) throws IllegalStateExceptio
*/
protected void shutdown()
{
if (State.PENDING.equals(state.get()) || State.RUNNING.equals(state.get())) {
if (State.PENDING.equals(state.get()) || State.RUNNING.equals(state.get()) || State.STOPPED.equals(state.get())) {

@kfaraz kfaraz Aug 24, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need to handle STOPPED state here? Wouldn't the job have already finished?

@georgew5656 georgew5656 Aug 24, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is saying thats its okay to shutdown if its stopped

kubernetesClient.deletePeonJob(taskId);
}
}
Expand Down Expand Up @@ -223,7 +221,7 @@ protected State getState()
*/
protected TaskLocation getTaskLocation()
{
if (!State.RUNNING.equals(state.get())) {
if (State.PENDING.equals(state.get()) || State.NOT_STARTED.equals(state.get())) {
log.debug("Can't get task location for non-running job. [%s]", taskId.getOriginalTaskId());
return TaskLocation.unknown();
}
Expand Down Expand Up @@ -251,7 +249,6 @@ protected TaskLocation getTaskLocation()
Boolean.parseBoolean(pod.getMetadata().getAnnotations().getOrDefault(DruidK8sConstants.TLS_ENABLED, "false")),
pod.getMetadata() != null ? pod.getMetadata().getName() : ""
);
log.info("K8s task %s is running at location %s", taskId.getOriginalTaskId(), taskLocation);
}

return taskLocation;
Expand Down
24 changes: 15 additions & 9 deletions ...overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,10 +182,6 @@ protected TaskStatus doTask(Task task, boolean run)
KubernetesWorkItem workItem = tasks.get(task.getId());

if (workItem == null) {
throw new ISE("Task [%s] disappeared", task.getId());
}

if (workItem.isShutdownRequested()) {
throw new ISE("Task [%s] has been shut down", task.getId());
}

Expand Down Expand Up @@ -213,11 +209,6 @@ protected TaskStatus doTask(Task task, boolean run)
log.error(e, "Task [%s] execution caught an exception", task.getId());
throw new RuntimeException(e);
}
finally {
synchronized (tasks) {
tasks.remove(task.getId());
}
}
}

@VisibleForTesting
Expand Down Expand Up @@ -271,6 +262,10 @@ public void shutdown(String taskid, String reason)
return;
}

synchronized (tasks) {

@YongGang YongGang Aug 25, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: the code can be simplified:

    KubernetesWorkItem workItem;
    synchronized (tasks) {
       workItem = tasks.remove(taskid);
    }
    if (workItem == null) {
      log.info("Ignoring request to cancel unknown task [%s]", taskid);
      return;
    }

tasks.remove(taskid);
}

workItem.shutdown();
}

Expand Down Expand Up @@ -440,6 +435,17 @@ public Collection<TaskRunnerWorkItem> getPendingTasks()
.collect(Collectors.toList());
}

@Override
public TaskLocation getTaskLocation(String taskId)
{
final KubernetesWorkItem workItem = tasks.get(taskId);
if (workItem == null) {
return TaskLocation.unknown();
} else {
return workItem.getLocation();
}
}

@Nullable
@Override
public RunnerTaskState getRunnerTaskState(String taskId)
Expand Down
9 changes: 0 additions & 9 deletions ...s-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesWorkItem.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,13 +30,10 @@
import org.apache.druid.java.util.common.ISE;

import java.io.InputStream;
import java.util.concurrent.atomic.AtomicBoolean;

public class KubernetesWorkItem extends TaskRunnerWorkItem
{
private final Task task;

private final AtomicBoolean shutdownRequested = new AtomicBoolean(false);
private KubernetesPeonLifecycle kubernetesPeonLifecycle = null;

public KubernetesWorkItem(Task task, ListenableFuture<TaskStatus> statusFuture)
Expand All @@ -53,19 +50,13 @@ protected synchronized void setKubernetesPeonLifecycle(KubernetesPeonLifecycle k

protected synchronized void shutdown()
{
this.shutdownRequested.set(true);

if (this.kubernetesPeonLifecycle != null) {
this.kubernetesPeonLifecycle.startWatchingLogs();
this.kubernetesPeonLifecycle.shutdown();
}
}

protected boolean isShutdownRequested()
{
return shutdownRequested.get();
}

protected boolean isPending()
{
return RunnerTaskState.PENDING.equals(getRunnerTaskState());
Expand Down
14 changes: 3 additions & 11 deletions ...d-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycleTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,9 +198,6 @@ protected synchronized TaskStatus join(long timeout)
EasyMock.anyLong(),
EasyMock.eq(TimeUnit.MILLISECONDS)
)).andReturn(null);
EasyMock.expect(kubernetesClient.deletePeonJob(
new K8sTaskId(ID)
)).andReturn(true);
Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());
stateListener.stateChanged(KubernetesPeonLifecycle.State.PENDING, ID);
EasyMock.expectLastCall().once();
Expand Down Expand Up @@ -245,7 +242,6 @@ public void test_join_withoutJob_returnsFailedTaskStatus() throws IOException
EasyMock.expectLastCall().once();
logWatch.close();
EasyMock.expectLastCall();
EasyMock.expect(kubernetesClient.deletePeonJob(k8sTaskId)).andReturn(true);

replayAll();

Expand Down Expand Up @@ -298,7 +294,6 @@ public void test_join() throws IOException
EasyMock.expectLastCall().once();
logWatch.close();
EasyMock.expectLastCall();
EasyMock.expect(kubernetesClient.deletePeonJob(k8sTaskId)).andReturn(true);

Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());

Expand Down Expand Up @@ -353,7 +348,6 @@ public void test_join_whenCalledMultipleTimes_raisesIllegalStateException() thro
EasyMock.expectLastCall().once();
logWatch.close();
EasyMock.expectLastCall();
EasyMock.expect(kubernetesClient.deletePeonJob(k8sTaskId)).andReturn(true);

Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());

Expand Down Expand Up @@ -408,7 +402,6 @@ public void test_join_withoutTaskStatus_returnsFailedTaskStatus() throws IOExcep
EasyMock.expectLastCall().once();
logWatch.close();
EasyMock.expectLastCall();
EasyMock.expect(kubernetesClient.deletePeonJob(k8sTaskId)).andReturn(true);

Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());

Expand Down Expand Up @@ -459,7 +452,6 @@ public void test_join_whenIOExceptionThrownWhileStreamingTaskStatus_returnsFaile
EasyMock.expectLastCall().once();
logWatch.close();
EasyMock.expectLastCall();
EasyMock.expect(kubernetesClient.deletePeonJob(k8sTaskId)).andReturn(true);

Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());

Expand Down Expand Up @@ -512,7 +504,6 @@ public void test_join_whenIOExceptionThrownWhileStreamingTaskLogs_isIgnored() th
EasyMock.expectLastCall().once();
logWatch.close();
EasyMock.expectLastCall();
EasyMock.expect(kubernetesClient.deletePeonJob(k8sTaskId)).andReturn(true);

Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());

Expand Down Expand Up @@ -554,8 +545,6 @@ public void test_join_whenRuntimeExceptionThrownWhileWaitingForKubernetesJob_thr
logWatch.close();
EasyMock.expectLastCall();

EasyMock.expect(kubernetesClient.deletePeonJob(k8sTaskId)).andReturn(true);

Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());

replayAll();
Expand Down Expand Up @@ -908,8 +897,11 @@ public void test_getTaskLocation_withStoppedTaskState_returnsUnknown()
stateListener
);
setPeonLifecycleState(peonLifecycle, KubernetesPeonLifecycle.State.STOPPED);
EasyMock.expect(kubernetesClient.getPeonPod(k8sTaskId.getK8sJobName())).andReturn(Optional.absent()).once();

replayAll();
Assert.assertEquals(TaskLocation.unknown(), peonLifecycle.getTaskLocation());
verifyAll();
}

private void setPeonLifecycleState(KubernetesPeonLifecycle peonLifecycle, KubernetesPeonLifecycle.State state)
Expand Down
62 changes: 31 additions & 31 deletions ...lord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,8 +152,6 @@ public void test_run_withoutExistingTask() throws IOException, ExecutionExceptio
Assert.assertEquals(taskStatus, future.get());

verifyAll();

Assert.assertFalse(runner.tasks.containsKey(task.getId()));
}

@Test
Expand Down Expand Up @@ -191,8 +189,6 @@ public void test_run_whenExceptionThrown_throwsRuntimeException() throws IOExcep
Assert.assertTrue(e.getCause() instanceof RuntimeException);

verifyAll();

Assert.assertFalse(runner.tasks.containsKey(task.getId()));
}

@Test
Expand All @@ -208,8 +204,6 @@ public void test_join_withoutExistingTask() throws ExecutionException, Interrupt
Assert.assertEquals(taskStatus, future.get());

verifyAll();

Assert.assertFalse(runner.tasks.containsKey(task.getId()));
}

@Test
Expand All @@ -236,28 +230,11 @@ public void test_join_whenExceptionThrown_throwsRuntimeException()
Assert.assertTrue(e.getCause() instanceof RuntimeException);

verifyAll();

Assert.assertFalse(runner.tasks.containsKey(task.getId()));
}

@Test
public void test_doTask_withoutWorkItem_throwsRuntimeException()
{
Assert.assertThrows(
"Task [id] disappeared",
RuntimeException.class,
() -> runner.doTask(task, true)
);
}

@Test
public void test_doTask_whenShutdownRequested_throwsRuntimeException()
{
KubernetesWorkItem workItem = new KubernetesWorkItem(task, null);
workItem.shutdown();

runner.tasks.put(task.getId(), workItem);

Assert.assertThrows(
"Task [id] has been shut down",
RuntimeException.class,
Expand All @@ -266,13 +243,7 @@ public void test_doTask_whenShutdownRequested_throwsRuntimeException()
}

@Test
public void test_shutdown_withoutExistingTask()
{
runner.shutdown(task.getId(), "");
}

@Test
public void test_shutdown_withExistingTask()
public void test_shutdown_withExistingTask_removesTaskFromMap()
{
KubernetesWorkItem workItem = new KubernetesWorkItem(task, null) {
@Override
Expand All @@ -282,7 +253,13 @@ protected synchronized void shutdown()
};

runner.tasks.put(task.getId(), workItem);
runner.shutdown(task.getId(), "");
Assert.assertTrue(runner.tasks.isEmpty());
}

@Test
public void test_shutdown_withoutExistingTask()
{
runner.shutdown(task.getId(), "");
}

Expand Down Expand Up @@ -629,6 +606,30 @@ public TaskLocation getLocation()
verifyAll();
}

@Test
public void test_getTaskLocation_withExistingTask()
{
KubernetesWorkItem workItem = new KubernetesWorkItem(task, null) {
@Override
public TaskLocation getLocation()
{
return TaskLocation.create("host", 0, 1, false);
}
};

runner.tasks.put(task.getId(), workItem);

TaskLocation taskLocation = runner.getTaskLocation(task.getId());
Assert.assertEquals(TaskLocation.create("host", 0, 1, false), taskLocation);
}

@Test
public void test_getTaskLocation_noTaskFound()
{
TaskLocation taskLocation = runner.getTaskLocation(task.getId());
Assert.assertEquals(TaskLocation.unknown(), taskLocation);
}

@Test
public void test_getTotalCapacity()
{
Expand All @@ -644,6 +645,5 @@ public void test_getUsedCapacity()
Assert.assertEquals(1, runner.getUsedCapacity());
runner.tasks.remove(task.getId());
Assert.assertEquals(0, runner.getUsedCapacity());

}
}
2 changes: 0 additions & 2 deletions ...erlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesWorkItemTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,6 @@ public void test_setKubernetesPeonLifecycleTwice_throwsIllegalStateException()
public void test_shutdown_withoutKubernetesPeonLifecycle()
{
workItem.shutdown();
Assert.assertTrue(workItem.isShutdownRequested());
}

@Test
Expand All @@ -91,7 +90,6 @@ public void test_shutdown_withKubernetesPeonLifecycle()

workItem.shutdown();
verifyAll();
Assert.assertTrue(workItem.isShutdownRequested());
}

@Test
Expand Down