Skip to content

Commit

Permalink
ITOAuth2Client: improve shutdown sequence
Browse files Browse the repository at this point in the history
Remarked while working on #7673: the call to close()
may disrupt ongoing token refreshes and / or refresh
schedule attempts; the client does close normally
eventually, but if DEBUG logs are enabled, many
strange errors may be logged. Also, if the executor
is scheduling a refresh when close() is called, this
may trigger the start of a thread, only to see it
being closed immediately after.

By guarding the shutdown sequence with a boolean flag,
we can achieve a much smoother shutdown sequence.
  • Loading branch information
adutra committed Oct 27, 2023
1 parent e14db90 commit 0b180ea
Showing 1 changed file with 34 additions and 19 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ class OAuth2Client implements OAuth2Authenticator, Closeable {
private final ObjectMapper objectMapper;
private final CompletableFuture<Void> started = new CompletableFuture<>();
/* Visible for testing. */ final AtomicBoolean sleeping = new AtomicBoolean();
private final AtomicBoolean closing = new AtomicBoolean();
private final Supplier<Instant> clock;

private volatile CompletionStage<Tokens> currentTokensStage;
Expand Down Expand Up @@ -156,27 +157,31 @@ public void start() {

@Override
public void close() {
try {
currentTokensStage.toCompletableFuture().cancel(true);
ScheduledFuture<?> tokenRefreshFuture = this.tokenRefreshFuture;
if (tokenRefreshFuture != null) {
tokenRefreshFuture.cancel(true);
}
if (shouldCloseExecutor) {
if (!executor.isShutdown()) {
executor.shutdown();
if (!executor.awaitTermination(10, TimeUnit.SECONDS)) {
executor.shutdownNow();
if (closing.compareAndSet(false, true)) {
LOGGER.debug("Closing...");
try {
currentTokensStage.toCompletableFuture().cancel(true);
ScheduledFuture<?> tokenRefreshFuture = this.tokenRefreshFuture;
if (tokenRefreshFuture != null) {
tokenRefreshFuture.cancel(true);
}
if (shouldCloseExecutor) {
if (!executor.isShutdown()) {
executor.shutdown();
if (!executor.awaitTermination(10, TimeUnit.SECONDS)) {
executor.shutdownNow();
}
}
}
if (password != null) {
Arrays.fill(password, (byte) 0);
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
} finally {
tokenRefreshFuture = null;
}
if (password != null) {
Arrays.fill(password, (byte) 0);
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
} finally {
tokenRefreshFuture = null;
LOGGER.debug("Closed");
}
}

Expand Down Expand Up @@ -205,11 +210,18 @@ private void maybeScheduleTokensRenewal(Tokens currentTokens) {
}

private void scheduleTokensRenewal(Duration delay) {
if (closing.get()) {
return;
}
LOGGER.debug("Scheduling token refresh in {}", delay);
try {
tokenRefreshFuture =
executor.schedule(this::renewTokens, delay.toMillis(), TimeUnit.MILLISECONDS);
} catch (RejectedExecutionException e) {
if (closing.get()) {
// We raced with close(), ignore
return;
}
LOGGER.warn("Failed to schedule next token renewal, forcibly sleeping", e);
sleeping.set(true);
}
Expand All @@ -230,7 +242,10 @@ private void renewTokens() {

private void log(Throwable error) {
if (error != null) {
LOGGER.error("Failed to renew tokens", error);
boolean tokensStageCancelled = error instanceof CancellationException && closing.get();
if (!tokensStageCancelled) {
LOGGER.error("Failed to renew tokens", error);
}
} else {
LOGGER.debug("Successfully renewed tokens");
}
Expand Down

0 comments on commit 0b180ea

Please sign in to comment.