NIFI-15885: Addressed issues that cause Offload to get stuck: when a … (#11184)

* NIFI-15885: Addressed issues that cause Offload to get stuck: when a processor extends AbstractSessionFactoryProcessor and stores a reference to a created ProcessSession but not its factory, garbage collection could cause framework to be unable to rollback sessions. When a processor is terminated, ensure we rollback retained sessions even if all threads are completed. Fixed related bug that caused background thread never to complete when waiting for processor thread count to reach 1 and terminate set it to 0.

* NIFI-15885: Fix to unit test

* NIFI-15885: Allow StandardProcessSession.migrate to unwrap delegating Session wrappers so FactoryRetainingProcessSession (used to keep an ActiveProcessSessionFactory reachable for offload/terminate) is recognized as a StandardProcessSession target. Adds a framework-internal DelegatingProcessSession contract implemented by the wrapper, and an IT covering the migrate-to-wrapper path that reproduces the MergeRecord CI failure.

Author: markap14

nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/StandardProcessorNode.java

@@ -1905,6 +1905,14 @@ public void run() {
                                     }
                                 }
                             }
+                        } else if (lifecycleState.isTerminated()) {
+                            // Termination was requested while the stop sequence was waiting for active threads to drain.
+                            // LifecycleState.terminate() reset the active thread count to zero, so the count==1
+                            // condition above will never be reached and rescheduling would loop forever. Complete the
+                            // stop action and exit. completeStopAction() is idempotent if procNode.terminate() already
+                            // invoked it.
+                            LOG.debug("Stop sequence for {} aborted because LifecycleState was terminated", this);
+                            completeStopAction();
                         } else {
                             // Not all of the active threads have finished. Try again in 100 milliseconds.
                             executor.schedule(this, 100, TimeUnit.MILLISECONDS);

nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/DelegatingProcessSession.java

@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.nifi.controller.repository;
+
+import org.apache.nifi.processor.ProcessSession;
+
+/**
+ * Framework-internal contract implemented by {@link ProcessSession} wrappers that forward all
+ * operations to an underlying delegate Session. This allows framework code that needs to operate
+ * on a concrete Session implementation (for example, {@code StandardProcessSession.migrate}, which
+ * requires its target to be another {@code StandardProcessSession}) to look through any number of
+ * wrapping layers and recover the underlying Session.
+ *
+ * Extensions must not implement this interface; it is intended for framework-internal Session
+ * wrappers only.
+ */
+public interface DelegatingProcessSession extends ProcessSession {
+
+    /**
+     * @return the Session that this wrapper forwards operations to. May itself be a
+     *         {@link DelegatingProcessSession} when wrappers are nested.
+     */
+    ProcessSession getDelegate();
+}

nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java

@@ -1501,11 +1501,23 @@ public void migrate(final ProcessSession newOwner, final Collection<FlowFile> fl
             throw new IllegalArgumentException("Must supply at least one FlowFile to migrate");
         }
 
-        if (!(newOwner instanceof StandardProcessSession)) {
+        // Look through any framework-internal Session wrappers (such as the one used to keep an
+        // ActiveProcessSessionFactory reachable for the offload/terminate path) so the underlying
+        // StandardProcessSession can be located.
+        ProcessSession resolvedOwner = newOwner;
+        while (resolvedOwner instanceof DelegatingProcessSession delegating) {
+            resolvedOwner = delegating.getDelegate();
+        }
+
+        if (!(resolvedOwner instanceof StandardProcessSession standardOwner)) {
             throw new IllegalArgumentException("Cannot migrate from a StandardProcessSession to a " + newOwner.getClass());
         }
 
-        migrate((StandardProcessSession) newOwner, flowFiles);
+        if (standardOwner == this) {
+            throw new IllegalArgumentException("Cannot migrate FlowFiles from a Process Session to itself");
+        }
+
+        migrate(standardOwner, flowFiles);
     }
 
     private synchronized void migrate(final StandardProcessSession newOwner, Collection<FlowFile> flowFiles) {

nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/WeakHashMapProcessSessionFactory.java

@@ -17,12 +17,32 @@
 
 package org.apache.nifi.controller.repository;
 
+import org.apache.nifi.components.state.Scope;
+import org.apache.nifi.components.state.StateMap;
+import org.apache.nifi.controller.queue.QueueSize;
+import org.apache.nifi.flowfile.FlowFile;
+import org.apache.nifi.processor.FlowFileFilter;
 import org.apache.nifi.processor.ProcessSession;
 import org.apache.nifi.processor.ProcessSessionFactory;
+import org.apache.nifi.processor.Relationship;
 import org.apache.nifi.processor.exception.TerminatedTaskException;
+import org.apache.nifi.processor.io.InputStreamCallback;
+import org.apache.nifi.processor.io.OutputStreamCallback;
+import org.apache.nifi.processor.io.StreamCallback;
+import org.apache.nifi.processor.metrics.CommitTiming;
+import org.apache.nifi.provenance.ProvenanceReporter;
 
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.file.Path;
+import java.util.Collection;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.WeakHashMap;
+import java.util.function.Consumer;
+import java.util.regex.Pattern;
 
 public class WeakHashMapProcessSessionFactory implements ActiveProcessSessionFactory {
     private final ProcessSessionFactory delegate;
@@ -39,9 +59,16 @@ public synchronized ProcessSession createSession() {
             throw new TerminatedTaskException();
         }
 
-        final ProcessSession session = delegate.createSession();
-        sessionMap.put(session, Boolean.TRUE);
-        return session;
+        final ProcessSession delegateSession = delegate.createSession();
+
+        // Wrap the delegate Session so that the returned Session strongly references this factory. The
+        // LifecycleState tracks ActiveProcessSessionFactory instances in a WeakHashMap, so without a
+        // back-reference from the Session a Processor that retains the Session but not the factory across
+        // onTrigger invocations would let the WeakHashMap entry clear, leaving terminate() unable to roll
+        // back the orphaned Session and preventing offload from completing.
+        final ProcessSession wrapper = new FactoryRetainingProcessSession(delegateSession, this);
+        sessionMap.put(wrapper, Boolean.TRUE);
+        return wrapper;
     }
 
     @Override
@@ -56,4 +83,261 @@ public synchronized void terminateActiveSessions() {
 
         sessionMap.clear();
     }
+
+    /**
+     * A delegating {@link ProcessSession} that strongly references the factory that produced it. This
+     * back-reference exists solely to keep the factory reachable so that the {@code LifecycleState}
+     * WeakHashMap entry tracking the factory is not cleared while the Session is still in use. All
+     * operations are forwarded to the underlying delegate Session.
+     */
+    private static final class FactoryRetainingProcessSession implements DelegatingProcessSession {
+        private final ProcessSession delegate;
+        @SuppressWarnings("unused") // Retained only to keep the factory reachable; see class Javadoc.
+        private final WeakHashMapProcessSessionFactory factoryRetention;
+
+        private FactoryRetainingProcessSession(final ProcessSession delegate, final WeakHashMapProcessSessionFactory factoryRetention) {
+            this.delegate = delegate;
+            this.factoryRetention = factoryRetention;
+        }
+
+        @Override
+        public ProcessSession getDelegate() {
+            return delegate;
+        }
+
+        @Override
+        public void commit() {
+            delegate.commit();
+        }
+
+        @Override
+        public void commitAsync() {
+            delegate.commitAsync();
+        }
+
+        @Override
+        public void commitAsync(final Runnable onSuccess, final Consumer<Throwable> onFailure) {
+            delegate.commitAsync(onSuccess, onFailure);
+        }
+
+        @Override
+        public void rollback() {
+            delegate.rollback();
+        }
+
+        @Override
+        public void rollback(final boolean penalize) {
+            delegate.rollback(penalize);
+        }
+
+        @Override
+        public void migrate(final ProcessSession newOwner, final Collection<FlowFile> flowFiles) {
+            delegate.migrate(newOwner, flowFiles);
+        }
+
+        @Override
+        public void migrate(final ProcessSession newOwner) {
+            delegate.migrate(newOwner);
+        }
+
+        @Override
+        public void adjustCounter(final String name, final long delta, final boolean immediate) {
+            delegate.adjustCounter(name, delta, immediate);
+        }
+
+        @Override
+        public void recordGauge(final String name, final double value, final CommitTiming commitTiming) {
+            delegate.recordGauge(name, value, commitTiming);
+        }
+
+        @Override
+        public FlowFile get() {
+            return delegate.get();
+        }
+
+        @Override
+        public List<FlowFile> get(final int maxResults) {
+            return delegate.get(maxResults);
+        }
+
+        @Override
+        public List<FlowFile> get(final FlowFileFilter filter) {
+            return delegate.get(filter);
+        }
+
+        @Override
+        public QueueSize getQueueSize() {
+            return delegate.getQueueSize();
+        }
+
+        @Override
+        public FlowFile create() {
+            return delegate.create();
+        }
+
+        @Override
+        public FlowFile create(final FlowFile parent) {
+            return delegate.create(parent);
+        }
+
+        @Override
+        public FlowFile create(final Collection<FlowFile> parents) {
+            return delegate.create(parents);
+        }
+
+        @Override
+        public FlowFile clone(final FlowFile example) {
+            return delegate.clone(example);
+        }
+
+        @Override
+        public FlowFile clone(final FlowFile example, final long offset, final long size) {
+            return delegate.clone(example, offset, size);
+        }
+
+        @Override
+        public FlowFile penalize(final FlowFile flowFile) {
+            return delegate.penalize(flowFile);
+        }
+
+        @Override
+        public FlowFile putAttribute(final FlowFile flowFile, final String key, final String value) {
+            return delegate.putAttribute(flowFile, key, value);
+        }
+
+        @Override
+        public FlowFile putAllAttributes(final FlowFile flowFile, final Map<String, String> attributes) {
+            return delegate.putAllAttributes(flowFile, attributes);
+        }
+
+        @Override
+        public FlowFile removeAttribute(final FlowFile flowFile, final String key) {
+            return delegate.removeAttribute(flowFile, key);
+        }
+
+        @Override
+        public FlowFile removeAllAttributes(final FlowFile flowFile, final Set<String> keys) {
+            return delegate.removeAllAttributes(flowFile, keys);
+        }
+
+        @Override
+        public FlowFile removeAllAttributes(final FlowFile flowFile, final Pattern keyPattern) {
+            return delegate.removeAllAttributes(flowFile, keyPattern);
+        }
+
+        @Override
+        public void transfer(final FlowFile flowFile, final Relationship relationship) {
+            delegate.transfer(flowFile, relationship);
+        }
+
+        @Override
+        public void transfer(final FlowFile flowFile) {
+            delegate.transfer(flowFile);
+        }
+
+        @Override
+        public void transfer(final Collection<FlowFile> flowFiles) {
+            delegate.transfer(flowFiles);
+        }
+
+        @Override
+        public void transfer(final Collection<FlowFile> flowFiles, final Relationship relationship) {
+            delegate.transfer(flowFiles, relationship);
+        }
+
+        @Override
+        public void remove(final FlowFile flowFile) {
+            delegate.remove(flowFile);
+        }
+
+        @Override
+        public void remove(final Collection<FlowFile> flowFiles) {
+            delegate.remove(flowFiles);
+        }
+
+        @Override
+        public void read(final FlowFile source, final InputStreamCallback reader) {
+            delegate.read(source, reader);
+        }
+
+        @Override
+        public InputStream read(final FlowFile flowFile) {
+            return delegate.read(flowFile);
+        }
+
+        @Override
+        public FlowFile merge(final Collection<FlowFile> sources, final FlowFile destination) {
+            return delegate.merge(sources, destination);
+        }
+
+        @Override
+        public FlowFile merge(final Collection<FlowFile> sources, final FlowFile destination, final byte[] header, final byte[] footer, final byte[] demarcator) {
+            return delegate.merge(sources, destination, header, footer, demarcator);
+        }
+
+        @Override
+        public FlowFile write(final FlowFile source, final OutputStreamCallback writer) {
+            return delegate.write(source, writer);
+        }
+
+        @Override
+        public FlowFile write(final FlowFile source, final StreamCallback writer) {
+            return delegate.write(source, writer);
+        }
+
+        @Override
+        public OutputStream write(final FlowFile source) {
+            return delegate.write(source);
+        }
+
+        @Override
+        public FlowFile append(final FlowFile source, final OutputStreamCallback writer) {
+            return delegate.append(source, writer);
+        }
+
+        @Override
+        public FlowFile importFrom(final Path source, final boolean keepSourceFile, final FlowFile destination) {
+            return delegate.importFrom(source, keepSourceFile, destination);
+        }
+
+        @Override
+        public FlowFile importFrom(final InputStream source, final FlowFile destination) {
+            return delegate.importFrom(source, destination);
+        }
+
+        @Override
+        public void exportTo(final FlowFile flowFile, final Path destination, final boolean append) {
+            delegate.exportTo(flowFile, destination, append);
+        }
+
+        @Override
+        public void exportTo(final FlowFile flowFile, final OutputStream destination) {
+            delegate.exportTo(flowFile, destination);
+        }
+
+        @Override
+        public ProvenanceReporter getProvenanceReporter() {
+            return delegate.getProvenanceReporter();
+        }
+
+        @Override
+        public void setState(final Map<String, String> state, final Scope scope) throws IOException {
+            delegate.setState(state, scope);
+        }
+
+        @Override
+        public StateMap getState(final Scope scope) throws IOException {
+            return delegate.getState(scope);
+        }
+
+        @Override
+        public boolean replaceState(final StateMap oldValue, final Map<String, String> newValue, final Scope scope) throws IOException {
+            return delegate.replaceState(oldValue, newValue, scope);
+        }
+
+        @Override
+        public void clearState(final Scope scope) throws IOException {
+            delegate.clearState(scope);
+        }
+    }
 }