org.openlca.proto.io.input.BatchImport Maven / Gradle / Ivy
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();
}
}
}