package com.findwise.hydra.memorydb;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.ConcurrentModificationException;
import java.util.Date;
import java.util.HashSet;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.findwise.hydra.DatabaseDocument;
import com.findwise.hydra.DatabaseQuery;
import com.findwise.hydra.Document;
import com.findwise.hydra.DocumentFile;
import com.findwise.hydra.DocumentID;
import com.findwise.hydra.DocumentReader;
import com.findwise.hydra.DocumentWriter;
import com.findwise.hydra.JsonException;
import com.findwise.hydra.TailableIterator;
import com.findwise.hydra.local.LocalDocumentID;
public class MemoryDocumentIO implements DocumentWriter<MemoryType>,
DocumentReader<MemoryType> {
private ConcurrentHashMap<MemoryDocument, Boolean> set;
private LinkedBlockingQueue<MemoryDocument> inactive;
private boolean[] b = new boolean[1];
private static Logger logger = LoggerFactory
.getLogger(MemoryDocumentIO.class);
private HashSet<DocumentFile<MemoryType>> files;
public static final int inactiveSize = 100;
public MemoryDocumentIO() {
set = new ConcurrentHashMap<MemoryDocument, Boolean>();
files = new HashSet<DocumentFile<MemoryType>>();
inactive = new LinkedBlockingQueue<MemoryDocument>(inactiveSize);
b[0] = false;
}
private void addInactive(MemoryDocument d) {
if (!inactive.offer(d)) {
inactive.poll();
b[0] = true;
inactive.offer(d);
}
}
@Override
public MemoryDocument getDocument(DatabaseQuery<MemoryType> q) {
List<DatabaseDocument<MemoryType>> list = getDocuments(q, 1);
if (list.size() == 0) {
return null;
}
return (MemoryDocument) list.get(0);
}
@Override
public MemoryDocument getDocumentById(DocumentID<MemoryType> id) {
return (MemoryDocument) getDocumentById(id, false);
}
@Override
public DatabaseDocument<MemoryType> getDocumentById(DocumentID<MemoryType> id,
boolean includeInactive) {
if (id == null) {
return null;
}
for (MemoryDocument d : set.keySet()) {
if (id.equals(d.getID())) {
return d;
}
}
if (includeInactive) {
for (MemoryDocument d : inactive) {
if (id.equals(d.getID())) {
return d;
}
}
}
return null;
}
@Override
public TailableIterator<MemoryType> getInactiveIterator() {
return getInactiveIterator(new MemoryQuery());
}
@SuppressWarnings("unchecked")
@Override
public List<DatabaseDocument<MemoryType>> getDocuments(
DatabaseQuery<MemoryType> q, int limit, int skip) {
ArrayList<MemoryDocument> list = new ArrayList<MemoryDocument>();
int matching = 0;
for (MemoryDocument doc : set.keySet()) {
if (list.size() >= limit)
break;
if (doc.matches((MemoryQuery) q)) {
if (matching >= skip) {
list.add(doc);
}
matching++;
}
}
return (List<DatabaseDocument<MemoryType>>) (Object) list;
}
@Override
public List<DatabaseDocument<MemoryType>> getDocuments(
DatabaseQuery<MemoryType> q, int limit) {
return getDocuments(q, limit, 0);
}
@Override
public long getNumberOfDocuments(DatabaseQuery<MemoryType> q) {
long matching = 0;
for (MemoryDocument doc : set.keySet()) {
if (doc.matches((MemoryQuery) q)) {
matching++;
}
}
return matching;
}
@Override
public DocumentFile<MemoryType> getDocumentFile(
DatabaseDocument<MemoryType> d, String fileName) {
for (DocumentFile<MemoryType> f : files) {
if (f.getDocumentId().getID().equals(d.getID().getID())
&& f.getFileName().equals(fileName)) {
try {
return copy(f);
} catch (IOException e) {
logger.error("Error copying the streams", e);
}
}
}
return null;
}
@Override
public boolean deleteDocumentFile(DatabaseDocument<MemoryType> d,
String fileName) {
for (DocumentFile<MemoryType> f : files) {
if (f.getDocumentId().equals(d.getID())
&& f.getFileName().equals(fileName)) {
files.remove(f);
return true;
}
}
return false;
}
@Override
public List<String> getDocumentFileNames(DatabaseDocument<MemoryType> d) {
ArrayList<String> list = new ArrayList<String>();
for (DocumentFile<MemoryType> f : files) {
if (f.getDocumentId().equals(d.getID())) {
list.add(f.getFileName());
}
}
return list;
}
@Override
public long getActiveDatabaseSize() {
return set.size();
}
@Override
public long getInactiveDatabaseSize() {
return inactive.size();
}
@Override
public MemoryDocument getAndTag(DatabaseQuery<MemoryType> query, String ... tag) {
for(String t : tag) {
((MemoryQuery) query).requireNotFetchedByStage(t);
}
MemoryDocument d = getDocument(query);
if (d != null) {
for(String t : tag) {
d.tag(Document.FETCHED_METADATA_TAG, t);
}
}
return d;
}
@Override
public Collection<DatabaseDocument<MemoryType>> getAndTag(
DatabaseQuery<MemoryType> query, int n, String ... tag) {
for(String t : tag) {
((MemoryQuery) query).requireNotFetchedByStage(t);
}
List<DatabaseDocument<MemoryType>> docs = getDocuments(query, n);
for (int i = 0; i < docs.size() && i < n; i++) {
DatabaseDocument<MemoryType> d = docs.get(i);
for(String t : tag) {
((MemoryDocument) d).tag(Document.FETCHED_METADATA_TAG, t);
}
}
return docs;
}
@Override
public boolean markTouched(DocumentID<MemoryType> id, String tag) {
MemoryDocument d = getDocumentById(id);
if (d == null) {
return false;
}
d.tag(Document.TOUCHED_METADATA_TAG, tag);
return true;
}
@Override
public boolean markProcessed(DatabaseDocument<MemoryType> d, String stage) {
return markDone(d, stage, Document.PROCESSED_METADATA_FLAG);
}
@Override
public boolean markDiscarded(DatabaseDocument<MemoryType> d, String stage) {
return markDone(d, stage, Document.DISCARDED_METADATA_FLAG);
}
@Override
public boolean markFailed(DatabaseDocument<MemoryType> d, String stage) {
return markDone(d, stage, Document.FAILED_METADATA_FLAG);
}
private boolean markDone(DatabaseDocument<MemoryType> d, String stage,
String flag) {
MemoryDocument temp = getDocumentById(d.getID());
if (temp == null) {
return false;
}
set.remove(temp);
((MemoryDocument) d).tag(flag, stage);
deleteAllFiles(d);
addInactive((MemoryDocument) d);
return true;
}
@Override
public boolean markPending(DatabaseDocument<MemoryType> d, String stage) {
MemoryDocument temp = getDocumentById(d.getID());
if (temp == null) {
return false;
}
temp.tag(Document.PENDING_METADATA_FLAG, stage);
return true;
}
@Override
public boolean insert(DatabaseDocument<MemoryType> d) {
MemoryDocument md = (MemoryDocument) d;
md.setID(new MemoryDocumentID(new LocalDocumentID(md.hashCode() + ""
+ System.currentTimeMillis())));
removeNullFields(md);
set.put(md, false);
md.markSynced();
return true;
}
@Override
public boolean insert(DatabaseDocument<MemoryType> d, List<DocumentFile<MemoryType>> attachments) {
if(attachments == null || attachments.isEmpty()) {
return insert(d);
}
d.putMetadataField(Document.COMMITTING_METADATA_FLAG, true);
if (!insert(d)) {
return false;
};
if (!writeAttachments(d, attachments)) {
delete(d);
return false;
}
d.putMetadataField(Document.COMMITTING_METADATA_FLAG, false);
return update(d);
}
private boolean writeAttachments(DatabaseDocument<MemoryType> d, List<DocumentFile<MemoryType>> attachments) {
for(DocumentFile<MemoryType> attachment: attachments) {
attachment.setDocumentId(d.getID());
try {
write(attachment);
} catch (IOException e) {
logger.error(
String.format(
"Exception while writing filename:%s for id:%s",
attachment.getFileName(),
d.getID()
),
e
);
return false;
}
}
return true;
}
private void removeNullFields(MemoryDocument md) {
HashSet<String> fields = new HashSet<String>();
for (String entry : md.getTouchedContent()) {
if (!md.hasContentField(entry)
|| md.getContentField(entry).equals(null)) {
fields.add(entry);
}
}
for (String field : fields) {
md.removeContentField(field);
}
}
private HashSet<String> getTouchedContentSnapshot(MemoryDocument md) {
while (true) {
try {
return new HashSet<String>(md.getTouchedContent());
} catch (ConcurrentModificationException e) {
logger.warn("Got concurrent modification in getTouchedContent");
}
}
}
private HashSet<String> getTouchedMetadataSnapshot(MemoryDocument md) {
while (true) {
try {
return new HashSet<String>(md.getTouchedMetadata());
} catch (ConcurrentModificationException e) {
logger.warn("Got concurrent modification in getTouchedContent");
}
}
}
@Override
public boolean update(DatabaseDocument<MemoryType> d) {
MemoryDocument md = (MemoryDocument) d;
MemoryDocument inDb = getDocumentById(d.getID());
if (inDb == null) {
set.put(md, false);
inDb = getDocumentById(d.getID());
}
if (inDb == null) {
logger.error("Can't find document with id '" + d.getID().getID()
+ "', unable to update.");
return false;
}
if (md.isTouchedAction()) {
inDb.setAction(md.getAction());
}
for (String s : getTouchedContentSnapshot(md)) {
if (md.getContentField(s) != null) {
inDb.putContentField(s, md.getContentField(s));
} else {
inDb.removeContentField(s);
}
}
for (String s : getTouchedMetadataSnapshot(md)) {
inDb.putMetadataField(s, md.getMetadataMap().get(s));
}
md.markSynced();
return true;
}
@Override
public void delete(DatabaseDocument<MemoryType> d) {
deleteAllFiles(d);
set.remove(getDocumentById(d.getID()));
}
private void deleteAllFiles(DatabaseDocument<MemoryType> d) {
for (String fileName : getDocumentFileNames(d)) {
deleteDocumentFile(d, fileName);
}
}
@Override
public void deleteAll() {
set.clear();
}
@Override
public void write(DocumentFile<MemoryType> df) throws IOException {
df.setUploadDate(new Date());
files.add(copy(df));
df.getStream().close();
}
private DocumentFile<MemoryType> copy(DocumentFile<MemoryType> df)
throws IOException {
String s = IOUtils.toString(df.getStream(), df.getEncoding());
df.getStream().close();
df.setStream(IOUtils.toInputStream(s, df.getEncoding()));
return new DocumentFile<MemoryType>(df.getDocumentId(),
df.getFileName(), IOUtils.toInputStream(s, df.getEncoding()),
df.getSavedByStage(), df.getUploadDate());
}
@Override
public void prepare() {
}
@Override
public DocumentID<MemoryType> toDocumentId(Object jsonPrimitive) {
return new MemoryDocumentID(new LocalDocumentID(jsonPrimitive));
}
@Override
public DocumentID<MemoryType> toDocumentIdFromJson(String json) {
try {
return new MemoryDocumentID(LocalDocumentID.getDocumentID(json));
} catch (JsonException e) {
logger.error("Unable to deserialize document id", e);
return null;
}
}
@Override
public TailableIterator<MemoryType> getInactiveIterator(
DatabaseQuery<MemoryType> query) {
return new MemoryTailableIterator(inactive, b, new MemoryQuery());
}
}