feat: restart interrupted data flows (#4612)

* feat: restart interrupted data flows

* PR remarks

* restart only push transfers

* pr remark

core/common/junit/src/main/java/org/eclipse/edc/junit/extensions/EmbeddedRuntime.java

@@ -31,6 +31,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicReference;
 
 import static java.util.concurrent.TimeUnit.SECONDS;
@@ -46,6 +47,7 @@ public class EmbeddedRuntime extends BaseRuntime {
     private final ExecutorService executorService = Executors.newSingleThreadExecutor();
     private final MultiSourceServiceLocator serviceLocator;
     private final URL[] classPathEntries;
+    private Future<?> runtimeThread;
 
     public EmbeddedRuntime(String name, Map<String, String> properties, String... additionalModules) {
         this(new MultiSourceServiceLocator(), name, properties, ClasspathReader.classpathFor(additionalModules));
@@ -76,7 +78,7 @@ public void boot(boolean addShutdownHook) {
             var runtimeException = new AtomicReference<Exception>();
             var latch = new CountDownLatch(1);
 
-            executorService.execute(() -> {
+            runtimeThread = executorService.submit(() -> {
                 try {
                     var classLoader = URLClassLoader.newInstance(classPathEntries);
 
@@ -108,6 +110,9 @@ public void boot(boolean addShutdownHook) {
     public void shutdown() {
         serviceLocator.clearSystemExtensions();
         super.shutdown();
+        if (runtimeThread != null && !runtimeThread.isDone()) {
+            runtimeThread.cancel(true);
+        }
     }
 
     @Override

core/common/lib/query-lib/src/main/java/org/eclipse/edc/query/CriterionOperatorRegistryImpl.java

@@ -45,6 +45,7 @@ public static CriterionOperatorRegistry ofDefaults() {
         registry.registerOperatorPredicate(ILIKE, new IlikeOperatorPredicate());
         registry.registerOperatorPredicate(CONTAINS, new ContainsOperatorPredicate());
         registry.registerOperatorPredicate(NOT_EQUAL, new NotEqualOperatorPredicate());
+        registry.registerOperatorPredicate(LESS_THAN, new LessThanOperatorPredicate());
         return registry;
     }
 

core/common/lib/query-lib/src/main/java/org/eclipse/edc/query/LessThanOperatorPredicate.java

@@ -0,0 +1,34 @@
+/*
+ *  Copyright (c) 2024 Cofinity-X
+ *
+ *  This program and the accompanying materials are made available under the
+ *  terms of the Apache License, Version 2.0 which is available at
+ *  https://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Contributors:
+ *       Cofinity-X - initial API and implementation
+ *
+ */
+
+package org.eclipse.edc.query;
+
+import org.eclipse.edc.spi.query.OperatorPredicate;
+
+import java.util.Comparator;
+
+public class LessThanOperatorPredicate implements OperatorPredicate {
+    @Override
+    public boolean test(Object value, Object comparedTo) {
+        if (value instanceof Number number1 && comparedTo instanceof Number number2) {
+            return Double.compare(number1.doubleValue(), number2.doubleValue()) < 0;
+        }
+
+        if (value instanceof String string1 && comparedTo instanceof String string2) {
+            return Comparator.<String>naturalOrder().compare(string1, string2) < 0;
+        }
+
+        return false;
+    }
+}

core/common/lib/query-lib/src/test/java/org/eclipse/edc/query/LessThanOperatorPredicateTest.java

@@ -0,0 +1,73 @@
+/*
+ *  Copyright (c) 2024 Cofinity-X
+ *
+ *  This program and the accompanying materials are made available under the
+ *  terms of the Apache License, Version 2.0 which is available at
+ *  https://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Contributors:
+ *       Cofinity-X - initial API and implementation
+ *
+ */
+
+package org.eclipse.edc.query;
+
+import org.eclipse.edc.spi.query.OperatorPredicate;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.ArgumentsProvider;
+import org.junit.jupiter.params.provider.ArgumentsSource;
+
+import java.util.stream.Stream;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.params.provider.Arguments.arguments;
+
+class LessThanOperatorPredicateTest {
+
+    private final OperatorPredicate predicate = new LessThanOperatorPredicate();
+
+    @ParameterizedTest
+    @ArgumentsSource(ValidValues.class)
+    void shouldReturnTrue_whenValueLessThanComparedOne(Object value, Object comparedTo) {
+        assertThat(predicate.test(value, comparedTo)).isTrue();
+    }
+
+    @ParameterizedTest
+    @ArgumentsSource(InvalidValues.class)
+    void shouldReturnFalse_whenValueNotLessThanComparedOne(Object value, Object comparedTo) {
+        assertThat(predicate.test(value, comparedTo)).isFalse();
+    }
+
+    private static class ValidValues implements ArgumentsProvider {
+
+        @Override
+        public Stream<? extends Arguments> provideArguments(ExtensionContext context) {
+            return Stream.of(
+                    arguments(1, 2),
+                    arguments(1, 2L),
+                    arguments(1, 1.01f),
+                    arguments(1, 1.01d),
+                    arguments("a", "b")
+            );
+        }
+    }
+
+    private static class InvalidValues implements ArgumentsProvider {
+
+        @Override
+        public Stream<? extends Arguments> provideArguments(ExtensionContext context) {
+            return Stream.of(
+                    arguments(1, 1),
+                    arguments(1, 1L),
+                    arguments(1, 1.0f),
+                    arguments(1, 1.0d),
+                    arguments("a", "a")
+            );
+        }
+    }
+
+}

core/common/lib/state-machine-lib/src/main/java/org/eclipse/edc/statemachine/AbstractStateEntityManager.java

@@ -55,7 +55,6 @@ public abstract class AbstractStateEntityManager<E extends StatefulEntity<E>, S
 
     @Override
     public void start() {
-        entityRetryProcessFactory = new EntityRetryProcessFactory(monitor, clock, entityRetryProcessConfiguration);
         var stateMachineManagerBuilder = StateMachineManager.Builder
                 .newInstance(getClass().getSimpleName(), monitor, executorInstrumentation, waitStrategy);
         stateMachineManager = configureStateMachineManager(stateMachineManagerBuilder).build();

core/data-plane/data-plane-core/src/main/java/org/eclipse/edc/connector/dataplane/framework/DataPlaneFrameworkExtension.java

@@ -143,6 +143,7 @@ public void initialize(ServiceExtensionContext context) {
 
     @Override
     public void start() {
+        dataPlaneManager.restartFlows();
         dataPlaneManager.start();
     }
 
@@ -151,6 +152,7 @@ public void shutdown() {
         if (dataPlaneManager != null) {
             dataPlaneManager.stop();
         }
+        pipelineService.closeAll();
     }
 
     @Provider

core/data-plane/data-plane-core/src/main/java/org/eclipse/edc/connector/dataplane/framework/manager/DataPlaneManagerImpl.java

@@ -35,6 +35,7 @@
 import org.eclipse.edc.statemachine.StateMachineManager;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.function.Function;
@@ -47,6 +48,8 @@
 import static org.eclipse.edc.spi.persistence.StateEntityStore.hasState;
 import static org.eclipse.edc.spi.response.ResponseStatus.FATAL_ERROR;
 import static org.eclipse.edc.spi.result.Result.success;
+import static org.eclipse.edc.spi.types.domain.transfer.FlowType.PULL;
+import static org.eclipse.edc.spi.types.domain.transfer.FlowType.PUSH;
 
 /**
  * Default data manager implementation.
@@ -65,7 +68,7 @@ private DataPlaneManagerImpl() {
     public Result<Boolean> validate(DataFlowStartMessage dataRequest) {
         // TODO for now no validation for pull scenario, since the transfer service registry
         //  is not applicable here. Probably validation only on the source part required.
-        if (FlowType.PULL.equals(dataRequest.getFlowType())) {
+        if (PULL.equals(dataRequest.getFlowType())) {
             return success(true);
         } else {
             var transferService = transferServiceRegistry.resolveTransferService(dataRequest);
@@ -121,6 +124,26 @@ public StatusResult<Void> terminate(String dataFlowId, @Nullable String reason)
                 });
     }
 
+    @Override
+    public StatusResult<Void> restartFlows() {
+        var now = clock.millis();
+        List<DataFlow> toBeRestarted;
+        do {
+            toBeRestarted = store.nextNotLeased(batchSize,
+                    hasState(STARTED.code()),
+                    new Criterion("stateTimestamp", "<", now),
+                    new Criterion("transferType.flowType", "=", PUSH.toString())
+            );
+
+            toBeRestarted.forEach(dataFlow -> {
+                dataFlow.transitToReceived();
+                processReceived(dataFlow);
+            });
+        } while (!toBeRestarted.isEmpty());
+
+        return StatusResult.success();
+    }
+
     @Override
     protected StateMachineManager.Builder configureStateMachineManager(StateMachineManager.Builder builder) {
         return builder
@@ -206,7 +229,7 @@ private boolean processReceived(DataFlow dataFlow) {
         store.save(dataFlow);
 
         return entityRetryProcessFactory.doAsyncProcess(dataFlow, () -> transferService.transfer(request))
-                .entityRetrieve(id -> store.findById(id))
+                .entityRetrieve(id -> store.findByIdAndLease(id).orElse(f -> null))
                 .onSuccess((f, r) -> {
                     if (f.getState() != STARTED.code()) {
                         return;
@@ -279,6 +302,7 @@ public Builder self() {
 
         @Override
         public DataPlaneManagerImpl build() {
+            super.build();
             Objects.requireNonNull(manager.transferProcessClient);
             return manager;
         }
@@ -297,6 +321,7 @@ public Builder authorizationService(DataPlaneAuthorizationService authorizationS
             manager.authorizationService = authorizationService;
             return this;
         }
+
     }
 
 }

core/data-plane/data-plane-core/src/main/java/org/eclipse/edc/connector/dataplane/framework/pipeline/PipelineServiceImpl.java

@@ -120,6 +120,11 @@ public StreamResult<Void> terminate(DataFlow dataFlow) {
         return terminate(dataFlow.getId());
     }
 
+    @Override
+    public void closeAll() {
+        sources.forEach((processId, source) -> terminate(processId));
+    }
+
     @Override
     public void registerFactory(DataSourceFactory factory) {
         sourceFactories.add(factory);

core/data-plane/data-plane-core/src/test/java/org/eclipse/edc/connector/dataplane/framework/DataPlaneFrameworkExtensionTest.java

@@ -15,6 +15,7 @@
 package org.eclipse.edc.connector.dataplane.framework;
 
 import org.eclipse.edc.connector.dataplane.framework.registry.TransferServiceRegistryImpl;
+import org.eclipse.edc.connector.dataplane.spi.pipeline.PipelineService;
 import org.eclipse.edc.connector.dataplane.spi.registry.TransferServiceRegistry;
 import org.eclipse.edc.junit.extensions.DependencyInjectionExtension;
 import org.eclipse.edc.spi.system.ExecutorInstrumentation;
@@ -24,12 +25,17 @@
 import org.junit.jupiter.api.extension.ExtendWith;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 
 @ExtendWith(DependencyInjectionExtension.class)
 class DataPlaneFrameworkExtensionTest {
 
+    private final PipelineService pipelineService = mock();
+
     @BeforeEach
     public void setUp(ServiceExtensionContext context) {
+        context.registerService(PipelineService.class, pipelineService);
         context.registerService(ExecutorInstrumentation.class, ExecutorInstrumentation.noop());
     }
 
@@ -40,4 +46,10 @@ void initialize_registers_transferService(ServiceExtensionContext context, DataP
         assertThat(context.getService(TransferServiceRegistry.class)).isInstanceOf(TransferServiceRegistryImpl.class);
     }
 
+    @Test
+    void shouldClosePipelineService_whenShutdown(DataPlaneFrameworkExtension extension) {
+        extension.shutdown();
+
+        verify(pipelineService).closeAll();
+    }
 }