[OOC] Adapt OOCStream to Forward QueueCallbacks

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

@@ -30,8 +30,12 @@ static <T> QueueCallback<T> eos(DMLRuntimeException e) {
 
 	void enqueue(T t);
 
+	void enqueue(QueueCallback<T> callback);
+
 	T dequeue();
 
+	QueueCallback<T> dequeueCB();
+
 	void closeInput();
 
 	void propagateFailure(DMLRuntimeException re);

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

@@ -52,18 +52,23 @@ public void enqueue(IndexedMatrixValue t) {
 		throw new DMLRuntimeException("Cannot enqueue to a playback stream");
 	}
 
+	@Override
+	public void enqueue(QueueCallback<IndexedMatrixValue> callback) {
+		throw new DMLRuntimeException("Cannot enqueue to a playback stream");
+	}
+
 	@Override
 	public void closeInput() {
 		throw new DMLRuntimeException("Cannot close a playback stream");
 	}
 
 	@Override
 	public synchronized IndexedMatrixValue dequeue() {
-		if (_subscriberSet.get())
+		if(_subscriberSet.get())
 			throw new IllegalStateException("Cannot dequeue from a playback stream if a subscriber has been set");
 
 		try {
-			if (_lastDequeue != null)
+			if(_lastDequeue != null)
 				_lastDequeue.close();
 			_lastDequeue = _streamCache.get(_streamIdx.getAndIncrement()).get();
 			return _lastDequeue.get();
@@ -72,6 +77,22 @@ public synchronized IndexedMatrixValue dequeue() {
 		}
 	}
 
+	@Override
+	public synchronized QueueCallback<IndexedMatrixValue> dequeueCB() {
+		if(_subscriberSet.get())
+			throw new IllegalStateException("Cannot dequeue from a playback stream if a subscriber has been set");
+
+		try {
+			if(_lastDequeue != null)
+				_lastDequeue.close();
+			_lastDequeue = _streamCache.get(_streamIdx.getAndIncrement()).get();
+			return _lastDequeue;
+		}
+		catch(InterruptedException | ExecutionException e) {
+			throw new DMLRuntimeException(e);
+		}
+	}
+
 	@Override
 	public OOCStream<IndexedMatrixValue> getReadStream() {
 		return _streamCache.getReadStream();
@@ -114,9 +135,9 @@ public void messageDownstream(OOCStreamMessage msg) {
 		if(msg.isCancelled())
 			return;
 		CopyOnWriteArrayList<Consumer<OOCStreamMessage>> relays = _downstreamRelays;
-		if (relays != null) {
-			for (Consumer<OOCStreamMessage> relay : relays) {
-				if (msg.isCancelled())
+		if(relays != null) {
+			for(Consumer<OOCStreamMessage> relay : relays) {
+				if(msg.isCancelled())
 					break;
 				relay.accept(msg);
 			}
@@ -125,7 +146,7 @@ public void messageDownstream(OOCStreamMessage msg) {
 
 	@Override
 	public void setSubscriber(Consumer<QueueCallback<IndexedMatrixValue>> subscriber) {
-		if (!_subscriberSet.compareAndSet(false, true))
+		if(!_subscriberSet.compareAndSet(false, true))
 			throw new IllegalArgumentException("Subscriber cannot be set multiple times");
 
 		_streamCache.setSubscriber(subscriber, false);
@@ -163,10 +184,10 @@ public void addUpstreamMessageRelay(Consumer<OOCStreamMessage> relay) {
 
 	@Override
 	public void addDownstreamMessageRelay(Consumer<OOCStreamMessage> relay) {
-		if (relay == null)
+		if(relay == null)
 			throw new IllegalArgumentException("Cannot set downstream relay to null");
 		CopyOnWriteArrayList<Consumer<OOCStreamMessage>> relays = _downstreamRelays;
-		if (relays == null) {
+		if(relays == null) {
 			synchronized(this) {
 				if (_downstreamRelays == null)
 					_downstreamRelays = new CopyOnWriteArrayList<>();

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

@@ -35,11 +35,12 @@
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 
-public class SubscribableTaskQueue<T> extends LocalTaskQueue<T> implements OOCStream<T> {
+public class SubscribableTaskQueue<T> extends LocalTaskQueue<OOCStream.QueueCallback<T>> implements OOCStream<T> {
 
 	private final AtomicInteger _availableCtr = new AtomicInteger(1);
 	private final AtomicBoolean _closed = new AtomicBoolean(false);
 	private final AtomicInteger _blockCount = new AtomicInteger(0);
+	private QueueCallback<T> _lastDequeued = null;
 	private CacheableData<?> _cdata;
 	private volatile Consumer<QueueCallback<T>> _subscriber = null;
 	private volatile CopyOnWriteArrayList<Consumer<OOCStreamMessage>> _upstreamMsgRelays = null;
@@ -48,7 +49,7 @@ public class SubscribableTaskQueue<T> extends LocalTaskQueue<T> implements OOCSt
 	private String _watchdogId;
 
 	public SubscribableTaskQueue() {
-		if (OOCWatchdog.WATCH) {
+		if(OOCWatchdog.WATCH) {
 			_watchdogId = "STQ-" + hashCode();
 			// Capture a short context to help identify origin
 			OOCWatchdog.registerOpen(_watchdogId, "SubscribableTaskQueue@" + hashCode(), getCtxMsg(), this);
@@ -71,12 +72,16 @@ private String getCtxMsg() {
 
 	@Override
 	public void enqueue(T t) {
-		if (t == NO_MORE_TASKS)
-			throw new DMLRuntimeException("Cannot enqueue NO_MORE_TASKS item");
+		enqueue(new SimpleQueueCallback<>(t, _failure));
+	}
 
+	@Override
+	public void enqueue(QueueCallback<T> cb) {
+		if(cb == NO_MORE_TASKS)
+			throw new DMLRuntimeException("Cannot enqueue NO_MORE_TASKS item");
 		int cnt = _availableCtr.incrementAndGet();
 
-		if (cnt <= 1) { // Then the queue was already closed and we disallow further enqueues
+		if(cnt <= 1) { // Then the queue was already closed and we disallow further enqueues
 			_availableCtr.decrementAndGet(); // Undo increment
 			throw new DMLRuntimeException("Cannot enqueue into closed SubscribableTaskQueue");
 		}
@@ -86,17 +91,17 @@ public void enqueue(T t) {
 		Consumer<QueueCallback<T>> s = _subscriber;
 		final Consumer<QueueCallback<T>> fS = s;
 
-		if (fS != null) {
-			fS.accept(new SimpleQueueCallback<>(t, _failure));
+		if(fS != null) {
+			fS.accept(cb);
 			onDeliveryFinished();
 			return;
 		}
 
-		synchronized (this) {
+		synchronized(this) {
 			// Re-check that subscriber is really null to avoid race conditions
-			if (_subscriber == null) {
+			if(_subscriber == null) {
 				try {
-					super.enqueueTask(t);
+					super.enqueueTask(cb);
 				}
 				catch(InterruptedException e) {
 					throw new DMLRuntimeException(e);
@@ -108,16 +113,16 @@ public void enqueue(T t) {
 		}
 
 		// Last case if due to race a subscriber has been set
-		s.accept(new SimpleQueueCallback<>(t, _failure));
+		s.accept(cb);
 		onDeliveryFinished();
 	}
 
 	protected boolean tryDeliverCallback(QueueCallback<T> cb, int blockCount) {
 		Consumer<QueueCallback<T>> s = _subscriber;
-		if (s == null)
+		if(s == null)
 			return false;
 		int cnt = _availableCtr.incrementAndGet();
-		if (cnt <= 1) { // Then the queue was already closed and we disallow further enqueues
+		if(cnt <= 1) { // Then the queue was already closed and we disallow further enqueues
 			_availableCtr.decrementAndGet(); // Undo increment
 			throw new DMLRuntimeException("Cannot enqueue into closed SubscribableTaskQueue");
 		}
@@ -128,49 +133,78 @@ protected boolean tryDeliverCallback(QueueCallback<T> cb, int blockCount) {
 	}
 
 	@Override
-	public synchronized void enqueueTask(T t) {
+	public synchronized void enqueueTask(OOCStream.QueueCallback<T> t) {
 		enqueue(t);
 	}
 
 	@Override
 	public T dequeue() {
 		try {
-			if (OOCWatchdog.WATCH)
+			if(OOCWatchdog.WATCH)
+				OOCWatchdog.addEvent(_watchdogId, "dequeue -- " + getCtxMsg());
+			if(_lastDequeued != null) {
+				_lastDequeued.close();
+				_lastDequeued = null;
+			}
+			OOCStream.QueueCallback<T> deq = super.dequeueTask();
+			if(deq != NO_MORE_TASKS) {
+				onDeliveryFinished();
+				_lastDequeued = deq;
+				return deq.get();
+			}
+			return null;
+		}
+		catch(InterruptedException e) {
+			throw new DMLRuntimeException(e);
+		}
+	}
+
+	@Override
+	public OOCStream.QueueCallback<T> dequeueCB() {
+		try {
+			if(OOCWatchdog.WATCH)
 				OOCWatchdog.addEvent(_watchdogId, "dequeue -- " + getCtxMsg());
-			T deq = super.dequeueTask();
-			if (deq != NO_MORE_TASKS)
+			if(_lastDequeued != null) {
+				_lastDequeued.close();
+				_lastDequeued = null;
+			}
+			OOCStream.QueueCallback<T> deq = super.dequeueTask();
+			if(deq != NO_MORE_TASKS) {
 				onDeliveryFinished();
-			return deq;
+				_lastDequeued = deq;
+			}
+			return deq == NO_MORE_TASKS ? null : deq;
 		}
 		catch(InterruptedException e) {
 			throw new DMLRuntimeException(e);
 		}
 	}
 
 	@Override
-	public synchronized T dequeueTask() {
-		return dequeue();
+	public synchronized OOCStream.QueueCallback<T> dequeueTask() {
+		return dequeueCB();
 	}
 
 	@Override
 	public synchronized void closeInput() {
-		if (_closed.compareAndSet(false, true)) {
+		if(_closed.compareAndSet(false, true)) {
 			super.closeInput();
 			onDeliveryFinished();
 			_upstreamMsgRelays = null;
 			_downstreamMsgRelays = null;
-		} else {
+		}
+		else {
 			throw new IllegalStateException("Multiple close input calls");
 		}
 	}
 
 	private void validateBlockCountOnClose() {
 		DataCharacteristics dc = getDataCharacteristics();
-		if (dc != null && dc.dimsKnown() && dc.getBlocksize() > 0) {
+		if(dc != null && dc.dimsKnown() && dc.getBlocksize() > 0) {
 			long expected = OOCUtils.getNumBlocks(dc);
-			if (expected >= 0 && _blockCount.get() != expected) {
-				throw new DMLRuntimeException("OOCStream block count mismatch: expected "
-					+ expected + " but saw " + _blockCount.get() + " (" + dc.getRows() + "x" + dc.getCols() + ")");
+			if(expected >= 0 && _blockCount.get() != expected) {
+				throw new DMLRuntimeException("OOCStream block count mismatch: expected " + expected + " but saw "
+					+ _blockCount.get() + " (" + dc.getRows() + "x" + dc.getCols() + ")");
 			}
 		}
 	}
@@ -180,7 +214,7 @@ public void setSubscriber(Consumer<QueueCallback<T>> subscriber) {
 		if(subscriber == null)
 			throw new IllegalArgumentException("Cannot set subscriber to null");
 
-		LinkedList<T> data;
+		LinkedList<QueueCallback<T>> data;
 		boolean needsEos;
 
 		synchronized(this) {
@@ -198,26 +232,25 @@ public void setSubscriber(Consumer<QueueCallback<T>> subscriber) {
 				_availableCtr.incrementAndGet(); // route terminal emission via onDeliveryFinished
 		}
 
-		for (T t : data) {
-			subscriber.accept(new SimpleQueueCallback<>(t, _failure));
+		for(QueueCallback<T> t : data) {
+			subscriber.accept(t);
 			onDeliveryFinished();
 		}
 
 		if(needsEos)
 			onDeliveryFinished();
 	}
 
-	@SuppressWarnings("unchecked")
 	private void onDeliveryFinished() {
 		int ctr = _availableCtr.decrementAndGet();
 
-		if (ctr == 0) {
+		if(ctr == 0) {
 			validateBlockCountOnClose();
 			Consumer<QueueCallback<T>> s = _subscriber;
-			if (s != null)
-				s.accept(new SimpleQueueCallback<>((T) LocalTaskQueue.NO_MORE_TASKS, _failure));
+			if(s != null)
+				s.accept(OOCStream.eos(_failure));
 
-			if (OOCWatchdog.WATCH)
+			if(OOCWatchdog.WATCH)
 				OOCWatchdog.registerClose(_watchdogId);
 		}
 	}
@@ -248,17 +281,17 @@ public void messageUpstream(OOCStreamMessage msg) {
 		if(msg.isCancelled())
 			return;
 		msg.addIXTransform(_ixTransform);
-		if (msg.isCancelled())
+		if(msg.isCancelled())
 			return;
-		if (msg instanceof OOCGetStreamTypeMessage) {
-			if (_cdata != null)
+		if(msg instanceof OOCGetStreamTypeMessage) {
+			if(_cdata != null)
 				((OOCGetStreamTypeMessage) msg).setInMemoryType();
 			return;
 		}
 		CopyOnWriteArrayList<Consumer<OOCStreamMessage>> relays = _upstreamMsgRelays;
 		if(relays != null) {
-			for (Consumer<OOCStreamMessage> relay : relays) {
-				if (msg.isCancelled())
+			for(Consumer<OOCStreamMessage> relay : relays) {
+				if(msg.isCancelled())
 					break;
 				relay.accept(msg);
 			}
@@ -272,8 +305,8 @@ public void messageDownstream(OOCStreamMessage msg) {
 		msg.addIXTransform(_ixTransform);
 		CopyOnWriteArrayList<Consumer<OOCStreamMessage>> relays = _downstreamMsgRelays;
 		if(relays != null) {
-			for (Consumer<OOCStreamMessage> relay : relays) {
-				if (msg.isCancelled())
+			for(Consumer<OOCStreamMessage> relay : relays) {
+				if(msg.isCancelled())
 					break;
 				relay.accept(msg);
 			}

src/main/java/org/apache/sysds/runtime/ooc/stream/FilteredOOCStream.java

@@ -35,6 +35,7 @@ public class FilteredOOCStream<T> implements OOCStream<T> {
 	private final OOCStream<T> _sourceStream;
 	private final Function<T, Boolean> _predicate;
 	private CacheableData<?> _data;
+	private QueueCallback<T> _last;
 
 	public FilteredOOCStream(OOCStream<T> sourceStream, Function<T, Boolean> predicate) {
 		_sourceStream = sourceStream;
@@ -46,12 +47,26 @@ public void enqueue(T t) {
 		_sourceStream.enqueue(t);
 	}
 
+	@Override
+	public void enqueue(QueueCallback<T> callback) {
+		_sourceStream.enqueue(callback);
+	}
+
 	@Override
 	public synchronized T dequeue() {
-		T next;
-		while((next = _sourceStream.dequeue()) != null) {
-			if(_predicate.apply(next))
-				return next;
+		QueueCallback<T> cb = dequeueCB();
+		return cb == null ? null : cb.get();
+	}
+
+	@Override
+	public synchronized QueueCallback<T> dequeueCB() {
+		if(_last != null)
+			_last.close();
+		while((_last = _sourceStream.dequeueCB()) != null) {
+			if(_predicate.apply(_last.get()))
+				return _last;
+			_last.close();
+			_last = null;
 		}
 		return null;
 	}
@@ -101,6 +116,8 @@ public void setSubscriber(Consumer<QueueCallback<T>> subscriber) {
 
 			if(_predicate.apply(cb.get()))
 				subscriber.accept(cb);
+			else
+				cb.close();
 		});
 	}