Fix lambda continuous testing failure

core/deployment/src/main/java/io/quarkus/deployment/QuarkusAugmentor.java

@@ -32,6 +32,7 @@
 import io.quarkus.deployment.builditem.LiveReloadBuildItem;
 import io.quarkus.deployment.builditem.QuarkusBuildCloseablesBuildItem;
 import io.quarkus.deployment.builditem.RawCommandLineArgumentsBuildItem;
+import io.quarkus.deployment.builditem.RuntimeApplicationShutdownBuildItem;
 import io.quarkus.deployment.builditem.ShutdownContextBuildItem;
 import io.quarkus.deployment.pkg.builditem.BuildSystemTargetBuildItem;
 import io.quarkus.dev.spi.DevModeType;
@@ -128,6 +129,9 @@ public BuildResult run() throws Exception {
             for (Consumer<BuildChainBuilder> i : buildChainCustomizers) {
                 i.accept(chainBuilder);
             }
+            if (launchMode.isDevOrTest()) {
+                chainBuilder.addFinal(RuntimeApplicationShutdownBuildItem.class);
+            }
 
             final ArchiveRootBuildItem.Builder rootBuilder = ArchiveRootBuildItem.builder();
             if (root != null) {

core/deployment/src/main/java/io/quarkus/deployment/builditem/RuntimeApplicationShutdownBuildItem.java

@@ -0,0 +1,29 @@
+package io.quarkus.deployment.builditem;
+
+import org.jboss.logging.Logger;
+
+import io.quarkus.builder.item.MultiBuildItem;
+
+/**
+ * Build Item that can be used to queue shutdown tasks that are run when the runtime application shuts down.
+ *
+ * This is similar to {@link ShutdownContextBuildItem} however it applies to tasks on the 'build' side, so if a processor
+ * wants to close something after the application has completed this item lets it do this.
+ *
+ * This has no effect for production applications, and is only useful in dev/test mode. The main use case for this is
+ * for shutting down deployment side test utilities at the end of a test run.
+ */
+public final class RuntimeApplicationShutdownBuildItem extends MultiBuildItem {
+
+    private static final Logger log = Logger.getLogger(RuntimeApplicationShutdownBuildItem.class);
+
+    private final Runnable closeTask;
+
+    public RuntimeApplicationShutdownBuildItem(Runnable closeTask) {
+        this.closeTask = closeTask;
+    }
+
+    public Runnable getCloseTask() {
+        return closeTask;
+    }
+}

core/deployment/src/main/java/io/quarkus/deployment/dev/testing/TestTracingProcessor.java

@@ -25,6 +25,7 @@
 import io.quarkus.deployment.builditem.ServiceStartBuildItem;
 import io.quarkus.deployment.logging.LogCleanupFilterBuildItem;
 import io.quarkus.dev.spi.DevModeType;
+import io.quarkus.dev.testing.ContinuousTestingSharedStateManager;
 import io.quarkus.dev.testing.TracingHandler;
 import io.quarkus.gizmo.Gizmo;
 
@@ -82,6 +83,8 @@ void startTesting(TestConfig config, LiveReloadBuildItem liveReloadBuildItem,
                 testSupport.stop();
             }
         }
+        QuarkusClassLoader cl = (QuarkusClassLoader) Thread.currentThread().getContextClassLoader();
+        ((QuarkusClassLoader) cl.parent()).addCloseTask(ContinuousTestingSharedStateManager::reset);
     }
 
     @BuildStep(onlyIf = IsTest.class)

core/deployment/src/main/java/io/quarkus/runner/bootstrap/StartupActionImpl.java

@@ -31,6 +31,7 @@
 import io.quarkus.deployment.builditem.GeneratedClassBuildItem;
 import io.quarkus.deployment.builditem.GeneratedResourceBuildItem;
 import io.quarkus.deployment.builditem.MainClassBuildItem;
