NIFI-4: Fixed issue with lifecycle of Controller Services

nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-cluster/src/main/java/org/apache/nifi/cluster/manager/impl/WebClusterManager.java

@@ -1320,6 +1320,18 @@ public class WebClusterManager implements HttpClusterManager, ProtocolHandler, C
         controllerServiceProvider.removeControllerService(serviceNode);
     }
     
+
+    @Override
+    public void enableControllerService(final ControllerServiceNode serviceNode) {
+        controllerServiceProvider.enableControllerService(serviceNode);
+    }
+    
+    @Override
+    public void disableControllerService(final ControllerServiceNode serviceNode) {
+        controllerServiceProvider.disableControllerService(serviceNode);
+    }
+    
+    
     /**
      * Handle a bulletins message.
      *

nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core-api/src/main/java/org/apache/nifi/controller/service/ControllerServiceProvider.java

@@ -53,4 +53,17 @@ public interface ControllerServiceProvider extends ControllerServiceLookup {
      * @throws IllegalStateException if the controller service is not disabled or is not a part of this flow
      */
     void removeControllerService(ControllerServiceNode serviceNode);
+    
+    /**
+     * Enables the given controller service that it can be used by other components
+     * @param serviceNode
+     */
+    void enableControllerService(ControllerServiceNode serviceNode);
+    
+    /**
+     * Disables the given controller service so that it cannot be used by other components. This allows
+     * configuration to be updated or allows service to be removed.
+     * @param serviceNode
+     */
+    void disableControllerService(ControllerServiceNode serviceNode);
 }

nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/FlowController.java

@@ -2569,16 +2569,16 @@ public class FlowController implements EventAccess, ControllerServiceProvider, H
         processScheduler.disableReportingTask(reportingTaskNode);
     }
     
+    @Override
     public void enableControllerService(final ControllerServiceNode serviceNode) {
         serviceNode.verifyCanEnable();
-        
+        controllerServiceProvider.enableControllerService(serviceNode);
-        processScheduler.enableControllerService(serviceNode);
     }
     
+    @Override
     public void disableControllerService(final ControllerServiceNode serviceNode) {
         serviceNode.verifyCanDisable();
-        
+        controllerServiceProvider.disableControllerService(serviceNode);
-        processScheduler.disableControllerService(serviceNode);
     }
 
     @Override

nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/scheduling/StandardProcessScheduler.java

@@ -581,37 +581,6 @@ public final class StandardProcessScheduler implements ProcessScheduler {
         }
     }
 
-    public synchronized void enableControllerService(final ControllerServiceNode serviceNode) {
-        if ( !serviceNode.isDisabled() ) {
-            throw new IllegalStateException("Controller Service cannot be enabled because it is not disabled");
-        }
-
-        // we set the service to enabled before invoking the @OnEnabled methods. We do this because it must be
-        // done in this order for disabling (serviceNode.setDisabled(true) will throw Exceptions if the service
-        // is currently known to be in use) and we want to be consistent with the ordering of calling setDisabled
-        // before annotated methods.
-        serviceNode.setDisabled(false);
-        
-        try (final NarCloseable x = NarCloseable.withNarLoader()) {
-            ReflectionUtils.quietlyInvokeMethodsWithAnnotation(OnEnabled.class, serviceNode.getControllerServiceImplementation());
-        }
-    }
-    
-    public synchronized void disableControllerService(final ControllerServiceNode serviceNode) {
-        if ( serviceNode.isDisabled() ) {
-            throw new IllegalStateException("Controller Service cannot be disabled because it is already disabled");
-        }
-
-        // We must set the service to disabled before we invoke the OnDisabled methods because the service node
-        // can throw Exceptions if we attempt to disable the service while it's known to be in use.
-        serviceNode.setDisabled(true);
-        
-        try (final NarCloseable x = NarCloseable.withNarLoader()) {
-            ReflectionUtils.quietlyInvokeMethodsWithAnnotation(OnDisabled.class, serviceNode.getControllerServiceImplementation());
-        }
-    }
-    
-    
     @Override
     public boolean isScheduled(final Object scheduled) {
         final ScheduleState scheduleState = scheduleStates.get(scheduled);

nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/service/ControllerServiceLoader.java

@@ -25,9 +25,7 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.StandardOpenOption;
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 
 import javax.xml.XMLConstants;
 import javax.xml.parsers.DocumentBuilder;
@@ -36,11 +34,10 @@ import javax.xml.parsers.ParserConfigurationException;
 import javax.xml.validation.Schema;
 import javax.xml.validation.SchemaFactory;
 
-import org.apache.nifi.util.file.FileUtils;
-import org.apache.nifi.util.DomUtils;
-
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import org.apache.nifi.util.DomUtils;
+import org.apache.nifi.util.file.FileUtils;
 import org.w3c.dom.Document;
 import org.w3c.dom.Element;
 import org.w3c.dom.NodeList;
@@ -139,7 +136,7 @@ public class ControllerServiceLoader {
                     }
 
                     services.add(serviceNode);
-                    serviceNode.setDisabled(false);
+                    provider.enableControllerService(serviceNode);
                 }
             }
         } catch (SAXException | ParserConfigurationException sxe) {

nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/service/StandardControllerServiceNode.java

@@ -154,6 +154,7 @@ public class StandardControllerServiceNode extends AbstractConfiguredComponent i
         return removed;
     }
     
