Produce contingency list as a flux to avoid 'DataBufferLimitException' for very large lists (#10)

Signed-off-by: Slimane AMAR <[email protected]>

Author: SlimaneAmar

src/main/java/org/gridsuite/securityanalysis/server/service/ActionsService.java

@@ -12,9 +12,8 @@
 import org.springframework.core.ParameterizedTypeReference;
 import org.springframework.stereotype.Service;
 import org.springframework.web.reactive.function.client.WebClient;
-import reactor.core.publisher.Mono;
+import reactor.core.publisher.Flux;
 
-import java.util.List;
 import java.util.Objects;
 import java.util.UUID;
 
@@ -39,7 +38,7 @@ public ActionsService(WebClient webClient) {
         this.webClient = Objects.requireNonNull(webClient);
     }
 
-    public Mono<List<Contingency>> getContingencyList(String name, UUID networkUuid) {
+    public Flux<Contingency> getContingencyList(String name, UUID networkUuid) {
         Objects.requireNonNull(name);
         Objects.requireNonNull(networkUuid);
         return webClient.get()
@@ -48,6 +47,7 @@ public Mono<List<Contingency>> getContingencyList(String name, UUID networkUuid)
                         .queryParam("networkUuid", networkUuid.toString())
                         .build(name))
                 .retrieve()
-                .bodyToMono(new ParameterizedTypeReference<>() { });
+                .bodyToFlux(new ParameterizedTypeReference<>() {
+                });
     }
 }

src/main/java/org/gridsuite/securityanalysis/server/service/SecurityAnalysisWorkerService.java