+import io.quarkus.deployment.builditem.RuntimeApplicationShutdownBuildItem;
 import io.quarkus.deployment.builditem.TransformedClassesBuildItem;
 import io.quarkus.deployment.configuration.RunTimeConfigurationGenerator;
 import io.quarkus.dev.appstate.ApplicationStateNotification;
@@ -106,6 +107,14 @@ public void run() {
                         if (ApplicationStateNotification.getState() == ApplicationStateNotification.State.INITIAL) {
                             ApplicationStateNotification.notifyStartupFailed(e);
                         }
+                    } finally {
+                        for (var i : buildResult.consumeMulti(RuntimeApplicationShutdownBuildItem.class)) {
+                            try {
+                                i.getCloseTask().run();
+                            } catch (Throwable t) {
+                                log.error("Failed to run close task", t);
+                            }
+                        }
                     }
                 }
             }, "Quarkus Main Thread");
@@ -173,6 +182,13 @@ public void accept(Integer integer) {
         } finally {
             runtimeClassLoader.close();
             Thread.currentThread().setContextClassLoader(old);
+            for (var i : buildResult.consumeMulti(RuntimeApplicationShutdownBuildItem.class)) {
+                try {
+                    i.getCloseTask().run();
+                } catch (Throwable t) {
+                    log.error("Failed to run close task", t);
+                }
+            }
         }
     }
 
