org.apache.flink.util.function.SupplierWithException Java Examples

private SupplierWithException<CheckpointStreamWithResultProvider, Exception> createCheckpointStreamSupplier(
	long checkpointId,
	CheckpointStreamFactory primaryStreamFactory,
	CheckpointOptions checkpointOptions) {

	return localRecoveryConfig.isLocalRecoveryEnabled() && !checkpointOptions.getCheckpointType().isSavepoint() ?

		() -> CheckpointStreamWithResultProvider.createDuplicatingStream(
			checkpointId,
			CheckpointedStateScope.EXCLUSIVE,
			primaryStreamFactory,
			localRecoveryConfig.getLocalStateDirectoryProvider()) :

		() -> CheckpointStreamWithResultProvider.createSimpleStream(
			CheckpointedStateScope.EXCLUSIVE,
			primaryStreamFactory);
}
 

@SuppressWarnings("unchecked")
DummyFlinkKafkaConsumer(
	SupplierWithException<AbstractFetcher<T, ?>, Exception> testFetcherSupplier,
	AbstractPartitionDiscoverer testPartitionDiscoverer,
	boolean isAutoCommitEnabled,
	long discoveryIntervalMillis,
	List<String> topics) {

	super(
		topics,
		null,
		(KeyedDeserializationSchema< T >) mock(KeyedDeserializationSchema.class),
		discoveryIntervalMillis,
		false);

	this.testFetcherSupplier = testFetcherSupplier;
	this.testPartitionDiscoverer = testPartitionDiscoverer;
	this.isAutoCommitEnabled = isAutoCommitEnabled;
}
 

/**
 * Runs the given runnable with the given ClassLoader as the thread's
 * {@link Thread#setContextClassLoader(ClassLoader) context class loader}.
 *
 * <p>The method will make sure to set the context class loader of the calling thread
 * back to what it was before after the runnable completed.
 */
public static <R, E extends Throwable> R withContextClassLoader(
		final ClassLoader cl,
		final SupplierWithException<R, E> s) throws E {

	final Thread currentThread = Thread.currentThread();
	final ClassLoader oldClassLoader = currentThread.getContextClassLoader();

	try {
		currentThread.setContextClassLoader(cl);
		return s.get();
	}
	finally {
		currentThread.setContextClassLoader(oldClassLoader);
	}
}
 

/**
 * Returns {@code requestValue} if it is not null, otherwise returns the query parameter value
 * if it is not null, otherwise returns the default value.
 */
public static <T> T fromRequestBodyOrQueryParameter(
		T requestValue,
		SupplierWithException<T, RestHandlerException> queryParameterExtractor,
		T defaultValue,
		Logger log) throws RestHandlerException {
	if (requestValue != null) {
		return requestValue;
	} else {
		T queryParameterValue = queryParameterExtractor.get();
		if (queryParameterValue != null) {
			log.warn("Configuring the job submission via query parameters is deprecated." +
				" Please migrate to submitting a JSON request instead.");
			return queryParameterValue;
		} else {
			return defaultValue;
		}
	}
}
 

TestingExecutor(
		List<SupplierWithException<TypedResult<List<Tuple2<Boolean, Row>>>, SqlExecutionException>> resultChanges,
		List<SupplierWithException<TypedResult<Integer>, SqlExecutionException>> snapshotResults,
		List<SupplierWithException<List<Row>, SqlExecutionException>> resultPages,
		BiConsumerWithException<String, String, SqlExecutionException> useCatalogConsumer,
		BiConsumerWithException<String, String, SqlExecutionException> useDatabaseConsumer,
		BiFunctionWithException<String, String, TableResult, SqlExecutionException> executeSqlConsumer,
		TriFunctionWithException<String, String, String, Void, SqlExecutionException> setSessionPropertyFunction,
		FunctionWithException<String, Void, SqlExecutionException> resetSessionPropertiesFunction) {
	this.resultChanges = resultChanges;
	this.snapshotResults = snapshotResults;
	this.resultPages = resultPages;
	this.useCatalogConsumer = useCatalogConsumer;
	this.useDatabaseConsumer = useDatabaseConsumer;
	this.executeSqlConsumer = executeSqlConsumer;
	this.setSessionPropertyFunction = setSessionPropertyFunction;
	this.resetSessionPropertiesFunction = resetSessionPropertiesFunction;
	helper = new SqlParserHelper();
	helper.registerTables();
}
 

