org.apache.flink.runtime.state.gemini.engine.snapshot.SnapshotOperation 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.state.gemini.engine.snapshot;

import org.apache.flink.api.common.typeutils.base.MapSerializer;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.core.fs.FileSystem;
import org.apache.flink.core.fs.Path;
import org.apache.flink.core.memory.DataOutputView;
import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.state.gemini.engine.GRegionID;
import org.apache.flink.runtime.state.gemini.engine.dbms.GContext;
import org.apache.flink.runtime.state.gemini.engine.fs.FileManager;
import org.apache.flink.util.Preconditions;

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

import javax.annotation.Nullable;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

/**
 * Snapshot operation.
 */
public abstract class SnapshotOperation {

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

	protected final GContext gContext;

	protected final SnapshotManager snapshotManager;

	SnapshotOperation(GContext gContext, SnapshotManager snapshotManager) {
		this.gContext = Preconditions.checkNotNull(gContext);
		this.snapshotManager = Preconditions.checkNotNull(snapshotManager);
	}

	public abstract SnapshotManager.PendingSnapshot createPendingSnapshot(
		BackendSnapshotMeta backendSnapshotMeta,
		long accessNumber);

	public abstract SnapshotManager.PendingSnapshot getPendingSnapshot();

	public abstract DBSnapshotResult getSnapshotResult() throws Exception;

	void writePageIndex(
		Map> copiedPageIndex,
		@Nullable SnapshotMetaFile.Writer localWriter,
		SnapshotMetaFile.Writer dfsWriter,
		@Nullable RegionSnapshot localRegionSnapshot,
		RegionSnapshot dfsRegionSnapshot,
		@Nullable Map>> localFileMapping,
		Map>> dfsFileMapping) throws IOException {

		// page index (file id + offset)
		writeInt(localWriter, dfsWriter, copiedPageIndex.size());
		for (Map.Entry> te : copiedPageIndex.entrySet()) {
			String tableName = te.getKey();
			Map regionMetas = te.getValue();
			writeUTF(localWriter, dfsWriter, tableName);
			writeInt(localWriter, dfsWriter, regionMetas.size());
			for (Map.Entry re : regionMetas.entrySet()) {
				writeInt(localWriter, dfsWriter, re.getKey().getRegionCode());
				SnapshotManager.GRegionSnapshotMeta meta = re.getValue();
				writeLong(localWriter, dfsWriter, meta.getLastSeqID());
				writeLong(localWriter, dfsWriter, meta.getRemoveAllSeqID());
				int groupId = re.getKey().getId();
				if (localRegionSnapshot != null) {
					localRegionSnapshot.reset(groupId);
				}
				dfsRegionSnapshot.reset(groupId);
				meta.getPageIndex().snapshot(localRegionSnapshot, dfsRegionSnapshot);
				if (localFileMapping != null) {
					updateFileMapping(localFileMapping, localRegionSnapshot.getRegionLocalFileMeta(), groupId);
				}
				updateFileMapping(dfsFileMapping, dfsRegionSnapshot.getRegionDfsFileMeta(), groupId);
			}
		}
	}

	void writeLocalAndDfsFileMapping(
		SnapshotManager.PendingSnapshot pendingSnapshot,
		SnapshotMetaFile.Writer writer,
		@Nullable FileManager localFileManager,
		FileManager dfsFileManager) throws IOException {

		// the offset of dfs file mapping in the meta file
		long dfsMappingOffset = writer.getPos();
		writeFileMapping(writer, dfsFileManager, pendingSnapshot.getSnapshotBasePath(),
			pendingSnapshot.getFileMapping(), false);

		// the offset of local file mapping in the meta file
		long localMappingOffset = writer.getPos();
		writeFileMapping(writer, localFileManager, pendingSnapshot.getLocalSnapshotBasePath(),
			pendingSnapshot.getLocalFileMapping(), true);

		// write the offset of dfs file mapping at the end of file
		writer.writeLong(dfsMappingOffset);
		// write the offset of local file mapping at the end of file
		writer.writeLong(localMappingOffset);
	}

	void closeAndDeleteWriterQuietly(SnapshotMetaFile.Writer writer) {
		if (writer != null) {
			Path path = writer.getFilePath();
			try {
				writer.close();
			} catch (IOException ioe) {
				LOG.warn("Failed to close meta writer, {}, {}.", path.toUri().toString(), ioe);
			}
			try {
				FileSystem fileSystem = path.getFileSystem();
				fileSystem.delete(path, false);
			} catch (IOException ioe) {
				LOG.warn("Failed to delete meta file, {}, {}.", path.toUri().toString(), ioe);
			}
		}
	}

	private void writeFileMapping(
		SnapshotMetaFile.Writer writer,
		FileManager fileManager,
		Path snapshotBasePath,
		@Nullable Map>> fileMapping,
		boolean isLocal) throws IOException {
		if (fileManager == null) {
			writer.writeBoolean(false);
		} else {
			writer.writeBoolean(true);
			writer.writeInt(fileMapping.size());
			writer.writeUTF(snapshotBasePath.toUri().toString());
			if (fileMapping.size() > 0) {
				for (Integer fileId : fileMapping.keySet()) {
					String path = fileManager.getFilePath(fileId);
					// for local file mapping, only record the final component of the path
					writer.writeUTF(isLocal ? new Path(path).getName() : path);
					writer.writeInt(fileId);
				}
				MapSerializer>> fileMappingSerializer = ((SnapshotManagerImpl) snapshotManager).getFileMappingSerializer();
				DataOutputView dataOutputView = new DataOutputViewStreamWrapper(writer);
				fileMappingSerializer.serialize(fileMapping, dataOutputView);
			}
		}
	}

	private void writeInt(
		SnapshotMetaFile.Writer localWriter,
		SnapshotMetaFile.Writer dfsWriter,
		int value) throws IOException {
		if (localWriter != null) {
			localWriter.writeInt(value);
		}
		if (dfsWriter != null) {
			dfsWriter.writeInt(value);
		}
	}

	private void writeLong(
		SnapshotMetaFile.Writer localWriter,
		SnapshotMetaFile.Writer dfsWriter,
		long value) throws IOException {
		if (localWriter != null) {
			localWriter.writeLong(value);
		}
		if (dfsWriter != null) {
			dfsWriter.writeLong(value);
		}
	}

	private void writeUTF(
		SnapshotMetaFile.Writer localWriter,
		SnapshotMetaFile.Writer dfsWriter,
		String value) throws IOException {
		if (localWriter != null) {
			localWriter.writeUTF(value);
		}
		if (dfsWriter != null) {
			dfsWriter.writeUTF(value);
		}
	}

	private void updateFileMapping(
		Map>> fileMapping,
		Map> regionFileMeta,
		int groupId) {
		Map> groupStat;
		for (Map.Entry> entry : regionFileMeta.entrySet()) {
			int fileId = entry.getKey();
			groupStat = fileMapping.computeIfAbsent(fileId, (nothing) -> new HashMap<>());
			Tuple2 tuple =
				groupStat.computeIfAbsent(groupId, (nothing) -> Tuple2.of(0, 0L));
			tuple.f0 += entry.getValue().f0;
			tuple.f1 += entry.getValue().f1;
		}
	}
}