Java Code Examples for org.apache.flink.runtime.state.KeyGroupRange#getEndKeyGroup()

/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange  Key group range to register
 * @param kvStateId      ID of the KvState instance at the key group index.
 * @param kvStateAddress Server address of the KvState instance at the key group index.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 */
public void registerKvState(KeyGroupRange keyGroupRange, KvStateID kvStateId, InetSocketAddress kvStateAddress) {

	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {

		if (kvStateIds[kgIdx] == null && kvStateAddresses[kgIdx] == null) {
			numRegisteredKeyGroups++;
		}

		kvStateIds[kgIdx] = kvStateId;
		kvStateAddresses[kgIdx] = kvStateAddress;
	}
}
 

/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange Key group range to unregister.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 * @throws IllegalArgumentException  If no location information registered for a key group index in the range.
 */
void unregisterKvState(KeyGroupRange keyGroupRange) {
	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {
		if (kvStateIds[kgIdx] == null || kvStateAddresses[kgIdx] == null) {
			throw new IllegalArgumentException("Not registered. Probably registration/unregistration race.");
		}

		numRegisteredKeyGroups--;

		kvStateIds[kgIdx] = null;
		kvStateAddresses[kgIdx] = null;
	}
}
 

/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange  Key group range to register
 * @param kvStateId      ID of the KvState instance at the key group index.
 * @param kvStateAddress Server address of the KvState instance at the key group index.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 */
public void registerKvState(KeyGroupRange keyGroupRange, KvStateID kvStateId, InetSocketAddress kvStateAddress) {

	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {

		if (kvStateIds[kgIdx] == null && kvStateAddresses[kgIdx] == null) {
			numRegisteredKeyGroups++;
		}

		kvStateIds[kgIdx] = kvStateId;
		kvStateAddresses[kgIdx] = kvStateAddress;
	}
}
 

/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange Key group range to unregister.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 * @throws IllegalArgumentException  If no location information registered for a key group index in the range.
 */
void unregisterKvState(KeyGroupRange keyGroupRange) {
	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {
		if (kvStateIds[kgIdx] == null || kvStateAddresses[kgIdx] == null) {
			throw new IllegalArgumentException("Not registered. Probably registration/unregistration race.");
		}

		numRegisteredKeyGroups--;

		kvStateIds[kgIdx] = null;
		kvStateAddresses[kgIdx] = null;
	}
}
 

/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange  Key group range to register
 * @param kvStateId      ID of the KvState instance at the key group index.
 * @param kvStateAddress Server address of the KvState instance at the key group index.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 */
public void registerKvState(KeyGroupRange keyGroupRange, KvStateID kvStateId, InetSocketAddress kvStateAddress) {

	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {

		if (kvStateIds[kgIdx] == null && kvStateAddresses[kgIdx] == null) {
			numRegisteredKeyGroups++;
		}

		kvStateIds[kgIdx] = kvStateId;
		kvStateAddresses[kgIdx] = kvStateAddress;
	}
}
 

/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange Key group range to unregister.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 * @throws IllegalArgumentException  If no location information registered for a key group index in the range.
 */
void unregisterKvState(KeyGroupRange keyGroupRange) {
	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {
		if (kvStateIds[kgIdx] == null || kvStateAddresses[kgIdx] == null) {
			throw new IllegalArgumentException("Not registered. Probably registration/unregistration race.");
		}

		numRegisteredKeyGroups--;

		kvStateIds[kgIdx] = null;
		kvStateAddresses[kgIdx] = null;
	}
}
 

/**
 * The method to clip the db instance according to the target key group range using
 * the {@link RocksDB#delete(ColumnFamilyHandle, byte[])}.
 *
 * @param db the RocksDB instance to be clipped.
 * @param columnFamilyHandles the column families in the db instance.
 * @param targetKeyGroupRange the target key group range.
 * @param currentKeyGroupRange the key group range of the db instance.
 * @param keyGroupPrefixBytes Number of bytes required to prefix the key groups.
 */
public static void clipDBWithKeyGroupRange(
	@Nonnull RocksDB db,
	@Nonnull List<ColumnFamilyHandle> columnFamilyHandles,
	@Nonnull KeyGroupRange targetKeyGroupRange,
	@Nonnull KeyGroupRange currentKeyGroupRange,
	@Nonnegative int keyGroupPrefixBytes) throws RocksDBException {

	final byte[] beginKeyGroupBytes = new byte[keyGroupPrefixBytes];
	final byte[] endKeyGroupBytes = new byte[keyGroupPrefixBytes];

	if (currentKeyGroupRange.getStartKeyGroup() < targetKeyGroupRange.getStartKeyGroup()) {
		RocksDBKeySerializationUtils.serializeKeyGroup(
			currentKeyGroupRange.getStartKeyGroup(), beginKeyGroupBytes);
		RocksDBKeySerializationUtils.serializeKeyGroup(
			targetKeyGroupRange.getStartKeyGroup(), endKeyGroupBytes);
		deleteRange(db, columnFamilyHandles, beginKeyGroupBytes, endKeyGroupBytes);
	}

	if (currentKeyGroupRange.getEndKeyGroup() > targetKeyGroupRange.getEndKeyGroup()) {
		RocksDBKeySerializationUtils.serializeKeyGroup(
			targetKeyGroupRange.getEndKeyGroup() + 1, beginKeyGroupBytes);
		RocksDBKeySerializationUtils.serializeKeyGroup(
			currentKeyGroupRange.getEndKeyGroup() + 1, endKeyGroupBytes);
		deleteRange(db, columnFamilyHandles, beginKeyGroupBytes, endKeyGroupBytes);
	}
}
 

