@Override
public void write(int b) throws IOException {
if (mShouldCacheCurrentBlock) {
try {
if (mCurrentBlockOutStream == null || mCurrentBlockOutStream.remaining() == 0) {
getNextBlock();
}
mCurrentBlockOutStream.write(b);
} catch (IOException ioe) {
handleCacheWriteException(ioe);
}
}
if (mUnderStorageType.isSyncPersist()) {
mUnderStorageOutputStream.write(b);
ClientContext.getClientMetrics().incBytesWrittenUfs(1);
}
mBytesWritten++;
}
protected void handleCacheWriteException(IOException ioe) throws IOException {
if (!mUnderStorageType.isSyncPersist()) {
throw new IOException(ExceptionMessage.FAILED_CACHE.getMessage(ioe.getMessage()), ioe);
}
LOG.warn("Failed to write into TachyonStore, canceling write attempt.", ioe);
if (mCurrentBlockOutStream != null) {
mShouldCacheCurrentBlock = false;
mCurrentBlockOutStream.cancel();
}
}
@Override
public void write(byte[] b, int off, int len) throws IOException {
Preconditions.checkArgument(b != null, PreconditionMessage.ERR_WRITE_BUFFER_NULL);
Preconditions.checkArgument(
off >= 0 && len >= 0 && len + off <= b.length,
PreconditionMessage.ERR_BUFFER_STATE,
b.length,
off,
len);
if (mShouldCacheCurrentBlock) {
try {
int tLen = len;
int tOff = off;
while (tLen > 0) {
if (mCurrentBlockOutStream == null || mCurrentBlockOutStream.remaining() == 0) {
getNextBlock();
}
long currentBlockLeftBytes = mCurrentBlockOutStream.remaining();
if (currentBlockLeftBytes >= tLen) {
mCurrentBlockOutStream.write(b, tOff, tLen);
tLen = 0;
} else {
mCurrentBlockOutStream.write(b, tOff, (int) currentBlockLeftBytes);
tOff += currentBlockLeftBytes;
tLen -= currentBlockLeftBytes;
}
}
} catch (IOException ioe) {
handleCacheWriteException(ioe);
}
}
if (mUnderStorageType.isSyncPersist()) {
mUnderStorageOutputStream.write(b, off, len);
ClientContext.getClientMetrics().incBytesWrittenUfs(len);
}
mBytesWritten += len;
}
private void getNextBlock() throws IOException {
if (mCurrentBlockOutStream != null) {
Preconditions.checkState(
mCurrentBlockOutStream.remaining() <= 0, PreconditionMessage.ERR_BLOCK_REMAINING);
mPreviousBlockOutStreams.add(mCurrentBlockOutStream);
}
if (mTachyonStorageType.isStore()) {
try {
NetAddress address =
mLocationPolicy.getWorkerForNextBlock(
mContext.getTachyonBlockStore().getWorkerInfoList(), mBlockSize);
mCurrentBlockOutStream =
mContext.getTachyonBlockStore().getOutStream(getNextBlockId(), mBlockSize, address);
mShouldCacheCurrentBlock = true;
} catch (TachyonException e) {
throw new IOException(e);
}
}
}
@Override
public void close() throws IOException {
if (mClosed) {
return;
}
if (mCurrentBlockOutStream != null) {
mPreviousBlockOutStreams.add(mCurrentBlockOutStream);
}
Boolean canComplete = false;
CompleteFileOptions options = CompleteFileOptions.defaults();
if (mUnderStorageType.isSyncPersist()) {
String tmpPath = PathUtils.temporaryFileName(mNonce, mUfsPath);
UnderFileSystem ufs = UnderFileSystem.get(tmpPath, ClientContext.getConf());
if (mCanceled) {
// TODO(yupeng): Handle this special case in under storage integrations.
mUnderStorageOutputStream.close();
if (!ufs.exists(tmpPath)) {
// Location of the temporary file has changed, recompute it.
updateUfsPath();
tmpPath = PathUtils.temporaryFileName(mNonce, mUfsPath);
}
ufs.delete(tmpPath, false);
} else {
mUnderStorageOutputStream.flush();
mUnderStorageOutputStream.close();
if (!ufs.exists(tmpPath)) {
// Location of the temporary file has changed, recompute it.
updateUfsPath();
tmpPath = PathUtils.temporaryFileName(mNonce, mUfsPath);
}
if (!ufs.rename(tmpPath, mUfsPath)) {
throw new IOException("Failed to rename " + tmpPath + " to " + mUfsPath);
}
options.setUfsLength(ufs.getFileSize(mUfsPath));
canComplete = true;
}
}
if (mTachyonStorageType.isStore()) {
try {
if (mCanceled) {
for (BufferedBlockOutStream bos : mPreviousBlockOutStreams) {
bos.cancel();
}
} else {
for (BufferedBlockOutStream bos : mPreviousBlockOutStreams) {
bos.close();
}
canComplete = true;
}
} catch (IOException ioe) {
handleCacheWriteException(ioe);
}
}
if (canComplete) {
FileSystemMasterClient masterClient = mContext.acquireMasterClient();
try {
masterClient.completeFile(mUri, options);
} catch (TachyonException e) {
throw new IOException(e);
} finally {
mContext.releaseMasterClient(masterClient);
}
}
if (mUnderStorageType.isAsyncPersist()) {
scheduleAsyncPersist();
}
mClosed = true;
}