Additional async handler fixes

impl/src/main/java/org/jboss/weld/invokable/AsyncHandlerRegistry.java

@@ -68,15 +68,15 @@ private void validateAndRegisterReturnType(AsyncHandler.ReturnType<?> handler) {
         Class<?> handlerClass = handler.getClass();
         validateDirectImplementation(handlerClass, AsyncHandler.ReturnType.class);
         Class<?> asyncType = extractAsyncType(handlerClass, AsyncHandler.ReturnType.class);
-        checkDuplicate(asyncType, handlerClass);
+        checkDuplicate(asyncType, handlerClass, true);
         handlers.put(asyncType, HandlerInfo.returnType(handler, asyncType));
     }
 
     private void validateAndRegisterParameterType(AsyncHandler.ParameterType<?> handler) {
         Class<?> handlerClass = handler.getClass();
         validateDirectImplementation(handlerClass, AsyncHandler.ParameterType.class);
         Class<?> asyncType = extractAsyncType(handlerClass, AsyncHandler.ParameterType.class);
-        checkDuplicate(asyncType, handlerClass);
+        checkDuplicate(asyncType, handlerClass, false);
         handlers.put(asyncType, HandlerInfo.parameterType(handler, asyncType));
     }
 
@@ -89,9 +89,12 @@ private void validateDirectImplementation(Class<?> handlerClass, Class<?> target
         throw InvokerLogger.LOG.asyncHandlerIndirectImplementation(handlerClass);
     }
 
