Skip to content

Introduce AbstractAsyncSelector #6097

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
jrhee17 merged 3 commits into from
Mar 26, 2025
Merged

Introduce AbstractAsyncSelector #6097

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6181416
introduce async selector
jrhee17 Feb 13, 2025
4102bbd
fix test
jrhee17 Feb 13, 2025
7ec56d7
macos flakiness
jrhee17 Feb 14, 2025
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
198 changes: 29 additions & 169 deletions core/src/main/java/com/linecorp/armeria/client/endpoint/AbstractEndpointSelector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,24 +20,16 @@
import java.util.List;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

import com.google.common.annotations.VisibleForTesting;
import com.google.errorprone.annotations.concurrent.GuardedBy;

import com.linecorp.armeria.client.ClientRequestContext;
import com.linecorp.armeria.client.Endpoint;
import com.linecorp.armeria.common.annotation.Nullable;
import com.linecorp.armeria.common.annotation.UnstableApi;
import com.linecorp.armeria.common.util.UnmodifiableFuture;
import com.linecorp.armeria.internal.client.AbstractAsyncSelector;
import com.linecorp.armeria.internal.client.ClientPendingThrowableUtil;
import com.linecorp.armeria.internal.common.util.IdentityHashStrategy;
import com.linecorp.armeria.internal.common.util.ReentrantShortLock;

import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenCustomHashSet;

/**
* A skeletal {@link EndpointSelector} implementation. This abstract class implements the
Expand All @@ -47,16 +39,14 @@
public abstract class AbstractEndpointSelector implements EndpointSelector {

private final EndpointGroup endpointGroup;
private final ReentrantShortLock lock = new ReentrantShortLock();
@GuardedBy("lock")
private final Set<ListeningFuture> pendingFutures =
new ObjectLinkedOpenCustomHashSet<>(IdentityHashStrategy.of());
private final EndpointAsyncSelector asyncSelector;

/**
* Creates a new instance that selects an {@link Endpoint} from the specified {@link EndpointGroup}.
*/
protected AbstractEndpointSelector(EndpointGroup endpointGroup) {
this.endpointGroup = requireNonNull(endpointGroup, "endpointGroup");
asyncSelector = new EndpointAsyncSelector(endpointGroup);
}

