/* Copyright (c) 2013-2014 Boundless and others.
* All rights reserved. This program and the accompanying materials
* are made available under the terms of the Eclipse Distribution License v1.0
* which accompanies this distribution, and is available at
* https://www.eclipse.org/org/documents/edl-v10.html
*
* Contributors:
* David Winslow (Boundless) - initial implementation
*/
package org.locationtech.geogig.storage.mongo;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicBoolean;
import org.locationtech.geogig.api.ObjectId;
import org.locationtech.geogig.api.RevCommit;
import org.locationtech.geogig.api.RevFeature;
import org.locationtech.geogig.api.RevFeatureType;
import org.locationtech.geogig.api.RevObject;
import org.locationtech.geogig.api.RevTag;
import org.locationtech.geogig.api.RevTree;
import org.locationtech.geogig.repository.RepositoryConnectionException;
import org.locationtech.geogig.storage.BulkOpListener;
import org.locationtech.geogig.storage.ConfigDatabase;
import org.locationtech.geogig.storage.ObjectDatabase;
import org.locationtech.geogig.storage.ObjectInserter;
import org.locationtech.geogig.storage.ObjectSerializingFactory;
import org.locationtech.geogig.storage.ObjectWriter;
import org.locationtech.geogig.storage.datastream.DataStreamSerializationFactoryV1;
import com.google.common.base.Functions;
import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.AbstractIterator;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.inject.Inject;
import com.mongodb.BasicDBObject;
import com.mongodb.BasicDBObjectBuilder;
import com.mongodb.BulkWriteOperation;
import com.mongodb.BulkWriteResult;
import com.mongodb.BulkWriteUpsert;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBCursor;
import com.mongodb.DBObject;
import com.mongodb.MongoClient;
import com.mongodb.WriteConcern;
import com.mongodb.WriteResult;
import com.ning.compress.lzf.LZFInputStream;
import com.ning.compress.lzf.LZFOutputStream;
/**
* An Object database that uses a MongoDB server for persistence.
*
* @see http://mongodb.com/
*/
public class MongoObjectDatabase implements ObjectDatabase {
private final MongoConnectionManager manager;
protected final ConfigDatabase config;
private MongoClient client = null;
protected DB db = null;
protected DBCollection collection = null;
protected ObjectSerializingFactory serializers = DataStreamSerializationFactoryV1.INSTANCE;
private String collectionName;
private ExecutorService executor;
@Inject
public MongoObjectDatabase(ConfigDatabase config, MongoConnectionManager manager,
ExecutorService executor) {
this(config, manager, "objects", executor);
}
MongoObjectDatabase(ConfigDatabase config, MongoConnectionManager manager,
String collectionName, ExecutorService executor) {
this.config = config;
this.manager = manager;
this.executor = executor;
this.collectionName = collectionName;
}
private RevObject fromBytes(ObjectId id, byte[] buffer) {
ByteArrayInputStream byteStream = new ByteArrayInputStream(buffer);
RevObject result;
try {
result = serializers.createObjectReader().read(id, new LZFInputStream(byteStream));
} catch (Exception e) {
throw Throwables.propagate(e);
}
return result;
}
private byte[] toBytes(RevObject object) {
ObjectWriter<RevObject> writer = serializers.createObjectWriter(object.getType());
ByteArrayOutputStream byteStream = new ByteArrayOutputStream();
LZFOutputStream cOut = new LZFOutputStream(byteStream);
try {
writer.write(object, cOut);
} catch (IOException e) {
throw new RuntimeException(e);
}
try {
cOut.close();
} catch (IOException e) {
throw Throwables.propagate(e);
}
return byteStream.toByteArray();
}
protected String getCollectionName() {
return collectionName;
}
@Override
public synchronized void open() {
if (client != null) {
return;
}
String uri = config.get("mongodb.uri").get();
String database = config.get("mongodb.database").get();
client = manager.acquire(new MongoAddress(uri));
db = client.getDB(database);
collection = db.getCollection(getCollectionName());
collection.ensureIndex("oid");
}
@Override
public synchronized boolean isOpen() {
return client != null;
}
@Override
public void configure() throws RepositoryConnectionException {
RepositoryConnectionException.StorageType.OBJECT.configure(config, "mongodb", "0.1");
String uri = config.get("mongodb.uri").or(config.getGlobal("mongodb.uri"))
.or("mongodb://localhost:27017/");
String database = config.get("mongodb.database").or(config.getGlobal("mongodb.database"))
.or("geogig");
config.put("mongodb.uri", uri);
config.put("mongodb.database", database);
}
@Override
public void checkConfig() throws RepositoryConnectionException {
RepositoryConnectionException.StorageType.OBJECT.verify(config, "mongodb", "0.1");
}
@Override
public synchronized void close() {
if (client != null) {
manager.release(client);
}
client = null;
db = null;
collection = null;
}
@Override
public boolean exists(ObjectId id) {
DBObject query = new BasicDBObject();
query.put("oid", id.toString());
return collection.find(query).hasNext();
}
@Override
public List<ObjectId> lookUp(final String partialId) {
if (partialId.matches("[a-fA-F0-9]+")) {
DBObject regex = new BasicDBObject();
regex.put("$regex", "^" + partialId);
DBObject query = new BasicDBObject();
query.put("oid", regex);
DBCursor cursor = collection.find(query);
List<ObjectId> ids = new ArrayList<ObjectId>();
while (cursor.hasNext()) {
DBObject elem = cursor.next();
String oid = (String) elem.get("oid");
ids.add(ObjectId.valueOf(oid));
}
return ids;
} else {
throw new IllegalArgumentException(
"Prefix query must be done with hexadecimal values only");
}
}
@Override
public RevObject get(ObjectId id) {
RevObject result = getIfPresent(id);
if (result != null) {
return result;
} else {
throw new NoSuchElementException("No object with id: " + id);
}
}
@Override
public <T extends RevObject> T get(ObjectId id, Class<T> clazz) {
return clazz.cast(get(id));
}
@Override
public RevObject getIfPresent(ObjectId id) {
DBObject query = new BasicDBObject();
query.put("oid", id.toString());
DBCursor results = collection.find(query);
if (results.hasNext()) {
DBObject result = results.next();
return fromBytes(id, (byte[]) result.get("serialized_object"));
} else {
return null;
}
}
@Override
public <T extends RevObject> T getIfPresent(ObjectId id, Class<T> clazz) {
return clazz.cast(getIfPresent(id));
}
@Override
public RevTree getTree(ObjectId id) {
return get(id, RevTree.class);
}
@Override
public RevFeature getFeature(ObjectId id) {
return get(id, RevFeature.class);
}
@Override
public RevFeatureType getFeatureType(ObjectId id) {
return get(id, RevFeatureType.class);
}
@Override
public RevCommit getCommit(ObjectId id) {
return get(id, RevCommit.class);
}
@Override
public RevTag getTag(ObjectId id) {
return get(id, RevTag.class);
}
private long deleteChunk(List<ObjectId> ids) {
List<String> idStrings = Lists.transform(ids, Functions.toStringFunction());
DBObject query = BasicDBObjectBuilder.start().push("oid").add("$in", idStrings).pop().get();
WriteResult result = collection.remove(query);
return result.getN();
}
@Override
public boolean delete(ObjectId id) {
DBObject query = new BasicDBObject();
query.put("oid", id.toString());
return collection.remove(query).getLastError().ok();
}
@Override
public long deleteAll(Iterator<ObjectId> ids) {
return deleteAll(ids, BulkOpListener.NOOP_LISTENER);
}
@Override
public long deleteAll(Iterator<ObjectId> ids, BulkOpListener listener) {
Iterator<List<ObjectId>> chunks = Iterators.partition(ids, 500);
long count = 0;
while (chunks.hasNext()) {
count += deleteChunk(chunks.next());
}
return count;
}
@Override
public boolean put(final RevObject object) {
DBObject query = new BasicDBObject();
query.put("oid", object.getId().toString());
DBObject record = toDocument(object);
return collection.update(query, record, true, false).getLastError().ok();
}
private DBObject toDocument(final RevObject object) {
DBObject record = new BasicDBObject();
record.put("oid", object.getId().toString());
record.put("serialized_object", toBytes(object));
return record;
}
@Override
public void putAll(final Iterator<? extends RevObject> objects) {
putAll(objects, BulkOpListener.NOOP_LISTENER);
}
@Override
public void putAll(Iterator<? extends RevObject> objects, BulkOpListener listener) {
Preconditions.checkNotNull(executor, "executor service not set");
if (!objects.hasNext()) {
return;
}
final int bulkSize = 1000;
final int maxRunningTasks = 10;
final AtomicBoolean cancelCondition = new AtomicBoolean();
List<ObjectId> ids = Lists.newArrayListWithCapacity(bulkSize);
List<Future<?>> runningTasks = new ArrayList<Future<?>>(maxRunningTasks);
BulkWriteOperation bulkOperation = collection.initializeOrderedBulkOperation();
try {
while (objects.hasNext()) {
RevObject object = objects.next();
bulkOperation.insert(toDocument(object));
ids.add(object.getId());
if (ids.size() == bulkSize || !objects.hasNext()) {
InsertTask task = new InsertTask(bulkOperation, listener, ids, cancelCondition);
runningTasks.add(executor.submit(task));
if (objects.hasNext()) {
bulkOperation = collection.initializeOrderedBulkOperation();
ids = Lists.newArrayListWithCapacity(bulkSize);
}
}
if (runningTasks.size() == maxRunningTasks) {
waitForTasks(runningTasks);
}
}
waitForTasks(runningTasks);
} catch (RuntimeException e) {
cancelCondition.set(true);
throw e;
}
}
private void waitForTasks(List<Future<?>> runningTasks) {
// wait...
for (Future<?> f : runningTasks) {
try {
f.get();
} catch (Exception e) {
throw Throwables.propagate(e);
}
}
runningTasks.clear();
}
private static class InsertTask implements Runnable {
private BulkWriteOperation bulkOperation;
private BulkOpListener listener;
private List<ObjectId> ids;
private AtomicBoolean cancelCondition;
public InsertTask(BulkWriteOperation bulkOperation, BulkOpListener listener,
List<ObjectId> ids, AtomicBoolean cancelCondition) {
this.bulkOperation = bulkOperation;
this.listener = listener;
this.ids = ids;
this.cancelCondition = cancelCondition;
}
@Override
public void run() {
if (cancelCondition.get()) {
return;
}
BulkWriteResult bulkResult = bulkOperation.execute(WriteConcern.ACKNOWLEDGED);
List<BulkWriteUpsert> upserts = bulkResult.getUpserts();
for (BulkWriteUpsert upsert : upserts) {
if (cancelCondition.get()) {
return;
}
int index = upsert.getIndex();
ObjectId existing = ids.set(index, null);
listener.found(existing, null);
}
for (ObjectId inserted : ids) {
if (cancelCondition.get()) {
return;
}
if (inserted != null) {
listener.inserted(inserted, null);
}
}
ids.clear();
}
}
@Override
public ObjectInserter newObjectInserter() {
return new ObjectInserter(this);
}
@Override
public Iterator<RevObject> getAll(Iterable<ObjectId> ids) {
return getAll(ids, BulkOpListener.NOOP_LISTENER);
}
@Override
public Iterator<RevObject> getAll(final Iterable<ObjectId> ids, final BulkOpListener listener) {
return new AbstractIterator<RevObject>() {
final Iterator<ObjectId> queryIds = ids.iterator();
@Override
protected RevObject computeNext() {
RevObject obj = null;
while (obj == null) {
if (!queryIds.hasNext()) {
return endOfData();
}
ObjectId id = queryIds.next();
obj = getIfPresent(id);
if (obj == null) {
listener.notFound(id);
} else {
listener.found(obj.getId(), null);
}
}
return obj == null ? endOfData() : obj;
}
};
}
public DBCollection getCollection(String name) {
return db.getCollection(name);
}
@Override
public String toString() {
return String.format("%s[db: %s, collection: %s]", getClass().getSimpleName(),
db == null ? "<unset>" : db, collectionName);
}
}