update TransferProcess after every check

Author: paullatzelsperger

extensions/transfer/transfer-core/src/main/java/com/microsoft/dagx/transfer/core/transfer/TransferProcessManagerImpl.java

@@ -139,10 +139,10 @@ private int checkProvisioned() {
                 } else {
                     process.transitionStreaming();
                 }
-                transferProcessStore.update(process);
             } else {
                 monitor.debug("Process " + process.getId() + " does not yet have provisioned resources, will stay in " + TransferProcessStates.REQUESTED_ACK);
             }
+            transferProcessStore.update(process);
         }
 
         return requestAcked.size();
@@ -164,7 +164,6 @@ private int checkCompleted() {
             // todo: maybe error out processes with uncheckable resources??
             final List<ProvisionedResource> resources = process.getProvisionedResourceSet().getResources().stream().filter(this::hasChecker).collect(Collectors.toList());
 
-
             //todo: comment this in if we want to error out uncheckable resources
 //            var resourcesWithNoChecker = resources.stream().filter(resource -> statusCheckerRegistry.resolve(resource) == null).collect(Collectors.toList());
 //            if (!resourcesWithNoChecker.isEmpty()) {
@@ -178,10 +177,9 @@ private int checkCompleted() {
             // update the process once ALL resources are completed
             if (resources.stream().allMatch(this::isComplete)) {
                 process.transitionCompleted();
-                transferProcessStore.update(process);
                 monitor.debug("Process " + process.getId() + " is now " + TransferProcessStates.COMPLETED);
             }
-
+            transferProcessStore.update(process);
         }
         return processesInProgress.size();
     }

extensions/transfer/transfer-core/src/test/java/com/microsoft/dagx/transfer/core/transfer/TransferManagerImplConsumerTest.java renamed to extensions/transfer/transfer-core/src/test/java/com/microsoft/dagx/transfer/core/transfer/TransferProcessManagerImplConsumerTest.java

@@ -24,7 +24,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.easymock.EasyMock.*;
 