@@ -31,12 +31,7 @@
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Schedulers;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-import java.util.Set;
-import java.util.UUID;
+import java.util.*;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
@@ -138,8 +133,7 @@ private Mono<SecurityAnalysisResult> run(SecurityAnalysisRunContext context, UUI
         Mono<Network> network = getNetwork(context.getNetworkUuid(), context.getOtherNetworkUuids());
 
         Mono<List<Contingency>> contingencies = Flux.fromIterable(context.getContingencyListNames())
-                .flatMap(contingencyListName -> actionsService.getContingencyList(contingencyListName, context.getNetworkUuid())
-                        .flatMapMany(Flux::fromIterable))
+                .flatMap(contingencyListName -> actionsService.getContingencyList(contingencyListName, context.getNetworkUuid()))
                 .collectList();
 
         return Mono.zip(network, contingencies)

src/test/java/org/gridsuite/securityanalysis/server/SecurityAnalysisControllerTest.java

@@ -11,9 +11,9 @@
 import com.powsybl.network.store.client.NetworkStoreService;
 import com.powsybl.network.store.client.PreloadingStrategy;
 import org.gridsuite.securityanalysis.server.service.ActionsService;
-import org.gridsuite.securityanalysis.server.service.UuidGeneratorService;
 import org.gridsuite.securityanalysis.server.service.SecurityAnalysisConfigService;
 import org.gridsuite.securityanalysis.server.service.SecurityAnalysisService;
+import org.gridsuite.securityanalysis.server.service.UuidGeneratorService;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -31,7 +31,7 @@
 import org.springframework.test.context.junit4.SpringRunner;
 import org.springframework.test.web.reactive.server.WebTestClient;
 import org.springframework.web.reactive.config.EnableWebFlux;
-import reactor.core.publisher.Mono;
+import reactor.core.publisher.Flux;
 
 import java.util.UUID;
 
@@ -93,9 +93,9 @@ public void setUp() {
 
         // action service mocking
         given(actionsService.getContingencyList(CONTINGENCY_LIST_NAME, NETWORK_UUID))
-                .willReturn(Mono.just(SecurityAnalysisFactoryMock.CONTINGENCIES));
+                .willReturn(Flux.fromIterable(SecurityAnalysisFactoryMock.CONTINGENCIES));
         given(actionsService.getContingencyList(CONTINGENCY_LIST2_NAME, NETWORK_UUID))
-                .willReturn(Mono.just(SecurityAnalysisFactoryMock.CONTINGENCIES));
+                .willReturn(Flux.fromIterable(SecurityAnalysisFactoryMock.CONTINGENCIES));
 
         // UUID service mocking to always generate the same result UUID
         given(uuidGeneratorService.generate()).willReturn(RESULT_UUID);

src/test/java/org/gridsuite/securityanalysis/server/service/ActionsServiceTest.java

@@ -6,7 +6,6 @@
  */
 package org.gridsuite.securityanalysis.server.service;
 
-import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.powsybl.contingency.BranchContingency;
 import com.powsybl.contingency.Contingency;
@@ -16,9 +15,11 @@
 import okhttp3.mockwebserver.MockWebServer;
 import okhttp3.mockwebserver.RecordedRequest;
 import org.gridsuite.securityanalysis.server.WebFluxConfig;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.codec.json.Jackson2JsonDecoder;
 import org.springframework.http.codec.json.Jackson2JsonEncoder;
@@ -27,30 +28,41 @@
 import org.springframework.web.reactive.function.client.WebClient;
 
 import java.io.IOException;
-import java.io.UncheckedIOException;
 import java.util.List;
+import java.util.Objects;
 import java.util.UUID;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import static org.junit.Assert.assertEquals;
 
 /**
  * @author Geoffroy Jamgotchian <geoffroy.jamgotchian at rte-france.com>
+ * @author Slimane Amar <slimane.amar at rte-france.com>
  */
 @RunWith(SpringRunner.class)
 public class ActionsServiceTest {
 
+    private static final int DATA_BUFFER_LIMIT = 256 * 1024; // AbstractJackson2Decoder.maxInMemorySize
+
     private static final String NETWORK_UUID = "7928181c-7977-4592-ba19-88027e4254e4";
 
     private static final String LIST_NAME = "myList";
 
+    private static final String VERY_LARGE_LIST_NAME = "veryLargelist";
+
     private static final Contingency CONTINGENCY = new Contingency("c1", new BranchContingency("b1"));
 
     private final ObjectMapper objectMapper = WebFluxConfig.createObjectMapper();
 
     private WebClient.Builder webClientBuilder;
 
+    private MockWebServer server;
+
+    private ActionsService actionsService;
+
     @Before
-    public void setUp() {
+    public void setUp() throws IOException {
         webClientBuilder = WebClient.builder();
         ExchangeStrategies strategies = ExchangeStrategies
                 .builder()
@@ -60,41 +72,66 @@ public void setUp() {
 
                 }).build();
         webClientBuilder.exchangeStrategies(strategies);
+
+        actionsService = new ActionsService(webClientBuilder, initMockWebServer());
     }
 
-    @Test
-    public void test() throws IOException {
-        MockWebServer server = new MockWebServer();
+    @After
+    public void tearDown() {
+        try {
+            server.shutdown();
+        } catch (Exception e) {
+            // Nothing to do
+        }
+    }
+
+    private String initMockWebServer() throws IOException {
+        server = new MockWebServer();
         server.start();
+
+        String jsonExpected = objectMapper.writeValueAsString(List.of(CONTINGENCY));
+        String veryLargeJsonExpected = objectMapper.writeValueAsString(createVeryLargeList());
+
         final Dispatcher dispatcher = new Dispatcher() {
             @Override
             public MockResponse dispatch(RecordedRequest request) {
-                if (("/" + ActionsService.ACTIONS_API_VERSION + "/contingency-lists/" + LIST_NAME + "/export?networkUuid=" + NETWORK_UUID)
-                        .equals(request.getPath())) {
-                    try {
-                        String json = objectMapper.writeValueAsString(List.of(CONTINGENCY));
-                        return new MockResponse()
-                                .setResponseCode(200)
-                                .setBody(json)
-                                .addHeader("Content-Type", "application/json; charset=utf-8");
-                    } catch (JsonProcessingException e) {
-                        throw new UncheckedIOException(e);
-                    }
+                String requestPath = Objects.requireNonNull(request.getPath());
+                if (requestPath.equals(String.format("/v1/contingency-lists/%s/export?networkUuid=%s", LIST_NAME, NETWORK_UUID))) {
+                    return new MockResponse().setResponseCode(HttpStatus.OK.value())
+                            .setBody(jsonExpected)
+                            .addHeader("Content-Type", "application/json; charset=utf-8");
+                } else if (requestPath.equals(String.format("/v1/contingency-lists/%s/export?networkUuid=%s", VERY_LARGE_LIST_NAME, NETWORK_UUID))) {
+                    return new MockResponse().setResponseCode(HttpStatus.OK.value())
+                            .setBody(veryLargeJsonExpected)
+                            .addHeader("Content-Type", "application/json; charset=utf-8");
+                } else {
+                    return new MockResponse().setResponseCode(HttpStatus.NOT_FOUND.value()).setBody("Path not supported: " + request.getPath());
                 }
-                return new MockResponse().setResponseCode(404);
             }
         };
+
         server.setDispatcher(dispatcher);
 
-        // get server base URL
+        // Ask the server for its URL. You'll need this to make HTTP requests.
         HttpUrl baseHttpUrl = server.url("");
-        String baseUrl = baseHttpUrl.toString().substring(0, baseHttpUrl.toString().length() - 1);
 
-        ActionsService actionsService = new ActionsService(webClientBuilder, baseUrl);
+        return baseHttpUrl.toString().substring(0, baseHttpUrl.toString().length() - 1);
+    }
+
+    private List<Contingency> createVeryLargeList() {
+        return IntStream.range(0, DATA_BUFFER_LIMIT).mapToObj(i -> new Contingency("l" + i, new BranchContingency("l" + i))).collect(Collectors.toList());
+    }
 
-        List<Contingency> list = actionsService.getContingencyList(LIST_NAME, UUID.fromString(NETWORK_UUID)).block();
+    @Test
+    public void test() {
+        List<Contingency> list = actionsService.getContingencyList(LIST_NAME, UUID.fromString(NETWORK_UUID)).collectList().block();
         assertEquals(List.of(CONTINGENCY), list);
+    }
 
-        server.shutdown();
+    @Test
+    public void testVeryLargeList() {
+        // DataBufferLimitException should not be thrown with this message : "Exceeded limit on max bytes to buffer : DATA_BUFFER_LIMIT"
+        List<Contingency> list = actionsService.getContingencyList(VERY_LARGE_LIST_NAME, UUID.fromString(NETWORK_UUID)).collectList().block();
+        assertEquals(createVeryLargeList(), list);
     }
 }