Introduce component decorators

config/src/main/java/org/axonframework/config/Component.java

@@ -21,6 +21,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.lang.invoke.MethodHandles;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
@@ -74,8 +75,9 @@ public Component(@Nonnull Supplier<Configuration> config,
 
     /**
      * Retrieves the object contained in this component, triggering the builder function if the component hasn't been
-     * built yet. Upon initiation of the instance the {@link LifecycleHandlerInspector#registerLifecycleHandlers(Configuration,
-     * Object)} methods will be called to resolve and register lifecycle methods.
+     * built yet. Upon initiation of the instance the
+     * {@link LifecycleHandlerInspector#registerLifecycleHandlers(Configuration, Object)} methods will be called to
+     * resolve and register lifecycle methods.
      *
      * @return the initialized component contained in this instance
      */
@@ -100,6 +102,26 @@ public void update(@Nonnull Function<Configuration, ? extends B> builderFunction
         this.builderFunction = builderFunction;
     }
 
+    /**
+     * Updates the builder function for this component by creating a new builderFunction that uses the
+     * {@code decoratingFunction} to change the original component. This can wrap the original component.
+     * <p>
+     * Lifecycle handlers are scanned for both the original component and the one created by the @code
+     * decoratingFunction}.
+     *
+     * @param decoratingFunction Function that takes the original component and can wrap or change it depending on
+     *                           requirements.
+     */
+    public void decorate(@Nonnull BiFunction<Configuration, B, ? extends B> decoratingFunction) {
+        Assert.state(instance == null, () -> "Cannot change " + name + ": it is already in use");
+        Function<Configuration, ? extends B> previous = builderFunction;
+        this.update((configuration) -> {
+            B wrappedComponent = previous.apply(configuration);
+            LifecycleHandlerInspector.registerLifecycleHandlers(configuration, wrappedComponent);
+            return decoratingFunction.apply(configuration, wrappedComponent);
+        });
+    }
+
     /**
      * Checks if the component is already initialized.
      *

config/src/main/java/org/axonframework/config/ComponentDecorator.java

@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) 2010-2022. Axon Framework
+ *
+ * Licensed 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.axonframework.config;
+
+/**
+ * Functional interface that is able to decorate a {@link Component}, wrapping the original component.
+ *
+ * @author Mitchell Herrijgers
+ * @since 4.7.0
+ * @param <T> The Component's type
+ */
+@FunctionalInterface
+public interface ComponentDecorator<T> {
+
+    T decorate(Configuration configuration, T component);
+}

config/src/main/java/org/axonframework/config/Configurer.java

@@ -212,6 +212,23 @@ Configurer configureCorrelationDataProviders(
     <C> Configurer registerComponent(@Nonnull Class<C> componentType,
                                      @Nonnull Function<Configuration, ? extends C> componentBuilder);
 
+
+    /**
+     * Registers a {@link ComponentDecorator decorator} for a component type. For any component that is created by the
+     * configuration that matches the {@code componentType}, the {@code decorator} will be called. It's up to the
+     * decorator to decide to return the original component or a wrapped version of it.
+     *
+     * @param componentType The declared type of the component, typically an interface
+     * @param decorator     The decorator function for this component
+     * @param <C>           The type of component
+     * @return the current instance of the Configurer, for chaining purposes
+     */
+    default <C> Configurer registerComponentDecorator(@Nonnull Class<C> componentType,
+                                              @Nonnull ComponentDecorator<C> decorator) {
+        // Default implementation for backwards compatibility
+        return this;
+    }
+
     /**
      * Registers a command handler bean with this {@link Configurer}. The bean may be of any type. The actual command
      * handler methods will be detected based on the annotations present on the bean's methods. Message handling

config/src/main/java/org/axonframework/config/DefaultConfigurer.java

@@ -159,6 +159,7 @@ public class DefaultConfigurer implements Configurer {
     );
 
     private final List<Consumer<Configuration>> initHandlers = new ArrayList<>();
+    private final List<Runnable> decoratorHandlers = new ArrayList<>();
     private final TreeMap<Integer, List<LifecycleHandler>> startHandlers = new TreeMap<>();
     private final TreeMap<Integer, List<LifecycleHandler>> shutdownHandlers = new TreeMap<>(Comparator.reverseOrder());
     private final List<ModuleConfiguration> modules = new ArrayList<>();
@@ -643,6 +644,20 @@ public <C> Configurer registerComponent(@Nonnull Class<C> componentType,
         return this;
     }
 
+    @SuppressWarnings("unchecked")
+    @Override
+    public <C> Configurer registerComponentDecorator(@Nonnull Class<C> componentType,
+                                                     @Nonnull ComponentDecorator<C> decorator) {
+        logger.debug("Registering decorator for component [{}]", componentType.getSimpleName());
+        decoratorHandlers.add(() -> {
+            Component<C> component = (Component<C>) components.getOrDefault(componentType, null);
+            if (component != null) {
+                component.decorate(decorator::decorate);
+            }
+        });
+        return this;
+    }
+
     @Override
     public Configurer registerCommandHandler(@Nonnull Function<Configuration, Object> commandHandlerBuilder) {
         messageHandlerRegistrars.add(new Component<>(
@@ -683,13 +698,13 @@ public Configurer registerQueryHandler(@Nonnull Function<Configuration, Object>
     public Configurer registerMessageHandler(@Nonnull Function<Configuration, Object> messageHandlerBuilder) {
         Component<Object> messageHandler = new Component<>(() -> config, "", messageHandlerBuilder);
         Class<?> handlerClass = messageHandler.get().getClass();
-        if (isCommandHandler(handlerClass)){
+        if (isCommandHandler(handlerClass)) {
             registerCommandHandler(c -> messageHandler.get());
         }
-        if (isEventHandler(handlerClass)){
+        if (isEventHandler(handlerClass)) {
             eventProcessing().registerEventHandler(c -> messageHandler.get());
         }
-        if (isQueryHandler(handlerClass)){
+        if (isQueryHandler(handlerClass)) {
             registerQueryHandler(c -> messageHandler.get());
         }
         return this;
@@ -749,6 +764,7 @@ public Configuration buildConfiguration() {
         if (!initialized) {
             verifyIdentifierFactory();
             prepareModules();
+            invokeDecoratorHandlers();
             prepareMessageHandlerRegistrars();
             invokeInitHandlers();
         }
@@ -792,6 +808,17 @@ protected void invokeInitHandlers() {
         initHandlers.forEach(h -> h.accept(config));
     }
 
+    /**
+     * Calls all registered decorators of this configuration and updates the components where appropriate. Decorating
+     * components is postponed until building to make sure all {@link Component} definitions are present and can be
+     * decorated.
+     * <p>
+     * Registration of decorates are ignore after initialization.
+     */
+    protected void invokeDecoratorHandlers() {
+        decoratorHandlers.forEach(Runnable::run);
+    }
+
     /**
      * Invokes all registered start handlers.
      */
@@ -865,7 +892,8 @@ private void invokeLifecycleHandlers(TreeMap<Integer, List<LifecycleHandler>> li
                         lifecycleState.description, currentLifecyclePhase
                 ));
             } catch (TimeoutException e) {
-            	final long lifecyclePhaseTimeoutInSeconds = TimeUnit.SECONDS.convert(lifecyclePhaseTimeout, lifecyclePhaseTimeunit);
+                final long lifecyclePhaseTimeoutInSeconds = TimeUnit.SECONDS.convert(lifecyclePhaseTimeout,
+                                                                                     lifecyclePhaseTimeunit);
                 logger.warn(String.format(
                         "Timed out during %s phase [%d] after %d second(s). Proceeding to following phase",
                         lifecycleState.description, currentLifecyclePhase, lifecyclePhaseTimeoutInSeconds

config/src/test/java/org/axonframework/config/DefaultConfigurerLifecycleOperationsTest.java

@@ -17,6 +17,7 @@
 package org.axonframework.config;
 
 import org.axonframework.common.AxonConfigurationException;
+import org.axonframework.lifecycle.Lifecycle;
 import org.axonframework.lifecycle.LifecycleHandlerInvocationException;
 import org.junit.jupiter.api.*;
 import org.mockito.*;
@@ -25,15 +26,16 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.ReentrantLock;
+import javax.annotation.Nonnull;
 
 import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.Mockito.*;
 
 /**
  * Test class validating the workings of the lifecycle operations registered and invoked on the {@link Configurer},
- * {@link Configuration} and the {@link DefaultConfigurer} implementation. As such, operations like the {@link
- * Configuration#onStart(int, LifecycleHandler)}, {@link Configuration#onShutdown(int, LifecycleHandler)}, {@link
- * Configurer#start()}, {@link Configuration#start()} and {@link Configuration#shutdown()} will be tested.
+ * {@link Configuration} and the {@link DefaultConfigurer} implementation. As such, operations like the
+ * {@link Configuration#onStart(int, LifecycleHandler)}, {@link Configuration#onShutdown(int, LifecycleHandler)},
+ * {@link Configurer#start()}, {@link Configuration#start()} and {@link Configuration#shutdown()} will be tested.
  *
  * @author Steven van Beelen
  */
@@ -468,6 +470,40 @@ void configureLifecyclePhaseTimeoutWithNullTimeUnitThrowsAxonConfigurationExcept
         );
     }
 
+    @Test
+    void decoratedComponentsHaveLifeCycleHandlersRegistered() {
+
+        LifeCycleComponent original = spy(new LifeCycleComponent());
+        LifeCycleComponent decorator1 = spy(new LifeCycleComponent());
+        LifeCycleComponent decorator2 = spy(new LifeCycleComponent());
+        LifeCycleComponent decorator3 = spy(new LifeCycleComponent());
+
+        Configuration testSubject = DefaultConfigurer
+                .defaultConfiguration()
+                .registerComponent(LifeCycleComponent.class, config -> original)
+                .registerComponentDecorator(LifeCycleComponent.class, (config, o) -> decorator1)
+                .registerComponentDecorator(LifeCycleComponent.class, (config, o) -> decorator2)
+                .registerComponentDecorator(LifeCycleComponent.class, (config, o) -> decorator3)
+                .buildConfiguration();
+
+        testSubject.getComponent(LifeCycleComponent.class);
+
+        testSubject.start();
+
+        InOrder lifecycleOrder =
+                inOrder(original, decorator1, decorator2, decorator3);
+        lifecycleOrder.verify(original).registerLifecycleHandlers(any());
+        lifecycleOrder.verify(decorator1).registerLifecycleHandlers(any());
+        lifecycleOrder.verify(decorator2).registerLifecycleHandlers(any());
+        lifecycleOrder.verify(decorator3).registerLifecycleHandlers(any());
+
+        assertTrue(original.isInvoked());
+        assertTrue(decorator1.isInvoked());
+        assertTrue(decorator2.isInvoked());
+        assertTrue(decorator3.isInvoked());
+    }
+
+
     private static class LifecycleManagedInstance {
 
         private final ReentrantLock lock;
@@ -534,9 +570,24 @@ public void failingStart() {
         }
     }
 
+    static class LifeCycleComponent implements Lifecycle {
+        private AtomicBoolean invoked = new AtomicBoolean(false);
+
+        @Override
+        public void registerLifecycleHandlers(@Nonnull LifecycleRegistry lifecycle) {
+            lifecycle.onStart(0, () -> {
+                invoked.set(true);
+            });
+        }
+
+        public boolean isInvoked() {
+            return invoked.get();
+        }
+    }
+
     @FunctionalInterface
     private interface LifecycleRegistration {
 
         void registerLifecycleHandler(Configuration configuration, int phase, Runnable lifecycleHandler);
     }
-}
+}

