/gpoptimizer-server/src/ru/ifmo/gpo/server/storage/MongoDBStorageManager.java

package ru.ifmo.gpo.server.storage;



import com.google.gson.Gson;

import com.google.gson.GsonBuilder;

import com.mongodb.*;

import com.mongodb.util.JSON;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;

import ru.ifmo.gpo.core.instructions.InstructionSequence;

import ru.ifmo.gpo.core.instructions.generic.IGenericInstruction;

import ru.ifmo.gpo.java.instructions.JVMInstruction;

import ru.ifmo.gpo.server.IStorageManager;

import ru.ifmo.gpo.util.CollectionUtils;

import ru.ifmo.gpo.util.ProbabilitySet;



import java.lang.reflect.Array;

import java.util.HashMap;

import java.util.Map;

import java.util.Random;

import java.util.Set;



/**

 * User: e_smirnov

 * Date: 11.09.2010

 * Time: 17:08:44

 * <p/>

 * Saves optimization results into Mongo databse.

 * Results are saved with following fields:

 * size - size of source sequence

 * success - flag showing if optimization was successfull

 * hash - source hash

 * source - source instruction sequence, serialized into JSON by gson library

 * optimized - optimization result, may be null if optimization has failed

 * instructionClass - fully qualified name for IGenericInstruction subclass that is used in these instruction sequences

 */

public class MongoDBStorageManager implements IStorageManager {



    private static final String collectionPrefix = "seq_";



    private DB mongoDb;

    private Logger logger = LoggerFactory.getLogger(getClass());



    private static final String sourceFieldName = "source";

    private static final String optimizedFieldName = "optimized";



    private static Gson jsonSerializer;



    @Override

    public void addOptimizationRecord(InstructionSequence source, InstructionSequence optimized) {

        if (source == null || optimized == null || source.isEmpty()) {

            logger.warn("Attempting to save null or empty sequences, ignoring");

            return;

        }

        logger.debug("Adding optimization record, source is {}, optimized is {}", source, optimized);



        DBCollection collection = mongoDb.getCollection(buildCollectionName(source));



        BasicDBObject newRecord = new BasicDBObject();

        newRecord.put("size", source.size());

        newRecord.put("success", true);

        newRecord.put("hash", source.hashCode());

        serializeInstructionSequence(newRecord, sourceFieldName, source);

        serializeInstructionSequence(newRecord, optimizedFieldName, optimized);



        newRecord.put("instructionClass", source.get(0).getClass().getCanonicalName());



        collection.insert(newRecord);

    }



    private DBObject findRecordForSource(InstructionSequence source) {

        DBCollection collection = mongoDb.getCollection(buildCollectionName(source));



        BasicDBObject query = new BasicDBObject();

        query.put("hash", source.hashCode());

        DBCursor cursor = collection.find(query);

        while (cursor.hasNext()) {

            DBObject obj = cursor.next();

            InstructionSequence sourceFromDb = deserialize(obj, sourceFieldName);

            if (sourceFromDb == null) {

                // record found but some error happened on deserialization

                return null;

            }

            if (sourceFromDb.equals(source)) {

                return obj;

            }

        }

        return null;

    }



    @Override

    public void addFailedOptimizationRecord(InstructionSequence source) {

        if (source == null || source.isEmpty()) {

            logger.warn("Attempting to save null or empty sequences, ignoring");

            return;

        }

        if (findRecordForSource(source) != null) {

            // database already contains record for this source

            return;

        }



        logger.debug("Adding failed optimization record, source is {}", source);



        DBCollection collection = mongoDb.getCollection(buildCollectionName(source));



        BasicDBObject newRecord = new BasicDBObject();

        newRecord.put("size", source.size());

        newRecord.put("success", false);

        newRecord.put("hash", source.hashCode());

        serializeInstructionSequence(newRecord, sourceFieldName, source);

        newRecord.put(optimizedFieldName, null);



        newRecord.put("instructionClass", source.get(0).getClass().getCanonicalName());



        collection.insert(newRecord);

    }



    @Override

    public InstructionSequence getOptimized(InstructionSequence source) {

        DBObject dbObj = findRecordForSource(source);

        if (dbObj == null) {

            return null;

        }

        return deserialize(dbObj, optimizedFieldName);

    }



    @Override

    public InstructionSequence getRandomRecordWithSolution() {

        return CollectionUtils.selectRandomElement(getAllResults().keySet()); // surely not very fast on big databases

    }



    @Override