<SE extends Throwable, CE extends Throwable, CLE extends Throwable, V> TtlValue<V> getWrappedWithTtlCheckAndUpdate(
	SupplierWithException<TtlValue<V>, SE> getter,
	ThrowingConsumer<TtlValue<V>, CE> updater,
	ThrowingRunnable<CLE> stateClear) throws SE, CE, CLE {
	TtlValue<V> ttlValue = getter.get();
	if (ttlValue == null) {
		return null;
	} else if (expired(ttlValue)) {
		stateClear.run();
		if (!returnExpired) {
			return null;
		}
	} else if (updateTsOnRead) {
		updater.accept(rewrapWithNewTs(ttlValue));
	}
	return ttlValue;
}
 

private SupplierWithException<CheckpointStreamWithResultProvider, Exception> createCheckpointStreamSupplier(
	long checkpointId,
	CheckpointStreamFactory primaryStreamFactory,
	CheckpointOptions checkpointOptions) {

	return localRecoveryConfig.isLocalRecoveryEnabled() &&
		(CheckpointType.SAVEPOINT != checkpointOptions.getCheckpointType()) ?

		() -> CheckpointStreamWithResultProvider.createDuplicatingStream(
			checkpointId,
			CheckpointedStateScope.EXCLUSIVE,
			primaryStreamFactory,
			localRecoveryConfig.getLocalStateDirectoryProvider()) :

		() -> CheckpointStreamWithResultProvider.createSimpleStream(
			CheckpointedStateScope.EXCLUSIVE,
			primaryStreamFactory);
}
 

@SuppressWarnings("unchecked")
DummyFlinkKafkaConsumer(
	SupplierWithException<AbstractFetcher<T, ?>, Exception> testFetcherSupplier,
	AbstractPartitionDiscoverer testPartitionDiscoverer,
	boolean isAutoCommitEnabled,
	long discoveryIntervalMillis,
	List<String> topics,
	KafkaDeserializationSchema<T> mock) {

	super(
		topics,
		null,
		mock,
		discoveryIntervalMillis,
		false);

	this.testFetcherSupplier = testFetcherSupplier;
	this.testPartitionDiscoverer = testPartitionDiscoverer;
	this.isAutoCommitEnabled = isAutoCommitEnabled;
}
 

@SuppressWarnings("unchecked")
DummyFlinkKafkaConsumer(
	SupplierWithException<AbstractFetcher<T, ?>, Exception> testFetcherSupplier,
	AbstractPartitionDiscoverer testPartitionDiscoverer,
	boolean isAutoCommitEnabled,
	long discoveryIntervalMillis,
	List<String> topics) {

	super(
		topics,
		null,
		(KeyedDeserializationSchema< T >) mock(KeyedDeserializationSchema.class),
		discoveryIntervalMillis,
		false);

	this.testFetcherSupplier = testFetcherSupplier;
	this.testPartitionDiscoverer = testPartitionDiscoverer;
	this.isAutoCommitEnabled = isAutoCommitEnabled;
}
 

private SupplierWithException<CheckpointStreamWithResultProvider, Exception> createCheckpointStreamSupplier(
	long checkpointId,
	CheckpointStreamFactory primaryStreamFactory,
	CheckpointOptions checkpointOptions) {

	return localRecoveryConfig.isLocalRecoveryEnabled() && !checkpointOptions.getCheckpointType().isSavepoint() ?

		() -> CheckpointStreamWithResultProvider.createDuplicatingStream(
			checkpointId,
			CheckpointedStateScope.EXCLUSIVE,
			primaryStreamFactory,
			localRecoveryConfig.getLocalStateDirectoryProvider()) :

		() -> CheckpointStreamWithResultProvider.createSimpleStream(
			CheckpointedStateScope.EXCLUSIVE,
			primaryStreamFactory);
}
 

