Merge pull request #6229 from pferraro/WFCORE-7039

WFCORE-7039 Add Installer.Builder methods to run arbitrary tasks on service start/stop/removal.

service/src/main/java/org/wildfly/service/Installer.java

@@ -5,14 +5,20 @@
 package org.wildfly.service;
 
 import java.util.Collection;
+import java.util.Collections;
+import java.util.EnumMap;
+import java.util.EnumSet;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
 import org.jboss.msc.Service;
+import org.jboss.msc.service.LifecycleEvent;
+import org.jboss.msc.service.LifecycleListener;
 import org.jboss.msc.service.ServiceBuilder;
 import org.jboss.msc.service.ServiceController;
 import org.jboss.msc.service.ServiceController.Mode;
@@ -80,6 +86,27 @@ public void accept(SB builder) {
          */
         B asActive();
 
+        /**
+         * Configures the specified task to be run after the installed service is started.
+         * @param task a task to execute upon service start
+         * @return a reference to this builder
+         */
+        B onStart(Runnable task);
+
+        /**
+         * Configures the specified task to be run after the installed service is stopped.
+         * @param task a task to execute upon service stop
+         * @return a reference to this builder
+         */
+        B onStop(Runnable task);
+
+        /**
+         * Configures the specified task to be run upon removal of the installed service.
+         * @param task a task to execute upon service removal
+         * @return a reference to this builder
+         */
+        B onRemove(Runnable task);
+
         /**
          * Builds a service installer.
          * @return a service installer
@@ -164,7 +191,14 @@ interface Configuration<SB extends DSB, DSB extends ServiceBuilder<?>> {
          * @return a service factory
          */
         Function<SB, Service> getServiceFactory();
-     }
+
+        /**
+         * Returns tasks to be run per lifecycle event.
+         * The returned map is either fully populated, or an empty map, if this service has no lifecycle tasks.
+         * @return a potentially empty map of tasks to be run per lifecycle event.
+         */
+        Map<LifecycleEvent, Collection<Runnable>> getLifecycleTasks();
+    }
 
     /**
      * Generic abstract installer implementation that installs a {@link UnaryService}.
@@ -180,25 +214,38 @@ class DefaultInstaller<ST extends ServiceTarget, SB extends DSB, DSB extends Ser
         private final ServiceController.Mode mode;
         private final Consumer<DSB> dependency;
         private final Function<SB, Service> serviceFactory;
+        private final Map<LifecycleEvent, Collection<Runnable>> lifecycleTasks;
 
         protected DefaultInstaller(Installer.Configuration<SB, DSB> config, Function<ST, SB> serviceBuilderFactory) {
             this.serviceBuilderFactory = serviceBuilderFactory;
             this.serviceFactory = config.getServiceFactory();
             this.mode = config.getInitialMode();
             this.dependency = config.getDependency();
+            this.lifecycleTasks = config.getLifecycleTasks();
         }
 
         @Override
         public ServiceController<?> install(ST target) {
             SB builder = this.serviceBuilderFactory.apply(target);
             this.dependency.accept(builder);
+            // N.B. map of tasks is either empty or fully populated
+            if (!this.lifecycleTasks.isEmpty()) {
+                Map<LifecycleEvent, Collection<Runnable>> tasks = this.lifecycleTasks;
+                builder.addListener(new LifecycleListener() {
+                    @Override
+                    public void handleEvent(ServiceController<?> controller, LifecycleEvent event) {
+                        tasks.get(event).forEach(Runnable::run);
+                    }
+                });
+            }
             return builder.setInstance(this.serviceFactory.apply(builder)).setInitialMode(this.mode).install();
         }
     }
 
     abstract class AbstractBuilder<B, I extends Installer<ST>, ST extends ServiceTarget, SB extends DSB, DSB extends ServiceBuilder<?>> implements Installer.Builder<B, I, ST, DSB>, Installer.Configuration<SB, DSB> {
         private volatile ServiceController.Mode mode = ServiceController.Mode.ON_DEMAND;
         private volatile Consumer<DSB> dependency = Functions.discardingConsumer();
+        private volatile Map<LifecycleEvent, List<Runnable>> lifecycleTasks = Map.of();
 
         protected abstract B builder();
 
@@ -220,6 +267,43 @@ public B requires(Consumer<DSB> dependency) {
             return this.builder();
         }
 
+        @Override
+        public B onStart(Runnable task) {
+            return this.onEvent(LifecycleEvent.UP, task);
+        }
+
+        @Override
+        public B onStop(Runnable task) {
+            return this.onEvent(LifecycleEvent.DOWN, task);
+        }
+
+        @Override
+        public B onRemove(Runnable task) {
+            return this.onEvent(LifecycleEvent.REMOVED, task);
+        }
+
+        private B onEvent(LifecycleEvent event, Runnable task) {
+            if (this.lifecycleTasks.isEmpty()) {
+                // Create EnumMap lazily, when needed
+                this.lifecycleTasks = new EnumMap<>(LifecycleEvent.class);
+                for (LifecycleEvent e : EnumSet.allOf(LifecycleEvent.class)) {
+                    this.lifecycleTasks.put(e, (e == event) ? List.of(task) : List.of());
+                }
+            } else {
+                List<Runnable> tasks = this.lifecycleTasks.get(event);
+                if (tasks.isEmpty()) {
+                    this.lifecycleTasks.put(event, List.of(task));
+                } else {
+                    if (tasks.size() == 1) {
+                        tasks = new LinkedList<>(tasks);
+                        this.lifecycleTasks.put(event, tasks);
+                    }
+                    tasks.add(task);
+                }
+            }
+            return this.builder();
+        }
+
         @Override
         public Mode getInitialMode() {
             return this.mode;
@@ -229,6 +313,17 @@ public Mode getInitialMode() {
         public Consumer<DSB> getDependency() {
             return this.dependency;
         }
+
+        @Override
+        public Map<LifecycleEvent, Collection<Runnable>> getLifecycleTasks() {
+            // Return empty map or fully unmodifiable copy
+            if (this.lifecycleTasks.isEmpty()) return Map.of();
+            Map<LifecycleEvent, Collection<Runnable>> result = new EnumMap<>(LifecycleEvent.class);
+            for (Map.Entry<LifecycleEvent, List<Runnable>> entry : this.lifecycleTasks.entrySet()) {
+                result.put(entry.getKey(), List.copyOf(entry.getValue()));
+            }
+            return Collections.unmodifiableMap(result);
+        }
     }
 
     abstract class AbstractNullaryBuilder<B, I extends Installer<ST>, ST extends ServiceTarget, SB extends DSB, DSB extends ServiceBuilder<?>> extends AbstractBuilder<B, I, ST, SB, DSB> implements Function<SB, Service> {

service/src/main/java/org/wildfly/service/capture/FunctionExecutor.java

@@ -32,7 +32,7 @@ public <R, E extends Exception> R execute(ExceptionFunction<V, R, E> function) t
     }
 
     /**
-     * Executes the given function.
+     * Executes the specified function, using a value provided by an associated {@link ValueRegistry}.
      * @param <R> the return type
      * @param <E> the exception type
      * @param function a function to execute

service/src/main/java/org/wildfly/service/capture/FunctionExecutorRegistry.java

@@ -12,9 +12,9 @@
  */
 public interface FunctionExecutorRegistry<K, V> {
     /**
-     * Returns the executor for the specified key.
+     * Returns the executor for the specified key, if one exists.
      * @param key a registry key
-     * @return an executor, or null, if no such executor exists in the registry
+     * @return the executor for the specified key, or null, if no such executor exists in the registry
      */
     FunctionExecutor<V> getExecutor(K key);
 }

