Skip to content

Commit

Permalink
Use deterministic ordering for inserted shutdown hooks
Browse files Browse the repository at this point in the history
Update `SpringApplicationShutdownHook` so the underlying set of
`Runnable` instances are stored in a `LinkedHashSet` rather than
a `Collections.newSetFromMap(new IdentityHashMap<>())`. This insures
that shutdown hooks are run in the order that they are added.

Fixes gh-43430
  • Loading branch information
philwebb committed Dec 17, 2024
1 parent e6e1274 commit 21203f0
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
package org.springframework.boot;

import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.LinkedHashSet;
import java.util.Set;
import java.util.WeakHashMap;
Expand Down Expand Up @@ -104,16 +103,16 @@ void deregisterFailedApplicationContext(ConfigurableApplicationContext applicati
public void run() {
Set<ConfigurableApplicationContext> contexts;
Set<ConfigurableApplicationContext> closedContexts;
Set<Runnable> actions;
Set<Handler> handlers;
synchronized (SpringApplicationShutdownHook.class) {
this.inProgress = true;
contexts = new LinkedHashSet<>(this.contexts);
closedContexts = new LinkedHashSet<>(this.closedContexts);
actions = new LinkedHashSet<>(this.handlers.getActions());
handlers = new LinkedHashSet<>(this.handlers.getActions());
}
contexts.forEach(this::closeAndWait);
closedContexts.forEach(this::closeAndWait);
actions.forEach(Runnable::run);
handlers.forEach(Handler::run);
}

boolean isApplicationContextRegistered(ConfigurableApplicationContext context) {
Expand Down Expand Up @@ -171,15 +170,15 @@ private void assertNotInProgress() {
*/
private final class Handlers implements SpringApplicationShutdownHandlers, Runnable {

private final Set<Runnable> actions = Collections.newSetFromMap(new IdentityHashMap<>());
private final Set<Handler> actions = new LinkedHashSet<>();

@Override
public void add(Runnable action) {
Assert.notNull(action, "Action must not be null");
addRuntimeShutdownHookIfNecessary();
synchronized (SpringApplicationShutdownHook.class) {
assertNotInProgress();
this.actions.add(action);
this.actions.add(new Handler(action));
}
}

Expand All @@ -188,11 +187,11 @@ public void remove(Runnable action) {
Assert.notNull(action, "Action must not be null");
synchronized (SpringApplicationShutdownHook.class) {
assertNotInProgress();
this.actions.remove(action);
this.actions.remove(new Handler(action));
}
}

Set<Runnable> getActions() {
Set<Handler> getActions() {
return this.actions;
}

Expand All @@ -204,6 +203,36 @@ public void run() {

}

/**
* A single handler that uses object identity for {@link #equals(Object)} and
* {@link #hashCode()}.
*
* @param runnable the handler runner
*/
record Handler(Runnable runnable) {

@Override
public int hashCode() {
return System.identityHashCode(this.runnable);
}

@Override
public boolean equals(Object obj) {
if (this == obj) {
return true;
}
if (obj == null || getClass() != obj.getClass()) {
return false;
}
return this.runnable == ((Handler) obj).runnable;
}

void run() {
this.runnable.run();
}

}

/**
* {@link ApplicationListener} to track closed contexts.
*/
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright 2012-2023 the original author or authors.
* Copyright 2012-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand All @@ -26,6 +26,7 @@

import org.awaitility.Awaitility;
import org.junit.jupiter.api.Test;
import org.mockito.InOrder;

import org.springframework.beans.factory.BeanCreationException;
import org.springframework.beans.factory.InitializingBean;
Expand All @@ -39,6 +40,8 @@
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;

/**
* Tests for {@link SpringApplicationShutdownHook}.
Expand Down Expand Up @@ -203,6 +206,23 @@ void deregistersFailedContext() {
assertThat(shutdownHook.isApplicationContextRegistered(context)).isFalse();
}

@Test
void handlersRunInDeterministicOrder() {
TestSpringApplicationShutdownHook shutdownHook = new TestSpringApplicationShutdownHook();
Runnable r1 = mock(Runnable.class);
Runnable r2 = mock(Runnable.class);
Runnable r3 = mock(Runnable.class);
shutdownHook.getHandlers().add(r2);
shutdownHook.getHandlers().add(r1);
shutdownHook.getHandlers().add(r3);
shutdownHook.run();
InOrder ordered = inOrder(r1, r2, r3);
ordered.verify(r2).run();
ordered.verify(r1).run();
ordered.verify(r3).run();
ordered.verifyNoMoreInteractions();
}

static class TestSpringApplicationShutdownHook extends SpringApplicationShutdownHook {

private boolean runtimeShutdownHookAdded;
Expand Down

0 comments on commit 21203f0

Please sign in to comment.