/**
Expand All @@ -66,76 +56,6 @@ protected final EndpointGroup group() {
return endpointGroup;
}

@SuppressWarnings("GuardedBy")
@VisibleForTesting
final Set<ListeningFuture> pendingFutures() {
return pendingFutures;
}

@Deprecated
@Override
public final CompletableFuture<Endpoint> select(ClientRequestContext ctx,
ScheduledExecutorService executor,
long timeoutMillis) {
return select(ctx, executor);
}

@Override
public final CompletableFuture<Endpoint> select(ClientRequestContext ctx,
ScheduledExecutorService executor) {
final Endpoint endpoint = selectNow(ctx);
if (endpoint != null) {
return UnmodifiableFuture.completedFuture(endpoint);
}

final ListeningFuture listeningFuture = new ListeningFuture(ctx, executor);
addPendingFuture(listeningFuture);

// The EndpointGroup have just been updated after adding ListeningFuture.
if (listeningFuture.isDone()) {
return listeningFuture;
}
if (endpointGroup.whenReady().isDone()) {
final Endpoint endpoint0 = selectNow(ctx);
if (endpoint0 != null) {
// The EndpointGroup have just been updated before adding ListeningFuture.
listeningFuture.complete(endpoint0);
return listeningFuture;
}
}

final long selectionTimeoutMillis = endpointGroup.selectionTimeoutMillis();
if (selectionTimeoutMillis == 0) {
// A static EndpointGroup.
return UnmodifiableFuture.completedFuture(null);
}

// Schedule the timeout task.
if (selectionTimeoutMillis < Long.MAX_VALUE) {
final ScheduledFuture<?> timeoutFuture = executor.schedule(() -> {
final EndpointSelectionTimeoutException ex =
EndpointSelectionTimeoutException.get(endpointGroup, selectionTimeoutMillis);
ClientPendingThrowableUtil.setPendingThrowable(ctx, ex);
// Don't complete exceptionally so that the throwable
// can be handled after executing the attached decorators
listeningFuture.complete(null);
}, selectionTimeoutMillis, TimeUnit.MILLISECONDS);
listeningFuture.timeoutFuture = timeoutFuture;

// Cancel the timeout task if listeningFuture is done already.
// This guards against the following race condition:
// 1) (Current thread) Timeout task is scheduled.
// 2) ( Other thread ) listeningFuture is completed, but the timeout task is not cancelled
// 3) (Current thread) timeoutFuture is assigned to listeningFuture.timeoutFuture, but it's too
// late.
if (listeningFuture.isDone()) {
timeoutFuture.cancel(false);
}
}

return listeningFuture;
}

/**
* Initialize this {@link EndpointSelector} to listen to the new endpoints emitted by the
* {@link EndpointGroup}. The new endpoints will be passed to {@link #updateNewEndpoints(List)}.
Expand All @@ -145,15 +65,21 @@ protected final void initialize() {
endpointGroup.addListener(this::refreshEndpoints, true);
}

@Override
public CompletableFuture<Endpoint> select(ClientRequestContext ctx, ScheduledExecutorService executor,
long timeoutMillis) {
return asyncSelector.select(ctx, executor, endpointGroup.selectionTimeoutMillis());
}

private void refreshEndpoints(List<Endpoint> endpoints) {
// Allow subclasses to update the endpoints first.
updateNewEndpoints(endpoints);
lock.lock();
try {
pendingFutures.removeIf(ListeningFuture::tryComplete);
} finally {
lock.unlock();
}
asyncSelector.refresh();
}

@VisibleForTesting
Set<? extends CompletableFuture<Endpoint>> pendingFutures() {
return asyncSelector.pendingFutures();
}

/**
Expand All @@ -162,97 +88,31 @@ private void refreshEndpoints(List<Endpoint> endpoints) {
@UnstableApi
protected void updateNewEndpoints(List<Endpoint> endpoints) {}

private void addPendingFuture(ListeningFuture future) {
lock.lock();
try {
pendingFutures.add(future);
} finally {
lock.unlock();
}
}

private void removePendingFuture(ListeningFuture future) {
lock.lock();
try {
pendingFutures.remove(future);
} finally {
lock.unlock();
}
}

@VisibleForTesting
final class ListeningFuture extends CompletableFuture<Endpoint> {
private final ClientRequestContext ctx;
private final Executor executor;
@Nullable
private volatile Endpoint selectedEndpoint;
@Nullable
private volatile ScheduledFuture<?> timeoutFuture;

ListeningFuture(ClientRequestContext ctx, Executor executor) {
this.ctx = ctx;
this.executor = executor;
}

/**
* Returns {@code true} if an {@link Endpoint} has been selected.
*/
public boolean tryComplete() {
if (selectedEndpoint != null || isDone()) {
return true;
}

try {
final Endpoint endpoint = selectNow(ctx);
if (endpoint == null) {
return false;
}
private class EndpointAsyncSelector extends AbstractAsyncSelector<Endpoint> {

cleanup(false);
// Complete with the selected endpoint.
selectedEndpoint = endpoint;
executor.execute(() -> super.complete(endpoint));
return true;
} catch (Throwable t) {
cleanup(false);
super.completeExceptionally(t);
return true;
}
}
private final EndpointGroup endpointGroup;

@Override
public boolean cancel(boolean mayInterruptIfRunning) {
cleanup(true);
return super.cancel(mayInterruptIfRunning);
EndpointAsyncSelector(EndpointGroup endpointGroup) {
this.endpointGroup = endpointGroup;
}

@Override
public boolean complete(@Nullable Endpoint value) {
cleanup(true);
return super.complete(value);
protected void onTimeout(ClientRequestContext ctx, long selectionTimeoutMillis) {
final EndpointSelectionTimeoutException ex =
EndpointSelectionTimeoutException.get(endpointGroup, selectionTimeoutMillis);
ClientPendingThrowableUtil.setPendingThrowable(ctx, ex);
}

@Nullable
@Override
public boolean completeExceptionally(Throwable ex) {
cleanup(true);
return super.completeExceptionally(ex);
}

private void cleanup(boolean removePendingFuture) {
if (removePendingFuture) {
removePendingFuture(this);
}
final ScheduledFuture<?> timeoutFuture = this.timeoutFuture;
if (timeoutFuture != null) {
this.timeoutFuture = null;
timeoutFuture.cancel(false);
}
protected Endpoint selectNow(ClientRequestContext ctx) {
return AbstractEndpointSelector.this.selectNow(ctx);
}

@Nullable
@VisibleForTesting
ScheduledFuture<?> timeoutFuture() {
return timeoutFuture;
@Override
protected Set<? extends CompletableFuture<Endpoint>> pendingFutures() {
return super.pendingFutures();
}
}
}
Loading