+    @SuppressWarnings("deprecation")
     private void onConfigured() {
         try (final NarCloseable x = NarCloseable.withNarLoader()) {
             final ConfigurationContext configContext = new StandardConfigurationContext(this, serviceProvider);

nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/service/StandardControllerServiceProvider.java

@@ -31,6 +31,8 @@ import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
 import org.apache.nifi.annotation.lifecycle.OnAdded;
+import org.apache.nifi.annotation.lifecycle.OnEnabled;
+import org.apache.nifi.annotation.lifecycle.OnDisabled;
 import org.apache.nifi.annotation.lifecycle.OnRemoved;
 import org.apache.nifi.controller.ConfigurationContext;
 import org.apache.nifi.controller.ControllerService;
@@ -165,6 +167,31 @@ public class StandardControllerServiceProvider implements ControllerServiceProvi
             }
         }
     }
+    
+    @Override
+    public void enableControllerService(final ControllerServiceNode serviceNode) {
+        serviceNode.verifyCanEnable();
+        
+        try (final NarCloseable x = NarCloseable.withNarLoader()) {
+            final ConfigurationContext configContext = new StandardConfigurationContext(serviceNode, this);
+            ReflectionUtils.quietlyInvokeMethodsWithAnnotation(OnEnabled.class, serviceNode.getControllerServiceImplementation(), configContext);
+        }
+        
+        serviceNode.setDisabled(false);
+    }
+    
+    @Override
+    public void disableControllerService(final ControllerServiceNode serviceNode) {
+        serviceNode.verifyCanDisable();
+
+        // We must set the service to disabled before we invoke the OnDisabled methods because the service node
+        // can throw Exceptions if we attempt to disable the service while it's known to be in use.
+        serviceNode.setDisabled(true);
+        
+        try (final NarCloseable x = NarCloseable.withNarLoader()) {
+            ReflectionUtils.quietlyInvokeMethodsWithAnnotation(OnDisabled.class, serviceNode.getControllerServiceImplementation());
+        }
+    }
 
     @Override
     public ControllerService getControllerService(final String serviceIdentifier) {

nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/controller/ControllerFacade.java

@@ -412,6 +412,16 @@ public class ControllerFacade implements ControllerServiceProvider {
         return flowController.isControllerServiceEnabled(serviceIdentifier);
     }
 
+    @Override
+    public void enableControllerService(final ControllerServiceNode serviceNode) {
+        flowController.enableControllerService(serviceNode);
+    }
+    
+    @Override
+    public void disableControllerService(ControllerServiceNode serviceNode) {
+        flowController.disableControllerService(serviceNode);
+    }
+    
     /**
      * Gets the status of this controller.
      *