D-Net

package eu.dnetlib.resolver.mdstore.plugin;

import com.mongodb.BasicDBObject;

import com.mongodb.BasicDBObjectBuilder;

import com.mongodb.DBObject;

import com.mongodb.QueryBuilder;

import com.mongodb.client.FindIterable;

import com.mongodb.client.MongoCollection;

import com.mongodb.client.MongoDatabase;

import com.mongodb.client.model.Filters;

import com.mongodb.client.model.UpdateOptions;

import eu.dnetlib.data.mdstore.modular.action.DoneCallback;

import eu.dnetlib.data.mdstore.modular.action.MDStorePlugin;

import eu.dnetlib.data.mdstore.modular.connector.MDStoreDao;

import eu.dnetlib.data.mdstore.modular.mongodb.MDStoreTransactionManagerImpl;

import eu.dnetlib.dli.resolver.PIDResolver;

import eu.dnetlib.dli.resolver.model.serializer.ResolverSerializer;

import eu.dnetlib.rmi.data.MDStoreServiceException;

import org.apache.commons.lang3.StringUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.bson.Document;

import org.bson.conversions.Bson;

import org.springframework.beans.factory.annotation.Autowired;

import java.util.ArrayList;

import java.util.List;

import java.util.Map;

import java.util.concurrent.*;

import static com.mongodb.client.model.Filters.eq;

/**

 * Created by sandro on 9/9/16.

 */

public class ResolverMDStorePlugin implements MDStorePlugin {

	private static final Log log = LogFactory.getLog(ResolverMDStorePlugin.class);

	public static DBObject DONE = new BasicDBObject();

	@Autowired

	private List<PIDResolver> pluginResolver;

	@Autowired

	private ResolverSerializer resolverSerializer;

	@Autowired

	private MDStoreTransactionManagerImpl transactionManager;

	public static void save(MongoCollection<DBObject> collection, DBObject document) {

		Object id = document.get("_id");

		if (id == null) {

			collection.insertOne(document);

		} else {

			collection.replaceOne(eq("_id", id), document, new UpdateOptions().upsert(true));

		}

	}

	@Override

	public void run(final MDStoreDao dao, final Map<String, String> params, final DoneCallback doneCallback) throws MDStoreServiceException {

		try {

			final String id = params.get("mdStoreId");

            int numberOfThreads = 4;

            final String numOfThreadsString = params.get("numberOfThreads");

            try {

                if (!StringUtils.isEmpty(numOfThreadsString)) {

                    numberOfThreads = Integer.parseInt(numOfThreadsString);

                }

            } catch (Throwable e) {

                log.error("Number of threads Param is not an int value it will apply by default 4");

            }

			final boolean refresh = params.get("refresh") != null && Boolean.parseBoolean(params.get("refresh"));

			final String internalId = transactionManager.readMdStore(id);

			final MongoDatabase db = transactionManager.getDb();

			final MongoCollection<DBObject> currentMdStoreCollection = db.getCollection(internalId, DBObject.class);

			final MongoCollection<DBObject> resolvedRecord = db.getCollection("resolved_" + StringUtils.substringBefore(id, "_"), DBObject.class);

            final BasicDBObject idx = new BasicDBObject();

            idx.put("resolved_ts", 1);

            resolvedRecord.createIndex(idx);

            upsertResolved(currentMdStoreCollection, resolvedRecord, 0);

			if (refresh) {

				resolvedRecord.drop();

			}

			final FindIterable<DBObject> mdstoreRecords = currentMdStoreCollection.find();

			final BlockingQueue<DBObject> queue = new ArrayBlockingQueue<>(100);

			final List<Future<Boolean>> responses = new ArrayList<>();

			final ExecutorService executor = Executors.newFixedThreadPool(100);

			final long total = currentMdStoreCollection.count();

			int previousPrintValue = -1;

			int currentPerc = 0;

            final long ts = System.currentTimeMillis();

            for (int i = 1; i < numberOfThreads; i++) {

                final RecordResolver resolver = new RecordResolver(ts);

                resolver.setInputQueue(queue);

				resolver.setOutputCollection(resolvedRecord);

				resolver.setSerializer(resolverSerializer);

				resolver.setPluginResolver(pluginResolver);

				responses.add(executor.submit(resolver));

			}

			int parsed = 0;

			for (DBObject currentMdStoreRecord : mdstoreRecords) {

				queue.put(currentMdStoreRecord);

				currentPerc = Math.round(((float) ++parsed / (float) total) * 100.0F);

				if (currentPerc != previousPrintValue) {

					log.info("Resolving process " + currentPerc + " %");

					previousPrintValue = currentPerc;

				}

			}

			queue.put(DONE);

			for (Future<Boolean> response : responses) {

				response.get();

			}

            upsertResolved(currentMdStoreCollection, resolvedRecord, ts - 1);

            doneCallback.call(params);

		} catch (Throwable e) {

            log.error(e);

            throw new MDStoreServiceException("Error on resolving records ", e);

		}

	}

    private void upsertResolved(MongoCollection<DBObject> currentMdStoreCollection, MongoCollection<DBObject> resolvedRecord, final long timestamp) {

        log.info("Updating resolved objects");

        final Bson queryByTs = Filters.gte("resolved_ts", Long.valueOf(timestamp));

        int i = 0;

        final FindIterable<DBObject> dbObjects = timestamp == 0 ? resolvedRecord.find() : resolvedRecord.find(queryByTs);

        for (DBObject object : dbObjects) {

			Bson query = Filters.eq("id", object.get("id").toString());

			final DBObject replacedObj = BasicDBObjectBuilder.start()

					.add("body", object.get("body").toString())

                    .add("resolved_ts", object.get("resolved_ts"))

                    .get();

			Bson newDocument = new Document("$set", replacedObj);

			currentMdStoreCollection.findOneAndUpdate(query, newDocument);

            i++;

		}

        log.info("Updated " + i);

    }

}