sila_java.library.manager.executor.BinaryDownloader Maven / Gradle / Ivy

Go to download
package sila_java.library.manager.executor;

import lombok.Getter;
import lombok.NonNull;
import org.apache.commons.io.IOUtils;
import sila2.org.silastandard.SiLABinaryTransfer;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.UUID;

/**
 * Binary downloader
 */
public class BinaryDownloader {
    static final int MAX_CHUNK_SIZE = 2097152;
    private final OutputStream outputStream;
    @Getter
    private final UUID binaryTransferUUID;

    @Getter
    private long binarySize = -1;
    @Getter
    private int chunkIndex = 0;
    @Getter
    private int downloadedBytes = 0;

    /**
     * Constructor
     * @param outputStream output stream to write chunks to
     * @param binaryTransferUUID the binary transfer UUID
     */
    public BinaryDownloader(
            @NonNull final OutputStream outputStream,
            @NonNull final UUID binaryTransferUUID
    ) {
        this.outputStream = outputStream;
        this.binaryTransferUUID = binaryTransferUUID;
    }

    /**
     * Get the next chunk download request
     * @param binaryInfo the binary info
     * @return the next chunk download request
     */
    public synchronized SiLABinaryTransfer.GetChunkRequest getNextChunkDownloadRequest(
            @NonNull final SiLABinaryTransfer.GetBinaryInfoResponse binaryInfo
    ) {
        this.binarySize = binaryInfo.getBinarySize();
        final int chunkSizeModulo = getChunkSizeModulo(binaryInfo.getBinarySize());
        final int chunkCount = getChunkCount(binaryInfo.getBinarySize());
        if (chunkCount == this.chunkIndex) {
            throw new RuntimeException("No more binary chunk to download");
        }
        final int lengthToRead = (this.chunkIndex == chunkCount - 1 && chunkSizeModulo > 0) ? (chunkSizeModulo) : (MAX_CHUNK_SIZE);
        final SiLABinaryTransfer.GetChunkRequest getChunkRequest = SiLABinaryTransfer.GetChunkRequest
                .newBuilder()
                .setLength(lengthToRead)
                .setOffset(((long) this.chunkIndex) * ((long)MAX_CHUNK_SIZE))
                .setBinaryTransferUUID(this.binaryTransferUUID.toString())
                .build();
        ++this.chunkIndex;
        return getChunkRequest;
    }

    /**
     * Get chunk size
     *
     * Size % {@link BinaryDownloader#MAX_CHUNK_SIZE}
     * @param binarySize the binary size in bytes
     * @return the chunk size
     */
    public static int getChunkSizeModulo(long binarySize) {
        return (int) (binarySize % MAX_CHUNK_SIZE);
    }

    /**
     * Get the number of chunk
     * @param binarySize the binary size in byte
     *
     * @return the number of chunk
     */
    public static int getChunkCount(long binarySize) {
        final int chunkSizeModulo = getChunkSizeModulo(binarySize);
        return (int) (binarySize / MAX_CHUNK_SIZE) + ((chunkSizeModulo > 0) ? 1 : 0);
    }

    /***
     * Write chunk into {@link BinaryDownloader#outputStream}
     *
     * @param chunkResponse the chunk response
     *
     * @return the number of bytes written
     * @throws IOException if enable to write
     */
    public synchronized int writeChunk(SiLABinaryTransfer.GetChunkResponse chunkResponse) throws IOException {
        if (this.binarySize < 0) {
            throw new RuntimeException("Invalid or unknown binary size!");
        }
        int receivedChunkIndex = (int) (chunkResponse.getOffset() / MAX_CHUNK_SIZE);
        if (this.chunkIndex - 1 != receivedChunkIndex) {
            throw new RuntimeException("Binary chunks must be written in sequential order");
        }
        try (final InputStream inputStream = chunkResponse.getPayload().newInput()) {
            final int copiedBytes = IOUtils.copy(inputStream, this.outputStream);
            this.downloadedBytes += copiedBytes;
            if (this.downloadedBytes > this.binarySize) {
                throw new RuntimeException("Binary Downloader downloaded too many bytes!");
            }
            return copiedBytes;
        }
    }

    public boolean isValidAndComplete() {
        return this.downloadedBytes == this.binarySize;
    }

    public void isValidAndCompleteOrThrow() {
        if (!isValidAndComplete()) {
            throw new RuntimeException(
                    "Binary Downloader expected to download " + this.binarySize + " bytes but received " + this.downloadedBytes
            );
        }
    }

    /**
     * Get binary info request
     * @return a new instance of {@link sila2.org.silastandard.SiLABinaryTransfer.GetBinaryInfoRequest}
     */
    public SiLABinaryTransfer.GetBinaryInfoRequest getBinaryInfoRequest() {
        return SiLABinaryTransfer.GetBinaryInfoRequest
                .newBuilder()
                .setBinaryTransferUUID(binaryTransferUUID.toString())
                .build();
    }
}