io.datakernel.stream.processor.StreamSorter Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of async-streams Show documentation
Composable asynchronous/reactive streams with powerful data processing capabilities.
The newest version!
/*
 * Copyright (C) 2015 SoftIndex LLC.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://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.
 */

package io.datakernel.stream.processor;

import com.google.common.base.Function;
import io.datakernel.async.CompletionCallback;
import io.datakernel.eventloop.Eventloop;
import io.datakernel.stream.*;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

/**
 * Represent {@link AbstractStreamTransformer_1_1} which receives data and saves it in collection, when it
 * receive end of stream it sorts it and streams to destination.
 *
 * @param  type of keys
 * @param  type of objects
 */
public class StreamSorter extends AbstractStreamConsumer implements StreamDataReceiver, StreamSorterMBean {

	private final StreamMergeSorterStorage storage;
	private final Function keyFunction;
	private final Comparator keyComparator;
	private final boolean deduplicate;
	protected final int itemsInMemorySize;

	private final Comparator itemComparator;

	protected List list;
	private List listOfPartitions;

	private StreamProducer saveProducer;

	private StreamForwarder result;

	protected long jmxItems;

	/**
	 * Creates a new instance of StreamSorter
	 *
	 * @param eventloop         event loop in which StreamSorter will run
	 * @param storage           storage for storing elements which was no placed to RAM
	 * @param keyFunction       function for searching key
	 * @param keyComparator     comparator for comparing key
	 * @param deduplicate       if it is true it means that in result will be not objects with same key
	 * @param itemsInMemorySize size of elements which can be saved in RAM before sorting
	 */
	public StreamSorter(Eventloop eventloop, StreamMergeSorterStorage storage,
	                    final Function keyFunction, final Comparator keyComparator, boolean deduplicate,
	                    int itemsInMemorySize) {
		super(eventloop);
		this.storage = checkNotNull(storage);
		this.keyComparator = checkNotNull(keyComparator);
		this.keyFunction = checkNotNull(keyFunction);
		this.deduplicate = deduplicate;
		checkArgument(itemsInMemorySize > 0, "itemsInMemorySize must be positive value, got %s", itemsInMemorySize);
		this.itemsInMemorySize = itemsInMemorySize;

		this.itemComparator = new Comparator() {
			private final Function _keyFunction = keyFunction;
			private final Comparator _keyComparator = keyComparator;

			@Override
			public int compare(T item1, T item2) {
				K key1 = _keyFunction.apply(item1);
				K key2 = _keyFunction.apply(item2);
				return _keyComparator.compare(key1, key2);
			}
		};
		this.list = new ArrayList<>(itemsInMemorySize + (itemsInMemorySize >> 4));
		this.listOfPartitions = new ArrayList<>();
		this.result = new StreamForwarder<>(eventloop);
		this.result.addCompletionCallback(new CompletionCallback() {
			@Override
			public void onComplete() {
				closeUpstream();
			}

			@Override
			public void onException(Exception exception) {
				closeUpstreamWithError(exception);
			}
		});
	}

	public StreamProducer getSortedStream() {
		return result;
	}

	@Override
	public StreamDataReceiver getDataReceiver() {
		return this;
	}

	/**
	 * Adds received data to storage, checks if its count bigger than itemsInMemorySize, if it is
	 * streams it to storage
	 *
	 * @param value received value
	 */
	@SuppressWarnings("AssertWithSideEffects")
	@Override
	public void onData(T value) {
		assert jmxItems != ++jmxItems;
		list.add(value);
		if (list.size() >= itemsInMemorySize) {
			nextState();
		}
	}

	protected void nextState() {
		if (result.getUpstream() != null) {
			return;
		}

		boolean bufferFull = list.size() >= itemsInMemorySize;

		if (saveProducer != null) {
			if (bufferFull) {
				suspendUpstream();
			}
			return;
		}

		if (getUpstreamStatus() == StreamProducer.END_OF_STREAM) {
			StreamMerger merger = StreamMerger.streamMerger(eventloop, keyFunction, keyComparator, deduplicate);

			Collections.sort(list, itemComparator);
			StreamProducer queueProducer = StreamProducers.ofIterable(eventloop, list);
			list = null;

			queueProducer.streamTo(merger.newInput());

			for (int partition : listOfPartitions) {
				storage.streamReader(partition).streamTo(merger.newInput());
			}

			merger.streamTo(result);
			return;
		}

		if (bufferFull) {
			Collections.sort(list, itemComparator);
			saveProducer = StreamProducers.ofIterable(eventloop, list);
			this.listOfPartitions.add(storage.nextPartition());
			StreamConsumer consumer = storage.streamWriter();
			saveProducer.streamTo(consumer);
			saveProducer.addCompletionCallback(new CompletionCallback() {
				@Override
				public void onComplete() {
					eventloop.post(new Runnable() {
						@Override
						public void run() {
							saveProducer = null;
							nextState();
						}
					});
				}

				@Override
				public void onException(Exception e) {
//					closeWithError(e);
					new StreamProducers.ClosingWithError(eventloop, e).streamTo(result);
				}
			});
			this.list = new ArrayList<>(list.size() + (list.size() >> 4));
			return;
		}

		resumeUpstream();
	}

	/**
	 * After producer end of stream it sorts elements and stream it to consumer
	 */
	@Override
	public void onEndOfStream() {
		nextState();
	}

	@Override
	public void onError(Exception e) {
		result.onError(e);
		upstreamProducer.closeWithError(e);
	}

	@Override
	public long getItems() {
		return jmxItems;
	}

	@SuppressWarnings("AssertWithSideEffects")
	@Override
	public String toString() {
		String items = "?";
		assert (items = "" + jmxItems) != null;
		return '{' + super.toString() + " items:" + items + '}';
	}
}