org.openlca.proto.io.input.BatchImport Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of olca-proto-io Show documentation
The newest version!
package org.openlca.proto.io.input;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

import org.openlca.core.model.ModelType;
import org.openlca.core.model.RootEntity;

class BatchImport {

	private final ProtoImport imp;
	private final Class clazz;
	private final ModelType type;
	private final int batchSize;
	private final Writer writer = new Writer();

	private final ArrayList inserts = new ArrayList<>();
	private final ArrayList updates = new ArrayList<>();

	BatchImport(ProtoImport imp, Class type, int batchSize) {
		this.imp = imp;
		this.type = imp.types.get(type);
		this.clazz = type;
		this.batchSize = batchSize;
	}

	static int batchSizeOf(ModelType type) {
		return switch (type) {
			case IMPACT_CATEGORY, PRODUCT_SYSTEM -> 1;
			case LOCATION, PROCESS, RESULT -> 100;
			default -> 1000;
		};
	}

	void run() {
		for (var refId : imp.reader.getIds(type)) {
			var item = imp.fetch(clazz, refId);
			if (item.isVisited() || item.isError())
				continue;
			if (item.isNew()) {
				insert(item.proto().read(imp));
			} else {
				T model = item.entity();
				item.proto().update(model, imp);
				update(model);
			}
		}
		if (inserts.size() > 0) {
			flushInserts();
		}
		if (updates.size() > 0) {
			flushUpdates();
		}
		writer.close();
	}

	private void insert(RootEntity entity) {
		inserts.add(entity);
		if (inserts.size() >= batchSize) {
			flushInserts();
		}
	}

	private void update(RootEntity entity) {
		updates.add(entity);
		if (updates.size() >= batchSize) {
			flushUpdates();
		}
	}

	private void flushInserts() {
		writer.insert(inserts);
		inserts.clear();
	}

	private void flushUpdates() {
		writer.update(updates);
		updates.clear();
	}

	private class Writer {

		private final ExecutorService exec = Executors.newFixedThreadPool(1);
		private final ArrayList buffer = new ArrayList<>();
		private volatile Future task;

		void insert(List batch) {
			next(batch);
			task = exec.submit(() -> {
				imp.db().transaction(em -> buffer.forEach(em::persist));
				buffer.forEach(imp::visited);
				buffer.clear();
			});
		}

		void update(List batch) {
			next(batch);
			task = exec.submit(() -> {
				imp.db().transaction(em -> buffer.replaceAll(em::merge));
				buffer.forEach(imp::visited);
				buffer.clear();
			});
		}

		private void next(List batch) {
			flush();
			buffer.clear();
			buffer.addAll(batch);
		}

		private void flush() {
			if (task != null) {
				try {
					task.get();
					task = null;
				} catch (Exception e) {
					throw new RuntimeException("failed to wait for worker", e);
				}
			}
		}

		private void close() {
			flush();
			exec.shutdown();
		}
	}

}