/**
 * Creates deep copy of the given {@link KeyedStateHandle}.
 */
public static KeyedStateHandle deepDummyCopy(KeyedStateHandle original) {

	if (original == null) {
		return null;
	}

	KeyGroupRange keyGroupRange = original.getKeyGroupRange();
	return new DummyKeyedStateHandle(
		new KeyGroupRange(keyGroupRange.getStartKeyGroup(), keyGroupRange.getEndKeyGroup()));
}
 

/**
 * The method to clip the db instance according to the target key group range using
 * the {@link RocksDB#delete(ColumnFamilyHandle, byte[])}.
 *
 * @param db the RocksDB instance to be clipped.
 * @param columnFamilyHandles the column families in the db instance.
 * @param targetKeyGroupRange the target key group range.
 * @param currentKeyGroupRange the key group range of the db instance.
 * @param keyGroupPrefixBytes Number of bytes required to prefix the key groups.
 */
public static void clipDBWithKeyGroupRange(
	@Nonnull RocksDB db,
	@Nonnull List<ColumnFamilyHandle> columnFamilyHandles,
	@Nonnull KeyGroupRange targetKeyGroupRange,
	@Nonnull KeyGroupRange currentKeyGroupRange,
	@Nonnegative int keyGroupPrefixBytes) throws RocksDBException {

	final byte[] beginKeyGroupBytes = new byte[keyGroupPrefixBytes];
	final byte[] endKeyGroupBytes = new byte[keyGroupPrefixBytes];

	if (currentKeyGroupRange.getStartKeyGroup() < targetKeyGroupRange.getStartKeyGroup()) {
		RocksDBKeySerializationUtils.serializeKeyGroup(
			currentKeyGroupRange.getStartKeyGroup(), beginKeyGroupBytes);
		RocksDBKeySerializationUtils.serializeKeyGroup(
			targetKeyGroupRange.getStartKeyGroup(), endKeyGroupBytes);
		deleteRange(db, columnFamilyHandles, beginKeyGroupBytes, endKeyGroupBytes);
	}

	if (currentKeyGroupRange.getEndKeyGroup() > targetKeyGroupRange.getEndKeyGroup()) {
		RocksDBKeySerializationUtils.serializeKeyGroup(
			targetKeyGroupRange.getEndKeyGroup() + 1, beginKeyGroupBytes);
		RocksDBKeySerializationUtils.serializeKeyGroup(
			currentKeyGroupRange.getEndKeyGroup() + 1, endKeyGroupBytes);
		deleteRange(db, columnFamilyHandles, beginKeyGroupBytes, endKeyGroupBytes);
	}
}
 

/**
 * Creates deep copy of the given {@link KeyedStateHandle}.
 */
public static KeyedStateHandle deepDummyCopy(KeyedStateHandle original) {

	if (original == null) {
		return null;
	}

	KeyGroupRange keyGroupRange = original.getKeyGroupRange();
	return new DummyKeyedStateHandle(
		new KeyGroupRange(keyGroupRange.getStartKeyGroup(), keyGroupRange.getEndKeyGroup()));
}
 

/**
 * The method to clip the db instance according to the target key group range using
 * the {@link RocksDB#delete(ColumnFamilyHandle, byte[])}.
 *
 * @param db the RocksDB instance to be clipped.
 * @param columnFamilyHandles the column families in the db instance.
 * @param targetKeyGroupRange the target key group range.
 * @param currentKeyGroupRange the key group range of the db instance.
 * @param keyGroupPrefixBytes Number of bytes required to prefix the key groups.
 */