@@ -230,6 +246,14 @@ public void close() throws IOException {
                         }
                     } finally {
                         ForkJoinClassLoading.setForkJoinClassLoader(ClassLoader.getSystemClassLoader());
+
+                        for (var i : buildResult.consumeMulti(RuntimeApplicationShutdownBuildItem.class)) {
+                            try {
+                                i.getCloseTask().run();
+                            } catch (Throwable t) {
+                                log.error("Failed to run close task", t);
+                            }
+                        }
                         if (curatedApplication.getQuarkusBootstrap().getMode() == QuarkusBootstrap.Mode.TEST &&
                                 !curatedApplication.getQuarkusBootstrap().isAuxiliaryApplication()) {
                             //for tests we just always shut down the curated application, as it is only used once

extensions/amazon-lambda-http/http-event-server/src/main/java/io/quarkus/amazon/lambda/runtime/MockHttpEventServer.java

@@ -80,6 +80,7 @@ public void handleHttpRequests(RoutingContext ctx) {
         try {
             byte[] mEvent = eventWriter.writeValueAsBytes(event);
             ctx.put(APIGatewayV2HTTPEvent.class.getName(), mEvent);
+            log.debugf("Putting message %s into the queue", requestId);
             queue.put(ctx);
         } catch (Exception e) {
             log.error("Publish failure", e);

extensions/amazon-lambda-rest/rest-event-server/src/main/java/io/quarkus/amazon/lambda/runtime/MockRestEventServer.java

@@ -104,6 +104,7 @@ public void handleHttpRequests(RoutingContext ctx) {
         try {
             byte[] mEvent = eventWriter.writeValueAsBytes(event);
             ctx.put(AwsProxyRequest.class.getName(), mEvent);
+            log.debugf("Putting message %s into the queue", requestId);
             queue.put(ctx);
         } catch (Exception e) {
             log.error("Publish failure", e);

extensions/amazon-lambda/common-deployment/src/main/java/io/quarkus/amazon/lambda/deployment/DevServicesLambdaProcessor.java

@@ -14,9 +14,11 @@
 import io.quarkus.deployment.IsNormal;
 import io.quarkus.deployment.annotations.BuildProducer;
 import io.quarkus.deployment.annotations.BuildStep;
+import io.quarkus.deployment.annotations.Produce;
 import io.quarkus.deployment.annotations.Record;
-import io.quarkus.deployment.builditem.DevServicesNativeConfigResultBuildItem;
+import io.quarkus.deployment.builditem.DevServicesConfigResultBuildItem;
 import io.quarkus.deployment.builditem.LaunchModeBuildItem;
+import io.quarkus.deployment.builditem.RuntimeApplicationShutdownBuildItem;
 import io.quarkus.deployment.builditem.ServiceStartBuildItem;
 import io.quarkus.runtime.LaunchMode;
 
@@ -47,12 +49,14 @@ private boolean legacyTestingEnabled() {
         }
     }
 
+    @Produce(ServiceStartBuildItem.class)
     @BuildStep(onlyIfNot = IsNormal.class)
     public void startEventServer(LaunchModeBuildItem launchMode,
             LambdaConfig config,
             Optional<EventServerOverrideBuildItem> override,
-            BuildProducer<DevServicesNativeConfigResultBuildItem> devServicePropertiesProducer,
-            BuildProducer<ServiceStartBuildItem> serviceStartBuildItemBuildProducer) throws Exception {
+            BuildProducer<DevServicesConfigResultBuildItem> devServicePropertiesProducer,
+            BuildProducer<RuntimeApplicationShutdownBuildItem> runtimeApplicationShutdownBuildItemBuildProducer)
+            throws Exception {
         if (!launchMode.getLaunchMode().isDevOrTest())
             return;
         if (legacyTestingEnabled())
@@ -73,9 +77,8 @@ public void startEventServer(LaunchModeBuildItem launchMode,
         startMode = launchMode.getLaunchMode();
         server.start(port);
         String baseUrl = "localhost:" + port + MockEventServer.BASE_PATH;
-        System.setProperty(AmazonLambdaApi.QUARKUS_INTERNAL_AWS_LAMBDA_TEST_API, baseUrl);
         devServicePropertiesProducer.produce(
-                new DevServicesNativeConfigResultBuildItem(AmazonLambdaApi.QUARKUS_INTERNAL_AWS_LAMBDA_TEST_API, baseUrl));
+                new DevServicesConfigResultBuildItem(AmazonLambdaApi.QUARKUS_INTERNAL_AWS_LAMBDA_TEST_API, baseUrl));
         Runnable closeTask = () -> {
             if (server != null) {
                 try {
@@ -91,5 +94,8 @@ public void startEventServer(LaunchModeBuildItem launchMode,
         };
         QuarkusClassLoader cl = (QuarkusClassLoader) Thread.currentThread().getContextClassLoader();
         ((QuarkusClassLoader) cl.parent()).addCloseTask(closeTask);
+        if (launchMode.isTest()) {
+            runtimeApplicationShutdownBuildItemBuildProducer.produce(new RuntimeApplicationShutdownBuildItem(closeTask));
+        }
     }
 }

extensions/amazon-lambda/common-runtime/src/main/java/io/quarkus/amazon/lambda/runtime/AbstractLambdaPollLoop.java

@@ -16,19 +16,23 @@
 import com.fasterxml.jackson.databind.ObjectReader;
 
 import io.quarkus.runtime.Application;
+import io.quarkus.runtime.LaunchMode;
 import io.quarkus.runtime.ShutdownContext;
 
 public abstract class AbstractLambdaPollLoop {
     private static final Logger log = Logger.getLogger(AbstractLambdaPollLoop.class);
 
-    private ObjectMapper objectMapper;
-    private ObjectReader cognitoIdReader;
-    private ObjectReader clientCtxReader;
+    private final ObjectMapper objectMapper;
+    private final ObjectReader cognitoIdReader;
+    private final ObjectReader clientCtxReader;
+    private final LaunchMode launchMode;
 
-    public AbstractLambdaPollLoop(ObjectMapper objectMapper, ObjectReader cognitoIdReader, ObjectReader clientCtxReader) {
+    public AbstractLambdaPollLoop(ObjectMapper objectMapper, ObjectReader cognitoIdReader, ObjectReader clientCtxReader,
+            LaunchMode launchMode) {
         this.objectMapper = objectMapper;
         this.cognitoIdReader = cognitoIdReader;
         this.clientCtxReader = clientCtxReader;
+        this.launchMode = launchMode;
     }
 
     protected abstract boolean isStream();
@@ -45,7 +49,7 @@ public void startPollLoop(ShutdownContext context) {
             @Override
             public void run() {
                 try {
-                    if (!LambdaHotReplacementRecorder.enabled) {
+                    if (!LambdaHotReplacementRecorder.enabled && launchMode == LaunchMode.DEVELOPMENT) {
                         // when running with continuous testing, this method fails
                         // because currentApplication is not set when running as an
                         // auxiliary application.  So, just skip it if hot replacement enabled.
@@ -81,7 +85,7 @@ public void run() {
                                 continue;
                             }
                             try {
-                                if (LambdaHotReplacementRecorder.enabled) {
+                                if (LambdaHotReplacementRecorder.enabled && launchMode == LaunchMode.DEVELOPMENT) {
                                     try {
                                         // do not interrupt during a hot replacement
                                         // as shutdown will abort and do nasty things.
@@ -125,7 +129,7 @@ public void run() {
                                 if (abortGracefully(e)) {
                                     return;
                                 }
-                                log.error("Failed to run lambda", e);
+                                log.error("Failed to run lambda (" + launchMode + ")", e);
 
                                 postError(AmazonLambdaApi.invocationError(baseUrl, requestId),
                                         new FunctionError(e.getClass().getName(), e.getMessage()));
@@ -134,7 +138,7 @@ public void run() {
 
                         } catch (Exception e) {
                             if (!abortGracefully(e))
-                                log.error("Error running lambda", e);
+                                log.error("Error running lambda (" + launchMode + ")", e);
                             Application app = Application.currentApplication();
                             if (app != null) {
                                 app.stop();
@@ -150,24 +154,24 @@ public void run() {
                     }
                 } catch (Exception e) {
                     try {
-                        log.error("Lambda init error", e);
+                        log.error("Lambda init error (" + launchMode + ")", e);
                         postError(AmazonLambdaApi.initError(baseUrl),
                                 new FunctionError(e.getClass().getName(), e.getMessage()));
                     } catch (Exception ex) {
-                        log.error("Failed to report init error", ex);
+                        log.error("Failed to report init error (" + launchMode + ")", ex);
                     } finally {
                         // our main loop is done, time to shutdown
                         Application app = Application.currentApplication();
                         if (app != null) {
-                            log.error("Shutting down Quarkus application because of error");
+                            log.error("Shutting down Quarkus application because of error (" + launchMode + ")");
                             app.stop();
                         }
                     }
                 } finally {
-                    log.info("Lambda polling thread complete");
+                    log.info("Lambda polling thread complete (" + launchMode + ")");
                 }
             }
-        }, "Lambda Thread");
+        }, "Lambda Thread (" + launchMode + ")");
         pollingThread.setDaemon(true);
         context.addShutdownTask(() -> {
             running.set(false);
@@ -261,7 +265,7 @@ boolean abortGracefully(Exception ex) {
         // if we are running in test mode, or native mode outside of the lambda container, then don't output stack trace for socket errors
 
         boolean lambdaEnv = System.getenv("AWS_LAMBDA_RUNTIME_API") != null;
-        boolean testEnv = System.getProperty(AmazonLambdaApi.QUARKUS_INTERNAL_AWS_LAMBDA_TEST_API) != null;
+        boolean testEnv = LaunchMode.current() == LaunchMode.TEST;
         boolean graceful = ((ex instanceof SocketException || ex instanceof ConnectException) && testEnv)
                 || (ex instanceof UnknownHostException && !lambdaEnv);
 

extensions/amazon-lambda/common-runtime/src/main/java/io/quarkus/amazon/lambda/runtime/AmazonLambdaApi.java

@@ -3,6 +3,10 @@
 import java.net.MalformedURLException;
 import java.net.URL;
 
+import org.eclipse.microprofile.config.ConfigProvider;
+
+import io.quarkus.runtime.LaunchMode;
+
 /**
  * Various constants and util methods used for communication with the AWS API.
  */
@@ -80,13 +84,15 @@ static String functionVersion() {
     }
 
     public static boolean isTestMode() {
-        return System.getProperty(AmazonLambdaApi.QUARKUS_INTERNAL_AWS_LAMBDA_TEST_API) != null;
+        //need this config check for native tests
+        return LaunchMode.current() == LaunchMode.TEST
+                || ConfigProvider.getConfig().getOptionalValue(QUARKUS_INTERNAL_AWS_LAMBDA_TEST_API, String.class).isPresent();
     }
 
     private static String runtimeApi() {
-        String testApi = System.getProperty(QUARKUS_INTERNAL_AWS_LAMBDA_TEST_API);
-        if (testApi != null) {
-            return testApi;
+        var testApi = ConfigProvider.getConfig().getOptionalValue(QUARKUS_INTERNAL_AWS_LAMBDA_TEST_API, String.class);
+        if (testApi.isPresent()) {
+            return testApi.get();
         }
         return System.getenv("AWS_LAMBDA_RUNTIME_API");
     }

extensions/amazon-lambda/deployment/src/main/java/io/quarkus/amazon/lambda/deployment/AmazonLambdaProcessor.java

@@ -240,9 +240,10 @@ public void recordHandlerClass(List<AmazonLambdaBuildItem> lambdas,
     @Record(value = ExecutionTime.RUNTIME_INIT)
     void startPoolLoop(AmazonLambdaRecorder recorder,
             ShutdownContextBuildItem shutdownContextBuildItem,
+            LaunchModeBuildItem launchModeBuildItem,
             List<ServiceStartBuildItem> orderServicesFirst // try to order this after service recorders
     ) {
-        recorder.startPollLoop(shutdownContextBuildItem);
+        recorder.startPollLoop(shutdownContextBuildItem, launchModeBuildItem.getLaunchMode());
     }
 
     @BuildStep
@@ -253,7 +254,7 @@ void startPoolLoopDevOrTest(AmazonLambdaRecorder recorder,
             LaunchModeBuildItem launchModeBuildItem) {
         LaunchMode mode = launchModeBuildItem.getLaunchMode();
         if (mode.isDevOrTest()) {
-            recorder.startPollLoop(shutdownContextBuildItem);
+            recorder.startPollLoop(shutdownContextBuildItem, mode);
         }
     }
 

extensions/amazon-lambda/deployment/src/test/java/io/quarkus/amazon/lambda/deployment/testing/LambdaDevServicesContinuousTestingTestCase.java

@@ -6,7 +6,7 @@
 import org.jboss.shrinkwrap.api.asset.StringAsset;
 import org.jboss.shrinkwrap.api.spec.JavaArchive;
 import org.junit.jupiter.api.Assertions;
-import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.RepeatedTest;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 import io.quarkus.test.ContinuousTestingTestUtils;
@@ -20,7 +20,9 @@ public class LambdaDevServicesContinuousTestingTestCase {
                 public JavaArchive get() {
                     return ShrinkWrap.create(JavaArchive.class)
                             .addClasses(GreetingLambda.class, Person.class)
-                            .addAsResource(new StringAsset(ContinuousTestingTestUtils.appProperties("")),
+                            .addAsResource(
+                                    new StringAsset(ContinuousTestingTestUtils.appProperties(
+                                            "quarkus.log.category.\"io.quarkus.amazon.lambda.runtime\".level=DEBUG")),
                                     "application.properties");
                 }
             }).setTestArchiveProducer(new Supplier<JavaArchive>() {
@@ -30,7 +32,8 @@ public JavaArchive get() {
                 }
             });
 
-    @Test
+    //run this twice, to make sure everything is cleaned up properly
+    @RepeatedTest(2)
     public void testLambda() throws Exception {
         ContinuousTestingTestUtils utils = new ContinuousTestingTestUtils();
         var result = utils.waitForNextCompletion();