From 668c278e3099598f7a42aaf6da1ce833c951fab2 Mon Sep 17 00:00:00 2001
From: Peter Gyori <peter.gyori.dev@gmail.com>
Date: Thu, 17 Sep 2020 15:51:32 +0200
Subject: [PATCH] NIFI-7815: Enhance the logging in MergeContent processor

Signed-off-by: Pierre Villard <pierre.villard.fr@gmail.com>

This closes #4534.
---
 .../processors/standard/MergeContent.java     | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/MergeContent.java b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/MergeContent.java
index 0d337d84bd..5d4506c311 100644
--- a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/MergeContent.java
+++ b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/MergeContent.java
@@ -577,6 +577,15 @@ public class MergeContent extends BinFiles {
         return NUMBER_PATTERN.matcher(value).matches();
     }
 
+    private void removeFlowFileFromSession(final ProcessSession session, final FlowFile flowFile, final ProcessContext context) {
+        try {
+            session.remove(flowFile);
+        } catch (final Exception e) {
+            getLogger().error("Failed to remove merged FlowFile from the session after merge failure during \""
+                    + context.getProperty(MERGE_FORMAT).getValue() + "\" merge.", e);
+        }
+    }
+
     private class BinaryConcatenationMerge implements MergeBin {
 
         private String mimeType = "application/octet-stream";
@@ -636,7 +645,7 @@ public class MergeContent extends BinFiles {
                     }
                 });
             } catch (final Exception e) {
-                session.remove(bundle);
+                removeFlowFileFromSession(session, bundle, context);
                 throw e;
             }
 
@@ -781,7 +790,7 @@ public class MergeContent extends BinFiles {
                     }
                 });
             } catch (final Exception e) {
-                session.remove(bundle);
+                removeFlowFileFromSession(session, bundle, context);
                 throw e;
             }
 
@@ -848,7 +857,7 @@ public class MergeContent extends BinFiles {
                     }
                 });
             } catch (final Exception e) {
-                session.remove(bundle);
+                removeFlowFileFromSession(session, bundle, context);
                 throw e;
             }
 
@@ -918,7 +927,7 @@ public class MergeContent extends BinFiles {
                     }
                 });
             } catch (final Exception e) {
-                session.remove(bundle);
+                removeFlowFileFromSession(session, bundle, context);
                 throw e;
             }
 
@@ -1042,7 +1051,7 @@ public class MergeContent extends BinFiles {
                     }
                 });
             } catch (final Exception e) {
-                session.remove(bundle);
+                removeFlowFileFromSession(session, bundle, context);
                 throw e;
             }