[MINOR] Fix Flaky Tests from Blocking Thread Pools (#2489)

Author: janniklinde

src/main/java/org/apache/sysds/runtime/instructions/ooc/AggregateUnaryOOCInstruction.java

@@ -35,7 +35,6 @@
 import org.apache.sysds.runtime.matrix.operators.AggregateUnaryOperator;
 import org.apache.sysds.runtime.matrix.operators.Operator;
 import org.apache.sysds.runtime.meta.DataCharacteristics;
-import org.apache.sysds.runtime.ooc.stream.StreamContext;
 import org.apache.sysds.runtime.util.IndexRange;
 
 import java.util.HashMap;
@@ -119,9 +118,10 @@ public void processInstruction( ExecutionContext ec ) {
 			});
 
 			// global reduce
-			submitOOCTask(() -> {
-				IndexedMatrixValue partial;
-				while ((partial = qLocal.dequeue()) != LocalTaskQueue.NO_MORE_TASKS) {
+			addOutStream(qOut);
+			submitOOCTasks(qLocal, callback -> {
+				IndexedMatrixValue partial = callback.get();
+				synchronized(aggTracker) {
 					long idx = aggun.isRowAggregate() ? partial.getIndexes().getRowIndex() : partial.getIndexes()
 						.getColumnIndex();
 
@@ -150,8 +150,7 @@ public void processInstruction( ExecutionContext ec ) {
 						corrs.remove(idx);
 					}
 				}
-				qOut.closeInput();
-			}, new StreamContext().addOutStream(qOut));
+			}).thenRun(qOut::closeInput);
 		}
 		// full aggregation
 		else {

src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCInstruction.java

@@ -1196,7 +1196,7 @@ protected <T> CompletableFuture<Void> submitOOCTasks(OOCStream<T> queue, Consume
 	}
 
 	protected CompletableFuture<Void> submitOOCTask(Runnable r, StreamContext ctx) {
-		ExecutorService pool = CommonThreadPool.get();
+		ExecutorService pool = CommonThreadPool.getDynamicPool();
 		final CompletableFuture<Void> future = new CompletableFuture<>();
 		try {
 			COMPUTE_IN_FLIGHT.incrementAndGet();
@@ -1220,9 +1220,6 @@ protected CompletableFuture<Void> submitOOCTask(Runnable r, StreamContext ctx) {
 			COMPUTE_IN_FLIGHT.decrementAndGet();
 			throw new DMLRuntimeException(ex);
 		}
-		finally {
-			pool.shutdown();
-		}
 
 		return future;
 	}