added a test to verify the non-starvation feature

Author: paullatzelsperger

extensions/transfer/transfer-core/build.gradle.kts

@@ -11,6 +11,7 @@ plugins {
 dependencies {
     api(project(":spi"))
 
+    testImplementation(project(":extensions:transfer:transfer-store-memory"))
 
 }
 

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

@@ -13,11 +13,14 @@
 import com.microsoft.dagx.spi.transfer.provision.ResourceManifestGenerator;
 import com.microsoft.dagx.spi.transfer.store.TransferProcessStore;
 import com.microsoft.dagx.spi.types.domain.transfer.*;
+import com.microsoft.dagx.transfer.store.memory.InMemoryTransferProcessStore;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
 
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
@@ -27,10 +30,12 @@
 class TransferProcessManagerImplConsumerTest {
 
     private static final long TIMEOUT = 5;
+    private final static int TRANSFER_MANAGER_BATCHSIZE = 10;
     private TransferProcessManagerImpl transferProcessManager;
     private ProvisionManager provisionManager;
     private RemoteMessageDispatcherRegistry dispatcherRegistry;
     private StatusCheckerRegistry statusCheckerRegistry;
+    private ExponentialWaitStrategy waitStrategyMock;
 
     @BeforeEach
     void setup() {
@@ -40,11 +45,17 @@ void setup() {
         ResourceManifestGenerator manifestGenerator = mock(ResourceManifestGenerator.class);
 
         statusCheckerRegistry = mock(StatusCheckerRegistry.class);
+
+        waitStrategyMock = partialMockBuilder(ExponentialWaitStrategy.class)
+                .withConstructor(1000L)
+                .addMockedMethod("success").strictMock();
+
+
         transferProcessManager = TransferProcessManagerImpl.Builder.newInstance()
                 .provisionManager(provisionManager)
                 .dataFlowManager(dataFlowManager)
-                .waitStrategy(new ExponentialWaitStrategy(1000))
-                .batchSize(10)
+                .waitStrategy(waitStrategyMock)
+                .batchSize(TRANSFER_MANAGER_BATCHSIZE)
                 .dispatcherRegistry(dispatcherRegistry)
                 .manifestGenerator(manifestGenerator)
                 .monitor(mock(Monitor.class))
@@ -303,17 +314,68 @@ void verifyCompleted_noCheckerForSomeResources() throws InterruptedException {
         assertThat(process.getState()).describedAs("State should be COMPLETED").isEqualTo(TransferProcessStates.COMPLETED.code());
     }
 
+    @Test
+    @DisplayName("Verify that no process 'starves' during two consecutive runs, when the batch size > number of processes")
+    void verifyProvision_shouldNotStarve() throws InterruptedException {
+        var numProcesses = TRANSFER_MANAGER_BATCHSIZE * 2;
+
+        //prepare process store
+        final TransferProcessStore inMemoryProcessStore = new InMemoryTransferProcessStore();
+
+        //create a few processes
+        var processes = new ArrayList<TransferProcess>();
+        for (int i = 0; i < numProcesses; i++) {
+            final TransferProcess process = createTransferProcess(TransferProcessStates.UNSAVED);
+            processes.add(process);
+            inMemoryProcessStore.create(process);
+        }
+
+        var processesToProvision = new CountDownLatch(numProcesses); //all processes should be provisioned
+
+        //prepare provision manager
+        provisionManager.provision(anyObject(TransferProcess.class));
+        expectLastCall().andAnswer(() -> {
+            processesToProvision.countDown();
+            return null;
+        }).anyTimes();
+        replay(provisionManager);
+
+        // use the waitstrategy to count the number of iterations by making sure "success" was called exactly twice
+        waitStrategyMock.success();
+        expectLastCall().times(2);
+        replay(waitStrategyMock);
+
+
+        //act
+        transferProcessManager.start(inMemoryProcessStore);
+
+        //assert
+        assertThat(processesToProvision.await(TIMEOUT, TimeUnit.SECONDS)).isTrue();
+        verify(provisionManager);
+        verify(waitStrategyMock);
+        assertThat(processes).describedAs("All transfer processes should be in PROVISIONING state").allSatisfy(p -> {
+            var id = p.getId();
+            var storedProcess = inMemoryProcessStore.find(id);
+            assertThat(storedProcess).describedAs("Should exist in the TransferProcessStore").isNotNull();
+            assertThat(storedProcess.getState()).isEqualTo(TransferProcessStates.PROVISIONING.code());
+        });
+    }
+
     private TransferProcess createTransferProcess(TransferProcessStates inState) {
         return createTransferProcess(inState, new TransferType());
     }
 
     private TransferProcess createTransferProcess(TransferProcessStates inState, TransferType type) {
+
+        String processId = UUID.randomUUID().toString();
+
         final DataRequest mock = niceMock(DataRequest.class);
         expect(mock.getTransferType()).andReturn(type).anyTimes();
+        expect(mock.getId()).andReturn(processId).anyTimes();
         replay(mock);
         return TransferProcess.Builder.newInstance()
                 .state(inState.code())
-                .id("test-process-id")
+                .id("test-process-" + processId)
                 .provisionedResourceSet(new ProvisionedResourceSet())
                 .type(TransferProcess.Type.CLIENT)
                 .dataRequest(mock)

extensions/transfer/transfer-demo-protocols/src/test/java/com/microsoft/dagx/transfer/demo/protocols/DemoProtocolsTransferExtensionTest.java

@@ -0,0 +1,100 @@
+package com.microsoft.dagx.transfer.demo.protocols;
+
+import com.microsoft.dagx.junit.DagxExtension;
+import com.microsoft.dagx.spi.monitor.Monitor;
+import com.microsoft.dagx.spi.security.Vault;
+import com.microsoft.dagx.spi.security.VaultResponse;
+import com.microsoft.dagx.spi.transfer.TransferProcessManager;
+import com.microsoft.dagx.spi.transfer.TransferWaitStrategy;
+import com.microsoft.dagx.spi.types.domain.metadata.DataEntry;
+import com.microsoft.dagx.spi.types.domain.transfer.DataAddress;
+import com.microsoft.dagx.spi.types.domain.transfer.DataRequest;
+import com.microsoft.dagx.transfer.demo.protocols.spi.DemoProtocols;
+import com.microsoft.dagx.transfer.demo.protocols.spi.stream.DestinationManager;
+import org.jetbrains.annotations.Nullable;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+/**
+ *
+ */
+@ExtendWith(DagxExtension.class)
+@Disabled
+class DemoProtocolsTransferExtensionTest {
+
+    /**
+     * Perform a push stream flow using the loopback protocol.
+     *
+     * @param processManager     the injected process manager
+     * @param destinationManager the injected destination manager
+     * @param monitor            the injected runtime monitor
+     */
+    @Test
+    void verifyPushStreamFlow(TransferProcessManager processManager, DestinationManager destinationManager, Monitor monitor) throws InterruptedException {
+        var latch = new CountDownLatch(1);
+
+        var destinationName = UUID.randomUUID().toString();
+        destinationManager.registerObserver((name, payload) -> {
+            monitor.info("Message: " + new String(payload));
+            latch.countDown();
+        });
+
+        var dataEntry = DataEntry.Builder.newInstance().id("test123").build();
+
+        var dataRequest = DataRequest.Builder.newInstance()
+                .id(UUID.randomUUID().toString())
+                .protocol("loopback")
+                .destinationType(DemoProtocols.PUSH_STREAM)
+                .dataEntry(dataEntry)
+                .dataDestination(DataAddress.Builder.newInstance().type(DemoProtocols.PUSH_STREAM)
+                        .property(DemoProtocols.DESTINATION_NAME, destinationName).build())
+                .connectorId("test").build();
+
+        processManager.initiateClientRequest(dataRequest);
+
+        latch.await(1, TimeUnit.MINUTES);
+    }
+
+    /**
+     * Fixture that obtains a reference to the runtime.
+     *
+     * @param extension the injected runtime instance
+     */
+    @BeforeEach
+    void before(DagxExtension extension) {
+        // register a mock Vault
+        extension.registerServiceMock(Vault.class, new MockVault());
+
+        // register a wait strategy of 1ms to speed up the interval between transfer manager iterations
+        extension.registerServiceMock(TransferWaitStrategy.class, () -> 1);
+    }
+
+    private static class MockVault implements Vault {
+        private final Map<String, String> secrets = new ConcurrentHashMap<>();
+
+        @Override
+        public @Nullable String resolveSecret(String key) {
+            return secrets.get(key);
+        }
+
+        @Override
+        public VaultResponse storeSecret(String key, String value) {
+            secrets.put(key, value);
+            return VaultResponse.OK;
+        }
+
+        @Override
+        public VaultResponse deleteSecret(String key) {
+            secrets.remove(key);
+            return VaultResponse.OK;
+        }
+    }
+}

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

@@ -250,7 +250,7 @@ public void updateStateTimestamp() {
     public String toString() {
         return "TransferProcess{" +
                 "id='" + id + '\'' +
-                ", state=" + state +
+                ", state=" + TransferProcessStates.from(state) +
                 ", stateTimestamp=" + Instant.ofEpochMilli(stateTimestamp) +
                 '}';
     }