Implement Purgeable on SubscriptionService (#620)

Author: Natchica

CHANGELOG.md

@@ -8,6 +8,7 @@ All notable changes to this project will be documented in this file.
 
 - Add workerpool address in configuration instead of reading from Scheduler `/workers/config` endpoint. (#607)
 - Set `0x0` as default value for Workerpool address and prevents startup if incorrectly configured. (#608)
+- Implement `Purgeable` on `SubscriptionService`. (#620)
 
 ### Bug fixes
 

src/main/java/com/iexec/worker/pubsub/SubscriptionService.java

@@ -16,6 +16,7 @@
 
 package com.iexec.worker.pubsub;
 
+import com.iexec.common.lifecycle.purge.Purgeable;
 import com.iexec.core.notification.TaskNotification;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
@@ -28,6 +29,7 @@
 import org.springframework.messaging.simp.stomp.StompSession.Subscription;
 import org.springframework.stereotype.Service;
 
+import javax.annotation.PreDestroy;
 import java.lang.reflect.Type;
 import java.util.Map;
 import java.util.Set;
@@ -38,8 +40,7 @@
 
 @Slf4j
 @Service
-// FIXME: implement `Purgeable`
-public class SubscriptionService {
+public class SubscriptionService implements Purgeable {
 
     private final Map<String, Subscription> chainTaskIdToSubscription = new ConcurrentHashMap<>();
     private final ApplicationEventPublisher eventPublisher;
@@ -64,14 +65,14 @@ public SubscriptionService(ApplicationEventPublisher applicationEventPublisher,
      *
      * @param chainTaskId id of the task to which to subscribe
      */
-    public void subscribeToTopic(String chainTaskId) {
-        String topic = getTaskTopicName(chainTaskId);
+    public void subscribeToTopic(final String chainTaskId) {
+        final String topic = getTaskTopicName(chainTaskId);
         if (this.chainTaskIdToSubscription.containsKey(chainTaskId)) {
             log.info("Already subscribed to topic [chainTaskId:{}, topic:{}]",
                     chainTaskId, topic);
             return;
         }
-        MessageHandler messageHandler = new MessageHandler(chainTaskId, this.workerWalletAddress);
+        final MessageHandler messageHandler = new MessageHandler(chainTaskId, this.workerWalletAddress);
         stompClientService.subscribeToTopic(topic, messageHandler).ifPresentOrElse(
                 subscription -> {
                     this.chainTaskIdToSubscription.put(chainTaskId, subscription);
@@ -84,16 +85,27 @@ public void subscribeToTopic(String chainTaskId) {
     /**
      * Unsubscribe from topic if already subscribed.
      *
-     * @param chainTaskId
+     * @param chainTaskId id of the task to unsubscribe and purge
+     * @return true if unsubscribed and purge topic successfully, false otherwise
      */
-    public void unsubscribeFromTopic(String chainTaskId) {
+    @Override
+    public boolean purgeTask(final String chainTaskId) {
+        log.debug("purgeTask [chainTaskId:{}]", chainTaskId);
         if (!isSubscribedToTopic(chainTaskId)) {
             log.error("Already unsubscribed from topic [chainTaskId:{}]", chainTaskId);
-            return;
+            return true;
         }
         this.chainTaskIdToSubscription.get(chainTaskId).unsubscribe();
         this.chainTaskIdToSubscription.remove(chainTaskId);
         log.info("Unsubscribed from topic [chainTaskId:{}]", chainTaskId);
+        return !isSubscribedToTopic(chainTaskId);
+    }
+
+    @Override
+    @PreDestroy
+    public void purgeAllTasksData() {
+        this.chainTaskIdToSubscription.keySet().forEach(this::purgeTask);
+        this.chainTaskIdToSubscription.clear();
     }
 
     /**
@@ -102,7 +114,7 @@ public void unsubscribeFromTopic(String chainTaskId) {
      * @param chainTaskId id of the task to check
      * @return true if subscribed, false otherwise
      */
-    public boolean isSubscribedToTopic(String chainTaskId) {
+    public boolean isSubscribedToTopic(final String chainTaskId) {
         return this.chainTaskIdToSubscription.containsKey(chainTaskId);
     }
 
@@ -113,7 +125,7 @@ public boolean isSubscribedToTopic(String chainTaskId) {
     @EventListener(SessionCreatedEvent.class)
     synchronized void reSubscribeToTopics() {
         log.debug("Received new SessionCreatedEvent");
-        Set<String> chainTaskIds = this.chainTaskIdToSubscription.keySet();
+        final Set<String> chainTaskIds = this.chainTaskIdToSubscription.keySet();
         if (chainTaskIds.isEmpty()) {
             log.info("No topic to resubscribe to");
         } else {
@@ -166,7 +178,7 @@ public void waitForSessionReady() throws InterruptedException {
         }
     }
 
-    private String getTaskTopicName(String chainTaskId) {
+    private String getTaskTopicName(final String chainTaskId) {
         return "/topic/task/" + chainTaskId;
     }
 
@@ -181,17 +193,17 @@ public class MessageHandler implements StompFrameHandler {
         private final String workerWalletAddress;
 
         @Override
-        public Type getPayloadType(StompHeaders headers) {
+        public Type getPayloadType(final StompHeaders headers) {
             return TaskNotification.class;
         }
 
         @Override
-        public void handleFrame(StompHeaders headers, @Nullable Object payload) {
+        public void handleFrame(final StompHeaders headers, @Nullable final Object payload) {
             if (payload == null) {
                 log.error("Payload of TaskNotification is null [chainTaskId:{}]", this.chainTaskId);
                 return;
             }
-            TaskNotification taskNotification = (TaskNotification) payload;
+            final TaskNotification taskNotification = (TaskNotification) payload;
             if (!isWorkerInvolved(taskNotification)) {
                 return;
             }
@@ -200,7 +212,7 @@ public void handleFrame(StompHeaders headers, @Nullable Object payload) {
             eventPublisher.publishEvent(taskNotification);
         }
 
-        private boolean isWorkerInvolved(TaskNotification notification) {
+        private boolean isWorkerInvolved(final TaskNotification notification) {
             return notification.getWorkersAddress() != null &&
                     (notification.getWorkersAddress().isEmpty() || // for all workers
                             notification.getWorkersAddress().contains(this.workerWalletAddress));

src/main/java/com/iexec/worker/task/TaskManagerService.java

@@ -35,7 +35,6 @@
 import com.iexec.worker.compute.post.PostComputeResponse;
 import com.iexec.worker.compute.pre.PreComputeResponse;
 import com.iexec.worker.dataset.DataService;
-import com.iexec.worker.pubsub.SubscriptionService;
 import com.iexec.worker.replicate.ReplicateActionResponse;
 import com.iexec.worker.result.ResultService;
 import com.iexec.worker.sms.SmsService;
@@ -68,7 +67,6 @@ public class TaskManagerService {
     private final DataService dataService;
     private final ResultService resultService;
     private final SmsService smsService;
-    private final SubscriptionService subscriptionService;
     private final PurgeService purgeService;
     private final String workerWalletAddress;
 
@@ -81,7 +79,6 @@ public TaskManagerService(
             DataService dataService,
             ResultService resultService,
             SmsService smsService,
-            SubscriptionService subscriptionService,
             PurgeService purgeService,
             String workerWalletAddress) {
         this.iexecHubService = iexecHubService;
@@ -92,7 +89,6 @@ public TaskManagerService(
         this.dataService = dataService;
         this.resultService = resultService;
         this.smsService = smsService;
-        this.subscriptionService = subscriptionService;
         this.purgeService = purgeService;
         this.workerWalletAddress = workerWalletAddress;
     }
@@ -468,7 +464,6 @@ ReplicateActionResponse complete(String chainTaskId) {
      */
     boolean abort(String chainTaskId) {
         log.info("Aborting task [chainTaskId:{}]", chainTaskId);
-        subscriptionService.unsubscribeFromTopic(chainTaskId);
         boolean allContainersStopped = computeManagerService.abort(chainTaskId);
         boolean allServicesPurged = purgeService.purgeAllServices(chainTaskId);
         final boolean isSuccess = allContainersStopped && allServicesPurged;

src/main/java/com/iexec/worker/task/TaskNotificationService.java

@@ -167,7 +167,6 @@ public void onTaskNotification(TaskNotification notification) {
             case PLEASE_COMPLETE:
                 updateStatusAndGetNextAction(chainTaskId, COMPLETING);
                 actionResponse = taskManagerService.complete(chainTaskId);
-                subscriptionService.unsubscribeFromTopic(chainTaskId);
                 nextStatus = actionResponse.isSuccess() ? COMPLETED : COMPLETE_FAILED;
                 nextAction = updateStatusAndGetNextAction(chainTaskId, nextStatus, actionResponse.getDetails());
                 break;

src/test/java/com/iexec/worker/pubsub/SubscriptionServiceTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2020 IEXEC BLOCKCHAIN TECH
+ * Copyright 2020-2024 IEXEC BLOCKCHAIN TECH
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,12 +16,12 @@
 
 package com.iexec.worker.pubsub;
 
-import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
 import org.mockito.Spy;
+import org.mockito.junit.jupiter.MockitoExtension;
 import org.springframework.context.ApplicationEventPublisher;
 import org.springframework.messaging.simp.stomp.StompSession.Subscription;
 import org.springframework.test.util.ReflectionTestUtils;
@@ -36,8 +36,11 @@
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.*;
 
+@ExtendWith(MockitoExtension.class)
 class SubscriptionServiceTests {
 
+    @Mock
+    private Subscription subscription;
     @Mock
     private ApplicationEventPublisher applicationEventPublisher;
     @Mock
@@ -47,19 +50,10 @@ class SubscriptionServiceTests {
     @InjectMocks
     SubscriptionService subscriptionService;
 
-    private static final String WORKER_WALLET_ADDRESS = "0x1234";
     private static final String CHAIN_TASK_ID = "chaintaskid";
     private static final String CHAIN_TASK_ID_2 = "chaintaskid2";
 
-    @Mock
-    private Subscription subscription;
-
-    @BeforeEach
-    void init() {
-        MockitoAnnotations.openMocks(this);
-        subscriptionService = spy(new SubscriptionService(applicationEventPublisher, stompClientService, WORKER_WALLET_ADDRESS));
-    }
-
+    // region subscribe
     @Test
     void shouldSubscribeToTopic() {
         when(stompClientService.subscribeToTopic(anyString(), any())).thenReturn(Optional.of(subscription));
@@ -79,24 +73,27 @@ void shouldNotSubscribeToExistingTopic() {
         assertThat(subscriptionService.isSubscribedToTopic(CHAIN_TASK_ID)).isTrue();
         verify(stompClientService, atMost(1)).subscribeToTopic(anyString(), any());
     }
+    // endregion
 
+    // region unsubscribe
     @Test
     void shouldUnsubscribeFromTopic() {
         when(stompClientService.subscribeToTopic(anyString(), any())).thenReturn(Optional.of(subscription));
 
         subscriptionService.subscribeToTopic(CHAIN_TASK_ID);
         assertThat(subscriptionService.isSubscribedToTopic(CHAIN_TASK_ID)).isTrue();
-        subscriptionService.unsubscribeFromTopic(CHAIN_TASK_ID);
+        assertThat(subscriptionService.purgeTask(CHAIN_TASK_ID)).isTrue();
+        verify(subscription).unsubscribe();
         assertThat(subscriptionService.isSubscribedToTopic(CHAIN_TASK_ID)).isFalse();
-        verify(subscription, times(1)).unsubscribe();
     }
 
     @Test
     void shouldNotUnsubscribeFromNonexistentTopic() {
         assertThat(subscriptionService.isSubscribedToTopic(CHAIN_TASK_ID)).isFalse();
-        subscriptionService.unsubscribeFromTopic(CHAIN_TASK_ID);
+        assertThat(subscriptionService.purgeTask(CHAIN_TASK_ID)).isTrue();
         verify(subscription, never()).unsubscribe();
     }
+    // endregion
 
     // region reSubscribeToTopics
     @Test
@@ -186,4 +183,26 @@ void shouldWaitForSessionReady() throws ExecutionException, InterruptedException
         future.get(10, TimeUnit.MILLISECONDS);
     }
     // endregion
+
+    // region purgeAllTasksData
+    @Test
+    void shouldPurgeMultipleTasks() {
+        when(stompClientService.subscribeToTopic(anyString(), any())).thenReturn(Optional.of(subscription));
+
+        subscriptionService.subscribeToTopic(CHAIN_TASK_ID);
+        subscriptionService.subscribeToTopic(CHAIN_TASK_ID_2);
+
+        subscriptionService.purgeAllTasksData();
+
+        verify(subscription, times(2)).unsubscribe();
+        assertThat(subscriptionService.isSubscribedToTopic(CHAIN_TASK_ID)).isFalse();
+        assertThat(subscriptionService.isSubscribedToTopic(CHAIN_TASK_ID_2)).isFalse();
+    }
+
+    @Test
+    void shouldHandleEmptyPurgeAllTasksData() {
+        subscriptionService.purgeAllTasksData();
+        verify(subscription, never()).unsubscribe();
+    }
+    // endregion
 }

src/test/java/com/iexec/worker/task/TaskManagerServiceTests.java

@@ -37,19 +37,18 @@
 import com.iexec.worker.compute.post.PostComputeResponse;
 import com.iexec.worker.compute.pre.PreComputeResponse;
 import com.iexec.worker.dataset.DataService;
-import com.iexec.worker.pubsub.SubscriptionService;
 import com.iexec.worker.replicate.ReplicateActionResponse;
 import com.iexec.worker.result.ResultService;
 import com.iexec.worker.sms.SmsService;
 import com.iexec.worker.tee.TeeService;
 import com.iexec.worker.tee.TeeServicesManager;
 import com.iexec.worker.utils.WorkflowException;
 import org.assertj.core.api.Assertions;
-import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.EnumSource;
+import org.mockito.InjectMocks;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 
@@ -75,6 +74,7 @@ class TaskManagerServiceTests {
                     .workerWallet(WORKER_ADDRESS)
                     .build();
 
+    @InjectMocks
     private TaskManagerService taskManagerService;
     @Mock
     private IexecHubService iexecHubService;
@@ -93,30 +93,11 @@ class TaskManagerServiceTests {
     @Mock
     private SmsService smsService;
     @Mock
-    private SubscriptionService subscriptionService;
-    @Mock
     private PurgeService purgeService;
 
     @Mock
     private TeeService teeMockedService;
 
-    @BeforeEach
-    void init() {
-        taskManagerService = new TaskManagerService(
-                iexecHubService,
-                contributionService,
-                revealService,
-                computeManagerService,
-                teeServicesManager,
-                dataService,
-                resultService,
-                smsService,
-                subscriptionService,
-                purgeService,
-                WORKER_ADDRESS
-        );
-    }
-
     TaskDescription.TaskDescriptionBuilder getTaskDescriptionBuilder(boolean isTeeTask) {
         final DealParams dealParams = DealParams.builder()
                 .iexecInputFiles(List.of("https://ab.cd/ef.jpeg"))
@@ -1364,7 +1345,6 @@ void shouldReturnFalseWhenRemainingContainers() {
         when(computeManagerService.abort(CHAIN_TASK_ID)).thenReturn(false);
         when(purgeService.purgeAllServices(CHAIN_TASK_ID)).thenReturn(true);
         assertThat(taskManagerService.abort(CHAIN_TASK_ID)).isFalse();
-        verify(subscriptionService).unsubscribeFromTopic(CHAIN_TASK_ID);
         verify(computeManagerService).abort(CHAIN_TASK_ID);
         verify(purgeService).purgeAllServices(CHAIN_TASK_ID);
     }
@@ -1374,7 +1354,6 @@ void shouldReturnFalseWhenRemainingService() {
         when(computeManagerService.abort(CHAIN_TASK_ID)).thenReturn(true);
         when(purgeService.purgeAllServices(CHAIN_TASK_ID)).thenReturn(false);
         assertThat(taskManagerService.abort(CHAIN_TASK_ID)).isFalse();
-        verify(subscriptionService).unsubscribeFromTopic(CHAIN_TASK_ID);
         verify(computeManagerService).abort(CHAIN_TASK_ID);
         verify(purgeService).purgeAllServices(CHAIN_TASK_ID);
     }
@@ -1384,7 +1363,6 @@ void shouldAbortTask() {
         when(computeManagerService.abort(CHAIN_TASK_ID)).thenReturn(true);
         when(purgeService.purgeAllServices(CHAIN_TASK_ID)).thenReturn(true);
         assertThat(taskManagerService.abort(CHAIN_TASK_ID)).isTrue();
-        verify(subscriptionService).unsubscribeFromTopic(CHAIN_TASK_ID);
         verify(computeManagerService).abort(CHAIN_TASK_ID);
         verify(purgeService).purgeAllServices(CHAIN_TASK_ID);
     }