From 2a7f135f1c8c9a31b1084860419383f4196e386c Mon Sep 17 00:00:00 2001
From: Mark Payne <markap14@hotmail.com>
Date: Tue, 19 Jul 2016 15:33:59 -0400
Subject: [PATCH] NIFI-2319: Ensure that when we set cluster node id's and node
 addresses, that we do so only if they are not already populated

This closes #680

Signed-off-by: jpercivall <joepercivall@yahoo.com>
---
 .../http/endpoints/ComponentStateEndpointMerger.java  |  7 +++++--
 .../http/endpoints/ListFlowFilesEndpointMerger.java   |  6 ++++--
 .../http/endpoints/ProvenanceEventEndpointMerger.java | 11 ++++++++---
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ComponentStateEndpointMerger.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ComponentStateEndpointMerger.java
index 5ab71c934e..3b03ed0b1e 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ComponentStateEndpointMerger.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ComponentStateEndpointMerger.java
@@ -74,8 +74,11 @@ public class ComponentStateEndpointMerger extends AbstractSingleDTOEndpoint<Comp
                 totalStateEntries += nodeLocalStateMap.getTotalEntryCount();
 
                 for (final StateEntryDTO nodeStateEntry : nodeLocalStateMap.getState()) {
-                    nodeStateEntry.setClusterNodeId(nodeId.getId());
-                    nodeStateEntry.setClusterNodeAddress(nodeAddress);
+                    if (nodeStateEntry.getClusterNodeId() == null || nodeStateEntry.getClusterNodeAddress() == null) {
+                        nodeStateEntry.setClusterNodeId(nodeId.getId());
+                        nodeStateEntry.setClusterNodeAddress(nodeAddress);
+                    }
+
                     localStateEntries.add(nodeStateEntry);
                 }
             }
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ListFlowFilesEndpointMerger.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ListFlowFilesEndpointMerger.java
index 2d17c4946e..72d42e2713 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ListFlowFilesEndpointMerger.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ListFlowFilesEndpointMerger.java
@@ -124,8 +124,10 @@ public class ListFlowFilesEndpointMerger extends AbstractSingleDTOEndpoint<Listi
 
             if (nodeRequest.getFlowFileSummaries() != null) {
                 for (final FlowFileSummaryDTO summaryDTO : nodeRequest.getFlowFileSummaries()) {
-                    summaryDTO.setClusterNodeId(nodeIdentifier.getId());
-                    summaryDTO.setClusterNodeAddress(nodeAddress);
+                    if (summaryDTO.getClusterNodeId() == null || summaryDTO.getClusterNodeAddress() == null) {
+                        summaryDTO.setClusterNodeId(nodeIdentifier.getId());
+                        summaryDTO.setClusterNodeAddress(nodeAddress);
+                    }
 
                     flowFileSummaries.add(summaryDTO);
 
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ProvenanceEventEndpointMerger.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ProvenanceEventEndpointMerger.java
index c05ca7ffee..e0372c0936 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ProvenanceEventEndpointMerger.java
+++ b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/endpoints/ProvenanceEventEndpointMerger.java
@@ -48,8 +48,13 @@ public class ProvenanceEventEndpointMerger extends AbstractSingleDTOEndpoint<Pro
     @Override
     protected void mergeResponses(ProvenanceEventDTO clientDto, Map<NodeIdentifier, ProvenanceEventDTO> dtoMap, Set<NodeResponse> successfulResponses, Set<NodeResponse> problematicResponses) {
         // The request for a Provenance Event is replicated to a single Node. We simply update its cluster node info.
-        final NodeIdentifier nodeId = successfulResponses.iterator().next().getNodeId();
-        clientDto.setClusterNodeId(nodeId.getId());
-        clientDto.setClusterNodeAddress(nodeId.getApiAddress() + ":" + nodeId.getApiPort());
+        // However, we only do this if the cluster node info isn't set, because if this is replicated across the cluster,
+        // the cluster coordinator will have already set it, and we will be receiving the response from the cluster
+        // coordinator. We do not want to overwrite this value on all DTO's with the cluster coordinator's information.
+        if (clientDto.getClusterNodeId() == null || clientDto.getClusterNodeAddress() == null) {
+            final NodeIdentifier nodeId = successfulResponses.iterator().next().getNodeId();
+            clientDto.setClusterNodeId(nodeId.getId());
+            clientDto.setClusterNodeAddress(nodeId.getApiAddress() + ":" + nodeId.getApiPort());
+        }
     }
 }