From 34fb909b411fb832eaaf8bb266f3f6e569e20c1e Mon Sep 17 00:00:00 2001
From: Marius Bogoevici <mbogoevici@pivotal.io>
Date: Mon, 7 Nov 2016 13:43:32 -0500
Subject: [PATCH] Support multiple root contexts in `Restarter`

Update devtools restarter to support multiple application contexts.

Fixes gh-7335
Closes gh-7336
---
 .../restart/RestartApplicationListener.java   |  2 +-
 .../boot/devtools/restart/Restarter.java      | 20 ++++++++++++++-----
 .../RestartApplicationListenerTests.java      | 10 ++++++----
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/spring-boot-devtools/src/main/java/org/springframework/boot/devtools/restart/RestartApplicationListener.java b/spring-boot-devtools/src/main/java/org/springframework/boot/devtools/restart/RestartApplicationListener.java
index 8cd7d81229..2966ac1456 100644
--- a/spring-boot-devtools/src/main/java/org/springframework/boot/devtools/restart/RestartApplicationListener.java
+++ b/spring-boot-devtools/src/main/java/org/springframework/boot/devtools/restart/RestartApplicationListener.java
@@ -52,7 +52,7 @@ public class RestartApplicationListener
 				|| event instanceof ApplicationFailedEvent) {
 			Restarter.getInstance().finish();
 			if (event instanceof ApplicationFailedEvent) {
-				Restarter.getInstance().prepare(null);
+				Restarter.getInstance().remove(((ApplicationFailedEvent) event).getApplicationContext());
 			}
 		}
 	}
diff --git a/spring-boot-devtools/src/main/java/org/springframework/boot/devtools/restart/Restarter.java b/spring-boot-devtools/src/main/java/org/springframework/boot/devtools/restart/Restarter.java
index 1c94bc5506..35b4b937c4 100644
--- a/spring-boot-devtools/src/main/java/org/springframework/boot/devtools/restart/Restarter.java
+++ b/spring-boot-devtools/src/main/java/org/springframework/boot/devtools/restart/Restarter.java
@@ -31,6 +31,7 @@ import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.BlockingDeque;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.LinkedBlockingDeque;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.locks.Lock;
@@ -51,6 +52,7 @@ import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.core.ResolvableType;
 import org.springframework.core.annotation.AnnotationUtils;
 import org.springframework.util.Assert;
+import org.springframework.util.CollectionUtils;
 import org.springframework.util.ReflectionUtils;
 
 /**
@@ -116,7 +118,7 @@ public class Restarter {
 
 	private boolean finished = false;
 
-	private volatile ConfigurableApplicationContext rootContext;
+	private final List<ConfigurableApplicationContext> rootContexts = new CopyOnWriteArrayList<ConfigurableApplicationContext>();
 
 	/**
 	 * Internal constructor to create a new {@link Restarter} instance.
@@ -314,9 +316,11 @@ public class Restarter {
 		this.logger.debug("Stopping application");
 		this.stopLock.lock();
 		try {
-			if (this.rootContext != null) {
-				this.rootContext.close();
-				this.rootContext = null;
+			if (!CollectionUtils.isEmpty(this.rootContexts)) {
+				for (ConfigurableApplicationContext rootContext : this.rootContexts) {
+					rootContext.close();
+				}
+				this.rootContexts.clear();
 			}
 			cleanupCaches();
 			if (this.forceReferenceCleanup) {
@@ -418,7 +422,13 @@ public class Restarter {
 		if (applicationContext != null && applicationContext.getParent() != null) {
 			return;
 		}
-		this.rootContext = applicationContext;
+		this.rootContexts.add(applicationContext);
+	}
+
+	void remove(ConfigurableApplicationContext applicationContext) {
+		if (applicationContext != null) {
+			this.rootContexts.remove(applicationContext);
+		}
 	}
 
 	private LeakSafeThread getLeakSafeThread() {
diff --git a/spring-boot-devtools/src/test/java/org/springframework/boot/devtools/restart/RestartApplicationListenerTests.java b/spring-boot-devtools/src/test/java/org/springframework/boot/devtools/restart/RestartApplicationListenerTests.java
index 31acad0f0e..34848adbc2 100644
--- a/spring-boot-devtools/src/test/java/org/springframework/boot/devtools/restart/RestartApplicationListenerTests.java
+++ b/spring-boot-devtools/src/test/java/org/springframework/boot/devtools/restart/RestartApplicationListenerTests.java
@@ -16,6 +16,8 @@
 
 package org.springframework.boot.devtools.restart;
 
+import java.util.List;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -64,8 +66,8 @@ public class RestartApplicationListenerTests {
 		assertThat(ReflectionTestUtils.getField(Restarter.getInstance(), "args"))
 				.isEqualTo(ARGS);
 		assertThat(Restarter.getInstance().isFinished()).isTrue();
-		assertThat(ReflectionTestUtils.getField(Restarter.getInstance(), "rootContext"))
-				.isNotNull();
+		assertThat((List<?>) ReflectionTestUtils.getField(Restarter.getInstance(),
+				"rootContexts")).isNotEmpty();
 	}
 
 	@Test
@@ -74,8 +76,8 @@ public class RestartApplicationListenerTests {
 		assertThat(ReflectionTestUtils.getField(Restarter.getInstance(), "args"))
 				.isEqualTo(ARGS);
 		assertThat(Restarter.getInstance().isFinished()).isTrue();
-		assertThat(ReflectionTestUtils.getField(Restarter.getInstance(), "rootContext"))
-				.isNull();
+		assertThat((List<?>) ReflectionTestUtils.getField(Restarter.getInstance(),
+				"rootContexts")).isEmpty();
 	}
 
 	@Test