org.apache.flink.runtime.checkpoint.AbstractAsyncSnapshotIOCallable Maven / Gradle / Ivy
/*
* 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.checkpoint;
import org.apache.commons.io.IOUtils;
import org.apache.flink.core.fs.CloseableRegistry;
import org.apache.flink.runtime.io.async.AbstractAsyncIOCallable;
import org.apache.flink.runtime.state.CheckpointStreamFactory;
import org.apache.flink.runtime.state.StateObject;
import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.util.Preconditions;
import java.io.IOException;
import java.util.concurrent.atomic.AtomicBoolean;
/**
* Abstract base class for async IO operations of snapshots against a
* {@link java.util.zip.CheckedOutputStream}. This includes participating in lifecycle management
* through a {@link CloseableRegistry}.
*/
public abstract class AbstractAsyncSnapshotIOCallable
extends AbstractAsyncIOCallable {
protected final long checkpointId;
protected final long timestamp;
protected final CheckpointStreamFactory streamFactory;
protected final CloseableRegistry closeStreamOnCancelRegistry;
protected final AtomicBoolean open;
public AbstractAsyncSnapshotIOCallable(
long checkpointId,
long timestamp,
CheckpointStreamFactory streamFactory,
CloseableRegistry closeStreamOnCancelRegistry) {
this.streamFactory = Preconditions.checkNotNull(streamFactory);
this.closeStreamOnCancelRegistry = Preconditions.checkNotNull(closeStreamOnCancelRegistry);
this.checkpointId = checkpointId;
this.timestamp = timestamp;
this.open = new AtomicBoolean(false);
}
@Override
public CheckpointStreamFactory.CheckpointStateOutputStream openIOHandle() throws Exception {
if (checkStreamClosedAndDoTransitionToOpen()) {
CheckpointStreamFactory.CheckpointStateOutputStream stream =
streamFactory.createCheckpointStateOutputStream(checkpointId, timestamp);
try {
closeStreamOnCancelRegistry.registerClosable(stream);
return stream;
} catch (Exception ex) {
open.set(false);
throw ex;
}
} else {
throw new IOException("Async snapshot: a checkpoint stream was already opened.");
}
}
@Override
public void done(boolean canceled) {
if (checkStreamOpenAndDoTransitionToClose()) {
CheckpointStreamFactory.CheckpointStateOutputStream stream = getIoHandle();
if (stream != null) {
closeStreamOnCancelRegistry.unregisterClosable(stream);
IOUtils.closeQuietly(stream);
}
}
}
protected boolean checkStreamClosedAndDoTransitionToOpen() {
return open.compareAndSet(false, true);
}
protected boolean checkStreamOpenAndDoTransitionToClose() {
return open.compareAndSet(true, false);
}
protected StreamStateHandle closeStreamAndGetStateHandle() throws IOException {
if (checkStreamOpenAndDoTransitionToClose()) {
final CheckpointStreamFactory.CheckpointStateOutputStream stream = getIoHandle();
try {
return stream.closeAndGetHandle();
} finally {
closeStreamOnCancelRegistry.unregisterClosable(stream);
}
} else {
throw new IOException("Checkpoint stream already closed.");
}
}
}