org.apache.flink.runtime.state.gemini.engine.fs.GeminiBufferedOutputStream Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of flink-runtime_2.11 Show documentation
There is a newer version: 1.5.1
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.runtime.state.gemini.engine.fs;

import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.util.Preconditions;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

/**
 * GeminiBufferedOutputStream.
 */
public class GeminiBufferedOutputStream extends FSDataOutputStream {

	private static final Logger LOG = LoggerFactory.getLogger(GeminiBufferedOutputStream.class);

	private final byte[] writeBuffer;

	private int pos;

	private FSDataOutputStream outStream;

	private volatile boolean closed;

	public GeminiBufferedOutputStream (
		FSDataOutputStream outStream,
		int bufferSize
	) {
		this.outStream = Preconditions.checkNotNull(outStream);
		this.writeBuffer = new byte[bufferSize];
	}

	@Override
	public synchronized void write(int b) throws IOException {
		if (pos >= writeBuffer.length) {
			flush();
		}
		writeBuffer[pos++] = (byte) b;
	}

	@Override
	public synchronized void write(byte[] b, int off, int len) throws IOException {
		if (len < writeBuffer.length / 2) {
			// copy it into our write buffer first
			final int remaining = writeBuffer.length - pos;
			if (len > remaining) {
				// copy as much as fits
				System.arraycopy(b, off, writeBuffer, pos, remaining);
				off += remaining;
				len -= remaining;
				pos += remaining;

				// flush the write buffer to make it clear again
				flush();
			}

			// copy what is in the buffer
			System.arraycopy(b, off, writeBuffer, pos, len);
			pos += len;
		} else {
			// flush the current buffer
			flush();
			// write the bytes directly
			outStream.write(b, off, len);
		}
	}

	@Override
	public synchronized long getPos() throws IOException {
		return pos + outStream.getPos();
	}

	@Override
	public synchronized void flush() throws IOException {
		if (!closed) {
			if (pos > 0) {
				outStream.write(writeBuffer, 0, pos);
				pos = 0;
			}
		} else {
			throw new IOException("closed");
		}
	}

	@Override
	public synchronized void sync() throws IOException {
		flush();
		outStream.sync();
	}

	/**
	 * Checks whether the stream is closed.
	 * @return True if the stream was closed, false if it is still open.
	 */
	public boolean isClosed() {
		return closed;
	}

	@Override
	public synchronized void close() throws IOException {
		if (!closed) {
			flush();
			outStream.close();
			closed = true;
		}
	}
}