public static void clipDBWithKeyGroupRange(
	@Nonnull RocksDB db,
	@Nonnull List<ColumnFamilyHandle> columnFamilyHandles,
	@Nonnull KeyGroupRange targetKeyGroupRange,
	@Nonnull KeyGroupRange currentKeyGroupRange,
	@Nonnegative int keyGroupPrefixBytes,
	@Nonnegative long writeBatchSize) throws RocksDBException {

	final byte[] beginKeyGroupBytes = new byte[keyGroupPrefixBytes];
	final byte[] endKeyGroupBytes = new byte[keyGroupPrefixBytes];

	if (currentKeyGroupRange.getStartKeyGroup() < targetKeyGroupRange.getStartKeyGroup()) {
		RocksDBKeySerializationUtils.serializeKeyGroup(
			currentKeyGroupRange.getStartKeyGroup(), beginKeyGroupBytes);
		RocksDBKeySerializationUtils.serializeKeyGroup(
			targetKeyGroupRange.getStartKeyGroup(), endKeyGroupBytes);
		deleteRange(db, columnFamilyHandles, beginKeyGroupBytes, endKeyGroupBytes, writeBatchSize);
	}

	if (currentKeyGroupRange.getEndKeyGroup() > targetKeyGroupRange.getEndKeyGroup()) {
		RocksDBKeySerializationUtils.serializeKeyGroup(
			targetKeyGroupRange.getEndKeyGroup() + 1, beginKeyGroupBytes);
		RocksDBKeySerializationUtils.serializeKeyGroup(
			currentKeyGroupRange.getEndKeyGroup() + 1, endKeyGroupBytes);
		deleteRange(db, columnFamilyHandles, beginKeyGroupBytes, endKeyGroupBytes, writeBatchSize);
	}
}
 

/**
 * Creates deep copy of the given {@link KeyedStateHandle}.
 */
public static KeyedStateHandle deepDummyCopy(KeyedStateHandle original) {

	if (original == null) {
		return null;
	}

	KeyGroupRange keyGroupRange = original.getKeyGroupRange();
	return new DummyKeyedStateHandle(
		new KeyGroupRange(keyGroupRange.getStartKeyGroup(), keyGroupRange.getEndKeyGroup()));
}
 

@Test
public void testStateAndTimerStateShufflingScalingDown() throws Exception {
	final int maxParallelism = 10;

	// first get two keys that will fall into different key-group ranges that go
	// to different operator subtasks when we restore

	// get two sub key-ranges so that we can restore two ranges separately
	KeyGroupRange subKeyGroupRange1 = new KeyGroupRange(0, (maxParallelism / 2) - 1);
	KeyGroupRange subKeyGroupRange2 = new KeyGroupRange(subKeyGroupRange1.getEndKeyGroup() + 1, maxParallelism - 1);

	// get two different keys, one per sub range
	int key1 = getKeyInKeyGroupRange(subKeyGroupRange1, maxParallelism);
	int key2 = getKeyInKeyGroupRange(subKeyGroupRange2, maxParallelism);

	OperatorSubtaskState snapshot1, snapshot2;
	// register some state with both instances and scale down to parallelism 1
	try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness1 =
			createTestHarness(maxParallelism, 2, 0)) {

		testHarness1.setup();
		testHarness1.open();

		testHarness1.processWatermark(0L);
		testHarness1.setProcessingTime(0L);

		testHarness1.processElement(new Tuple2<>(key1, "SET_EVENT_TIME_TIMER:30"), 0);
		testHarness1.processElement(new Tuple2<>(key1, "SET_PROC_TIME_TIMER:30"), 0);
		testHarness1.processElement(new Tuple2<>(key1, "SET_STATE:HELLO"), 0);

		snapshot1 = testHarness1.snapshot(0, 0);
	}

	try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness2 =
			createTestHarness(maxParallelism, 2, 1)) {
		testHarness2.setup();
		testHarness2.open();

		testHarness2.processWatermark(0L);
		testHarness2.setProcessingTime(0L);

		testHarness2.processElement(new Tuple2<>(key2, "SET_EVENT_TIME_TIMER:40"), 0);
		testHarness2.processElement(new Tuple2<>(key2, "SET_PROC_TIME_TIMER:40"), 0);
		testHarness2.processElement(new Tuple2<>(key2, "SET_STATE:CIAO"), 0);

		snapshot2 = testHarness2.snapshot(0, 0);
	}
	// take a snapshot from each one of the "parallel" instances of the operator
	// and combine them into one so that we can scale down

	OperatorSubtaskState repackagedState =
		AbstractStreamOperatorTestHarness.repackageState(snapshot1, snapshot2);

	OperatorSubtaskState initSubTaskState =
		AbstractStreamOperatorTestHarness.repartitionOperatorState(repackagedState, maxParallelism, 2, 1, 0);

	try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness3 =
			createTestHarness(maxParallelism, 1, 0)) {
		testHarness3.setup();
		testHarness3.initializeState(initSubTaskState);
		testHarness3.open();

		testHarness3.processWatermark(30L);
		assertThat(extractResult(testHarness3), contains("ON_EVENT_TIME:HELLO"));
		assertTrue(extractResult(testHarness3).isEmpty());

		testHarness3.processWatermark(40L);
		assertThat(extractResult(testHarness3), contains("ON_EVENT_TIME:CIAO"));
		assertTrue(extractResult(testHarness3).isEmpty());

		testHarness3.setProcessingTime(30L);
		assertThat(extractResult(testHarness3), contains("ON_PROC_TIME:HELLO"));
		assertTrue(extractResult(testHarness3).isEmpty());

		testHarness3.setProcessingTime(40L);
		assertThat(extractResult(testHarness3), contains("ON_PROC_TIME:CIAO"));
		assertTrue(extractResult(testHarness3).isEmpty());
	}
}