io.bdeploy.bhive.op.ObjectWriteOperation Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of api Show documentation
Public API including dependencies, ready to be used for integrations and plugins.
There is a newer version: 7.4.0
package io.bdeploy.bhive.op;

import java.io.BufferedOutputStream;
import java.io.DataOutputStream;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;
import java.util.zip.GZIPOutputStream;

import com.google.common.io.CountingOutputStream;

import io.bdeploy.bhive.BHive;
import io.bdeploy.bhive.ReadOnlyOperation;
import io.bdeploy.bhive.audit.AuditParameterExtractor.AuditStrategy;
import io.bdeploy.bhive.audit.AuditParameterExtractor.AuditWith;
import io.bdeploy.bhive.model.Manifest;
import io.bdeploy.bhive.model.Manifest.Key;
import io.bdeploy.bhive.model.ObjectId;
import io.bdeploy.bhive.util.StorageHelper;
import io.bdeploy.common.ActivityReporter.Activity;
import io.bdeploy.common.util.RuntimeAssert;
import io.bdeploy.common.util.StreamHelper;

/**
 * Writes one or more objects to a stream.
 */
@ReadOnlyOperation
public class ObjectWriteOperation extends BHive.Operation {

    public static final int BUFFER_SIZE = 8192;

    @AuditWith(AuditStrategy.COLLECTION_PEEK)
    private final Set objects = new LinkedHashSet<>();

    private final Set manifests = new LinkedHashSet<>();

    private OutputStream output;

    @Override
    public Long call() throws Exception {
        RuntimeAssert.assertNotNull(output);

        for (Manifest.Key key : new ArrayList<>(manifests)) {
            Collection refs = execute(new ManifestRefScanOperation().setAllowMissingObjects(true).setManifest(key)).values();
            manifests.addAll(refs);
        }

        // Collect the entire size that we are going to transfer.
        // ATTENTION: insertion order is important.
        long totalSize = 0;
        Map object2FileSize = new LinkedHashMap<>();
        for (ObjectId object : objects) {
            long size = getObjectManager().db(db -> db.getObjectSize(object));
            object2FileSize.put(object, size);
            totalSize += size;
        }

        try (Activity activity = getActivityReporter().start("Sending", totalSize);
                CountingOutputStream countingOut = new CountingOutputStream(output);
                GZIPOutputStream zipOut = new GZIPOutputStream(countingOut, BUFFER_SIZE);
                BufferedOutputStream buffOut = new BufferedOutputStream(zipOut, BUFFER_SIZE * 2);
                DataOutputStream dataOut = new DataOutputStream(zipOut)) {

            // First we send the total size so that the client can display a progress bar
            dataOut.writeLong(totalSize);

            // Stream all manifests
            dataOut.writeLong(manifests.size());
            for (Manifest.Key key : manifests) {
                Manifest mf = getManifestDatabase().getManifest(key);
                byte[] bytes = StorageHelper.toRawBytes(mf);
                dataOut.writeLong(bytes.length);
                dataOut.write(bytes);
            }

            // Stream all objects
            dataOut.writeLong(objects.size());
            for (Map.Entry entry : object2FileSize.entrySet()) {
                ObjectId objectId = entry.getKey();
                long size = entry.getValue();
                getObjectManager().db(db -> {
                    dataOut.writeLong(size);
                    try (InputStream input = db.getStream(objectId)) {
                        StreamHelper.copy(input, dataOut);
                    }
                    return null;
                });
                activity.worked(size);
            }
            return countingOut.getCount();
        } finally {
            StreamHelper.close(output);
        }
    }

    /**
     * The stream to write the objects to
     */
    public ObjectWriteOperation stream(OutputStream output) {
        this.output = output;
        return this;
    }

    /**
     * The {@link ObjectId object} to fetch
     */
    public ObjectWriteOperation objects(ObjectId obj) {
        this.objects.add(obj);
        return this;
    }

    /**
     * The {@link ObjectId objects} to fetch
     */
    public ObjectWriteOperation objects(Collection objects) {
        this.objects.addAll(objects);
        return this;
    }

    /**
     * The {@link Key manifests} to fetch
     */
    public ObjectWriteOperation manifest(Manifest.Key obj) {
        this.manifests.add(obj);
        return this;
    }

    /**
     * The {@link Key manifests} to fetch
     */
    public ObjectWriteOperation manifests(Collection manifests) {
        this.manifests.addAll(manifests);
        return this;
    }

}