NIFI-12731: Ensure state is updated in GetHBase whenever the session is committed

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

This closes #8346.

nifi-nar-bundles/nifi-hbase-bundle/nifi-hbase-processors/src/main/java/org/apache/nifi/hbase/GetHBase.java

@@ -140,7 +140,7 @@ public class GetHBase extends AbstractProcessor implements VisibilityFetchSuppor
             .build();
 
     private final AtomicReference<ScanResult> lastResult = new AtomicReference<>();
-    private volatile List<Column> columns = new ArrayList<>();
+    private final List<Column> columns = new ArrayList<>();
     private volatile String previousTable = null;
 
     @Override
@@ -201,11 +201,11 @@ public class GetHBase extends AbstractProcessor implements VisibilityFetchSuppor
         for (final String column : columns) {
             if (column.contains(":"))  {
                 final String[] parts = column.split(":");
-                final byte[] cf = parts[0].getBytes(Charset.forName("UTF-8"));
+                final byte[] cf = parts[0].getBytes(StandardCharsets.UTF_8);
-                final byte[] cq = parts[1].getBytes(Charset.forName("UTF-8"));
+                final byte[] cq = parts[1].getBytes(StandardCharsets.UTF_8);
                 this.columns.add(new Column(cf, cq));
             } else {
-                final byte[] cf = column.getBytes(Charset.forName("UTF-8"));
+                final byte[] cf = column.getBytes(StandardCharsets.UTF_8);
                 this.columns.add(new Column(cf, null));
             }
         }
@@ -307,11 +307,7 @@ public class GetHBase extends AbstractProcessor implements VisibilityFetchSuppor
                             final byte[] cellValue = Arrays.copyOfRange(cell.getValueArray(), cell.getValueOffset(), cell.getValueLength() + cell.getValueOffset());
 
                             final String rowHash = new String(rowValue, StandardCharsets.UTF_8);
-                            Set<String> cellHashes = cellsMatchingTimestamp.get(rowHash);
+                            Set<String> cellHashes = cellsMatchingTimestamp.computeIfAbsent(rowHash, k -> new HashSet<>());
-                            if (cellHashes == null) {
-                                cellHashes = new HashSet<>();
-                                cellsMatchingTimestamp.put(rowHash, cellHashes);
-                            }
                             cellHashes.add(new String(cellValue, StandardCharsets.UTF_8));
                         }
                     }
@@ -336,11 +332,23 @@ public class GetHBase extends AbstractProcessor implements VisibilityFetchSuppor
                 rowsPulledHolder.set(++rowsPulled);
 
                 if (++rowsPulled % getBatchSize() == 0) {
-                    session.commitAsync();
+                    updateStateAndCommit(session, latestTimestampHolder.get(), cellsMatchingTimestamp);
                 }
             });
 
-            final ScanResult scanResults = new ScanResult(latestTimestampHolder.get(), cellsMatchingTimestamp);
+            updateStateAndCommit(session, latestTimestampHolder.get(), cellsMatchingTimestamp);
+        } catch (final IOException e) {
+            getLogger().error("Failed to receive data from HBase due to {}", e);
+            session.rollback();
+        } finally {
+            // if we failed, we want to yield so that we don't hammer hbase. If we succeed, then we have
+            // pulled all of the records, so we want to wait a bit before hitting hbase again anyway.
+            context.yield();
+        }
+    }
+
+    private void updateStateAndCommit(final ProcessSession session, final long latestTimestamp, final Map<String, Set<String>> cellsMatchingTimestamp) throws IOException {
+        final ScanResult scanResults = new ScanResult(latestTimestamp, cellsMatchingTimestamp);
 
         final ScanResult latestResult = lastResult.get();
         if (latestResult == null || scanResults.getTimestamp() > latestResult.getTimestamp()) {
@@ -370,14 +378,6 @@ public class GetHBase extends AbstractProcessor implements VisibilityFetchSuppor
 
             session.commitAsync(() -> updateScanResultsIfNewer(scanResult));
         }
-        } catch (final IOException e) {
-            getLogger().error("Failed to receive data from HBase due to {}", e);
-            session.rollback();
-        } finally {
-            // if we failed, we want to yield so that we don't hammer hbase. If we succeed, then we have
-            // pulled all of the records, so we want to wait a bit before hitting hbase again anyway.
-            context.yield();
-        }
     }
 
     private void updateScanResultsIfNewer(final ScanResult scanResult) {
@@ -495,11 +495,7 @@ public class GetHBase extends AbstractProcessor implements VisibilityFetchSuppor
                 final String rowIndex = matcher.group(1);
                 final String cellIndex = matcher.group(3);
 
-                Set<String> cellHashes = rowIndexToMatchingCellHashes.get(rowIndex);
+                Set<String> cellHashes = rowIndexToMatchingCellHashes.computeIfAbsent(rowIndex, k -> new HashSet<>());
-                if (cellHashes == null) {
-                    cellHashes = new HashSet<>();
-                    rowIndexToMatchingCellHashes.put(rowIndex, cellHashes);
-                }
 
                 if (cellIndex == null) {
                     // this provides a Row ID.

nifi-nar-bundles/nifi-standard-services/nifi-hbase-client-service-api/src/main/java/org/apache/nifi/hbase/scan/ResultHandler.java

@@ -16,11 +16,13 @@
  */
 package org.apache.nifi.hbase.scan;
 
+import java.io.IOException;
+
 /**
  * Handles a single row from an HBase scan.
  */
 public interface ResultHandler {
 
-    void handle(byte[] row, ResultCell[] resultCells);
+    void handle(byte[] row, ResultCell[] resultCells) throws IOException;
 
 }