FileOutputOutputStream.java

package org.embulk.spi.util;

import java.io.OutputStream;
import org.embulk.spi.Buffer;
import org.embulk.spi.BufferAllocator;
import org.embulk.spi.FileOutput;

public class FileOutputOutputStream
        extends OutputStream
{
    private final FileOutput out;
    private final BufferAllocator allocator;
    private final CloseMode closeMode;
    private int pos;
    private Buffer buffer;

    public static enum CloseMode {
        FLUSH,
        FLUSH_FINISH,
        FLUSH_FINISH_CLOSE,
        CLOSE;
    }

    public FileOutputOutputStream(FileOutput out, BufferAllocator allocator, CloseMode closeMode)
    {
        this.out = out;
        this.allocator = allocator;
        this.buffer = allocator.allocate();
        this.closeMode = closeMode;
    }

    public void nextFile()
    {
        out.nextFile();
    }

    public void finish()
    {
        doFlush();
        out.finish();
    }

    @Override
    public void write(int b)
    {
        buffer.array()[buffer.offset() + pos] = (byte) b;
        pos++;
        if (pos >= buffer.capacity()) {
            flush();
        }
    }

    @Override
    public void write(byte[] b, int off, int len)
    {
        while (true) {
            int available = buffer.capacity() - pos;
            if (available < len) {
                buffer.setBytes(pos, b, off, available);
                pos += available;
                len -= available;
                off += available;
                flush();
            } else {
                buffer.setBytes(pos, b, off, len);
                pos += len;
                if (available <= len) {
                    flush();
                }
                break;
            }
        }
    }

    private boolean doFlush()
    {
        if (pos > 0) {
            buffer.limit(pos);
            out.add(buffer);
            buffer = Buffer.EMPTY;
            pos = 0;
            return true;
        }
        return false;
    }

    @Override
    public void flush()
    {
        if (doFlush()) {
            buffer = allocator.allocate();
        }
    }

    @Override
    public void close()
    {
        switch (closeMode) {
        case FLUSH:
            doFlush();
            break;
        case FLUSH_FINISH:
            doFlush();
            out.finish();
            break;
        case FLUSH_FINISH_CLOSE:
            doFlush();
            out.finish();
            out.close();
            break;
        case CLOSE:
            out.close();
            break;
        }
        buffer.release();
        buffer = Buffer.EMPTY;
        pos = 0;
    }
}