Fixed the bug that lambda function is blocked (#34)

* Fixed the bug that lambda function is blocked

Author: yaozhaoy

src/main/java/software/amazon/cloudwatchlogs/emf/environment/EnvironmentProvider.java

@@ -16,6 +16,8 @@
 
 package software.amazon.cloudwatchlogs.emf.environment;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import lombok.AllArgsConstructor;
@@ -49,18 +51,15 @@ public CompletableFuture<Environment> resolveEnvironment() {
             return CompletableFuture.completedFuture(cachedEnvironment);
         }
 
-        CompletableFuture<Optional<EnvironmentResolveResult>> resolvedEnv =
-                discoverEnvironmentAsync();
+        CompletableFuture<Optional<Environment>> resolvedEnv = discoverEnvironmentAsync();
 
         return resolvedEnv.thenApply(
                 optionalEnv ->
-                        optionalEnv
-                                .map(EnvironmentResolveResult::getEnvironment)
-                                .orElseGet(
-                                        () -> {
-                                            cachedEnvironment = defaultEnvironment;
-                                            return cachedEnvironment;
-                                        }));
+                        optionalEnv.orElseGet(
+                                () -> {
+                                    cachedEnvironment = defaultEnvironment;
+                                    return cachedEnvironment;
+                                }));
     }
 
     public Environment getDefaultEnvironment() {
@@ -72,28 +71,30 @@ void cleanResolvedEnvironment() {
         cachedEnvironment = null;
     }
 
-    private CompletableFuture<Optional<EnvironmentResolveResult>> discoverEnvironmentAsync() {
+    private CompletableFuture<Optional<Environment>> discoverEnvironmentAsync() {
 
-        CompletableFuture<Optional<EnvironmentResolveResult>> ans =
-                CompletableFuture.completedFuture(Optional.empty());
+        CompletableFuture<Optional<Environment>> ans = new CompletableFuture<>();
+
+        List<CompletableFuture<EnvironmentResolveResult>> futures = new ArrayList<>();
         for (Environment env : environments) {
             CompletableFuture<EnvironmentResolveResult> future =
                     CompletableFuture.supplyAsync(
                             () -> new EnvironmentResolveResult(env.probe(), env));
-            ans =
-                    ans.thenCombine(
-                            future,
-                            (optionalEnv, envResult) -> {
-                                if (optionalEnv.isPresent()) {
-                                    return optionalEnv;
-                                }
-                                if (envResult.isCandidate) {
-                                    return Optional.of(envResult);
-                                }
-                                return Optional.empty();
-                            });
+            futures.add(future);
         }
 
+        CompletableFuture.runAsync(
+                () -> {
+                    for (CompletableFuture<EnvironmentResolveResult> future : futures) {
+                        EnvironmentResolveResult result = future.join();
+                        if (result.isCandidate) {
+                            ans.complete(Optional.of(result.environment));
+                            return;
+                        }
+                    }
+                    ans.complete(Optional.empty());
+                });
+
         return ans;
     }
 

src/test/java/software/amazon/cloudwatchlogs/emf/environment/EnvironmentProviderTest.java

@@ -28,6 +28,8 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.internal.util.reflection.FieldSetter;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 import org.powermock.api.mockito.PowerMockito;
 import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
@@ -168,4 +170,33 @@ public void testResolveEnvironmentEC2AndECSEnvs() throws Exception {
         assertSame(expectedEnv, mockedEC2Env);
         environmentProvider.cleanResolvedEnvironment();
     }
+
+    @Test
+    public void testResolveEnvironmentReturnFirstDetectedEnvironment() throws Exception {
+
+        long startTime = System.currentTimeMillis();
+        LambdaEnvironment mockedLambdaEnv = mock(LambdaEnvironment.class);
+        when(mockedLambdaEnv.probe()).thenReturn(true);
+        EC2Environment mockedEC2Env = mock(EC2Environment.class);
+        when(mockedEC2Env.probe()).thenReturn(true);
+        DefaultEnvironment mockedDefaultEnv = mock(DefaultEnvironment.class);
+        when(mockedDefaultEnv.probe())
+                .thenAnswer(
+                        new Answer<Boolean>() {
+                            @Override
+                            public Boolean answer(InvocationOnMock invocation) throws Throwable {
+                                Thread.sleep(5_000);
+                                return true;
+                            }
+                        });
+        Environment[] envs = new Environment[] {mockedLambdaEnv, mockedDefaultEnv, mockedEC2Env};
+
+        FieldSetter.setField(
+                environmentProvider,
+                EnvironmentProvider.class.getDeclaredField("environments"),
+                envs);
+        Environment env = environmentProvider.resolveEnvironment().join();
+        assertSame(env, mockedLambdaEnv);
+        assertTrue(System.currentTimeMillis() - startTime < 3_000);
+    }
 }