ClassUploader.java

package edu.harvard.mcb.leschziner.load;

import java.io.File;
import java.io.IOException;
import java.security.NoSuchAlgorithmException;
import java.util.Map;
import java.util.UUID;
import java.util.Vector;

import org.apache.commons.compress.archivers.zip.ZipArchiveEntry;
import org.apache.commons.compress.archivers.zip.ZipArchiveOutputStream;
import org.jets3t.service.S3Service;
import org.jets3t.service.S3ServiceException;
import org.jets3t.service.ServiceException;
import org.jets3t.service.impl.rest.httpclient.RestS3Service;
import org.jets3t.service.model.S3Object;
import org.jets3t.service.multi.SimpleThreadedStorageService;
import org.vertx.java.core.json.JsonArray;
import org.vertx.java.core.json.JsonObject;

import com.hazelcast.core.MultiMap;

import edu.harvard.mcb.leschziner.aws.DefaultCredentials;
import edu.harvard.mcb.leschziner.core.Particle;
import edu.harvard.mcb.leschziner.core.ParticleClassifier;
import edu.harvard.mcb.leschziner.storage.DefaultStorageEngine;

public class ClassUploader {
    private final ParticleClassifier     classifier;

    private final String                 targetBucket;
    private final String                 pipelineId;
    private final String                 classMapName;
    private final MultiMap<Long, String> uploadedClasses;
    private final String                 averageMapName;
    private final Map<Long, String>      uploadedAverages;

    private S3Service                    s3Service;

    public ClassUploader(ParticleClassifier classifier,
                         String bucket,
                         String pipelineId) {
        super();
        this.classifier = classifier;
        classMapName = "UploadedClasses_" + this.hashCode();
        uploadedClasses = DefaultStorageEngine.getStorageEngine()
                                              .getMultiMap(classMapName);
        averageMapName = "UploadedAverages_" + this.hashCode();
        uploadedAverages = DefaultStorageEngine.getStorageEngine()
                                               .getMap(averageMapName);
        targetBucket = bucket;
        this.pipelineId = pipelineId;

        initS3Connection();
    }

    private boolean initS3Connection() {
        try {
            s3Service = new RestS3Service(DefaultCredentials.getAwsCredentials());
            return true;
        } catch (S3ServiceException e) {
            System.err.println("["
                               + this.getClass().getSimpleName()
                               + "]: Failed to connect to S3");
            s3Service = null;
            return false;
        }
    }

    public void uploadAll() {
        if (s3Service == null && !initS3Connection()) {
            return;
        }
        SimpleThreadedStorageService simpleMulti = new SimpleThreadedStorageService(s3Service);
        Vector<S3Object> objects = new Vector<S3Object>(classifier.getClassIds()
                                                                  .size());
        // For each class
        for (Long classId : classifier.getClassIds()) {
            System.out.println("["
                               + this.getClass().getSimpleName()
                               + "]: Uploading class "
                               + classId);

            UUID uuid = UUID.randomUUID();

            // Prep the class average
            String filename = pipelineId + "/" + uuid.toString() + ".png";
            String url = "https://s3.amazonaws.com/"
                         + targetBucket
                         + "/"
                         + filename;
            S3Object object;
            try {
                object = new S3Object(filename,
                                      classifier.getClassAverage(classId)
                                                .toPng());
                objects.add(object);

                uploadedAverages.put(classId, url);
                System.out.println("["
                                   + this.getClass().getSimpleName()
                                   + "]: Put average, "
                                   + url
                                   + " in S3");
            } catch (NoSuchAlgorithmException | IOException e) {
                e.printStackTrace();
            }

        }
        try {
            simpleMulti.putObjects(targetBucket,
                                   objects.toArray(new S3Object[objects.size()]));
        } catch (ServiceException e) {
            e.printStackTrace();
        }

        // Reuse the vector
        objects.clear();
        Vector<File> tempFiles = new Vector<>();
        for (Long classId : classifier.getClassIds()) {
            String filename = classId + ".zip";
            File zipFile = new File("download/" + filename);

            try (ZipArchiveOutputStream zos = new ZipArchiveOutputStream(zipFile)) {
                for (Particle target : classifier.getClass(classId)) {
                    // Give the particle a filename
                    UUID uuid = UUID.randomUUID();
                    String path = uuid + ".png";

                    ZipArchiveEntry zipEntry = new ZipArchiveEntry(path);
                    byte[] buffer = target.toPng();
                    zipEntry.setSize(buffer.length);
                    zos.putArchiveEntry(zipEntry);
                    zos.write(buffer);
                    zos.closeArchiveEntry();
                }
                zos.finish();

                // Keep track of the temporary file
                tempFiles.add(zipFile);
                // Upload the entire class
                S3Object object = new S3Object(zipFile);
                String key = pipelineId + "/" + filename;
                String url = "https://s3.amazonaws.com/"
                             + targetBucket
                             + "/"
                             + key;
                object.setKey(key);
                objects.add(object);

                uploadedClasses.put(classId, url);
            } catch (IOException | NoSuchAlgorithmException e) {
                e.printStackTrace();
            }
        }

        try {
            simpleMulti.putObjects(targetBucket,
                                   objects.toArray(new S3Object[objects.size()]));
        } catch (ServiceException e) {
            e.printStackTrace();
        }

        // Clean up the files
        for (File file : tempFiles) {
            file.delete();
        }
    }

    private String s3Put(Particle target) {
        // Generate a UUID for the object
        UUID uuid = UUID.randomUUID();

        // Prep the particle for putting
        try {
            String filename = pipelineId + "/" + uuid.toString() + ".png";
            String url = "https://s3.amazonaws.com/"
                         + targetBucket
                         + "/"
                         + filename;
            S3Object object = new S3Object(filename, target.toPng());
            // Connect to S3
            s3Service.putObject(targetBucket, object);
            return url;
        } catch (NoSuchAlgorithmException | IOException | S3ServiceException e) {
            e.printStackTrace();
        }
        return null;
    }

    public String getTargetBucket() {
        return targetBucket;
    }

    public String getClassMapName() {
        return classMapName;
    }

    public MultiMap<Long, String> getUploadedClasses() {
        return uploadedClasses;
    }

    public String getAverageMapName() {
        return averageMapName;
    }

    public Map<Long, String> getUploadedAverages() {
        return uploadedAverages;
    }

    public String getResultsJson() {
        JsonObject results = new JsonObject();

        for (Long classId : uploadedClasses.keySet()) {
            JsonObject classResults = new JsonObject();
            JsonArray classParticles = new JsonArray();
            for (String url : uploadedClasses.get(classId)) {
                classParticles.addString(url);
            }
            classResults.putArray("particles", classParticles);
            classResults.putString("average", uploadedAverages.get(classId));
            classResults.putNumber("count", classifier.getClass(classId).size());
            results.putObject(classId.toString(), classResults);
        }

        return results.encode();
    }
}