ai.h2o.mojos.runtime.api.backend.MemoryReaderBackend Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of mojo2-runtime-api Show documentation
MOJO2 Runtime API
There is a newer version: 2.8.7.1
package ai.h2o.mojos.runtime.api.backend;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.security.NoSuchAlgorithmException;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Access to in-memory fileset; primarily for testing */
public class MemoryReaderBackend implements ReaderBackend {
    private static final Logger log = LoggerFactory.getLogger(MemoryReaderBackend.class);
    private final Map files;
    private final ResourceInfo.Cache cache = new ResourceInfo.Cache();

    private MemoryReaderBackend(Map files) {
        this.files = files;
    }

    public static ReaderBackend open(Map files) {
        return new MemoryReaderBackend(files);
    }

    /**
     * Read zipped files from inputstream. Useful when the container is nested in a java resource.
     */
    public static ReaderBackend fromZipStream(final InputStream is) throws IOException {
        log.info("Opening mojo stream: {}", is);
        final HashMap files = new HashMap<>();
        try (final ZipInputStream zis = new ZipInputStream(is)) {
            while (true) {
                final ZipEntry entry = zis.getNextEntry();
                if (entry == null) {
                    break;
                }
                if (entry.getSize() > Integer.MAX_VALUE) {
                    throw new IOException("File is too large: " + entry.getName());
                }
                if (!entry.isDirectory()) {
                    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
                    ReaderBackendUtils.copy(zis, baos, 8 * 1024);
                    files.put(entry.getName(), baos.toByteArray());
                }
                zis.closeEntry();
            }
        }
        return MemoryReaderBackend.open(files);
    }

    public void put(final String resourceName, final byte[] bytes) {
        files.put(resourceName, bytes);
    }

    @Override
    public ResourceInfo getResourceInfo(final String resourceName) throws IOException {
        ResourceInfo info = cache.get(resourceName);
        if (info == null) {
            final byte[] bytes = files.get(resourceName);
            if (bytes == null) {
                throw new FileNotFoundException(resourceName);
            }
            try {
                final String md5 = ResourceInfo.computeMD5(bytes); // here we compute strong hash instead, as the bytes are already in memory anyway
                info = new ResourceInfo(bytes.length, "MD5:" + md5);
                cache.put(resourceName, info);
            } catch (NoSuchAlgorithmException e) {
                throw new IOException(e);
            }
        }
        return info;
    }

    @Override
    public InputStream getInputStream(String resourceName) throws FileNotFoundException {
        final byte[] bytes = files.get(resourceName);
        if (bytes == null) {
            throw new FileNotFoundException(resourceName);
        }
        return new ByteArrayInputStream(bytes);
    }

    @Override
    public boolean exists(String resourceName) {
        return files.containsKey(resourceName);
    }

    @Override
    public Collection list() {
        return files.keySet();
    }

    @Override
    public void close() {}

    @Override
    public String toString() {
        return String.format("%s[%s]", getClass().getSimpleName(), files.size());
    }
}