From eb25c8547a8fec278296ab7b846deac2295f4af4 Mon Sep 17 00:00:00 2001
From: Mark Payne <markap14@hotmail.com>
Date: Wed, 24 May 2017 16:15:18 -0400
Subject: [PATCH] NIFI-3972: Ensure that we wait until service state becomes
 enabled before triggering completable future that says that it's enabled

---
 .../service/StandardControllerServiceNode.java           | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/service/StandardControllerServiceNode.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/service/StandardControllerServiceNode.java
index a18fd1f592..e399e7ffbb 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/service/StandardControllerServiceNode.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/service/StandardControllerServiceNode.java
@@ -407,12 +407,13 @@ public class StandardControllerServiceNode extends AbstractConfiguredComponent i
                             ReflectionUtils.invokeMethodsWithAnnotation(OnEnabled.class, getControllerServiceImplementation(), configContext);
                         }
 
-                        future.complete(null);
-
                         boolean shouldEnable = false;
                         synchronized (active) {
                             shouldEnable = active.get() && stateRef.compareAndSet(ControllerServiceState.ENABLING, ControllerServiceState.ENABLED);
                         }
+
+                        future.complete(null);
+
                         if (!shouldEnable) {
                             LOG.debug("Disabling service " + this + " after it has been enabled due to disable action being initiated.");
                             // Can only happen if user initiated DISABLE operation before service finished enabling. It's state will be
@@ -421,6 +422,8 @@ public class StandardControllerServiceNode extends AbstractConfiguredComponent i
                             stateRef.set(ControllerServiceState.DISABLED);
                         }
                     } catch (Exception e) {
+                        future.completeExceptionally(e);
+
                         final Throwable cause = e instanceof InvocationTargetException ? e.getCause() : e;
                         final ComponentLog componentLog = new SimpleProcessLogger(getIdentifier(), StandardControllerServiceNode.this);
                         componentLog.error("Failed to invoke @OnEnabled method due to {}", cause);
@@ -438,6 +441,8 @@ public class StandardControllerServiceNode extends AbstractConfiguredComponent i
                     }
                 }
             });
+        } else {
+            future.complete(null);
         }
 
         return future;