From 27ceda38bf77f5e8403a4427cdda2eaa1838520c Mon Sep 17 00:00:00 2001
From: greyp9 <greyp9@users.noreply.github.com>
Date: Mon, 31 Jan 2022 09:20:54 -0500
Subject: [PATCH] NIFI-9233 - Improve reliability of system integration tests
 (#5720)

---
 .../state/server/ZooKeeperStateServer.java    | 32 ++++++++++++-------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/server/ZooKeeperStateServer.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/server/ZooKeeperStateServer.java
index 354750392f..4caf58d213 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/server/ZooKeeperStateServer.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/server/ZooKeeperStateServer.java
@@ -175,6 +175,26 @@ public class ZooKeeperStateServer extends ZooKeeperServerMain {
         if (started) {
             started = false;
 
+            if (quorumPeer != null && quorumPeer.isRunning()) {
+                quorumPeer.shutdown();
+            }
+
+            if (connectionFactory != null) {
+                try {
+                    connectionFactory.shutdown();
+                } catch (Exception e) {
+                    logger.warn("Failed to shutdown Connection Factory", e);
+                }
+            }
+
+            if (embeddedZkServer != null && embeddedZkServer.isRunning()) {
+                try {
+                    embeddedZkServer.shutdown();
+                } catch (Exception e) {
+                    logger.warn("Failed to shutdown Embedded Zookeeper", e);
+                }
+            }
+
             if (transactionLog != null) {
                 try {
                     transactionLog.close();
@@ -183,18 +203,6 @@ public class ZooKeeperStateServer extends ZooKeeperServerMain {
                 }
             }
 
-            if (connectionFactory != null) {
-                connectionFactory.shutdown();
-            }
-
-            if (quorumPeer != null && quorumPeer.isRunning()) {
-                quorumPeer.shutdown();
-            }
-
-            if (embeddedZkServer != null && embeddedZkServer.isRunning()) {
-                embeddedZkServer.shutdown();
-            }
-
             if (datadirCleanupManager != null) {
                 datadirCleanupManager.shutdown();
             }