org.apache.druid.rpc.indexing.SpecificTaskServiceLocator Maven / Gradle / Ivy
/*
* 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.druid.rpc.indexing;
import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import com.google.errorprone.annotations.concurrent.GuardedBy;
import org.apache.druid.client.indexing.TaskStatusResponse;
import org.apache.druid.indexer.TaskLocation;
import org.apache.druid.indexer.TaskState;
import org.apache.druid.indexer.TaskStatus;
import org.apache.druid.indexer.TaskStatusPlus;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.concurrent.Execs;
import org.apache.druid.rpc.ServiceLocation;
import org.apache.druid.rpc.ServiceLocations;
import org.apache.druid.rpc.ServiceLocator;
import java.util.Collections;
import java.util.Map;
import java.util.Set;
/**
* Service locator for a specific task. Uses the {@link OverlordClient#taskStatuses(Set)} API to locate tasks.
*
* This locator has an internal cache that is updated if the last check has been over {@link #LOCATION_CACHE_MS} ago.
*
* This locator is Closeable, like all ServiceLocators, but it is not essential that you actually close it. Closing
* does not free any resources: it merely makes future calls to {@link #locate()} return
* {@link ServiceLocations#closed()}.
*/
public class SpecificTaskServiceLocator implements ServiceLocator
{
private static final String BASE_PATH = "/druid/worker/v1/chat";
private static final long LOCATION_CACHE_MS = 30_000;
private final String taskId;
private final OverlordClient overlordClient;
private final Object lock = new Object();
@GuardedBy("lock")
private TaskState lastKnownState = TaskState.RUNNING; // Assume task starts out running.
@GuardedBy("lock")
private ServiceLocation lastKnownLocation;
@GuardedBy("lock")
private boolean closed = false;
@GuardedBy("lock")
private long lastUpdateTime = -1;
@GuardedBy("lock")
private SettableFuture pendingFuture = null;
public SpecificTaskServiceLocator(final String taskId, final OverlordClient overlordClient)
{
this.taskId = taskId;
this.overlordClient = overlordClient;
}
@Override
public ListenableFuture locate()
{
synchronized (lock) {
if (pendingFuture != null) {
return Futures.nonCancellationPropagating(pendingFuture);
} else if (closed || lastKnownState != TaskState.RUNNING) {
return Futures.immediateFuture(ServiceLocations.closed());
} else if (lastKnownLocation == null || lastUpdateTime + LOCATION_CACHE_MS < System.currentTimeMillis()) {
final ListenableFuture> taskStatusFuture;
try {
taskStatusFuture = overlordClient.taskStatuses(ImmutableSet.of(taskId));
}
catch (Exception e) {
throw new RuntimeException(e);
}
// Use shared future for concurrent calls to "locate"; don't want multiple calls out to the Overlord at once.
// Alias pendingFuture to retVal in case taskStatusFuture is already resolved. (This will make the callback
// below execute immediately, firing and nulling out pendingFuture.)
final SettableFuture retVal = (pendingFuture = SettableFuture.create());
pendingFuture.addListener(
() -> {
if (!taskStatusFuture.isDone()) {
// pendingFuture may resolve without taskStatusFuture due to close().
taskStatusFuture.cancel(true);
}
},
Execs.directExecutor()
);
Futures.addCallback(
taskStatusFuture,
new FutureCallback>()
{
@Override
public void onSuccess(final Map taskStatusMap)
{
synchronized (lock) {
if (pendingFuture != null) {
lastUpdateTime = System.currentTimeMillis();
final TaskStatus status = taskStatusMap.get(taskId);
if (status == null) {
// If the task status is unknown, we'll treat it as closed.
resolvePendingFuture(null, null);
} else if (TaskLocation.unknown().equals(status.getLocation())) {
// Do not resolve the future just yet, try the fallback API instead
fetchFallbackTaskLocation();
} else {
resolvePendingFuture(status.getStatusCode(), status.getLocation());
}
}
}
}
@Override
public void onFailure(Throwable t)
{
resolvePendingFutureOnException(t);
}
},
Execs.directExecutor()
);
return Futures.nonCancellationPropagating(retVal);
} else {
return Futures.immediateFuture(ServiceLocations.forLocation(lastKnownLocation));
}
}
}
@Override
public void close()
{
// Class-level Javadocs promise that this method does not actually free resources: it only alters behavior
// for future calls to locate(). This is exploited in TaskServiceClients.makeClient.
synchronized (lock) {
// Idempotent: can call close() multiple times so long as start() has already been called.
if (!closed) {
if (pendingFuture != null) {
pendingFuture.set(ServiceLocations.closed());
// Clear pendingFuture once it has been set.
pendingFuture = null;
}
closed = true;
}
}
}
private void resolvePendingFuture(TaskState state, TaskLocation location)
{
synchronized (lock) {
if (pendingFuture != null) {
lastKnownState = state;
lastKnownLocation = location == null ? null : new ServiceLocation(
location.getHost(),
location.getPort(),
location.getTlsPort(),
StringUtils.format("%s/%s", BASE_PATH, StringUtils.urlEncode(taskId))
);
if (lastKnownState != TaskState.RUNNING) {
pendingFuture.set(ServiceLocations.closed());
} else if (lastKnownLocation == null) {
pendingFuture.set(ServiceLocations.forLocations(Collections.emptySet()));
} else {
pendingFuture.set(ServiceLocations.forLocation(lastKnownLocation));
}
// Clear pendingFuture once it has been set.
pendingFuture = null;
}
}
}
private void resolvePendingFutureOnException(Throwable t)
{
synchronized (lock) {
if (pendingFuture != null) {
pendingFuture.setException(t);
// Clear pendingFuture once it has been set.
pendingFuture = null;
}
}
}
/**
* Invokes the single task status API {@link OverlordClient#taskStatus} if the
* multi-task status API returns an unknown location (this can happen if the
* Overlord is running on a version older than Druid 30.0.0 (pre #15724)).
*/
private void fetchFallbackTaskLocation()
{
synchronized (lock) {
if (pendingFuture != null) {
final ListenableFuture taskStatusFuture;
try {
taskStatusFuture = overlordClient.taskStatus(taskId);
}
catch (Exception e) {
resolvePendingFutureOnException(e);
return;
}
pendingFuture.addListener(
() -> {
if (!taskStatusFuture.isDone()) {
// pendingFuture may resolve without taskStatusFuture due to close().
taskStatusFuture.cancel(true);
}
},
Execs.directExecutor()
);
Futures.addCallback(
taskStatusFuture,
new FutureCallback()
{
@Override
public void onSuccess(final TaskStatusResponse taskStatusResponse)
{
synchronized (lock) {
if (pendingFuture != null) {
lastUpdateTime = System.currentTimeMillis();
final TaskStatusPlus status = taskStatusResponse.getStatus();
if (status == null) {
// If the task status is unknown, we'll treat it as closed.
resolvePendingFuture(null, null);
} else if (TaskLocation.unknown().equals(status.getLocation())) {
resolvePendingFuture(status.getStatusCode(), null);
} else {
resolvePendingFuture(status.getStatusCode(), status.getLocation());
}
}
}
}
@Override
public void onFailure(Throwable t)
{
resolvePendingFutureOnException(t);
}
},
Execs.directExecutor()
);
}
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy