org.apache.flink.runtime.state.InternalPriorityQueue Java Examples

@SuppressWarnings("unchecked")
public KeyGroupPartitionedPriorityQueue(
	@Nonnull KeyExtractorFunction<T> keyExtractor,
	@Nonnull PriorityComparator<T> elementPriorityComparator,
	@Nonnull PartitionQueueSetFactory<T, PQ> orderedCacheFactory,
	@Nonnull KeyGroupRange keyGroupRange,
	@Nonnegative int totalKeyGroups) {

	this.keyExtractor = keyExtractor;
	this.totalKeyGroups = totalKeyGroups;
	this.firstKeyGroup = keyGroupRange.getStartKeyGroup();
	this.keyGroupedHeaps = (PQ[]) new InternalPriorityQueue[keyGroupRange.getNumberOfKeyGroups()];
	this.heapOfKeyGroupedHeaps = new HeapPriorityQueue<>(
		new InternalPriorityQueueComparator<>(elementPriorityComparator),
		keyGroupRange.getNumberOfKeyGroups());
	for (int i = 0; i < keyGroupedHeaps.length; i++) {
		final PQ keyGroupSubHeap =
			orderedCacheFactory.create(firstKeyGroup + i, totalKeyGroups, keyExtractor, elementPriorityComparator);
		keyGroupedHeaps[i] = keyGroupSubHeap;
		heapOfKeyGroupedHeaps.add(keyGroupSubHeap);
	}
}
 

@SuppressWarnings("unchecked")
public KeyGroupPartitionedPriorityQueue(
	@Nonnull KeyExtractorFunction<T> keyExtractor,
	@Nonnull PriorityComparator<T> elementPriorityComparator,
	@Nonnull PartitionQueueSetFactory<T, PQ> orderedCacheFactory,
	@Nonnull KeyGroupRange keyGroupRange,
	@Nonnegative int totalKeyGroups) {

	this.keyExtractor = keyExtractor;
	this.totalKeyGroups = totalKeyGroups;
	this.firstKeyGroup = keyGroupRange.getStartKeyGroup();
	this.keyGroupedHeaps = (PQ[]) new InternalPriorityQueue[keyGroupRange.getNumberOfKeyGroups()];
	this.heapOfKeyGroupedHeaps = new HeapPriorityQueue<>(
		new InternalPriorityQueueComparator<>(elementPriorityComparator),
		keyGroupRange.getNumberOfKeyGroups());
	for (int i = 0; i < keyGroupedHeaps.length; i++) {
		final PQ keyGroupSubHeap =
			orderedCacheFactory.create(firstKeyGroup + i, totalKeyGroups, keyExtractor, elementPriorityComparator);
		keyGroupedHeaps[i] = keyGroupSubHeap;
		heapOfKeyGroupedHeaps.add(keyGroupSubHeap);
	}
}
 

@SuppressWarnings("all")
public static InternalPriorityQueue<InternalTimer<Object, TimerInternals.TimerData>>
    retrieveInternalProcessingTimerQueue(
        InternalTimerService<TimerInternals.TimerData> timerService) {
  Field internalProcessingTimerQueue = null;
  try {
    internalProcessingTimerQueue =
        InternalTimerServiceImpl.class.getDeclaredField("processingTimeTimersQueue");
    internalProcessingTimerQueue.setAccessible(true);
    return (InternalPriorityQueue) internalProcessingTimerQueue.get(timerService);
  } catch (Exception e) {
    throw new IllegalStateException("Failed to retrieve processing timer queue.", e);
  } finally {
    if (internalProcessingTimerQueue != null) {
      internalProcessingTimerQueue.setAccessible(false);
    }
  }
}
 

@SuppressWarnings("unchecked")
public KeyGroupPartitionedPriorityQueue(
	@Nonnull KeyExtractorFunction<T> keyExtractor,
	@Nonnull PriorityComparator<T> elementPriorityComparator,
	@Nonnull PartitionQueueSetFactory<T, PQ> orderedCacheFactory,
	@Nonnull KeyGroupRange keyGroupRange,
	@Nonnegative int totalKeyGroups) {

	this.keyExtractor = keyExtractor;
	this.totalKeyGroups = totalKeyGroups;
	this.firstKeyGroup = keyGroupRange.getStartKeyGroup();
	this.keyGroupedHeaps = (PQ[]) new InternalPriorityQueue[keyGroupRange.getNumberOfKeyGroups()];
	this.heapOfKeyGroupedHeaps = new HeapPriorityQueue<>(
		new InternalPriorityQueueComparator<>(elementPriorityComparator),
		keyGroupRange.getNumberOfKeyGroups());
	for (int i = 0; i < keyGroupedHeaps.length; i++) {
		final PQ keyGroupSubHeap =
			orderedCacheFactory.create(firstKeyGroup + i, totalKeyGroups, keyExtractor, elementPriorityComparator);
		keyGroupedHeaps[i] = keyGroupSubHeap;
		heapOfKeyGroupedHeaps.add(keyGroupSubHeap);
	}
}
 

@Override
protected InternalPriorityQueue<TestElement> newPriorityQueue(int initialCapacity) {
	return new KeyGroupPartitionedPriorityQueue<>(
		KEY_EXTRACTOR_FUNCTION,
		TEST_ELEMENT_PRIORITY_COMPARATOR,
		newFactory(),
		KEY_GROUP_RANGE, KEY_GROUP_RANGE.getNumberOfKeyGroups());
}
 

@Override
protected InternalPriorityQueue<TestElement> newPriorityQueue(int initialCapacity) {
	return new KeyGroupPartitionedPriorityQueue<>(
		KEY_EXTRACTOR_FUNCTION,
		TEST_ELEMENT_PRIORITY_COMPARATOR,
		newFactory(initialCapacity),
		KEY_GROUP_RANGE, KEY_GROUP_RANGE.getNumberOfKeyGroups());
}
 

private int countTimersInNamespaceInternal(N namespace, InternalPriorityQueue<TimerHeapInternalTimer<K, N>> queue) {
	int count = 0;
	try (final CloseableIterator<TimerHeapInternalTimer<K, N>> iterator = queue.iterator()) {
		while (iterator.hasNext()) {
			final TimerHeapInternalTimer<K, N> timer = iterator.next();
			if (timer.getNamespace().equals(namespace)) {
				count++;
			}
		}
	} catch (Exception e) {
		throw new FlinkRuntimeException("Exception when closing iterator.", e);
	}
	return count;
}
 

@Override
protected InternalPriorityQueue<TestElement> newPriorityQueue(int initialCapacity) {
	return new KeyGroupPartitionedPriorityQueue<>(
		KEY_EXTRACTOR_FUNCTION,
		TEST_ELEMENT_PRIORITY_COMPARATOR,
		newFactory(),
		KEY_GROUP_RANGE, KEY_GROUP_RANGE.getNumberOfKeyGroups());
}
 

@Override
protected InternalPriorityQueue<TestElement> newPriorityQueue(int initialCapacity) {
	return new KeyGroupPartitionedPriorityQueue<>(
		KEY_EXTRACTOR_FUNCTION,
		TEST_ELEMENT_PRIORITY_COMPARATOR,
		newFactory(initialCapacity),
		KEY_GROUP_RANGE, KEY_GROUP_RANGE.getNumberOfKeyGroups());
}
 

private int countTimersInNamespaceInternal(N namespace, InternalPriorityQueue<TimerHeapInternalTimer<K, N>> queue) {
	int count = 0;
	try (final CloseableIterator<TimerHeapInternalTimer<K, N>> iterator = queue.iterator()) {
		while (iterator.hasNext()) {
			final TimerHeapInternalTimer<K, N> timer = iterator.next();
			if (timer.getNamespace().equals(namespace)) {
				count++;
			}
		}
	} catch (Exception e) {
		throw new FlinkRuntimeException("Exception when closing iterator.", e);
	}
	return count;
}
 

/**
 * Processes all pending processing timers. This is intended for use during shutdown. From Flink
 * 1.10 on, processing timer execution is stopped when the operator is closed. This leads to
 * problems for applications which assume all pending timers will be completed. Although Flink
 * does drain the remaining timers after close(), this is not sufficient because no new timers
 * are allowed to be scheduled anymore. This breaks Beam pipelines which rely on all processing
 * timers to be scheduled and executed.
 */
void processPendingProcessingTimeTimers() {
  final KeyedStateBackend<Object> keyedStateBackend = getKeyedStateBackend();
  final InternalPriorityQueue<InternalTimer<Object, TimerData>> processingTimeTimersQueue =
      Workarounds.retrieveInternalProcessingTimerQueue(timerService);

  InternalTimer<Object, TimerData> internalTimer;
  while ((internalTimer = processingTimeTimersQueue.poll()) != null) {
    keyedStateBackend.setCurrentKey(internalTimer.getKey());
    TimerData timer = internalTimer.getNamespace();
    checkInvokeStartBundle();
    fireTimer(timer);
  }
}
 

@Override
protected InternalPriorityQueue<TestElement> newPriorityQueue(int initialCapacity) {
	return new KeyGroupPartitionedPriorityQueue<>(
		KEY_EXTRACTOR_FUNCTION,
		TEST_ELEMENT_PRIORITY_COMPARATOR,
		newFactory(),
		KEY_GROUP_RANGE, KEY_GROUP_RANGE.getNumberOfKeyGroups());
}
 

@Override
protected InternalPriorityQueue<TestElement> newPriorityQueue(int initialCapacity) {
	return new KeyGroupPartitionedPriorityQueue<>(
		KEY_EXTRACTOR_FUNCTION,
		TEST_ELEMENT_PRIORITY_COMPARATOR,
		newFactory(initialCapacity),
		KEY_GROUP_RANGE, KEY_GROUP_RANGE.getNumberOfKeyGroups());
}
 

private int countTimersInNamespaceInternal(N namespace, InternalPriorityQueue<TimerHeapInternalTimer<K, N>> queue) {
	int count = 0;
	try (final CloseableIterator<TimerHeapInternalTimer<K, N>> iterator = queue.iterator()) {
		while (iterator.hasNext()) {
			final TimerHeapInternalTimer<K, N> timer = iterator.next();
			if (timer.getNamespace().equals(namespace)) {
				count++;
			}
		}
	} catch (Exception e) {
		throw new FlinkRuntimeException("Exception when closing iterator.", e);
	}
	return count;
}