How to use

TaskMaster

in

org.apache.druid.indexing.overlord

@Override
public Map<String, Long> getSuccessfulTaskCount()
{
 Optional<TaskQueue> taskQueue = getTaskQueue();
 if (taskQueue.isPresent()) {
  return taskQueue.get().getSuccessfulTaskCount();
 } else {
  return null;
 }
}

@GET
@Path("/leader")
@ResourceFilters(StateResourceFilter.class)
@Produces(MediaType.APPLICATION_JSON)
public Response getLeader()
{
 return Response.ok(taskMaster.getCurrentLeader()).build();
}

private Response asLeaderWithSupervisorManager(Function<SupervisorManager, Response> f)
{
 Optional<SupervisorManager> supervisorManager = taskMaster.getSupervisorManager();
 if (supervisorManager.isPresent()) {
  return f.apply(supervisorManager.get());
 } else {
  // Encourage client to try again soon, when we'll likely have a redirect set up
  return Response.status(Response.Status.SERVICE_UNAVAILABLE).build();
 }
}

@Test
public void testShutdownTask()
{
 // This is disabled since OverlordResource.doShutdown is annotated with TaskResourceFilter
 // This should be fixed in https://github.com/apache/incubator-druid/issues/6685.
 // expectAuthorizationTokenCheck();
 TaskQueue mockQueue = EasyMock.createMock(TaskQueue.class);
 EasyMock.expect(taskMaster.isLeader()).andReturn(true).anyTimes();
 EasyMock.expect(taskMaster.getTaskRunner()).andReturn(
   Optional.of(taskRunner)
 ).anyTimes();
 EasyMock.expect(taskMaster.getTaskQueue()).andReturn(
   Optional.of(mockQueue)
 ).anyTimes();
 mockQueue.shutdown("id_1", "Shutdown request from user");
 EasyMock.expectLastCall();
 EasyMock.replay(taskRunner, taskMaster, taskStorageQueryAdapter, indexerMetadataStorageAdapter, req, mockQueue);
 final Map<String, Integer> response = (Map<String, Integer>) overlordResource
   .doShutdown("id_1")
   .getEntity();
 Assert.assertEquals("id_1", response.get("task"));
}

@Override
public boolean doLocal(String requestURI)
{
 return (requestURI != null && LOCAL_PATHS.contains(requestURI)) || taskMaster.isLeader();
}

 private HttpRemoteTaskRunner getHttpRemoteTaskRunner()
 {
  Optional<TaskRunner> taskRunnerOpt = taskMaster.getTaskRunner();
  if (taskRunnerOpt.isPresent() && taskRunnerOpt.get() instanceof HttpRemoteTaskRunner) {
   return (HttpRemoteTaskRunner) taskRunnerOpt.get();
  } else {
   return null;
  }
 }
}

taskMaster.start();
announcementLatch.await();
while (!taskMaster.isLeader()) {
Assert.assertEquals(taskMaster.getCurrentLeader(), druidNode.getHostAndPort());
taskMaster.stop();
Assert.assertFalse(taskMaster.isLeader());
EasyMock.verify(taskLockbox, taskActionClientFactory);

@Test
public void testShutdownAllTasksForNonExistingDataSource()
{
 final TaskQueue taskQueue = EasyMock.createMock(TaskQueue.class);
 EasyMock.expect(taskMaster.isLeader()).andReturn(true).anyTimes();
 EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes();
 EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo(EasyMock.anyString())).andReturn(Collections.emptyList());
 EasyMock.replay(taskRunner, taskMaster, taskStorageQueryAdapter, indexerMetadataStorageAdapter, req);
 final Response response = overlordResource.shutdownTasksForDataSource("notExisting");
 Assert.assertEquals(Status.NOT_FOUND.getStatusCode(), response.getStatus());
}

@GET
@Path("/scaling")
@Produces(MediaType.APPLICATION_JSON)
@ResourceFilters(StateResourceFilter.class)
public Response getScalingState()
{
 // Don't use asLeaderWith, since we want to return 200 instead of 503 when missing an autoscaler.
 final Optional<ScalingStats> rms = taskMaster.getScalingStats();
 if (rms.isPresent()) {
  return Response.ok(rms.get()).build();
 } else {
  return Response.ok().build();
 }
}

taskMaster.getTaskActionClient(holder.getTask()),
new Function<TaskActionClient, Response>()

/**
 * Stops forever (not just this particular leadership session). Should only be called once throughout the life of
 * the program.
 */
@LifecycleStop
public void stop()
{
 giant.lock();
 try {
  gracefulStopLeaderLifecycle();
  overlordLeaderSelector.unregisterListener();
 }
 finally {
  giant.unlock();
 }
}

druidNode = new DruidNode("hey", "what", false, 1234, null, true, false);
ServiceEmitter serviceEmitter = new NoopServiceEmitter();
taskMaster = new TaskMaster(
  new TaskQueueConfig(null, new Period(1), null, new Period(10)),
  taskLockbox,

EasyMock.expect(taskMaster.isLeader()).andReturn(true).anyTimes();
EasyMock.expect(taskMaster.getTaskRunner()).andReturn(
  Optional.of(taskRunner)
).anyTimes();
EasyMock.expect(taskMaster.getTaskQueue()).andReturn(
  Optional.of(mockQueue)
).anyTimes();

public Optional<SupervisorManager> getSupervisorManager()
{
 if (isLeader()) {
  return Optional.of(supervisorManager);
 } else {
  return Optional.absent();
 }
}

 @Override
 public Optional<ByteSource> streamTaskLog(String taskid, long offset) throws IOException
 {
  final TaskRunner runner = taskMaster.getTaskRunner().orNull();
  if (runner instanceof TaskLogStreamer) {
   return ((TaskLogStreamer) runner).streamTaskLog(taskid, offset);
  } else {
   return Optional.absent();
  }
 }
}

@GET
@Path("/scaling")
@Produces(MediaType.APPLICATION_JSON)
@ResourceFilters(StateResourceFilter.class)
public Response getScalingState()
{
 // Don't use asLeaderWith, since we want to return 200 instead of 503 when missing an autoscaler.
 final Optional<ScalingStats> rms = taskMaster.getScalingStats();
 if (rms.isPresent()) {
  return Response.ok(rms.get()).build();
 } else {
  return Response.ok().build();
 }
}

taskMaster.getTaskActionClient(holder.getTask()),
new Function<TaskActionClient, Response>()

@Override
public Map<String, Long> getFailedTaskCount()
{
 Optional<TaskQueue> taskQueue = getTaskQueue();
 if (taskQueue.isPresent()) {
  return taskQueue.get().getFailedTaskCount();
 } else {
  return null;
 }
}

public Optional<TaskQueue> getTaskQueue()
{
 if (isLeader()) {
  return Optional.of(taskQueue);
 } else {
  return Optional.absent();
 }
}

@Test
public void testTerminateAll()
{
 EasyMock.expect(taskMaster.getSupervisorManager()).andReturn(Optional.of(supervisorManager));
 supervisorManager.stopAndRemoveAllSupervisors();
 EasyMock.expectLastCall();
 replayAll();
 Response response = supervisorResource.terminateAll();
 Assert.assertEquals(200, response.getStatus());
 Assert.assertEquals(ImmutableMap.of("status", "success"), response.getEntity());
 verifyAll();
}