<SE extends Throwable, CE extends Throwable, CLE extends Throwable, V> TtlValue<V> getWrappedWithTtlCheckAndUpdate(
	SupplierWithException<TtlValue<V>, SE> getter,
	ThrowingConsumer<TtlValue<V>, CE> updater,
	ThrowingRunnable<CLE> stateClear) throws SE, CE, CLE {
	TtlValue<V> ttlValue = getter.get();
	if (ttlValue == null) {
		return null;
	} else if (expired(ttlValue)) {
		stateClear.run();
		if (!returnExpired) {
			return null;
		}
	} else if (updateTsOnRead) {
		updater.accept(rewrapWithNewTs(ttlValue));
	}
	return ttlValue;
}
 

/**
 * Returns {@code requestValue} if it is not null, otherwise returns the query parameter value
 * if it is not null, otherwise returns the default value.
 */
public static <T> T fromRequestBodyOrQueryParameter(
		T requestValue,
		SupplierWithException<T, RestHandlerException> queryParameterExtractor,
		T defaultValue,
		Logger log) throws RestHandlerException {
	if (requestValue != null) {
		return requestValue;
	} else {
		T queryParameterValue = queryParameterExtractor.get();
		if (queryParameterValue != null) {
			log.warn("Configuring the job submission via query parameters is deprecated." +
				" Please migrate to submitting a JSON request instead.");
			return queryParameterValue;
		} else {
			return defaultValue;
		}
	}
}
 

/**
 * Returns {@code requestValue} if it is not null, otherwise returns the query parameter value
 * if it is not null, otherwise returns the default value.
 */
public static <T> T fromRequestBodyOrQueryParameter(
		T requestValue,
		SupplierWithException<T, RestHandlerException> queryParameterExtractor,
		T defaultValue,
		Logger log) throws RestHandlerException {
	if (requestValue != null) {
		return requestValue;
	} else {
		T queryParameterValue = queryParameterExtractor.get();
		if (queryParameterValue != null) {
			log.warn("Configuring the job submission via query parameters is deprecated." +
				" Please migrate to submitting a JSON request instead.");
			return queryParameterValue;
		} else {
			return defaultValue;
		}
	}
}
 

<SE extends Throwable, CE extends Throwable, CLE extends Throwable, V> TtlValue<V> getWrappedWithTtlCheckAndUpdate(
	SupplierWithException<TtlValue<V>, SE> getter,
	ThrowingConsumer<TtlValue<V>, CE> updater,
	ThrowingRunnable<CLE> stateClear) throws SE, CE, CLE {
	TtlValue<V> ttlValue = getter.get();
	if (ttlValue == null) {
		return null;
	} else if (expired(ttlValue)) {
		stateClear.run();
		if (!returnExpired) {
			return null;
		}
	} else if (updateTsOnRead) {
		updater.accept(rewrapWithNewTs(ttlValue));
	}
	return ttlValue;
}
 

private SupplierWithException<Boolean, Exception> jobIsRunning(Supplier<CompletableFuture<? extends AccessExecutionGraph>> executionGraphFutureSupplier) {
	final Predicate<AccessExecution> runningOrFinished = ExecutionGraphTestUtils.isInExecutionState(ExecutionState.RUNNING).or(ExecutionGraphTestUtils.isInExecutionState(ExecutionState.FINISHED));
	final Predicate<AccessExecutionGraph> allExecutionsRunning = ExecutionGraphTestUtils.allExecutionsPredicate(runningOrFinished);

	return () -> {
		final AccessExecutionGraph executionGraph = executionGraphFutureSupplier.get().join();
		return allExecutionsRunning.test(executionGraph);
	};
}
 

private SupplierWithException<Boolean, Exception> isJobVertexBackPressured(final JobVertex sourceJobVertex) {
	return () -> {
		final OperatorBackPressureStatsResponse backPressureStatsResponse = dispatcherGateway
			.requestOperatorBackPressureStats(TEST_JOB_ID, sourceJobVertex.getID())
			.get();

		return backPressureStatsResponse.getOperatorBackPressureStats()
			.map(backPressureStats -> isBackPressured(backPressureStats))
			.orElse(false);
	};
}
 

private static <X> X fetchMetric(final SupplierWithException<CompletableFuture<X>, IOException> clientOperation, final Predicate<X> predicate) throws InterruptedException, ExecutionException, TimeoutException {
	final CompletableFuture<X> responseFuture = FutureUtils.retrySuccessfulWithDelay(() -> {
			try {
				return clientOperation.get();
			} catch (IOException e) {
				throw new RuntimeException(e);
			}
		},
		Time.seconds(1),
		Deadline.fromNow(Duration.ofSeconds(5)),
		predicate,
		new ScheduledExecutorServiceAdapter(scheduledExecutorService));

	return responseFuture.get(30, TimeUnit.SECONDS);
}
 

/**
 * Creates an input gate and all of its input channels.
 */
public SingleInputGate create(
		@Nonnull String owningTaskName,
		int gateIndex,
		@Nonnull InputGateDeploymentDescriptor igdd,
		@Nonnull PartitionProducerStateProvider partitionProducerStateProvider,
		@Nonnull InputChannelMetrics metrics) {
	SupplierWithException<BufferPool, IOException> bufferPoolFactory = createBufferPoolFactory(
		networkBufferPool,
		networkBuffersPerChannel,
		floatingNetworkBuffersPerGate,
		igdd.getShuffleDescriptors().length,
		igdd.getConsumedPartitionType());

	BufferDecompressor bufferDecompressor = null;
	if (igdd.getConsumedPartitionType().isBlocking() && blockingShuffleCompressionEnabled) {
		bufferDecompressor = new BufferDecompressor(networkBufferSize, compressionCodec);
	}

	SingleInputGate inputGate = new SingleInputGate(
		owningTaskName,
		gateIndex,
		igdd.getConsumedResultId(),
		igdd.getConsumedPartitionType(),
		igdd.getConsumedSubpartitionIndex(),
		igdd.getShuffleDescriptors().length,
		partitionProducerStateProvider,
		bufferPoolFactory,
		bufferDecompressor,
		networkBufferPool);

	createInputChannels(owningTaskName, igdd, inputGate, metrics);
	return inputGate;
}
 

@VisibleForTesting
static SupplierWithException<BufferPool, IOException> createBufferPoolFactory(
		BufferPoolFactory bufferPoolFactory,
		int networkBuffersPerChannel,
		int floatingNetworkBuffersPerGate,
		int size,
		ResultPartitionType type) {
	// Note that we should guarantee at-least one floating buffer for local channel state recovery.
	return () -> bufferPoolFactory.createBufferPool(1, floatingNetworkBuffersPerGate);
}
 

public SingleInputGate(
	String owningTaskName,
	int gateIndex,
	IntermediateDataSetID consumedResultId,
	final ResultPartitionType consumedPartitionType,
	int consumedSubpartitionIndex,
	int numberOfInputChannels,
	PartitionProducerStateProvider partitionProducerStateProvider,
	SupplierWithException<BufferPool, IOException> bufferPoolFactory,
	@Nullable BufferDecompressor bufferDecompressor,
	MemorySegmentProvider memorySegmentProvider) {

	this.owningTaskName = checkNotNull(owningTaskName);
	Preconditions.checkArgument(0 <= gateIndex, "The gate index must be positive.");
	this.gateIndex = gateIndex;

	this.consumedResultId = checkNotNull(consumedResultId);
	this.consumedPartitionType = checkNotNull(consumedPartitionType);
	this.bufferPoolFactory = checkNotNull(bufferPoolFactory);

	checkArgument(consumedSubpartitionIndex >= 0);
	this.consumedSubpartitionIndex = consumedSubpartitionIndex;

	checkArgument(numberOfInputChannels > 0);
	this.numberOfInputChannels = numberOfInputChannels;

	this.inputChannels = new HashMap<>(numberOfInputChannels);
	this.channels = new InputChannel[numberOfInputChannels];
	this.channelsWithEndOfPartitionEvents = new BitSet(numberOfInputChannels);
	this.enqueuedInputChannelsWithData = new BitSet(numberOfInputChannels);

	this.partitionProducerStateProvider = checkNotNull(partitionProducerStateProvider);

	this.bufferDecompressor = bufferDecompressor;
	this.memorySegmentProvider = checkNotNull(memorySegmentProvider);

	this.closeFuture = new CompletableFuture<>();
}
 

public static void waitUntilCondition(SupplierWithException<Boolean, Exception> condition, Deadline timeout, long retryIntervalMillis) throws Exception {
	while (timeout.hasTimeLeft() && !condition.get()) {
		final long timeLeft = Math.max(0, timeout.timeLeft().toMillis());
		Thread.sleep(Math.min(retryIntervalMillis, timeLeft));
	}

	if (!timeout.hasTimeLeft()) {
		throw new TimeoutException("Condition was not met in given timeout.");
	}
}
 

SnapshotAsynchronousPartCallable(
	@Nonnull SupplierWithException<CheckpointStreamWithResultProvider, Exception> checkpointStreamSupplier,
	@Nonnull ResourceGuard.Lease dbLease,
	@Nonnull Snapshot snapshot,
	@Nonnull List<StateMetaInfoSnapshot> stateMetaInfoSnapshots,
	@Nonnull List<RocksDbKvStateInfo> metaDataCopy,
	@Nonnull String logPathString) {

	this.checkpointStreamSupplier = checkpointStreamSupplier;
	this.dbLease = dbLease;
	this.snapshot = snapshot;
	this.stateMetaInfoSnapshots = stateMetaInfoSnapshots;
	this.metaData = fillMetaData(metaDataCopy);
	this.logPathString = logPathString;
}
 

@SuppressWarnings("unchecked")
DummyFlinkKafkaConsumer(
		SupplierWithException<AbstractFetcher<T, ?>, Exception> testFetcherSupplier,
		AbstractPartitionDiscoverer testPartitionDiscoverer,
		boolean isAutoCommitEnabled,
		long discoveryIntervalMillis) {
	this(
		testFetcherSupplier,
		testPartitionDiscoverer,
		isAutoCommitEnabled,
		discoveryIntervalMillis,
		Collections.singletonList("dummy-topic"),
		(KeyedDeserializationSchema<T>) mock(KeyedDeserializationSchema.class)
	);
}
 

@SuppressWarnings("unchecked")
DummyFlinkKafkaConsumer(
		SupplierWithException<AbstractFetcher<T, ?>, Exception> abstractFetcherSupplier,
		AbstractPartitionDiscoverer abstractPartitionDiscoverer,
		long discoveryIntervalMillis) {
	this(
			abstractFetcherSupplier,
			abstractPartitionDiscoverer,
			false,
			discoveryIntervalMillis);
}
 

private static <X> X fetchMetric(final SupplierWithException<CompletableFuture<X>, IOException> clientOperation, final Predicate<X> predicate) throws InterruptedException, ExecutionException, TimeoutException {
	final CompletableFuture<X> responseFuture = FutureUtils.retrySuccessfulWithDelay(() -> {
			try {
				return clientOperation.get();
			} catch (IOException e) {
				throw new RuntimeException(e);
			}
		},
		Time.seconds(1),
		Deadline.fromNow(Duration.ofSeconds(5)),
		predicate,
		new ScheduledExecutorServiceAdapter(scheduledExecutorService));

	return responseFuture.get(30, TimeUnit.SECONDS);
}
 

public static void waitUntilCondition(SupplierWithException<Boolean, Exception> condition, Deadline timeout, long retryIntervalMillis) throws Exception {
	while (timeout.hasTimeLeft() && !condition.get()) {
		final long timeLeft = Math.max(0, timeout.timeLeft().toMillis());
		Thread.sleep(Math.min(retryIntervalMillis, timeLeft));
	}

	if (!timeout.hasTimeLeft()) {
		throw new TimeoutException("Condition was not met in given timeout.");
	}
}
 

/**
 * Executes the given supplier with the main thread executor until completion, returns the result or a exception.
 * This method blocks until the execution is complete.
 */
public <U> U execute(@Nonnull SupplierWithException<U, Throwable> supplierWithException) {
	return CompletableFuture.supplyAsync(
		FunctionUtils.uncheckedSupplier(supplierWithException),
		mainThreadExecutor)
		.join();
}
 

private SupplierWithException<Boolean, Exception> jobIsRunning(Supplier<CompletableFuture<? extends AccessExecutionGraph>> executionGraphFutureSupplier) {
	final Predicate<AccessExecution> runningOrFinished = ExecutionGraphTestUtils.isInExecutionState(ExecutionState.RUNNING).or(ExecutionGraphTestUtils.isInExecutionState(ExecutionState.FINISHED));
	final Predicate<AccessExecutionGraph> allExecutionsRunning = ExecutionGraphTestUtils.allExecutionsPredicate(runningOrFinished);

	return () -> {
		final AccessExecutionGraph executionGraph = executionGraphFutureSupplier.get().join();
		return allExecutionsRunning.test(executionGraph);
	};
}
 

public SingleInputGate(
	String owningTaskName,
	IntermediateDataSetID consumedResultId,
	final ResultPartitionType consumedPartitionType,
	int consumedSubpartitionIndex,
	int numberOfInputChannels,
	PartitionProducerStateProvider partitionProducerStateProvider,
	boolean isCreditBased,
	SupplierWithException<BufferPool, IOException> bufferPoolFactory) {

	this.owningTaskName = checkNotNull(owningTaskName);

	this.consumedResultId = checkNotNull(consumedResultId);
	this.consumedPartitionType = checkNotNull(consumedPartitionType);
	this.bufferPoolFactory = checkNotNull(bufferPoolFactory);

	checkArgument(consumedSubpartitionIndex >= 0);
	this.consumedSubpartitionIndex = consumedSubpartitionIndex;

	checkArgument(numberOfInputChannels > 0);
	this.numberOfInputChannels = numberOfInputChannels;

	this.inputChannels = new HashMap<>(numberOfInputChannels);
	this.channelsWithEndOfPartitionEvents = new BitSet(numberOfInputChannels);
	this.enqueuedInputChannelsWithData = new BitSet(numberOfInputChannels);

	this.partitionProducerStateProvider = checkNotNull(partitionProducerStateProvider);

	this.isCreditBased = isCreditBased;

	this.closeFuture = new CompletableFuture<>();
}
 

/**
 * Creates an input gate and all of its input channels.
 */
public SingleInputGate create(
		@Nonnull String owningTaskName,
		@Nonnull InputGateDeploymentDescriptor igdd,
		@Nonnull PartitionProducerStateProvider partitionProducerStateProvider,
		@Nonnull InputChannelMetrics metrics) {
	SupplierWithException<BufferPool, IOException> bufferPoolFactory = createBufferPoolFactory(
		networkBufferPool,
		isCreditBased,
		networkBuffersPerChannel,
		floatingNetworkBuffersPerGate,
		igdd.getShuffleDescriptors().length,
		igdd.getConsumedPartitionType());

	SingleInputGate inputGate = new SingleInputGate(
		owningTaskName,
		igdd.getConsumedResultId(),
		igdd.getConsumedPartitionType(),
		igdd.getConsumedSubpartitionIndex(),
		igdd.getShuffleDescriptors().length,
		partitionProducerStateProvider,
		isCreditBased,
		bufferPoolFactory);

	createInputChannels(owningTaskName, igdd, inputGate, metrics);
	return inputGate;
}