All Downloads are FREE. Search and download functionalities are using the official Maven repository.

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

There is a newer version: 1.5.1
Show newest version
/*
 * 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.annotation.VisibleForTesting;
import org.apache.flink.util.Preconditions;

import java.util.BitSet;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Implementation of file id generator. We have three parts to compose the whole id: flag, sub-task-index, unique increased id in a circular range.
 * For total tasks number in range [1, 1024], we have distribution as: |1_bit|10_bit|21_bit|
 * FOr total tasks number in range [1025, 32768], we have distribution as: |2_bit|15_bit|15_bit|
 */
public class FileIDGenerator {

	private final FileIdMode fileIdMode;
	private final int initValue;
	private final int subTaskIndex;
	private final int totalTaskNum;

	/** Unlike the flag and sub-task-index, this value would increase and stay unique.*/
	private final AtomicInteger uniqueId;

	private final FileIDRecorder fileIDRecorder;

	public FileIDGenerator(int subTaskIndex, int totalTaskNum) {
		this(subTaskIndex, totalTaskNum, 0);
	}

	@VisibleForTesting
	FileIDGenerator(int subTaskIndex, int totalTaskNum, int initUniqueIdValue) {
		Preconditions.checkArgument(subTaskIndex >= 0, "Illegal sub task index: " + subTaskIndex);
		Preconditions.checkArgument(totalTaskNum > 0, "Illegal total task num: " + totalTaskNum);
		Preconditions.checkArgument(subTaskIndex < totalTaskNum,
			"Illegal subTaskIndex v.s totalTaskNum: " + totalTaskNum + " v.s " + totalTaskNum);

		if (totalTaskNum <= FileIdMode.MODE1.getMaxAllowedTaskNum()) {
			fileIdMode = FileIdMode.MODE1;
		} else if (totalTaskNum <= FileIdMode.MODE2.getMaxAllowedTaskNum()) {
			fileIdMode = FileIdMode.MODE2;
		} else {
			throw new IllegalArgumentException("Unexpected total task number: " + totalTaskNum + ", cannot initialize any file id mode.");
		}
		this.totalTaskNum = totalTaskNum;
		this.subTaskIndex = subTaskIndex;
		this.initValue = fileIdMode.getInitValue() | (this.subTaskIndex << fileIdMode.getSubtaskIndexOffset());
		this.uniqueId = new AtomicInteger(initUniqueIdValue);
		this.fileIDRecorder = new FileIDRecorder(fileIdMode.maxAllowedUniqueID);
	}

	public FileID generate() {
		if (fileIDRecorder.recordedFileIDNum() >= getMaxAllowedUniqueID()) {
			throw new IllegalStateException("Useful file unique ids have been consumed completely, not able to create new file writer.");
		}

		int currentUniqueId = uniqueId.get();
		while (fileIDRecorder.isIDRecorded(currentUniqueId)) {
			currentUniqueId = uniqueId.incrementAndGet();
			// if we reach the max allowed uniqueId, we should return from the scratch.
			if (currentUniqueId > fileIdMode.getMaxAllowedUniqueID()) {
				// Note. This method is not thread safe, please ensure to use this method in thread-safe situation.
				uniqueId.set(0);
				currentUniqueId = 0;
			}
		}

		fileIDRecorder.recordUniqueID(currentUniqueId);
		return new FileIDImpl(initValue | currentUniqueId);
	}

	public void recycleFileID(FileID fileID) {
		recycleUniqueID(fileID.getUniqueID());
	}

	@VisibleForTesting
	void recycleUniqueID(int uniqueID) {
		fileIDRecorder.recycleFileID(uniqueID);
	}

	/**
	 * Record the specific filID is occupied and would not generate this unless recycled.
	 */
	void restoreFileID(FileID fileID) {
		if (fileID.getMaxAllowedUniqueID() == getMaxAllowedUniqueID()) {
			fileIDRecorder.recordUniqueID(fileID.getUniqueID());
		}
	}

	/**
	 * Record the specific filIDs are occupied and would not generate them unless recycled.
	 */
	void restoreFileIDs(Set fileIDs) {
		if (!fileIDs.isEmpty()) {
			fileIDs.forEach(f -> {
				if (f.getMaxAllowedUniqueID() == getMaxAllowedUniqueID()) {
					fileIDRecorder.recordUniqueID(f.getUniqueID());
				}
			});
		}
	}

	@VisibleForTesting
	FileID get() {
		return new FileIDImpl(initValue | uniqueId.get());
	}

	@VisibleForTesting
	int getUniqueID() {
		return uniqueId.get();
	}

	@VisibleForTesting
	FileIDRecorder getFileIDRecorder() {
		return fileIDRecorder;
	}

	@VisibleForTesting
	void setUniqueID(int uniqueID) {
		Preconditions.checkArgument(uniqueID <= fileIdMode.getMaxAllowedUniqueID(),
			"UniqueID to set: " + uniqueID + " is larger than max allowed uniqueID.");
		this.uniqueId.set(uniqueID);
	}

	public int getMaxAllowedUniqueID() {
		return fileIdMode.getMaxAllowedUniqueID();
	}

	@VisibleForTesting
	FileIdMode getFileIdMode() {
		return fileIdMode;
	}

