com.powsybl.ws.commons.computation.service.AbstractWorkerService Maven / Gradle / Ivy
/**
* Copyright (c) 2024, 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 com.powsybl.ws.commons.computation.service;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.powsybl.commons.PowsyblException;
import com.powsybl.commons.report.ReportNode;
import com.powsybl.iidm.network.Network;
import com.powsybl.iidm.network.VariantManagerConstants;
import com.powsybl.network.store.client.NetworkStoreService;
import com.powsybl.network.store.client.PreloadingStrategy;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
import org.springframework.messaging.Message;
import org.springframework.web.server.ResponseStatusException;
import java.util.Map;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Consumer;
/**
* @author Mathieu Deharbe
* @param powsybl Result class specific to the computation
* @param Run context specific to a computation, including parameters
* @param powsybl and gridsuite Parameters specifics to the computation
* @param result service specific to the computation
*/
public abstract class AbstractWorkerService, P, S extends AbstractComputationResultService> {
private static final Logger LOGGER = LoggerFactory.getLogger(AbstractWorkerService.class);
protected final Lock lockRunAndCancel = new ReentrantLock();
protected final ObjectMapper objectMapper;
protected final NetworkStoreService networkStoreService;
protected final ReportService reportService;
protected final ExecutionService executionService;
protected final NotificationService notificationService;
protected final AbstractComputationObserver observer;
protected final Map> futures = new ConcurrentHashMap<>();
protected final Map cancelComputationRequests = new ConcurrentHashMap<>();
protected final S resultService;
protected AbstractWorkerService(NetworkStoreService networkStoreService,
NotificationService notificationService,
ReportService reportService,
S resultService,
ExecutionService executionService,
AbstractComputationObserver observer,
ObjectMapper objectMapper) {
this.networkStoreService = networkStoreService;
this.notificationService = notificationService;
this.reportService = reportService;
this.resultService = resultService;
this.executionService = executionService;
this.observer = observer;
this.objectMapper = objectMapper;
}
protected PreloadingStrategy getNetworkPreloadingStrategy() {
return PreloadingStrategy.COLLECTION;
}
protected Network getNetwork(UUID networkUuid, String variantId) {
try {
Network network = networkStoreService.getNetwork(networkUuid, getNetworkPreloadingStrategy());
String variant = StringUtils.isBlank(variantId) ? VariantManagerConstants.INITIAL_VARIANT_ID : variantId;
network.getVariantManager().setWorkingVariant(variant);
return network;
} catch (PowsyblException e) {
throw new ResponseStatusException(HttpStatus.NOT_FOUND, e.getMessage());
}
}
protected void cleanResultsAndPublishCancel(UUID resultUuid, String receiver) {
resultService.delete(resultUuid);
notificationService.publishStop(resultUuid, receiver, getComputationType());
if (LOGGER.isInfoEnabled()) {
LOGGER.info("{} (resultUuid='{}')",
NotificationService.getCancelMessage(getComputationType()),
resultUuid);
}
}
private boolean cancelAsync(CancelContext cancelContext) {
lockRunAndCancel.lock();
boolean isCanceled = false;
try {
cancelComputationRequests.put(cancelContext.resultUuid(), cancelContext);
// find the completableFuture associated with result uuid
CompletableFuture future = futures.get(cancelContext.resultUuid());
if (future != null) {
isCanceled = future.cancel(true); // cancel computation in progress
if (isCanceled) {
cleanResultsAndPublishCancel(cancelContext.resultUuid(), cancelContext.receiver());
}
}
} finally {
lockRunAndCancel.unlock();
}
return isCanceled;
}
protected abstract AbstractResultContext fromMessage(Message message);
protected boolean resultCanBeSaved(R result) {
return result != null;
}
public Consumer> consumeRun() {
return message -> {
AbstractResultContext resultContext = fromMessage(message);
AtomicReference rootReporter = new AtomicReference<>(ReportNode.NO_OP);
try {
long startTime = System.nanoTime();
Network network = getNetwork(resultContext.getRunContext().getNetworkUuid(),
resultContext.getRunContext().getVariantId());
resultContext.getRunContext().setNetwork(network);
R result = run(resultContext.getRunContext(), resultContext.getResultUuid(), rootReporter);
LOGGER.info("Just run in {}s", TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() - startTime));
if (resultCanBeSaved(result)) {
startTime = System.nanoTime();
observer.observe("results.save", resultContext.getRunContext(), () -> saveResult(network, resultContext, result));
LOGGER.info("Stored in {}s", TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() - startTime));
sendResultMessage(resultContext, result);
LOGGER.info("{} complete (resultUuid='{}')", getComputationType(), resultContext.getResultUuid());
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
} catch (Exception e) {
if (!(e instanceof CancellationException)) {
LOGGER.error(NotificationService.getFailedMessage(getComputationType()), e);
publishFail(resultContext, e.getMessage());
resultService.delete(resultContext.getResultUuid());
this.handleNonCancellationException(resultContext, e, rootReporter);
}
} finally {
futures.remove(resultContext.getResultUuid());
cancelComputationRequests.remove(resultContext.getResultUuid());
}
};
}
/**
* Handle exception in consumeRun that is not a CancellationException
* @param resultContext The context of the computation
* @param exception The exception to handle
*/
protected void handleNonCancellationException(AbstractResultContext resultContext, Exception exception, AtomicReference rootReporter) {
}
public Consumer> consumeCancel() {
return message -> {
CancelContext cancelContext = CancelContext.fromMessage(message);
boolean isCancelled = cancelAsync(cancelContext);
if (!isCancelled) {
notificationService.publishCancelFailed(cancelContext.resultUuid(), cancelContext.receiver(), getComputationType(), cancelContext.userId());
}
};
}
protected abstract void saveResult(Network network, AbstractResultContext resultContext, R result);
protected void sendResultMessage(AbstractResultContext resultContext, R ignoredResult) {
notificationService.sendResultMessage(resultContext.getResultUuid(), resultContext.getRunContext().getReceiver(),
resultContext.getRunContext().getUserId(), null);
}
protected void publishFail(AbstractResultContext resultContext, String message) {
notificationService.publishFail(resultContext.getResultUuid(), resultContext.getRunContext().getReceiver(),
message, resultContext.getRunContext().getUserId(), getComputationType(), null);
}
/**
* Do some extra task before running the computation, e.g. print log or init extra data for the run context
* @param ignoredRunContext This context may be used for further computation in overriding classes
*/
protected void preRun(C ignoredRunContext) {
LOGGER.info("Run {} computation...", getComputationType());
}
protected R run(C runContext, UUID resultUuid, AtomicReference rootReporter) throws Exception {
String provider = runContext.getProvider();
ReportNode reportNode = ReportNode.NO_OP;
if (runContext.getReportInfos() != null && runContext.getReportInfos().reportUuid() != null) {
final String reportType = runContext.getReportInfos().computationType();
String rootReporterId = runContext.getReportInfos().reporterId();
rootReporter.set(ReportNode.newRootReportNode().withMessageTemplate(rootReporterId, rootReporterId).build());
reportNode = rootReporter.get().newReportNode().withMessageTemplate(reportType, reportType + (provider != null ? " (" + provider + ")" : ""))
.withUntypedValue("providerToUse", Objects.requireNonNullElse(provider, "")).add();
// Delete any previous computation logs
observer.observe("report.delete",
runContext, () -> reportService.deleteReport(runContext.getReportInfos().reportUuid()));
}
runContext.setReportNode(reportNode);
preRun(runContext);
CompletableFuture future = runAsync(runContext, provider, resultUuid);
R result = future == null ? null : observer.observeRun("run", runContext, future::get);
postRun(runContext, rootReporter, result);
return result;
}
/**
* Do some extra task after running the computation
* @param runContext This context may be used for extra task in overriding classes
* @param rootReportNode root of the reporter tree
* @param ignoredResult The result of the computation
*/
protected void postRun(C runContext, AtomicReference rootReportNode, R ignoredResult) {
if (runContext.getReportInfos().reportUuid() != null) {
observer.observe("report.send", runContext, () -> reportService.sendReport(runContext.getReportInfos().reportUuid(), rootReportNode.get()));
}
}
protected CompletableFuture runAsync(
C runContext,
String provider,
UUID resultUuid) {
lockRunAndCancel.lock();
try {
if (resultUuid != null && cancelComputationRequests.get(resultUuid) != null) {
return null;
}
CompletableFuture future = getCompletableFuture(runContext, provider, resultUuid);
if (resultUuid != null) {
futures.put(resultUuid, future);
}
return future;
} finally {
lockRunAndCancel.unlock();
}
}
protected abstract String getComputationType();
protected abstract CompletableFuture getCompletableFuture(C runContext, String provider, UUID resultUuid);
}