Stop a security analysis computation in progress (#8)

FranckLecuyer · web-flow · commit e6944e20cb13 · 2021-02-17T17:06:11.000+01:00
Stop a security analysis computation in progress
diff --git a/src/main/java/org/gridsuite/securityanalysis/server/SecurityAnalysisController.java b/src/main/java/org/gridsuite/securityanalysis/server/SecurityAnalysisController.java
@@ -35,6 +35,7 @@
 
 /**
  * @author Geoffroy Jamgotchian <geoffroy.jamgotchian at rte-france.com>
+ * @author Franck Lecuyer <franck.lecuyer at rte-france.com>
  */
 @RestController
 @RequestMapping(value = "/" + SecurityAnalysisApi.API_VERSION)
@@ -134,4 +135,13 @@ public ResponseEntity<Mono<Void>> invalidateStatus(@Parameter(description = "Res
         Mono<Void> result = service.setStatus(resultUuid, SecurityAnalysisStatus.NOT_DONE.name());
         return ResponseEntity.ok().body(result);
     }
+
+    @PutMapping(value = "/results/{resultUuid}/stop", produces = APPLICATION_JSON_VALUE)
+    @Operation(summary = "Stop a security analysis computation")
+    @ApiResponses(value = {@ApiResponse(responseCode = "200", description = "The security analysis has been stopped")})
+    public ResponseEntity<Mono<Void>> stop(@Parameter(description = "Result UUID") @PathVariable("resultUuid") UUID resultUuid,
+                                           @Parameter(description = "Result receiver") @RequestParam(name = "receiver", required = false) String receiver) {
+        Mono<Void> result = service.stop(resultUuid, receiver);
+        return ResponseEntity.ok().body(result);
+    }
 }
diff --git a/src/main/java/org/gridsuite/securityanalysis/server/service/SecurityAnalysisCancelContext.java b/src/main/java/org/gridsuite/securityanalysis/server/service/SecurityAnalysisCancelContext.java
@@ -0,0 +1,61 @@
+/**
+ * Copyright (c) 2021, RTE (http://www.rte-france.com)
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ */
+package org.gridsuite.securityanalysis.server.service;
+
+import com.powsybl.commons.PowsyblException;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageHeaders;
+import org.springframework.messaging.support.MessageBuilder;
+
+import java.util.Objects;
+import java.util.UUID;
+
+/**
+ * @author Franck Lecuyer <franck.lecuyer at rte-france.com>
+ */
+public class SecurityAnalysisCancelContext {
+
+    private final UUID resultUuid;
+
+    private final String receiver;
+
+    public SecurityAnalysisCancelContext(UUID resultUuid, String receiver) {
+        this.resultUuid = Objects.requireNonNull(resultUuid);
+        this.receiver = Objects.requireNonNull(receiver);
+    }
+
+    public UUID getResultUuid() {
+        return resultUuid;
+    }
+
+    public String getReceiver() {
+        return receiver;
+    }
+
+    private static String getNonNullHeader(MessageHeaders headers, String name) {
+        String header = (String) headers.get(name);
+        if (header == null) {
+            throw new PowsyblException("Header '" + name + "' not found");
+        }
+        return header;
+    }
+
+    public static SecurityAnalysisCancelContext fromMessage(Message<String> message) {
+        Objects.requireNonNull(message);
+        MessageHeaders headers = message.getHeaders();
+        UUID resultUuid = UUID.fromString(getNonNullHeader(headers, "resultUuid"));
+        String receiver = (String) headers.get("receiver");
+        return new SecurityAnalysisCancelContext(resultUuid, receiver);
+    }
+
+    public Message<String> toMessage() {
+        return MessageBuilder.withPayload("")
+                .setHeader("resultUuid", resultUuid.toString())
+                .setHeader("receiver", receiver)
+                .build();
+    }
+}
diff --git a/src/main/java/org/gridsuite/securityanalysis/server/service/SecurityAnalysisCancelPublisherService.java b/src/main/java/org/gridsuite/securityanalysis/server/service/SecurityAnalysisCancelPublisherService.java
@@ -0,0 +1,39 @@
+/**
+ * Copyright (c) 2021, RTE (http://www.rte-france.com)
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ */
+package org.gridsuite.securityanalysis.server.service;
+
+import org.springframework.context.annotation.Bean;
+import org.springframework.messaging.Message;
+import org.springframework.stereotype.Service;
+import reactor.core.publisher.EmitterProcessor;
+import reactor.core.publisher.Flux;
+
+import java.util.Objects;
+import java.util.function.Supplier;
+import java.util.logging.Level;
+
+/**
+ * @author Franck Lecuyer <franck.lecuyer at rte-france.com>
+ */
+@Service
+public class SecurityAnalysisCancelPublisherService {
+
+    private static final String CATEGORY_BROKER_OUTPUT = SecurityAnalysisCancelPublisherService.class.getName()
+            + ".output-broker-messages";
+
+    private final EmitterProcessor<Message<String>> cancelMessagePublisher = EmitterProcessor.create();
+
+    @Bean
+    public Supplier<Flux<Message<String>>> publishCancel() {
+        return () -> cancelMessagePublisher.log(CATEGORY_BROKER_OUTPUT, Level.FINE);
+    }
+
+    public void publish(SecurityAnalysisCancelContext cancelContext) {
+        Objects.requireNonNull(cancelContext);
+        cancelMessagePublisher.onNext(cancelContext.toMessage());
+    }
+}
diff --git a/src/main/java/org/gridsuite/securityanalysis/server/service/SecurityAnalysisService.java b/src/main/java/org/gridsuite/securityanalysis/server/service/SecurityAnalysisService.java
@@ -19,20 +19,25 @@
 
 /**
  * @author Geoffroy Jamgotchian <geoffroy.jamgotchian at rte-france.com>
+ * @author Franck Lecuyer <franck.lecuyer at rte-france.com>
  */
 @Service
 public class SecurityAnalysisService {
     private SecurityAnalysisResultRepository resultRepository;
 
     private SecurityAnalysisRunPublisherService runPublisherService;
 
+    private SecurityAnalysisCancelPublisherService cancelPublisherService;
+
     private UuidGeneratorService uuidGeneratorService;
 
     public SecurityAnalysisService(SecurityAnalysisResultRepository resultRepository,
                                    SecurityAnalysisRunPublisherService runPublisherService,
+                                   SecurityAnalysisCancelPublisherService cancelPublisherService,
                                    UuidGeneratorService uuidGeneratorService) {
         this.resultRepository = Objects.requireNonNull(resultRepository);
         this.runPublisherService = Objects.requireNonNull(runPublisherService);
+        this.cancelPublisherService = Objects.requireNonNull(cancelPublisherService);
         this.uuidGeneratorService = Objects.requireNonNull(uuidGeneratorService);
     }
 
@@ -66,4 +71,9 @@ public Mono<String> getStatus(UUID resultUuid) {
     public Mono<Void> setStatus(UUID resultUuid, String status) {
         return resultRepository.insertStatus(resultUuid, status);
     }
+
+    public Mono<Void> stop(UUID resultUuid, String receiver) {
+        return Mono.fromRunnable(() ->
+                cancelPublisherService.publish(new SecurityAnalysisCancelContext(resultUuid, receiver))).then();
+    }
 }
diff --git a/src/main/java/org/gridsuite/securityanalysis/server/service/SecurityAnalysisStoppedPublisherService.java b/src/main/java/org/gridsuite/securityanalysis/server/service/SecurityAnalysisStoppedPublisherService.java
@@ -0,0 +1,43 @@
+/**
+ * Copyright (c) 2021, RTE (http://www.rte-france.com)
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ */
+package org.gridsuite.securityanalysis.server.service;
+
+import org.springframework.context.annotation.Bean;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.support.MessageBuilder;
+import org.springframework.stereotype.Service;
+import reactor.core.publisher.EmitterProcessor;
+import reactor.core.publisher.Flux;
+
+import java.util.UUID;
+import java.util.function.Supplier;
+import java.util.logging.Level;
+
+/**
+ * @author Franck Lecuyer <franck.lecuyer at rte-france.com>
+ */
+@Service
+public class SecurityAnalysisStoppedPublisherService {
+
+    private static final String CATEGORY_BROKER_OUTPUT = SecurityAnalysisStoppedPublisherService.class.getName()
+            + ".output-broker-messages";
+
+    private final EmitterProcessor<Message<String>> stoppedMessagePublisher = EmitterProcessor.create();
+
+    @Bean
+    public Supplier<Flux<Message<String>>> publishStopped() {
+        return () -> stoppedMessagePublisher.log(CATEGORY_BROKER_OUTPUT, Level.FINE);
+    }
+
+    public void publish(UUID resultUuid, String receiver) {
+        stoppedMessagePublisher.onNext(MessageBuilder
+                .withPayload("")
+                .setHeader("resultUuid", resultUuid.toString())
+                .setHeader("receiver", receiver)
+                .build());
+    }
+}
diff --git a/src/main/java/org/gridsuite/securityanalysis/server/service/SecurityAnalysisWorkerService.java b/src/main/java/org/gridsuite/securityanalysis/server/service/SecurityAnalysisWorkerService.java
@@ -7,6 +7,7 @@
 package org.gridsuite.securityanalysis.server.service;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.Sets;
 import com.powsybl.commons.PowsyblException;
 import com.powsybl.computation.local.LocalComputationManager;
 import com.powsybl.contingency.Contingency;
@@ -32,15 +33,19 @@
 
 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.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
-import java.util.logging.Level;
 import java.util.stream.Collectors;
 
 /**
  * @author Geoffroy Jamgotchian <geoffroy.jamgotchian at rte-france.com>
+ * @author Franck Lecuyer <franck.lecuyer at rte-france.com>
  */
 @Service
 public class SecurityAnalysisWorkerService {
@@ -62,15 +67,26 @@ public class SecurityAnalysisWorkerService {
 
     private SecurityAnalysisResultPublisherService resultPublisherService;
 
+    private SecurityAnalysisStoppedPublisherService stoppedPublisherService;
+
+    private Map<UUID, CompletableFuture<SecurityAnalysisResult>> futures = new ConcurrentHashMap<>();
+
+    private Map<UUID, SecurityAnalysisCancelContext> cancelComputationRequests = new ConcurrentHashMap<>();
+
+    private Set<UUID> runRequests = Sets.newConcurrentHashSet();
+
     public SecurityAnalysisWorkerService(NetworkStoreService networkStoreService, ActionsService actionsService,
                                          SecurityAnalysisResultRepository resultRepository, ObjectMapper objectMapper,
-                                         SecurityAnalysisConfigService configService, SecurityAnalysisResultPublisherService resultPublisherService) {
+                                         SecurityAnalysisConfigService configService,
+                                         SecurityAnalysisResultPublisherService resultPublisherService,
+                                         SecurityAnalysisStoppedPublisherService stoppedPublisherService) {
         this.networkStoreService = Objects.requireNonNull(networkStoreService);
         this.actionsService = Objects.requireNonNull(actionsService);
         this.resultRepository = Objects.requireNonNull(resultRepository);
         this.objectMapper = Objects.requireNonNull(objectMapper);
         this.configService = Objects.requireNonNull(configService);
         this.resultPublisherService = Objects.requireNonNull(resultPublisherService);
+        this.stoppedPublisherService = Objects.requireNonNull(stoppedPublisherService);
     }
 
     private static String sanitizeParam(String param) {
@@ -111,6 +127,10 @@ private Mono<Network> getNetwork(UUID networkUuid, List<UUID> otherNetworkUuids)
     }
 
     public Mono<SecurityAnalysisResult> run(SecurityAnalysisRunContext context) {
+        return run(context, null);
+    }
+
+    private Mono<SecurityAnalysisResult> run(SecurityAnalysisRunContext context, UUID resultUuid) {
         Objects.requireNonNull(context);
 
         LOGGER.info("Run security analysis on contingency lists: {}", context.getContingencyListNames().stream().map(SecurityAnalysisWorkerService::sanitizeParam).collect(Collectors.toList()));
@@ -125,26 +145,75 @@ public Mono<SecurityAnalysisResult> run(SecurityAnalysisRunContext context) {
         return Mono.zip(network, contingencies)
                 .flatMap(tuple -> {
                     SecurityAnalysis securityAnalysis = configService.getSecurityAnalysisFactory().create(tuple.getT1(), LocalComputationManager.getDefault(), 0);
-                    CompletableFuture<SecurityAnalysisResult> result = securityAnalysis.run(VariantManagerConstants.INITIAL_VARIANT_ID, context.getParameters(), n -> tuple.getT2());
-                    return Mono.fromCompletionStage(result);
+                    CompletableFuture<SecurityAnalysisResult> future = securityAnalysis.run(VariantManagerConstants.INITIAL_VARIANT_ID, context.getParameters(), n -> tuple.getT2());
+                    if (resultUuid != null) {
+                        futures.put(resultUuid, future);
+                    }
+                    if (resultUuid != null && cancelComputationRequests.get(resultUuid) != null) {
+                        return Mono.empty();
+                    } else {
+                        return Mono.fromCompletionStage(future);
+                    }
                 });
     }
 
     @Bean
-    public Consumer<Flux<Message<String>>> consumeRun() {
-        return f -> f.log(CATEGORY_BROKER_INPUT, Level.FINE)
-                .flatMap(message -> {
-                    SecurityAnalysisResultContext resultContext = SecurityAnalysisResultContext.fromMessage(message, objectMapper);
-
-                    return run(resultContext.getRunContext())
-                            .flatMap(result -> resultRepository.insert(resultContext.getResultUuid(), result)
-                                    .then(resultRepository.insertStatus(resultContext.getResultUuid(), SecurityAnalysisStatus.COMPLETED.name())))
-                            .doOnSuccess(unused -> {
-                                resultPublisherService.publish(resultContext.getResultUuid(), resultContext.getRunContext().getReceiver());
-                                LOGGER.info("Security analysis complete (resultUuid='{}')", resultContext.getResultUuid());
-                            });
-                })
-                .doOnError(throwable -> LOGGER.error(throwable.toString(), throwable))
-                .subscribe();
+    public Consumer<Message<String>> consumeRun() {
+        return message -> {
+
+            SecurityAnalysisResultContext resultContext = SecurityAnalysisResultContext.fromMessage(message, objectMapper);
+            runRequests.add(resultContext.getResultUuid());
+
+            run(resultContext.getRunContext(), resultContext.getResultUuid())
+                    .flatMap(result -> resultRepository.insert(resultContext.getResultUuid(), result)
+                            .then(resultRepository.insertStatus(resultContext.getResultUuid(), SecurityAnalysisStatus.COMPLETED.name()))
+                            .then(Mono.just(result)))
+                    .doOnSuccess(result -> {
+                        if (result != null) {  // result available
+                            resultPublisherService.publish(resultContext.getResultUuid(), resultContext.getRunContext().getReceiver());
+                            LOGGER.info("Security analysis complete (resultUuid='{}')", resultContext.getResultUuid());
+                        } else {  // result not available : stop computation request
+                            if (cancelComputationRequests.get(resultContext.getResultUuid()) != null) {
+                                stoppedPublisherService.publish(resultContext.getResultUuid(), cancelComputationRequests.get(resultContext.getResultUuid()).getReceiver());
+                            }
+                        }
+                    })
+                    .doOnError(throwable -> {
+                        if (!(throwable instanceof CancellationException)) {
+                            LOGGER.error(throwable.toString(), throwable);
+                        }
+                    })
+                    .doFinally(s -> {
+                        futures.remove(resultContext.getResultUuid());
+                        cancelComputationRequests.remove(resultContext.getResultUuid());
+                        runRequests.remove(resultContext.getResultUuid());
+                    })
+                    .subscribe();
+        };
+    }
+
+    @Bean
+    public Consumer<Message<String>> consumeCancel() {
+        return message -> {
+            SecurityAnalysisCancelContext cancelContext = SecurityAnalysisCancelContext.fromMessage(message);
+
+            if (runRequests.contains(cancelContext.getResultUuid())) {
+                cancelComputationRequests.put(cancelContext.getResultUuid(), cancelContext);
+            }
+
+            // find the completableFuture associated with result uuid
+            CompletableFuture<SecurityAnalysisResult> future = futures.get(cancelContext.getResultUuid());
+            if (future != null) {
+                future.cancel(true);  // cancel computation in progress
+
+                resultRepository.delete(cancelContext.getResultUuid())
+                        .doOnSuccess(unused -> {
+                            stoppedPublisherService.publish(cancelContext.getResultUuid(), cancelContext.getReceiver());
+                            LOGGER.info("Security analysis stopped (resultUuid='{}')", cancelContext.getResultUuid());
+                        })
+                        .doOnError(throwable -> LOGGER.error(throwable.toString(), throwable))
+                        .subscribe();
+            }
+        };
     }
 }
diff --git a/src/main/resources/application.yaml b/src/main/resources/application.yaml
@@ -5,7 +5,7 @@ spring:
 
   cloud:
     function:
-      definition: consumeRun;publishRun;publishResult
+      definition: consumeRun;publishRun;publishResult;consumeCancel;publishCancel;publishStopped
     stream:
       bindings:
         consumeRun-in-0:
@@ -15,6 +15,12 @@ spring:
           destination: sa.run
         publishResult-out-0:
           destination: sa.result
+        consumeCancel-in-0:
+            destination: sa.cancel
+        publishCancel-out-0:
+            destination: sa.cancel
+        publishStopped-out-0:
+          destination: sa.stopped
 
 springdoc:
   swagger-ui:
diff --git a/src/test/java/org/gridsuite/securityanalysis/server/SecurityAnalysisControllerTest.java b/src/test/java/org/gridsuite/securityanalysis/server/SecurityAnalysisControllerTest.java