-    private void checkDuplicate(Class<?> asyncType, Class<?> handlerClass) {
+    private void checkDuplicate(Class<?> asyncType, Class<?> handlerClass, boolean isReturnType) {
         HandlerInfo existing = handlers.get(asyncType);
         if (existing != null && !existing.isBuiltin()) {
+            if (existing.getHandlerClass() == handlerClass && existing.isReturnType() != isReturnType) {
+                throw InvokerLogger.LOG.asyncHandlerBothKinds(handlerClass, asyncType);
+            }
             throw InvokerLogger.LOG.asyncHandlerDuplicate(asyncType, handlerClass);
         }
     }
@@ -186,23 +189,25 @@ public static class HandlerInfo {
         private final AsyncHandler.ReturnType<?> returnTypeHandler;
         private final AsyncHandler.ParameterType<?> parameterTypeHandler;
         private final Class<?> asyncType;
+        private final Class<?> handlerClass;
         private final boolean isReturnType;
         private boolean builtin;
 
         static HandlerInfo returnType(AsyncHandler.ReturnType<?> handler, Class<?> asyncType) {
-            return new HandlerInfo(handler, null, asyncType, true);
+            return new HandlerInfo(handler, null, asyncType, handler.getClass(), true);
         }
 
         static HandlerInfo parameterType(AsyncHandler.ParameterType<?> handler, Class<?> asyncType) {
-            return new HandlerInfo(null, handler, asyncType, false);
+            return new HandlerInfo(null, handler, asyncType, handler.getClass(), false);
         }
 
         private HandlerInfo(AsyncHandler.ReturnType<?> returnTypeHandler,
                 AsyncHandler.ParameterType<?> parameterTypeHandler,
-                Class<?> asyncType, boolean isReturnType) {
+                Class<?> asyncType, Class<?> handlerClass, boolean isReturnType) {
             this.returnTypeHandler = returnTypeHandler;
             this.parameterTypeHandler = parameterTypeHandler;
             this.asyncType = asyncType;
+            this.handlerClass = handlerClass;
             this.isReturnType = isReturnType;
         }
 
@@ -220,6 +225,10 @@ public Class<?> getAsyncType() {
             return asyncType;
         }
 
+        public Class<?> getHandlerClass() {
+            return handlerClass;
+        }
+
         public boolean isReturnType() {
             return isReturnType;
         }

impl/src/main/java/org/jboss/weld/invokable/AsyncInvokerImpl.java

@@ -1,6 +1,7 @@
 package org.jboss.weld.invokable;
 
 import java.lang.invoke.MethodHandle;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import jakarta.enterprise.inject.build.compatible.spi.InvokerInfo;
 import jakarta.enterprise.invoke.AsyncHandler;
@@ -18,9 +19,11 @@
  * ({@code foldArguments} is skipped). Instead, this invoker creates
  * {@code CleanupActions} externally and passes it as the first parameter to
  * {@code mh.invoke()}. This allows {@code transformArgument} to receive
- * {@code ca::cleanup} as the completion callback before the method runs, which is
- * required for correct behavior when the async parameter completes synchronously
- * during the method body.
+ * the completion callback before the method runs.
+ * <p>
+ * Cleanup of dependent beans is deferred via {@link DeferredCleanup} so that
+ * it only happens once both the method has returned and the completion callback
+ * has been called. See {@link DeferredCleanup} for details.
  * <p>
  * When no cleanup is needed, the chain has no {@code CleanupActions} at all and
  * the handler receives a no-op completion callback.
@@ -44,11 +47,14 @@ class AsyncInvokerImpl<T, R> implements Invoker<T, R>, InvokerInfo {
     public R invoke(T instance, Object[] arguments) throws Exception {
         Runnable completion;
         CleanupActions ca;
+        DeferredCleanup deferred;
         if (requiresCleanup) {
             ca = new CleanupActions();
-            completion = ca::cleanup;
+            deferred = new DeferredCleanup(ca);
+            completion = deferred;
         } else {
             ca = null;
+            deferred = null;
             completion = () -> {
             };
         }
@@ -63,11 +69,92 @@ public R invoke(T instance, Object[] arguments) throws Exception {
             } else {
                 result = (R) mh.invoke(instance, arguments);
             }
+            if (deferred != null) {
+                deferred.methodReturned();
+            }
+            // completion is still valid here; the state machine handles repeated calls
             return (R) parameterTypeHandler.transformReturnValue(result, completion);
         } catch (ValueCarryingException e) {
+            if (deferred != null) {
+                deferred.forceCleanup();
+            }
             return (R) e.getMethodReturnValue();
         } catch (Throwable e) {
+            if (deferred != null) {
+                deferred.forceCleanup();
+            }
             throw SneakyThrow.sneakyThrow(e);
         }
     }
+
+    /**
+     * Coordinates cleanup of dependent beans between the completion callback
+     * (fired by the async handler) and the method return.
+     * <p>
+     * Cleanup must only run when <em>both</em> the completion callback has been
+     * called and the target method has returned. This prevents premature destruction
+     * of dependent beans when the completion callback fires synchronously during the
+     * method body (e.g., the method calls {@code asyncParam.resume()} directly).
+     * <p>
+     * Uses a four-state machine driven by {@link AtomicInteger} CAS operations:
+     *
+     * <pre>
+     *   PENDING (0)  ──run()──►  COMPLETION_SIGNALED (1)
+     *       │                             │
+     *   methodReturned()              methodReturned()
+     *       │                             │
+     *       ▼                             ▼
+     *   METHOD_RETURNED (2)  ──run()──►  DONE (3) → cleanup runs
+     * </pre>
+     * <ul>
+     * <li><b>Sync case</b> (completion fires during method): transitions
+     * 0→1, then 1→3 on method return; cleanup runs after the method.</li>
+     * <li><b>Async case</b> (completion fires after method): transitions
+     * 0→2 on method return, then 2→3 on completion; cleanup runs on completion.</li>
+     * <li><b>Exception case</b>: {@link #forceCleanup()} sets state to 3 directly.
+     * The MH chain's {@code tryFinally/runExceptionOnly} may also call
+     * {@code ca.cleanup()}, but that is safe because {@link CleanupActions#cleanup()}
+     * clears its internal lists on first call.</li>
+     * </ul>
+     * Thread safety is ensured by CAS; the async completion callback may fire
+     * on a different thread than the invoker.
+     */
+    private static class DeferredCleanup implements Runnable {
+        private static final int PENDING = 0;
+        private static final int COMPLETION_SIGNALED = 1;
+        private static final int METHOD_RETURNED = 2;
+        private static final int DONE = 3;
+
+        private final CleanupActions ca;
+        private final AtomicInteger state = new AtomicInteger(PENDING);
+
+        DeferredCleanup(CleanupActions ca) {
+            this.ca = ca;
+        }
+
+        @Override
+        public void run() {
+            if (state.compareAndSet(PENDING, COMPLETION_SIGNALED)) {
+                return;
+            }
+            if (state.compareAndSet(METHOD_RETURNED, DONE)) {
+                ca.cleanup();
+            }
+        }
+
+        void methodReturned() {
+            if (state.compareAndSet(PENDING, METHOD_RETURNED)) {
+                return;
+            }
+            if (state.compareAndSet(COMPLETION_SIGNALED, DONE)) {
+                ca.cleanup();
+            }
+        }
+
+        void forceCleanup() {
+            if (state.getAndSet(DONE) != DONE) {
+                ca.cleanup();
+            }
+        }
+    }
 }

impl/src/main/java/org/jboss/weld/logging/InvokerLogger.java

@@ -114,4 +114,7 @@ public interface InvokerLogger extends WeldLogger {
 
     @Message(id = 2029, value = "Unhandled primitive type: {0}", format = Format.MESSAGE_FORMAT)
     RuntimeException unhandledPrimitiveType(Object primitive);
+
+    @Message(id = 2030, value = "AsyncHandler {0} implements both ReturnType and ParameterType for the same async type {1}", format = Format.MESSAGE_FORMAT)
+    DefinitionException asyncHandlerBothKinds(Object handlerClass, Object asyncType);
 }

tests-arquillian/src/test/java/org/jboss/weld/tests/invokable/async/broken/AnotherParamTypeHandler.java

@@ -0,0 +1,10 @@
+package org.jboss.weld.tests.invokable.async.broken;
+
+import jakarta.enterprise.invoke.AsyncHandler;
+
+public class AnotherParamTypeHandler<T> implements AsyncHandler.ParameterType<MyAsyncType<T>> {
+    @Override
+    public MyAsyncType<T> transformArgument(MyAsyncType<T> original, Runnable completion) {
+        return original;
+    }
+}

tests-arquillian/src/test/java/org/jboss/weld/tests/invokable/async/broken/AnotherReturnTypeHandler.java

@@ -0,0 +1,10 @@
+package org.jboss.weld.tests.invokable.async.broken;
+
+import jakarta.enterprise.invoke.AsyncHandler;
+
+public class AnotherReturnTypeHandler<T> implements AsyncHandler.ReturnType<MyAsyncType<T>> {
+    @Override
+    public MyAsyncType<T> transform(MyAsyncType<T> original, Runnable completion) {
+        return original;
+    }
+}

tests-arquillian/src/test/java/org/jboss/weld/tests/invokable/async/broken/BothInterfacesHandler.java

@@ -2,14 +2,15 @@
 
 import jakarta.enterprise.invoke.AsyncHandler;
 
-public class BothInterfacesHandler<T> implements AsyncHandler.ReturnType<T>, AsyncHandler.ParameterType<T> {
+public class BothInterfacesHandler<T>
+        implements AsyncHandler.ReturnType<MyAsyncType<T>>, AsyncHandler.ParameterType<MyAsyncType<T>> {
     @Override
-    public T transform(T original, Runnable completion) {
+    public MyAsyncType<T> transform(MyAsyncType<T> original, Runnable completion) {
         return original;
     }
 
     @Override
-    public T transformArgument(T original, Runnable completion) {
+    public MyAsyncType<T> transformArgument(MyAsyncType<T> original, Runnable completion) {
         return original;
     }
 }

tests-arquillian/src/test/java/org/jboss/weld/tests/invokable/async/broken/BothInterfacesHandlerTest.java

@@ -21,7 +21,7 @@ public class BothInterfacesHandlerTest {
     public static Archive<?> deploy() {
         return ShrinkWrap.create(BeanArchive.class,
                 Utils.getDeploymentNameAsHash(BothInterfacesHandlerTest.class))
-                .addClass(BothInterfacesHandler.class)
+                .addClasses(BothInterfacesHandler.class, MyAsyncType.class)
                 .addAsServiceProvider(AsyncHandler.ReturnType.class, BothInterfacesHandler.class)
                 .addAsServiceProvider(AsyncHandler.ParameterType.class, BothInterfacesHandler.class);
     }

tests-arquillian/src/test/java/org/jboss/weld/tests/invokable/async/broken/DuplicateParamTypeHandler.java

@@ -0,0 +1,10 @@
+package org.jboss.weld.tests.invokable.async.broken;
+
+import jakarta.enterprise.invoke.AsyncHandler;
+
+public class DuplicateParamTypeHandler<T> implements AsyncHandler.ParameterType<MyAsyncType<T>> {
+    @Override
+    public MyAsyncType<T> transformArgument(MyAsyncType<T> original, Runnable completion) {
+        return original;
+    }
+}

tests-arquillian/src/test/java/org/jboss/weld/tests/invokable/async/broken/DuplicateParamTypeHandlerTest.java

@@ -0,0 +1,32 @@
+package org.jboss.weld.tests.invokable.async.broken;
+
+import jakarta.enterprise.inject.spi.DeploymentException;
+import jakarta.enterprise.invoke.AsyncHandler;
+
+import org.jboss.arquillian.container.test.api.Deployment;
+import org.jboss.arquillian.container.test.api.ShouldThrowException;
+import org.jboss.arquillian.junit.Arquillian;
+import org.jboss.shrinkwrap.api.Archive;
+import org.jboss.shrinkwrap.api.BeanArchive;
+import org.jboss.shrinkwrap.api.ShrinkWrap;
+import org.jboss.weld.test.util.Utils;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(Arquillian.class)
+public class DuplicateParamTypeHandlerTest {
+
+    @Deployment
+    @ShouldThrowException(DeploymentException.class)
+    public static Archive<?> deploy() {
+        return ShrinkWrap.create(BeanArchive.class,
+                Utils.getDeploymentNameAsHash(DuplicateParamTypeHandlerTest.class))
+                .addClasses(DuplicateParamTypeHandler.class, AnotherParamTypeHandler.class, MyAsyncType.class)
+                .addAsServiceProvider(AsyncHandler.ParameterType.class,
+                        DuplicateParamTypeHandler.class, AnotherParamTypeHandler.class);
+    }
+
+    @Test
+    public void trigger() {
+    }
+}

tests-arquillian/src/test/java/org/jboss/weld/tests/invokable/async/broken/DuplicateReturnTypeHandler.java

@@ -0,0 +1,10 @@
+package org.jboss.weld.tests.invokable.async.broken;
+
+import jakarta.enterprise.invoke.AsyncHandler;
+
+public class DuplicateReturnTypeHandler<T> implements AsyncHandler.ReturnType<MyAsyncType<T>> {
+    @Override
+    public MyAsyncType<T> transform(MyAsyncType<T> original, Runnable completion) {
+        return original;
+    }
+}

tests-arquillian/src/test/java/org/jboss/weld/tests/invokable/async/broken/DuplicateReturnTypeHandlerTest.java

@@ -0,0 +1,32 @@
+package org.jboss.weld.tests.invokable.async.broken;
+
+import jakarta.enterprise.inject.spi.DeploymentException;
+import jakarta.enterprise.invoke.AsyncHandler;
+
+import org.jboss.arquillian.container.test.api.Deployment;
+import org.jboss.arquillian.container.test.api.ShouldThrowException;
+import org.jboss.arquillian.junit.Arquillian;
+import org.jboss.shrinkwrap.api.Archive;
+import org.jboss.shrinkwrap.api.BeanArchive;
+import org.jboss.shrinkwrap.api.ShrinkWrap;
+import org.jboss.weld.test.util.Utils;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(Arquillian.class)
+public class DuplicateReturnTypeHandlerTest {
+
+    @Deployment
+    @ShouldThrowException(DeploymentException.class)
+    public static Archive<?> deploy() {
+        return ShrinkWrap.create(BeanArchive.class,
+                Utils.getDeploymentNameAsHash(DuplicateReturnTypeHandlerTest.class))
+                .addClasses(DuplicateReturnTypeHandler.class, AnotherReturnTypeHandler.class, MyAsyncType.class)
+                .addAsServiceProvider(AsyncHandler.ReturnType.class,
+                        DuplicateReturnTypeHandler.class, AnotherReturnTypeHandler.class);
+    }
+
+    @Test
+    public void trigger() {
+    }
+}

tests-arquillian/src/test/java/org/jboss/weld/tests/invokable/async/broken/MyAsyncType.java

@@ -0,0 +1,4 @@
+package org.jboss.weld.tests.invokable.async.broken;
+
+public interface MyAsyncType<T> {
+}

tests-arquillian/src/test/java/org/jboss/weld/tests/invokable/async/broken/TypeVariableAsyncHandler.java

@@ -0,0 +1,10 @@
+package org.jboss.weld.tests.invokable.async.broken;
+
+import jakarta.enterprise.invoke.AsyncHandler;
+
+public class TypeVariableAsyncHandler<T> implements AsyncHandler.ReturnType<T> {
+    @Override
+    public T transform(T original, Runnable completion) {
+        return original;
+    }
+}

tests-arquillian/src/test/java/org/jboss/weld/tests/invokable/async/broken/TypeVariableAsyncHandlerTest.java

@@ -0,0 +1,31 @@
+package org.jboss.weld.tests.invokable.async.broken;
+
+import jakarta.enterprise.inject.spi.DefinitionException;
+import jakarta.enterprise.invoke.AsyncHandler;
+
+import org.jboss.arquillian.container.test.api.Deployment;
+import org.jboss.arquillian.container.test.api.ShouldThrowException;
+import org.jboss.arquillian.junit.Arquillian;
+import org.jboss.shrinkwrap.api.Archive;
+import org.jboss.shrinkwrap.api.BeanArchive;
+import org.jboss.shrinkwrap.api.ShrinkWrap;
+import org.jboss.weld.test.util.Utils;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(Arquillian.class)
+public class TypeVariableAsyncHandlerTest {
+
+    @Deployment
+    @ShouldThrowException(DefinitionException.class)
+    public static Archive<?> deploy() {
+        return ShrinkWrap.create(BeanArchive.class,
+                Utils.getDeploymentNameAsHash(TypeVariableAsyncHandlerTest.class))
+                .addClass(TypeVariableAsyncHandler.class)
+                .addAsServiceProvider(AsyncHandler.ReturnType.class, TypeVariableAsyncHandler.class);
+    }
+
+    @Test
+    public void trigger() {
+    }
+}