From 02f828d0a948603a94c7835d1e7307f173ed141d Mon Sep 17 00:00:00 2001
From: Thang PHAM <phamthang37@gmail.com>
Date: Wed, 30 Oct 2024 12:03:15 +0100
Subject: [PATCH] default implementation for clean method

Signed-off-by: Thang PHAM <phamthang37@gmail.com>
---
 .../computation/service/AbstractWorkerService.java   | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/powsybl/ws/commons/computation/service/AbstractWorkerService.java b/src/main/java/com/powsybl/ws/commons/computation/service/AbstractWorkerService.java
index 3bd1155..4787ad4 100644
--- a/src/main/java/com/powsybl/ws/commons/computation/service/AbstractWorkerService.java
+++ b/src/main/java/com/powsybl/ws/commons/computation/service/AbstractWorkerService.java
@@ -153,12 +153,20 @@ public Consumer<Message<String>> consumeRun() {
                     this.handleNonCancellationException(resultContext, e, rootReporter);
                 }
             } finally {
-                futures.remove(resultContext.getResultUuid());
-                cancelComputationRequests.remove(resultContext.getResultUuid());
+                clean(resultContext);
             }
         };
     }
 
+    /**
+     * Perform cleaning
+     * @param resultContext The context of the computation
+     */
+    protected void clean(AbstractResultContext<C> resultContext) {
+        futures.remove(resultContext.getResultUuid());
+        cancelComputationRequests.remove(resultContext.getResultUuid());
+    }
+
     /**
      * Handle exception in consumeRun that is not a CancellationException
      * @param resultContext The context of the computation