service/src/test/java/org/wildfly/service/DefaultServiceInstallerTestCase.java

@@ -7,11 +7,18 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 
+import java.util.EnumMap;
+import java.util.EnumSet;
+import java.util.Map;
 import java.util.function.Consumer;
 import java.util.function.Function;
 
 import org.jboss.msc.Service;
+import org.jboss.msc.service.LifecycleEvent;
+import org.jboss.msc.service.LifecycleListener;
 import org.jboss.msc.service.ServiceBuilder;
 import org.jboss.msc.service.ServiceController;
 import org.jboss.msc.service.ServiceName;
@@ -38,12 +45,29 @@ public void test() throws StartException {
         ServiceDependency<Object> dependency = mock(ServiceDependency.class);
         Consumer<Object> captor = mock(Consumer.class);
         Consumer<Object> combinedCaptor = mock(Consumer.class);
-        Consumer<Object> startTask = mock(Consumer.class);
-        Consumer<Object> stopTask = mock(Consumer.class);
+        Consumer<Object> startConsumer = mock(Consumer.class);
+        Consumer<Object> stopConsumer = mock(Consumer.class);
         ArgumentCaptor<Service> service = ArgumentCaptor.forClass(Service.class);
         ArgumentCaptor<ServiceController.Mode> mode = ArgumentCaptor.forClass(ServiceController.Mode.class);
+        Map<LifecycleEvent, Runnable> tasks = new EnumMap<>(LifecycleEvent.class);
+        for (LifecycleEvent event : EnumSet.of(LifecycleEvent.UP, LifecycleEvent.DOWN, LifecycleEvent.REMOVED)) {
+            tasks.put(event, mock(Runnable.class));
+        }
 
-        ServiceInstaller installer = ServiceInstaller.builder(mapper, dependency).provides(name).requires(dependency).withCaptor(captor).onStart(startTask).onStop(stopTask).build();
+        ArgumentCaptor<LifecycleListener> capturedListener = ArgumentCaptor.forClass(LifecycleListener.class);
+
+        doReturn(builder).when(builder).addListener(capturedListener.capture());
+
+        ServiceInstaller installer = ServiceInstaller.builder(mapper, dependency)
+                .provides(name)
+                .requires(dependency)
+                .withCaptor(captor)
+                .onStart(startConsumer)
+                .onStop(stopConsumer)
+                .onStart(tasks.get(LifecycleEvent.UP))
+                .onStop(tasks.get(LifecycleEvent.DOWN))
+                .onRemove(tasks.get(LifecycleEvent.REMOVED))
+                .build();
 
         doReturn(builder).when(target).addService();
         doReturn(injector).when(builder).provides(name);
@@ -59,7 +83,26 @@ public void test() throws StartException {
 
         verify(dependency).accept(builder);
         verify(builder).provides(name);
+        for (Runnable task : tasks.values()) {
+            verifyNoInteractions(task);
+        }
+
+        LifecycleListener listener = capturedListener.getValue();
+
+        Assert.assertNotNull(listener);
+
+        for (LifecycleEvent event : EnumSet.allOf(LifecycleEvent.class)) {
+            listener.handleEvent(controller, event);
+
+            for (Map.Entry<LifecycleEvent, Runnable> entry : tasks.entrySet()) {
+                if (event == entry.getKey()) {
+                    verify(entry.getValue()).run();
+                } else {
+                    verifyNoMoreInteractions(entry.getValue());
+                }
+            }
+        }
 
-        DefaultServiceTestCase.test(service.getValue(), "value", "mappedValue", combinedCaptor, mapper, dependency, startTask, stopTask);
+        DefaultServiceTestCase.test(service.getValue(), "value", "mappedValue", combinedCaptor, mapper, dependency, startConsumer, stopConsumer);
     }
 }