From 2ceb35383880b389c685508213b16eabaa890349 Mon Sep 17 00:00:00 2001
From: Michael Minella <mminella@pivotal.io>
Date: Fri, 22 Jan 2016 13:03:21 -0600
Subject: [PATCH] Use spring Boot events to get exit code

Spring Boot now provides the result of an ExitCodeExceptionMapper via an
event called ExitCodeEvent.  This commit takes advantage of this new
event.

Resolves spring-cloud/spring-cloud-task#48
---
 pom.xml                                       |  2 +-
 .../task/listener/TaskLifecycleListener.java  | 54 +++++++++----------
 .../listener/TaskLifecycleListenerTests.java  |  1 +
 spring-cloud-task-samples/timestamp/pom.xml   |  2 +-
 4 files changed, 28 insertions(+), 31 deletions(-)

diff --git a/pom.xml b/pom.xml
index f2bb2da0..7e75be69 100755
--- a/pom.xml
+++ b/pom.xml
@@ -8,7 +8,7 @@
 				<!-- Import dependency management from Spring Boot -->
 				<groupId>org.springframework.boot</groupId>
 				<artifactId>spring-boot-dependencies</artifactId>
-				<version>1.3.2.BUILD-SNAPSHOT</version>
+				<version>1.3.2.RELEASE</version>
 				<type>pom</type>
 				<scope>import</scope>
 			</dependency>
diff --git a/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/listener/TaskLifecycleListener.java b/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/listener/TaskLifecycleListener.java
index 7d171ef5..ae2be9ef 100644
--- a/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/listener/TaskLifecycleListener.java
+++ b/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/listener/TaskLifecycleListener.java
@@ -26,6 +26,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import org.springframework.boot.ApplicationArguments;
+import org.springframework.boot.ExitCodeEvent;
 import org.springframework.boot.context.event.ApplicationFailedEvent;
 import org.springframework.cloud.task.repository.TaskExecution;
 import org.springframework.cloud.task.repository.TaskNameResolver;
@@ -68,6 +69,10 @@ public class TaskLifecycleListener implements ApplicationListener<ApplicationEve
 
 	private ApplicationArguments applicationArguments;
 
+	private ApplicationFailedEvent applicationFailedEvent;
+
+	private ExitCodeEvent exitCodeEvent;
+
 	/**
 	 * @param taskRepository The repository to record executions in.
 	 */
@@ -99,28 +104,15 @@ public class TaskLifecycleListener implements ApplicationListener<ApplicationEve
 			doTaskStart();
 			started = true;
 		}
+		else if(applicationEvent instanceof ApplicationFailedEvent) {
+			this.applicationFailedEvent = (ApplicationFailedEvent) applicationEvent;
+		}
+		else if(applicationEvent instanceof ExitCodeEvent){
+			this.exitCodeEvent = (ExitCodeEvent) applicationEvent;
+		}
 		else if(applicationEvent instanceof ContextClosedEvent) {
 			doTaskEnd();
 		}
-		else if(applicationEvent instanceof ApplicationFailedEvent) {
-			doTaskFailed(((ApplicationFailedEvent) applicationEvent).getException());
-		}
-	}
-
-	private void doTaskFailed(Throwable exception) {
-
-		if(started) {
-			this.taskExecution.setEndTime(new Date());
-			//TODO: get exit code from new event thrown by Boot.
-			this.taskExecution.setExitCode(1);
-			this.taskExecution.setExitMessage(stackTraceToString(exception));
-
-			taskRepository.update(taskExecution);
-		}
-		else  {
-			logger.error("An event to fail a task has been received for a task that has " +
-					"not yet started.");
-		}
 	}
 
 	private String stackTraceToString(Throwable exception) {
@@ -135,8 +127,20 @@ public class TaskLifecycleListener implements ApplicationListener<ApplicationEve
 	private void doTaskEnd() {
 		if(started) {
 			this.taskExecution.setEndTime(new Date());
-			//TODO: get exit code from new event thrown by Boot.
-			this.taskExecution.setExitCode(0);
+
+			if(this.exitCodeEvent != null) {
+				this.taskExecution.setExitCode(exitCodeEvent.getExitCode());
+			}
+			else if(this.applicationFailedEvent != null){
+				this.taskExecution.setExitCode(1);
+			}
+			else{
+				this.taskExecution.setExitCode(0);
+			}
+
+			if(this.applicationFailedEvent != null) {
+				this.taskExecution.setExitMessage(stackTraceToString(this.applicationFailedEvent.getException()));
+			}
 
 			taskRepository.update(taskExecution);
 		}
@@ -166,12 +170,4 @@ public class TaskLifecycleListener implements ApplicationListener<ApplicationEve
 					"recorded.");
 		}
 	}
-
-	/**
-	 * Used for testing purposes
-	 * @return the current {@link TaskExecution}
-	 */
-	TaskExecution getTaskExecution() {
-		return this.taskExecution;
-	}
 }
diff --git a/spring-cloud-task-core/src/test/java/org/springframework/cloud/task/listener/TaskLifecycleListenerTests.java b/spring-cloud-task-core/src/test/java/org/springframework/cloud/task/listener/TaskLifecycleListenerTests.java
index b9c057ef..ecc5b065 100644
--- a/spring-cloud-task-core/src/test/java/org/springframework/cloud/task/listener/TaskLifecycleListenerTests.java
+++ b/spring-cloud-task-core/src/test/java/org/springframework/cloud/task/listener/TaskLifecycleListenerTests.java
@@ -99,6 +99,7 @@ public class TaskLifecycleListenerTests {
 		context.refresh();
 		RuntimeException exception = new RuntimeException("This was expected");
 		context.publishEvent(new ApplicationFailedEvent(new SpringApplication(), new String[0], context, exception));
+		context.publishEvent(new ContextClosedEvent(context));
 
 		verifyTaskExecution(0, true, 1, exception);
 	}
diff --git a/spring-cloud-task-samples/timestamp/pom.xml b/spring-cloud-task-samples/timestamp/pom.xml
index 261caf13..cd1caa1a 100644
--- a/spring-cloud-task-samples/timestamp/pom.xml
+++ b/spring-cloud-task-samples/timestamp/pom.xml
@@ -15,7 +15,7 @@
 	<parent>
 		<groupId>org.springframework.boot</groupId>
 		<artifactId>spring-boot-starter-parent</artifactId>
-		<version>1.3.0.RELEASE</version>
+		<version>1.3.2.RELEASE</version>
 	</parent>
 
 	<properties>