config/src/test/java/org/axonframework/config/DefaultConfigurerTest.java

@@ -620,13 +620,31 @@ void customConfiguredSpanFactory() {
         SpanFactory custom = mock(SpanFactory.class);
 
         SpanFactory result = DefaultConfigurer.defaultConfiguration()
-                        .configureSpanFactory((config) -> custom)
-                                 .buildConfiguration()
-                                 .spanFactory();
+                                              .configureSpanFactory((config) -> custom)
+                                              .buildConfiguration()
+                                              .spanFactory();
 
         assertSame(custom, result);
     }
 
+    @Test
+    void canDecorateComponentsSuchAsSpanFactory() {
+        SpanFactory custom = mock(SpanFactory.class);
+        SpanFactory decorated = mock(SpanFactory.class);
+
+        SpanFactory result = DefaultConfigurer.defaultConfiguration()
+                                              .configureSpanFactory((config) -> custom)
+                                              .registerComponentDecorator(SpanFactory.class,
+                                                                          (configuration, component) -> {
+                                                                              assertSame(custom, component);
+                                                                              return decorated;
+                                                                          })
+                                              .buildConfiguration()
+                                              .spanFactory();
+
+        assertSame(decorated, result);
+    }
+
     @Test
     void defaultConfiguredScopeAwareProvider() {
         ScopeAwareProvider result = DefaultConfigurer.defaultConfiguration()
@@ -637,8 +655,8 @@ void defaultConfiguredScopeAwareProvider() {
     }
 
     @Test
-    void whenStubAggregateRegisteredWithRegisterMessageHandler_thenRightThingsCalled(){
-        Configurer configurer =  spy(DefaultConfigurer.defaultConfiguration());
+    void whenStubAggregateRegisteredWithRegisterMessageHandler_thenRightThingsCalled() {
+        Configurer configurer = spy(DefaultConfigurer.defaultConfiguration());
         configurer.registerMessageHandler(c -> new StubAggregate());
 
         verify(configurer, times(1)).registerCommandHandler(any());
@@ -647,8 +665,8 @@ void whenStubAggregateRegisteredWithRegisterMessageHandler_thenRightThingsCalled
     }
 
     @Test
-    void whenQueryHandlerRegisteredWithRegisterMessageHandler_thenRightThingsCalled(){
-        Configurer configurer =  spy(DefaultConfigurer.defaultConfiguration());
+    void whenQueryHandlerRegisteredWithRegisterMessageHandler_thenRightThingsCalled() {
+        Configurer configurer = spy(DefaultConfigurer.defaultConfiguration());
         configurer.registerMessageHandler(c -> new StubQueryHandler());
 
         verify(configurer, never()).registerCommandHandler(any());

messaging/src/main/java/org/axonframework/tracing/SpanFactory.java

@@ -132,11 +132,11 @@ Span createHandlerSpan(Supplier<String> operationNameSupplier, Message<?> parent
      * {@link SpanUtils#determineMessageName(Message)}.
      *
      * @param operationNameSupplier Supplier of the operation's name.
-     * @param parentMessage  The message that is being handled.
+     * @param message  The message that is being sent.
      * @param linkedSiblings Optional parameter, providing this will link the provided messages to the current.
      * @return The created {@link Span}.
      */
-    Span createDispatchSpan(Supplier<String> operationNameSupplier, Message<?> parentMessage, Message<?>... linkedSiblings);
+    Span createDispatchSpan(Supplier<String> operationNameSupplier, Message<?> message, Message<?>... linkedSiblings);
 
     /**
      * Creates a new {@link Span} linked to the currently active span. This is useful for tracing different parts of