-public class TransferManagerImplConsumerTest {
+class TransferProcessManagerImplConsumerTest {
 
     private static final long TIMEOUT = 5;
     private TransferProcessManagerImpl transferProcessManager;
@@ -243,6 +243,8 @@ void verifyCompleted_notAllYetCompleted() throws InterruptedException {
             cdl.countDown();
             return Collections.emptyList();
         }).anyTimes();
+        processStoreMock.update(eq(process));
+        expectLastCall().anyTimes();
         replay(processStoreMock);
 
         // prepare statuschecker registry

extensions/transfer/transfer-store-memory/src/main/java/com/microsoft/dagx/transfer/store/memory/InMemoryTransferProcessStore.java

@@ -47,7 +47,11 @@ public String processIdForTransferId(String id) {
     public @NotNull List<TransferProcess> nextForState(int state, int max) {
         return readLock(() -> {
             var set = stateCache.get(state);
-            return set == null ? Collections.emptyList() : set.stream().limit(max).map(TransferProcess::copy).collect(toList());
+            return set == null ? Collections.emptyList() : set.stream()
+                    .sorted(Comparator.comparingLong(TransferProcess::getStateTimestamp)) //order by state timestamp, oldest first
+                    .limit(max)
+                    .map(TransferProcess::copy)
+                    .collect(toList());
         });
     }
 
@@ -67,6 +71,7 @@ public void create(TransferProcess process) {
     @Override
     public void update(TransferProcess process) {
         writeLock(() -> {
+            process.updateStateTimestamp();
             delete(process.getId());
             TransferProcess internalCopy = process.copy();
             processesByExternalId.put(process.getDataRequest().getId(), internalCopy);

extensions/transfer/transfer-store-memory/src/test/java/com/microsoft/dagx/transfer/store/memory/InMemoryTransferProcessStoreTest.java

@@ -5,22 +5,18 @@
 
 package com.microsoft.dagx.transfer.store.memory;
 
-import com.microsoft.dagx.spi.types.domain.transfer.DataRequest;
-import com.microsoft.dagx.spi.types.domain.transfer.ResourceManifest;
-import com.microsoft.dagx.spi.types.domain.transfer.TransferProcess;
-import com.microsoft.dagx.spi.types.domain.transfer.TransferProcessStates;
+import com.microsoft.dagx.spi.types.domain.transfer.*;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import java.util.List;
 import java.util.UUID;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNotNull;
-import static org.junit.jupiter.api.Assertions.assertNotSame;
-import static org.junit.jupiter.api.Assertions.assertNull;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.easymock.EasyMock.niceMock;
+import static org.easymock.EasyMock.replay;
+import static org.junit.jupiter.api.Assertions.*;
 
 /**
  *
@@ -108,6 +104,46 @@ void verifyMutlipleRequets() {
 
     }
 
+    @Test
+    void verifyOrderingByTimestamp() {
+        for (int i = 0; i < 100; i++) {
+            final TransferProcess process = createProcess("test-process-" + i);
+            store.create(process);
+        }
+
+        final List<TransferProcess> processes = store.nextForState(TransferProcessStates.INITIAL.code(), 50);
+
+        assertThat(processes).hasSize(50);
+        assertThat(processes).allMatch(p -> p.getStateTimestamp() > 0);
+    }
+
+    @Test
+    void verifyNextForState_avoidsStarvation() throws InterruptedException {
+        for (int i = 0; i < 10; i++) {
+            final TransferProcess process = createProcess("test-process-" + i);
+            store.create(process);
+        }
+
+        var list1 = store.nextForState(TransferProcessStates.INITIAL.code(), 5);
+        Thread.sleep(50); //simulate a short delay to generate different timestamps
+        list1.forEach(tp -> store.update(tp));
+        var list2 = store.nextForState(TransferProcessStates.INITIAL.code(), 5);
+        assertThat(list1).isNotEqualTo(list2).doesNotContainAnyElementsOf(list2);
+    }
+
+    private TransferProcess createProcess(String name) {
+        final DataRequest mock = niceMock(DataRequest.class);
+        replay(mock);
+        return TransferProcess.Builder.newInstance()
+                .type(TransferProcess.Type.CLIENT)
+                .id(name)
+                .stateTimestamp(0)
+                .state(TransferProcessStates.UNSAVED.code())
+                .provisionedResourceSet(new ProvisionedResourceSet())
+                .dataRequest(mock)
+                .build();
+    }
+
 
     @BeforeEach
     void setUp() {

spi/src/main/java/com/microsoft/dagx/spi/transfer/store/TransferProcessStore.java

@@ -22,6 +22,20 @@ public interface TransferProcessStore {
     @Nullable
     String processIdForTransferId(String id);
 
+    /**
+     * Returns a list of TransferProcesses that are in a specific state.
+     * <p>
+     * Implementors MUST handle the starvation scenario, i.e. when the number of processes is greater than the number
+     * passedin via {@code max}.
+     * E.g. database-based implementations should perform a query along the lines of {@code SELECT ... ORDER BY TransferProcess#stateTimestamp}.
+     * Then, after the check, users of this method must update the {@code TransferProcess#stateTimestamp} even if the process
+     * remains unchanged.
+     * Some database frameworks such as Spring have automatic lastChanged columns.
+     *
+     * @param state The state that the processes of interest should be in.
+     * @param max   The maximum amount of result items.
+     * @return A list of TransferProcesses (at most _max_) that are in the desired state.
+     */
     @NotNull
     List<TransferProcess> nextForState(int state, int max);
 

spi/src/main/java/com/microsoft/dagx/spi/types/domain/transfer/TransferProcess.java

@@ -197,14 +197,14 @@ public void transitionError(@Nullable String errorDetail) {
         state = TransferProcessStates.ERROR.code();
         this.errorDetail = errorDetail;
         stateCount = 1;
-        stateTimestamp = Instant.now().toEpochMilli();
+        updateStateTimestamp();
     }
 
 
     public void rollbackState(TransferProcessStates state) {
         this.state = state.code();
         stateCount = 1;
-        stateTimestamp = Instant.now().toEpochMilli();
+        updateStateTimestamp();
     }
 
     public TransferProcess copy() {
@@ -222,7 +222,7 @@ private void transition(TransferProcessStates end, TransferProcessStates... star
         }
         stateCount = state == end.code() ? stateCount + 1 : 1;
         state = end.code();
-        stateTimestamp = Instant.now().toEpochMilli();
+        updateStateTimestamp();
     }
 
     @Override
@@ -242,6 +242,18 @@ public int hashCode() {
         return Objects.hash(id);
     }
 
+    public void updateStateTimestamp() {
+        stateTimestamp = Instant.now().toEpochMilli();
+    }
+
+    @Override
+    public String toString() {
+        return "TransferProcess{" +
+                "id='" + id + '\'' +
+                ", state=" + state +
+                ", stateTimestamp=" + Instant.ofEpochMilli(stateTimestamp) +
+                '}';
+    }
 
     public enum Type {
         CLIENT, PROVIDER