From 16d1446aa417a660a34c36e140361dee94758cf3 Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Singh Date: Fri, 22 May 2026 18:00:50 +0530 Subject: [PATCH] NIFI-15883: Fix Wait/Notify race condition with CAS-aware replace and complete - WaitNotifyProtocol.replace(Signal): Changed from returning boolean to void. Throws ConcurrentModificationException when the underlying cache.replace() returns false (CAS failure due to concurrent Notify), making the API consistent with complete(Signal) which also throws on concurrent modification. - WaitNotifyProtocol.notify(): Updated retry loop to catch ConcurrentModificationException from replace() instead of checking the boolean return value. - WaitNotifyProtocol.complete(Signal): Replaces raw cache.remove() with a version-aware approach that re-fetches the cache entry and compares revisions before removing. Throws ConcurrentModificationException if a concurrent Notify updated the signal between Wait's initial read and the remove call, preventing silent data loss. - Wait.java: Simplified call sites - protocol.replace(signal) and protocol.complete(signal) both now throw ConcurrentModificationException on concurrent modification. The catch block handles both IOException and ConcurrentModificationException with session rollback and ProcessException. - TestWaitNotifyProtocol: Add three tests for complete(Signal): happy-path removal, concurrent-modification detection, and no-op when the signal is already absent. - TestWait: Add two integration tests that verify session rollback and ProcessException wrapping when replace() throws on CAS failure (waitProgressed path) and when complete() detects a stale revision (waitCompleted path). --- .../apache/nifi/processors/standard/Wait.java | 7 +- .../standard/WaitNotifyProtocol.java | 48 +++++++- .../nifi/processors/standard/TestWait.java | 112 ++++++++++++++++++ .../standard/TestWaitNotifyProtocol.java | 60 ++++++++++ 4 files changed, 220 insertions(+), 7 deletions(-) diff --git a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/Wait.java b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/Wait.java index f653f31f3a49..1c282ce9c3cb 100644 --- a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/Wait.java +++ b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/Wait.java @@ -46,6 +46,7 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.ConcurrentModificationException; import java.util.HashMap; import java.util.Iterator; import java.util.List; @@ -558,7 +559,7 @@ public void onTrigger(final ProcessContext context, final ProcessSession session // Update signal if needed. try { if (waitCompleted) { - protocol.complete(signalId); + protocol.complete(signal); if (logger.isDebugEnabled()) { logger.debug("Completed wait for signalId='{}' and removed signal from cache", signalId); } @@ -570,6 +571,10 @@ public void onTrigger(final ProcessContext context, final ProcessSession session } } + } catch (final ConcurrentModificationException e) { + logger.warn("Concurrent modification detected for signal [{}], rolling back session to retry: {}", signalId, e.getMessage()); + session.rollback(); + throw new ProcessException(String.format("Concurrent modification detected while updating signal %s", signalId), e); } catch (final IOException e) { session.rollback(); throw new ProcessException(String.format("Unable to communicate with cache while updating %s due to %s", signalId, e), e); diff --git a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/WaitNotifyProtocol.java b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/WaitNotifyProtocol.java index 1e684a8a0d45..7cd154258d56 100644 --- a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/WaitNotifyProtocol.java +++ b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/WaitNotifyProtocol.java @@ -199,8 +199,11 @@ public Signal notify(final String signalId, final Map deltas, f signal.counts.put(counterName, count); }); - if (replace(signal)) { + try { + replace(signal); return signal; + } catch (final ConcurrentModificationException ignored) { + // CAS failed; retry with fresh signal state after backoff. } long waitMillis = REPLACE_RETRY_WAIT_MILLIS * (i + 1); @@ -238,7 +241,7 @@ public Signal notify(final String signalId, final String counterName, final int /** * Retrieve a stored Signal in the cache engine. - * If a caller gets satisfied with the returned Signal state and finish waiting, it should call {@link #complete(String)} + * If a caller gets satisfied with the returned Signal state and finish waiting, it should call {@link #complete(Signal)} * to complete the Wait Notify protocol. * @param signalId a key in the underlying cache engine * @return A Signal instance @@ -281,14 +284,44 @@ public Signal getSignal(final String signalId) throws IOException, Deserializati /** * Finish protocol and remove the cache entry. - * @param signalId a key in the underlying cache engine + * + *

This method performs a best-effort version check before removing the entry. If the signal + * was concurrently modified by a Notify processor after the caller last read it, a + * {@link ConcurrentModificationException} is thrown so the caller can roll back and retry + * rather than silently discarding the concurrent notification.

+ * + *

Note: there is a small inherent TOCTOU window between the version re-fetch and the + * remove call. A {@link AtomicDistributedMapCacheClient} API extension for atomic + * compare-and-delete would eliminate this entirely, but this approach covers the common case.

