W-16941297: Scatter Gather timeout exception

core/src/main/java/org/mule/runtime/core/internal/event/AbstractEventContext.java

@@ -52,7 +52,7 @@
  *
  * @since 4.0
  */
-abstract class AbstractEventContext implements SpanContextAware, BaseEventContext {
+public abstract class AbstractEventContext implements SpanContextAware, BaseEventContext {
 
   private static final int STATE_READY = 0;
   private static final int STATE_RESPONSE_RECEIVED = 1;

modules/core-components/src/main/java/org/mule/runtime/core/internal/routing/forkjoin/AbstractForkJoinStrategyFactory.java

@@ -9,10 +9,10 @@
 import static org.mule.runtime.core.api.event.CoreEvent.builder;
 import static org.mule.runtime.core.internal.exception.ErrorHandlerContextManager.ERROR_HANDLER_CONTEXT;
 import static org.mule.runtime.core.internal.routing.ForkJoinStrategy.RoutingPair.of;
-import static org.mule.runtime.core.privileged.processor.MessageProcessors.processWithChildContextDontComplete;
 
+import static java.lang.Boolean.parseBoolean;
 import static java.lang.Long.MAX_VALUE;
-import static java.util.Optional.empty;
+import static java.lang.System.getProperty;
 import static java.util.Optional.ofNullable;
 import static java.util.stream.Collectors.toList;
 
@@ -74,6 +74,10 @@
  */
 public abstract class AbstractForkJoinStrategyFactory implements ForkJoinStrategyFactory {
 
+  // TODO W-17814249: Remove kill switch
+  private static final boolean COMPLETE_CHILDREN_ON_TIMEOUT =
+      parseBoolean(getProperty("mule.forkJoin.completeChildContextsOnTimeout", "true"));
+
   public static final String TIMEOUT_EXCEPTION_DESCRIPTION = "Route Timeout";
   public static final String TIMEOUT_EXCEPTION_DETAILED_DESCRIPTION_PREFIX = "Timeout while processing route/part:";
   private final boolean mergeVariables;
@@ -171,30 +175,34 @@ private Function<RoutingPair, Publisher<Pair<CoreEvent, EventProcessingException
     return pair -> {
       ReactiveProcessor route = publisher -> from(publisher)
           .transform(pair.getRoute());
-      return from(processWithChildContextDontComplete(pair.getEvent(),
-                                                      applyProcessingStrategy(processingStrategy, route, maxConcurrency),
-                                                      empty()))
-                                                          .timeout(timeout,
-                                                                   onTimeout(processingStrategy, delayErrors, timeoutErrorType,
-                                                                             pair),
-                                                                   timeoutScheduler)
-                                                          .map(coreEvent -> new Pair<CoreEvent, EventProcessingException>(
-                                                                                                                          ((DefaultEventBuilder) CoreEvent
-                                                                                                                              .builder(coreEvent))
-                                                                                                                                  .removeInternalParameter(ERROR_HANDLER_CONTEXT)
-                                                                                                                                  .build(),
-                                                                                                                          null))
-                                                          .onErrorResume(MessagingException.class,
-                                                                         me -> getPublisher(delayErrors, me));
+      route = applyProcessingStrategy(processingStrategy, route, maxConcurrency);
+
+      // TODO W-17814249: Remove kill switch
+      RoutePairPublisherAssemblyHelper routePairPublisherAssemblyHelper = COMPLETE_CHILDREN_ON_TIMEOUT
+          ? new DefaultRoutePairPublisherAssemblyHelper(pair.getEvent(), route)
+          : new LegacyRoutePairPublisherAssemblyHelper(pair.getEvent(), route);
+
+      return from(routePairPublisherAssemblyHelper.getPublisherOnChildContext())
+          .timeout(timeout,
+                   routePairPublisherAssemblyHelper
+                       .decorateTimeoutPublisher(onTimeout(processingStrategy, delayErrors, timeoutErrorType, pair)),
+                   timeoutScheduler)
+          .map(this::eventToPair)
+          .onErrorResume(MessagingException.class, me -> getPublisher(delayErrors, me));
     };
   }
 
+  private Pair<CoreEvent, EventProcessingException> eventToPair(CoreEvent coreEvent) {
+    return new Pair<>(((DefaultEventBuilder) CoreEvent.builder(coreEvent))
+        .removeInternalParameter(ERROR_HANDLER_CONTEXT)
+        .build(), null);
+  }
+
   private Publisher<Pair<CoreEvent, EventProcessingException>> getPublisher(boolean delayErrors, EventProcessingException me) {
     Pair<CoreEvent, EventProcessingException> pair = new Pair<>(me.getEvent(), me);
     return delayErrors ? just(pair) : error(me);
   }
 
-
   private Mono<CoreEvent> onTimeout(ProcessingStrategy processingStrategy, boolean delayErrors, ErrorType timeoutErrorType,
                                     RoutingPair pair) {
     return defer(() -> delayErrors ? just(createTimeoutErrorEvent(timeoutErrorType, pair))

modules/core-components/src/main/java/org/mule/runtime/core/internal/routing/forkjoin/DefaultRoutePairPublisherAssemblyHelper.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2023 Salesforce, Inc. All rights reserved.
+ * The software in this package is published under the terms of the CPAL v1.0
+ * license, a copy of which has been included with this distribution in the
+ * LICENSE.txt file.
+ */
+package org.mule.runtime.core.internal.routing.forkjoin;
+
+import static org.mule.runtime.core.internal.event.EventQuickCopy.quickCopy;
+import static org.mule.runtime.core.privileged.processor.MessageProcessors.newChildContext;
+import static org.mule.runtime.core.privileged.processor.MessageProcessors.processWithChildContext;
+
+import static java.util.Optional.empty;
+
+import org.mule.runtime.core.api.event.CoreEvent;
+import org.mule.runtime.core.api.processor.ReactiveProcessor;
+import org.mule.runtime.core.internal.event.AbstractEventContext;
+import org.mule.runtime.core.privileged.event.BaseEventContext;
+import org.mule.runtime.core.privileged.exception.MessagingException;
+
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Mono;
+
+/**
+ * Default implementation of {@link RoutePairPublisherAssemblyHelper} which takes care of completing the child contexts in case of
+ * timeouts.
+ */
+class DefaultRoutePairPublisherAssemblyHelper implements RoutePairPublisherAssemblyHelper {
+
+  private final Publisher<CoreEvent> publisherWithChildContext;
+  private final BaseEventContext childContext;
+
+  DefaultRoutePairPublisherAssemblyHelper(CoreEvent routeEvent, ReactiveProcessor chain) {
+    // This sequence is equivalent to processWithChildContextDontComplete (used in the legacy version), only this way we can keep
+    // track of the child context, which we will need later on.
+    this.childContext = newChildContext(routeEvent, empty());
+    this.publisherWithChildContext = processWithChildContext(routeEvent, chain, childContext);
+  }
+
+  @Override
+  public Publisher<CoreEvent> getPublisherOnChildContext() {
+    return publisherWithChildContext;
+  }
+
+  @Override
+  public Publisher<CoreEvent> decorateTimeoutPublisher(Publisher<CoreEvent> timeoutPublisher) {
+    // When the timeout happens, the subscription to the original publisher is cancelled, so the inner MessageProcessorChains
+    // never finishes and the child contexts are never completed, hence we have to complete them manually on timeout
+    return Mono.from(timeoutPublisher)
+        .doOnSuccess(completeRecursively(childContext, BaseEventContext::error));
+  }
+
+  private Consumer<CoreEvent> completeRecursively(BaseEventContext eventContext,
+                                                  BiConsumer<BaseEventContext, MessagingException> forEachChild) {
+    return some -> {
+      // Tracks the child contexts first and then completes them, that way we are not retaining the read lock all the time.
+      Deque<BaseEventContext> allContexts = new ArrayDeque<>();
+      allContexts.push(eventContext);
+      ((AbstractEventContext) eventContext).forEachChild(allContexts::push);
+
+      while (!allContexts.isEmpty()) {
+        BaseEventContext ctx = allContexts.pop();
+        // Some context may already be terminated, as the children can propagate the termination up to their parents.
+        if (!ctx.isTerminated()) {
+          // It is important to swap the context so it has the right flow stack among other things.
+          forEachChild.accept(ctx, new MessagingException(quickCopy(ctx, some), some.getError().get().getCause()));
+        }
+      }
+    };
+  }
+}

modules/core-components/src/main/java/org/mule/runtime/core/internal/routing/forkjoin/LegacyRoutePairPublisherAssemblyHelper.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2023 Salesforce, Inc. All rights reserved.
+ * The software in this package is published under the terms of the CPAL v1.0
+ * license, a copy of which has been included with this distribution in the
+ * LICENSE.txt file.
+ */
+package org.mule.runtime.core.internal.routing.forkjoin;
+
+import static org.mule.runtime.core.privileged.processor.MessageProcessors.processWithChildContextDontComplete;
+
+import static java.util.Optional.empty;
+
+import org.mule.runtime.core.api.event.CoreEvent;
+import org.mule.runtime.core.api.processor.ReactiveProcessor;
+
+import org.reactivestreams.Publisher;
+
+/**
+ * Implementation of {@link RoutePairPublisherAssemblyHelper} which does not complete the child contexts on timeout.
+ * <p>
+ * Exists as a fallback for the kill switch COMPLETE_CHILDREN_ON_TIMEOUT.
+ */
+// TODO W-17814249: Remove once the kill switch is removed
+class LegacyRoutePairPublisherAssemblyHelper implements RoutePairPublisherAssemblyHelper {
+
+  private final Publisher<CoreEvent> publisherWithChildContext;
+
+  LegacyRoutePairPublisherAssemblyHelper(CoreEvent routeEvent, ReactiveProcessor chain) {
+    this.publisherWithChildContext = processWithChildContextDontComplete(routeEvent, chain, empty());
+  }
+
+  @Override
+  public Publisher<CoreEvent> getPublisherOnChildContext() {
+    return publisherWithChildContext;
+  }
+
+  @Override
+  public Publisher<CoreEvent> decorateTimeoutPublisher(Publisher<CoreEvent> timeoutPublisher) {
+    // Adds nothing to the timeout publisher
+    return timeoutPublisher;
+  }
+}

modules/core-components/src/main/java/org/mule/runtime/core/internal/routing/forkjoin/RoutePairPublisherAssemblyHelper.java

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2023 Salesforce, Inc. All rights reserved.
+ * The software in this package is published under the terms of the CPAL v1.0
+ * license, a copy of which has been included with this distribution in the
+ * LICENSE.txt file.
+ */
+package org.mule.runtime.core.internal.routing.forkjoin;
+
+import org.mule.runtime.core.api.event.CoreEvent;
+
+import org.reactivestreams.Publisher;
+
+/**
+ * Helper class to assist in assembling the routing pair publisher.
+ */
+interface RoutePairPublisherAssemblyHelper {
+
+  Publisher<CoreEvent> getPublisherOnChildContext();
+
+  Publisher<CoreEvent> decorateTimeoutPublisher(Publisher<CoreEvent> timeoutPublisher);
+}