cavis/deeplearning4j/deeplearning4j-core/src/main/java/org/deeplearning4j/core/parallelism/AsyncIterator.java

/*******************************************************************************
 * Copyright (c) 2015-2018 Skymind, Inc.
 *
 * 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.
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 ******************************************************************************/

package org.deeplearning4j.core.parallelism;

import lombok.Getter;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;

import java.util.Iterator;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Asynchronous Iterator for better performance of iterators in dl4j-nn & dl4j-nlp
 *
 * @author raver119@gmail.com
 */
@Slf4j
public class AsyncIterator<T extends Object> implements Iterator<T> {
    @Getter
    protected BlockingQueue<T> buffer;
    protected ReaderThread<T> thread;
    protected Iterator<T> iterator;
    @Getter
    protected T terminator = (T) new Object();
    protected T nextElement;
    protected AtomicBoolean shouldWork = new AtomicBoolean(true);

    public AsyncIterator(@NonNull Iterator<T> iterator, int bufferSize) {
        this.buffer = new LinkedBlockingQueue<>(bufferSize);
        this.iterator = iterator;

        thread = new ReaderThread<>(iterator, this.buffer, terminator);
        thread.start();
    }

    public AsyncIterator(@NonNull Iterator<T> iterator) {
        this(iterator, 1024);
    }

    @Override
    public boolean hasNext() {
        try {
            if (nextElement != null && nextElement != terminator) {
                return true;
            }

            // if on previous run we've got terminator - just return false
            if (nextElement == terminator)
                return false;

            nextElement = buffer.take();

            // same on this run
            return (nextElement != terminator);
        } catch (Exception e) {
            log.error("Premature end of loop!");
            return false;
        }
    }

    @Override
    public T next() {
        T temp = nextElement;
        nextElement = temp == terminator ? terminator : null;
        return temp;
    }

    @Override
    public void remove() {
        // no-op
    }

    public void shutdown() {
        if (shouldWork.get()) {
            shouldWork.set(false);
            thread.interrupt();
            try {
                // Shutdown() should be a synchronous operation since the iterator is reset after shutdown() is
                // called in AsyncLabelAwareIterator.reset().
                thread.join();
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
            nextElement = terminator;
            buffer.clear();
        }
    }


    private class ReaderThread<T> extends Thread implements Runnable {
        private BlockingQueue<T> buffer;
        private Iterator<T> iterator;
        private T terminator;

        public ReaderThread(Iterator<T> iterator, BlockingQueue<T> buffer, T terminator) {
            this.buffer = buffer;
            this.iterator = iterator;
            this.terminator = terminator;

            setDaemon(true);
            setName("AsyncIterator Reader thread");
        }

        @Override
        public void run() {
            //log.info("AsyncReader [{}] started", Thread.currentThread().getId());
            try {
                while (iterator.hasNext() && shouldWork.get()) {
                    T smth = iterator.next();

                    if (smth != null)
                        buffer.put(smth);
                }
                buffer.put(terminator);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                // do nothing
                shouldWork.set(false);
            } catch (Exception e) {
                // TODO: pass that forward
                throw new RuntimeException(e);
            } finally {
                //log.info("AsyncReader [{}] stopped", Thread.currentThread().getId());
            }
        }
    }
}