+ * + * @param signal the Signal obtained from the most recent {@link #getSignal(String)} call; + * its cached revision is used to detect concurrent modifications * @throws IOException thrown when it failed interacting with the cache engine + * @throws ConcurrentModificationException thrown if the signal was concurrently modified + * or removed since the caller last read it */ - public void complete(final String signalId) throws IOException { + public void complete(final Signal signal) throws IOException, ConcurrentModificationException { + final String signalId = signal.identifier; + + // Re-fetch to detect concurrent updates since the signal was last read. + final Signal current = getSignal(signalId); + if (current == null) { + throw new ConcurrentModificationException(String.format( + "Failed to complete signal [%s]: signal was concurrently removed.", signalId)); + } + + final Object expectedRevision = signal.cachedEntry != null ? signal.cachedEntry.getRevision().orElse(null) : null; + final Object actualRevision = current.cachedEntry.getRevision().orElse(null); + if (expectedRevision != null && !expectedRevision.equals(actualRevision)) { + throw new ConcurrentModificationException(String.format( + "Failed to complete signal [%s]: signal was concurrently modified (expected revision %s, found %s).", + signalId, expectedRevision, actualRevision)); + } + cache.remove(signalId, stringSerializer); } - public boolean replace(final Signal signal) throws IOException { + public void replace(final Signal signal) throws IOException, ConcurrentModificationException { final String signalJson = objectMapper.writeValueAsString(signal); if (signal.cachedEntry == null) { @@ -296,7 +329,10 @@ public boolean replace(final Signal signal) throws IOException { } else { signal.cachedEntry.setValue(signalJson); } - return cache.replace(signal.cachedEntry, stringSerializer, stringSerializer); + if (!cache.replace(signal.cachedEntry, stringSerializer, stringSerializer)) { + throw new ConcurrentModificationException(String.format( + "Failed to update signal [%s] in cache due to concurrent modification.", signal.identifier)); + } } } diff --git a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestWait.java b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestWait.java index 8c58f07312bb..a2be90d10d41 100644 --- a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestWait.java +++ b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestWait.java @@ -16,6 +16,9 @@ */ package org.apache.nifi.processors.standard; +import org.apache.nifi.distributed.cache.client.AtomicCacheEntry; +import org.apache.nifi.distributed.cache.client.Deserializer; +import org.apache.nifi.distributed.cache.client.Serializer; import org.apache.nifi.processor.exception.ProcessException; import org.apache.nifi.processors.standard.TestNotify.MockCacheClient; import org.apache.nifi.reporting.InitializationException; @@ -785,6 +788,115 @@ public void testOpenGate() throws IOException { assertEquals(0, signal.getReleasableCount()); } + /** + * A {@link MockCacheClient} extension that can simulate concurrent-modification scenarios + * by returning {@code false} from {@link #replace} on demand, or by returning a stale + * (bumped) revision from {@link #fetch} on a specific call number. + */ + private static class ControllableCacheClient extends MockCacheClient { + private volatile boolean failNextReplace = false; + private volatile int bumpRevisionOnFetchCall = -1; + private volatile int fetchCallCount = 0; + + void setFailNextReplace() { + failNextReplace = true; + } + + void setBumpRevisionOnFetchCall(final int callNumber) { + fetchCallCount = 0; + bumpRevisionOnFetchCall = callNumber; + } + + @Override + public boolean replace(final AtomicCacheEntry entry, + final Serializer keySerializer, final Serializer valueSerializer) throws IOException { + if (failNextReplace) { + failNextReplace = false; + return false; + } + return super.replace(entry, keySerializer, valueSerializer); + } + + @Override + @SuppressWarnings("unchecked") + public AtomicCacheEntry fetch(final K key, + final Serializer keySerializer, final Deserializer valueDeserializer) throws IOException { + final AtomicCacheEntry entry = super.fetch(key, keySerializer, valueDeserializer); + final int call = ++fetchCallCount; + if (bumpRevisionOnFetchCall == call && entry != null) { + return new AtomicCacheEntry<>(entry.getKey(), entry.getValue(), + entry.getRevision().orElse(0L) + 999L); + } + return entry; + } + } + + /** + * When a concurrent Notify wins the CAS slot, {@code protocol.replace(signal)} + * throws {@link ConcurrentModificationException}. Wait must roll back the session and + * throw a {@link ProcessException} so FlowFiles are retried rather than silently lost. + */ + @Test + public void testWaitProgressedRollsBackOnConcurrentReplaceFailure() throws Exception { + final ControllableCacheClient controllableService = new ControllableCacheClient(); + runner.addControllerService("controllable-service", controllableService); + runner.enableControllerService(controllableService); + runner.setProperty(Wait.DISTRIBUTED_CACHE_SERVICE, "controllable-service"); + + // Signal count=2, target=1, releasable=1 (default): after consuming 1 signal to release + // the FlowFile, releasableCount becomes 1 (leftover ticket). That means waitCompleted=false + // and waitProgressed=true → Wait calls protocol.replace(signal). + runner.setProperty(Wait.RELEASE_SIGNAL_IDENTIFIER, "${releaseSignalAttribute}"); + runner.setProperty(Wait.TARGET_SIGNAL_COUNT, "1"); + + final WaitNotifyProtocol protocol = new WaitNotifyProtocol(controllableService); + protocol.notify("key", WaitNotifyProtocol.DEFAULT_COUNT_NAME, 2, null); + + final Map attrs = new HashMap<>(); + attrs.put("releaseSignalAttribute", "key"); + runner.enqueue(new byte[]{}, attrs); + + // Simulate a concurrent Notify winning the CAS slot: replace() will return false. + controllableService.setFailNextReplace(); + + final AssertionError e = assertThrows(AssertionError.class, () -> runner.run()); + assertInstanceOf(ProcessException.class, e.getCause()); + assertInstanceOf(ConcurrentModificationException.class, e.getCause().getCause()); + } + + /** + * When {@code protocol.complete(signal)} detects that the cached entry was + * modified by a concurrent Notify between Wait's initial fetch and the remove call, Wait + * must roll back the session and throw a {@link ProcessException} to avoid silently + * discarding the Notify update. + */ + @Test + public void testWaitCompletedRollsBackOnConcurrentSignalModification() throws Exception { + final ControllableCacheClient controllableService = new ControllableCacheClient(); + runner.addControllerService("controllable-service", controllableService); + runner.enableControllerService(controllableService); + runner.setProperty(Wait.DISTRIBUTED_CACHE_SERVICE, "controllable-service"); + + runner.setProperty(Wait.RELEASE_SIGNAL_IDENTIFIER, "${releaseSignalAttribute}"); + runner.setProperty(Wait.TARGET_SIGNAL_COUNT, "1"); + + final WaitNotifyProtocol protocol = new WaitNotifyProtocol(controllableService); + protocol.notify("key", WaitNotifyProtocol.DEFAULT_COUNT_NAME, 1, null); + + final Map attrs = new HashMap<>(); + attrs.put("releaseSignalAttribute", "key"); + runner.enqueue(new byte[]{}, attrs); + + // After notify(), reset the counter. During runner.run(): + // call 1 = Wait's own getSignal() → returns normally (revision N) + // call 2 = complete()'s internal getSignal() → returns bumped revision (simulating concurrent Notify) + controllableService.setBumpRevisionOnFetchCall(2); + + final AssertionError e = assertThrows(AssertionError.class, () -> runner.run()); + assertInstanceOf(ProcessException.class, e.getCause()); + assertInstanceOf(ConcurrentModificationException.class, e.getCause().getCause()); + } + @Test void testMigrateProperties() { final Map expectedRenamed = Map.ofEntries( diff --git a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestWaitNotifyProtocol.java b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestWaitNotifyProtocol.java index c32f14a026bc..d2c2149f668b 100644 --- a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestWaitNotifyProtocol.java +++ b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestWaitNotifyProtocol.java @@ -424,6 +424,66 @@ public void testReleaseCandidateTotal() throws Exception { } + @Test + public void testCompleteRemovesSignalFromCache() throws Exception { + doAnswer(successfulReplace).when(cache).replace(any(), any(), any()); + doAnswer(invocation -> { + cacheEntries.remove(invocation.getArguments()[0]); + return true; + }).when(cache).remove(any(), any()); + + final WaitNotifyProtocol protocol = new WaitNotifyProtocol(cache); + final String signalId = "signal-id"; + + protocol.notify(signalId, "a", 1, null); + assertTrue(cacheEntries.containsKey(signalId)); + + final Signal signal = protocol.getSignal(signalId); + assertNotNull(signal); + + protocol.complete(signal); + assertFalse(cacheEntries.containsKey(signalId)); + } + + @Test + public void testCompleteThrowsOnConcurrentModification() throws Exception { + doAnswer(successfulReplace).when(cache).replace(any(), any(), any()); + + final WaitNotifyProtocol protocol = new WaitNotifyProtocol(cache); + final String signalId = "signal-id"; + + // Notify creates the signal at revision 1. + protocol.notify(signalId, "a", 1, null); + final Signal signalBeforeRace = protocol.getSignal(signalId); + assertNotNull(signalBeforeRace); + + // Simulate a concurrent Notify that updates the signal, bumping the revision. + protocol.notify(signalId, "a", 1, null); + + // complete() with a stale signal should detect the version mismatch and throw. + assertThrows(ConcurrentModificationException.class, () -> protocol.complete(signalBeforeRace)); + + // The entry must still be present — complete() must NOT have removed it. + assertTrue(cacheEntries.containsKey(signalId)); + } + + @Test + public void testCompleteThrowsWhenAlreadyRemoved() throws Exception { + doAnswer(successfulReplace).when(cache).replace(any(), any(), any()); + + final WaitNotifyProtocol protocol = new WaitNotifyProtocol(cache); + final String signalId = "signal-id"; + + protocol.notify(signalId, "a", 1, null); + final Signal signal = protocol.getSignal(signalId); + + // Remove the entry from the cache directly, simulating concurrent removal by another process. + cacheEntries.remove(signalId); + + // complete() on a signal that was concurrently removed should throw ConcurrentModificationException. + assertThrows(ConcurrentModificationException.class, () -> protocol.complete(signal)); + } + public void assertValueEquals(String expected, String value) throws Exception { assertEquals(mapper.readTree(expected), mapper.readTree(value)); }