    public InstructionSequence getRandomRecordWithoutSolution() {

        Random random = new Random(System.currentTimeMillis());

        BasicDBObject query = new BasicDBObject();

        query.put("success", false);

        ProbabilitySet<DBCollection> collectionSet = new ProbabilitySet<DBCollection>(random);

        for (String s : mongoDb.getCollectionNames()) {

            DBCollection coll = mongoDb.getCollection(s);

            collectionSet.put(coll, (double) coll.count(query));

        }

        if (collectionSet.isEmpty()) {

            return null;

        }

        DBCollection selectedCollection = collectionSet.getRandom();

        int recordIdx = random.nextInt(collectionSet.get(selectedCollection).intValue());

        DBCursor cur = selectedCollection.find(query);

        for (int idx = 0; cur.hasNext(); idx++) {

            DBObject obj = cur.next();

            if (idx == recordIdx) {

                return deserialize(obj, sourceFieldName);

            }

        }

        return null;

    }



    @Override

    public int getRecordsNum() {

        Set<String> collections = mongoDb.getCollectionNames();

        int rz = 0;

        for (String s : collections) {

            DBCollection coll = mongoDb.getCollection(s);

            rz += coll.count();

        }

        return rz;

    }



    @Override

    public boolean init(String dbName) {

        try {

            Mongo mongo = new Mongo("localhost");

            mongoDb = mongo.getDB(dbName);

            GsonBuilder gsonBuilder = new GsonBuilder();

            gsonBuilder.registerTypeAdapter(JVMInstruction.class, new JVMInstructionSerializer());

            jsonSerializer = gsonBuilder.create();

            return true;

        } catch (Exception ex) {

            logger.error("Mongo storage manager init failed", ex);

            return false;

        }

    }



    @Override

    public void term() {

    }



    @Override

    public void flush() {

    }



    @Override

    public void clear() {

        mongoDb.dropDatabase();

    }



    private String buildCollectionName(InstructionSequence sequence) {

        return collectionPrefix + sequence.size();

    }



    private void serializeInstructionSequence(DBObject parentObject, String fieldName, InstructionSequence instructionSequence) {

        if (instructionSequence == null) {

            parentObject.put(fieldName, null);

            return;

        }

        final String serializedJSON = jsonSerializer.toJson(instructionSequence);

        final Object objectToPut = JSON.parse(serializedJSON);

        parentObject.put(fieldName, objectToPut);

    }



    private InstructionSequence deserialize(DBObject parentObject, String fieldName) {

        String className = "";

        try {

            className = (String) parentObject.get("instructionClass");

            Class<? extends IGenericInstruction> instructionClass = (Class<? extends IGenericInstruction>) Class.forName(className);

            return deserialize0(parentObject, fieldName, instructionClass);

        } catch (ClassNotFoundException e) {

            logger.error("Can not deserialize InstructionSequence as its instruction class " + className + " can not be loaded:", e);

            return null;

        } catch (Exception e) {

            logger.error("Error while deserializing InstructionSequence", e);

            return null;

        }

    }



    private <T extends IGenericInstruction> InstructionSequence deserialize0(DBObject parentObject, String fieldName, Class<T> instructionClass) {

        Object obj = parentObject.get(fieldName);

        if (obj == null) {

            return null;

        }

        final String jsonString = JSON.serialize(obj);



        T[] coll = (T[]) jsonSerializer.fromJson(jsonString, Array.newInstance(instructionClass, 0).getClass());

        return new InstructionSequence(coll);

    }



    @Override

    public boolean containsFailedOptimizationRecord(InstructionSequence source) {

        DBObject obj = findRecordForSource(source);

        return obj != null && obj.get("success").equals(false);

    }



    /**

     * Returns pair of source-optimized for all successfull optimizations that exist in database

     *

     * @return All succesfull optimization results

     */

    public Map<InstructionSequence, InstructionSequence> getAllResults() {

        Map<InstructionSequence, InstructionSequence> result = new HashMap<InstructionSequence, InstructionSequence>();

        Set<String> collections = mongoDb.getCollectionNames();

        for (String s : collections) {

            DBCollection coll = mongoDb.getCollection(s);

            BasicDBObject query = new BasicDBObject();

            query.put("success", true);

            DBCursor cur = coll.find(query);

            for (; cur.hasNext();) {

                DBObject obj = cur.next();

                final InstructionSequence source = deserialize(obj, sourceFieldName);

                final InstructionSequence optimized = deserialize(obj, optimizedFieldName);

                result.put(source, optimized);

            }

        }

        return result;

    }

}