From a931e7278719f0b6030ca5c7da66cae22f4d08dd Mon Sep 17 00:00:00 2001
From: Aldrin Piri <aldrin@apache.org>
Date: Mon, 1 Feb 2016 15:50:37 -0500
Subject: [PATCH] NIFI-259:  Expose the ZK exception code in the event of a
 KeeperException within the ZooKeeperStateProvider.

---
 .../zookeeper/ZooKeeperStateProvider.java     | 33 ++++++++++---------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/zookeeper/ZooKeeperStateProvider.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/zookeeper/ZooKeeperStateProvider.java
index 5c8b4c4eda..239b68a759 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/zookeeper/ZooKeeperStateProvider.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/state/providers/zookeeper/ZooKeeperStateProvider.java
@@ -244,16 +244,17 @@ public class ZooKeeperStateProvider extends AbstractStateProvider {
             ZKUtil.deleteRecursive(getZooKeeper(), getComponentPath(componentId));
         } catch (final KeeperException ke) {
             // Node doesn't exist so just ignore
-            if (Code.NONODE == ke.code()) {
+            final Code exceptionCode = ke.code();
+            if (Code.NONODE == exceptionCode) {
                 return;
             }
-            if (Code.SESSIONEXPIRED == ke.code()) {
+            if (Code.SESSIONEXPIRED == exceptionCode) {
                 invalidateClient();
                 onComponentRemoved(componentId);
                 return;
             }
 
-            throw new IOException("Unable to remove state for component with ID '" + componentId + "' from ZooKeeper", ke);
+            throw new IOException("Unable to remove state for component with ID '" + componentId + " with exception code " + exceptionCode, ke);
         } catch (final InterruptedException e) {
             Thread.currentThread().interrupt();
             throw new IOException("Failed to remove state for component with ID '" + componentId + "' from ZooKeeper due to being interrupted", e);
@@ -331,7 +332,8 @@ public class ZooKeeperStateProvider extends AbstractStateProvider {
             try {
                 keeper.setData(path, data, version);
             } catch (final KeeperException ke) {
-                if (ke.code() == Code.NONODE) {
+                final Code exceptionCode = ke.code();
+                if (exceptionCode == Code.NONODE) {
                     createNode(path, data);
                     return;
                 } else {
@@ -342,17 +344,18 @@ public class ZooKeeperStateProvider extends AbstractStateProvider {
             Thread.currentThread().interrupt();
             throw new IOException("Failed to set cluster-wide state in ZooKeeper for component with ID " + componentId + " due to interruption", e);
         } catch (final KeeperException ke) {
-            if (Code.SESSIONEXPIRED == ke.code()) {
+            final Code exceptionCode = ke.code();
+            if (Code.SESSIONEXPIRED == exceptionCode) {
                 invalidateClient();
                 setState(stateValues, version, componentId);
                 return;
             }
-            if (Code.NODEEXISTS == ke.code()) {
+            if (Code.NODEEXISTS == exceptionCode) {
                 setState(stateValues, version, componentId);
                 return;
             }
 
-            throw new IOException("Failed to set cluster-wide state in ZooKeeper for component with ID " + componentId, ke);
+            throw new IOException("Failed to set cluster-wide state in ZooKeeper for component with ID " + componentId + " with exception code " + exceptionCode, ke);
         } catch (final IOException ioe) {
             throw new IOException("Failed to set cluster-wide state in ZooKeeper for component with ID " + componentId, ioe);
         }
@@ -365,20 +368,21 @@ public class ZooKeeperStateProvider extends AbstractStateProvider {
         } catch (final InterruptedException ie) {
             throw new IOException("Failed to update cluster-wide state due to interruption", ie);
         } catch (final KeeperException ke) {
-            if (ke.code() == Code.NONODE) {
+            final Code exceptionCode = ke.code();
+            if (Code.NONODE == exceptionCode) {
                 final String parentPath = StringUtils.substringBeforeLast(path, "/");
                 createNode(parentPath, null);
                 createNode(path, data);
                 return;
             }
-            if (Code.SESSIONEXPIRED == ke.code()) {
+            if (Code.SESSIONEXPIRED == exceptionCode) {
                 invalidateClient();
                 createNode(path, data);
                 return;
             }
 
             // Node already exists. Node must have been created by "someone else". Just set the data.
-            if (ke.code() == Code.NODEEXISTS) {
+            if (Code.NODEEXISTS == exceptionCode) {
                 try {
                     getZooKeeper().setData(path, data, -1);
                     return;
@@ -392,8 +396,6 @@ public class ZooKeeperStateProvider extends AbstractStateProvider {
                     throw new IOException("Failed to update cluster-wide state due to interruption", ie);
                 }
             }
-
-
             throw ke;
         }
     }
@@ -413,15 +415,16 @@ public class ZooKeeperStateProvider extends AbstractStateProvider {
             Thread.currentThread().interrupt();
             throw new IOException("Failed to obtain value from ZooKeeper for component with ID " + componentId + ", due to interruption", e);
         } catch (final KeeperException ke) {
-            if (ke.code() == Code.NONODE) {
+            final Code exceptionCode = ke.code();
+            if (Code.NONODE == exceptionCode) {
                 return new StandardStateMap(null, -1L);
             }
-            if (Code.SESSIONEXPIRED == ke.code()) {
+            if (Code.SESSIONEXPIRED == exceptionCode) {
                 invalidateClient();
                 return getState(componentId);
             }
 
-            throw new IOException("Failed to obtain value from ZooKeeper for component with ID " + componentId, ke);
+            throw new IOException("Failed to obtain value from ZooKeeper for component with ID " + componentId + " with exception code " + exceptionCode, ke);
         } catch (final IOException ioe) {
             // provide more context in the error message
             throw new IOException("Failed to obtain value from ZooKeeper for component with ID " + componentId, ioe);