	int recordedFileIDNum() {
		return fileIDRecorder.recordedFileIDNum();
	}

	boolean isFileIDRecorded(FileID fileID) {
		return fileIDRecorder.isFileIDRecorded(fileID);
	}

	public int getUniqueId() {
		return uniqueId.get();
	}

	public int getSubTaskIndex() {
		return subTaskIndex;
	}

	public int getTotalTaskNum() {
		return totalTaskNum;
	}

	public static FileIDGenerator fromFileID(FileIDImpl fileID) {
		return new FileIDGenerator(fileID.getSubTaskIndex(), fileID.getFileIdMode().getMaxAllowedTaskNum(), fileID.getUniqueID());
	}

	@Override
	public String toString() {
		return "FileIDGenerator{" +
			"fileIdMode=" + fileIdMode +
			", subTaskIndex=" + subTaskIndex +
			", totalTaskNum=" + totalTaskNum +
			", uniqueId=" + uniqueId +
			'}';
	}

	/**
	 * Mode of file ID.
	 */
	public enum FileIdMode {
		/**
		 * Mode1 for total task number in range [1, 1024]
		 */
		MODE1(
			0,
			0x8000_0000,
			31,
			0x7FE0_0000,
			21,
			0x0001_FFFF),
		/**
		 * Mode2 for total task number in range [1025, 32768]
		 */
		MODE2(
			2,
			0xC000_0000,
			30,
			0X3FFF_8000,
			15,
			0X0000_7FFF);

		/** The flag represented in decimal value. */
		private final int flagValue;
		private final int maxAllowedTaskNum;
		private final int maxAllowedUniqueID;
		private final int flagMask;
		private final int flagOffset;
		private final int subtaskIndexMask;
		private final int subtaskIndexOffset;
		private final int uniqueIDMask;
		private final int initValue;
		private static final int TOTAL_LENGTH = 32;

		FileIdMode(
			int flagValue,
			int flagMask,
			int flagOffset,
			int subTaskMask,
			int subTaskIndexOffset,
			int uniqueIDMask) {

			this.flagValue = flagValue;
			this.flagMask = flagMask;
			this.flagOffset = flagOffset;
			this.initValue = flagValue << flagOffset;
			this.subtaskIndexMask = subTaskMask;
			this.subtaskIndexOffset = subTaskIndexOffset;
			this.maxAllowedTaskNum = (subTaskMask >> subTaskIndexOffset) + 1;
			this.uniqueIDMask = uniqueIDMask;
			this.maxAllowedUniqueID = uniqueIDMask;
		}


		public int getFlagValue() {
			return flagValue;
		}

		public int getFlagOffset() {
			return flagOffset;
		}

		public int getMaxAllowedTaskNum() {
			return maxAllowedTaskNum;
		}

		public int getFlagMask() {
			return flagMask;
		}

		public int getSubtaskIndexOffset() {
			return subtaskIndexOffset;
		}

		public int getSubtaskIndexMask() {
			return subtaskIndexMask;
		}

		public int getUniqueIDMask() {
			return uniqueIDMask;
		}

		public int getMaxAllowedUniqueID() {
			return maxAllowedUniqueID;
		}

		public int getInitValue() {
			return initValue;
		}

		public int getTotalLength() {
			return TOTAL_LENGTH;
		}
	}

	/**
	 * Recorde all used file ids.
	 */
	public class FileIDRecorder {
		/**
		 * Bit set for used file unique id, refer to {@link FileIDImpl} to know more details of unique id.
		 */
		private final BitSet usedFileUniqueIDs;

		/**
		 * maximum of allowed unique id, this is also useful to judge whether the file id to record is acceptable in is recorder.
		 */
		private final int maxAllowedUniqueID;

		public FileIDRecorder(int maxAllowedUniqueID) {
			this.maxAllowedUniqueID = maxAllowedUniqueID;
			this.usedFileUniqueIDs = new BitSet(maxAllowedUniqueID);
		}

		public int recordedFileIDNum() {
			return usedFileUniqueIDs.cardinality();
		}

		public boolean isFileIDRecorded(FileID fileID) {
			return usedFileUniqueIDs.get(fileID.getUniqueID());
		}

		boolean isIDRecorded(int uniqueID) {
			return usedFileUniqueIDs.get(uniqueID);
		}

		void recordFileID(FileID fileID) {
			if (fileID.getMaxAllowedUniqueID() == maxAllowedUniqueID) {
				usedFileUniqueIDs.set(fileID.getUniqueID());
			}
		}

		void recordUniqueID(int uniqueID) {
			usedFileUniqueIDs.set(uniqueID);
		}

		void recycleFileID(int uniqueID) {
			usedFileUniqueIDs.clear(uniqueID);
		}

		public void recordFileIDs(Set fileIDs) {
			if (!fileIDs.isEmpty() & fileIDs.iterator().next().getMaxAllowedUniqueID() == maxAllowedUniqueID) {
				fileIDs.forEach(f -> usedFileUniqueIDs.set(f.getUniqueID()));
			}
		}

		@VisibleForTesting
		BitSet getUsedFileUniqueIDs() {
			return BitSet.valueOf(usedFileUniqueIDs.toLongArray());
		}
	}

}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy