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

TestingTaskExecutorGateway(
		String address,
		String hostname,
		BiConsumer<ResourceID, AllocatedSlotReport> heartbeatJobManagerConsumer,
		BiConsumer<JobID, Throwable> disconnectJobManagerConsumer,
		BiFunction<TaskDeploymentDescriptor, JobMasterId, CompletableFuture<Acknowledge>> submitTaskConsumer,
		Function<Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId>, CompletableFuture<Acknowledge>> requestSlotFunction,
		BiFunction<AllocationID, Throwable, CompletableFuture<Acknowledge>> freeSlotFunction,
		Consumer<ResourceID> heartbeatResourceManagerConsumer,
		Consumer<Exception> disconnectResourceManagerConsumer,
		Function<ExecutionAttemptID, CompletableFuture<Acknowledge>> cancelTaskFunction,
		Supplier<CompletableFuture<Boolean>> canBeReleasedSupplier,
		TriConsumer<JobID, Set<ResultPartitionID>, Set<ResultPartitionID>> releaseOrPromotePartitionsConsumer,
		Consumer<Collection<IntermediateDataSetID>> releaseClusterPartitionsConsumer,
		TriFunction<ExecutionAttemptID, OperatorID, SerializedValue<OperatorEvent>, CompletableFuture<Acknowledge>> operatorEventHandler,
		Supplier<CompletableFuture<ThreadDumpInfo>> requestThreadDumpSupplier) {

	this.address = Preconditions.checkNotNull(address);
	this.hostname = Preconditions.checkNotNull(hostname);
	this.heartbeatJobManagerConsumer = Preconditions.checkNotNull(heartbeatJobManagerConsumer);
	this.disconnectJobManagerConsumer = Preconditions.checkNotNull(disconnectJobManagerConsumer);
	this.submitTaskConsumer = Preconditions.checkNotNull(submitTaskConsumer);
	this.requestSlotFunction = Preconditions.checkNotNull(requestSlotFunction);
	this.freeSlotFunction = Preconditions.checkNotNull(freeSlotFunction);
	this.heartbeatResourceManagerConsumer = heartbeatResourceManagerConsumer;
	this.disconnectResourceManagerConsumer = disconnectResourceManagerConsumer;
	this.cancelTaskFunction = cancelTaskFunction;
	this.canBeReleasedSupplier = canBeReleasedSupplier;
	this.releaseOrPromotePartitionsConsumer = releaseOrPromotePartitionsConsumer;
	this.releaseClusterPartitionsConsumer = releaseClusterPartitionsConsumer;
	this.operatorEventHandler = operatorEventHandler;
	this.requestThreadDumpSupplier = requestThreadDumpSupplier;
}
 

public TestingJobMasterGatewayBuilder setFailSlotConsumer(TriConsumer<ResourceID, AllocationID, Throwable> failSlotConsumer) {
	this.failSlotConsumer = failSlotConsumer;
	return this;
}
 

public TestingJobMasterGateway(
		@Nonnull String address,
		@Nonnull String hostname,
		@Nonnull Supplier<CompletableFuture<Acknowledge>> cancelFunction,
		@Nonnull Supplier<CompletableFuture<Acknowledge>> stopFunction,
		@Nonnull BiFunction<Integer, RescalingBehaviour, CompletableFuture<Acknowledge>> rescalingJobFunction,
		@Nonnull TriFunction<Collection<JobVertexID>, Integer, RescalingBehaviour, CompletableFuture<Acknowledge>> rescalingOperatorsFunction,
		@Nonnull Function<TaskExecutionState, CompletableFuture<Acknowledge>> updateTaskExecutionStateFunction,
		@Nonnull BiFunction<JobVertexID, ExecutionAttemptID, CompletableFuture<SerializedInputSplit>> requestNextInputSplitFunction,
		@Nonnull BiFunction<IntermediateDataSetID, ResultPartitionID, CompletableFuture<ExecutionState>> requestPartitionStateFunction,
		@Nonnull Function<ResultPartitionID, CompletableFuture<Acknowledge>> scheduleOrUpdateConsumersFunction,
		@Nonnull Function<ResourceID, CompletableFuture<Acknowledge>> disconnectTaskManagerFunction,
		@Nonnull Consumer<ResourceManagerId> disconnectResourceManagerConsumer,
		@Nonnull Supplier<CompletableFuture<ClassloadingProps>> classloadingPropsSupplier,
		@Nonnull BiFunction<ResourceID, Collection<SlotOffer>, CompletableFuture<Collection<SlotOffer>>> offerSlotsFunction,
		@Nonnull TriConsumer<ResourceID, AllocationID, Throwable> failSlotConsumer,
		@Nonnull BiFunction<String, TaskManagerLocation, CompletableFuture<RegistrationResponse>> registerTaskManagerFunction,
		@Nonnull BiConsumer<ResourceID, AccumulatorReport> taskManagerHeartbeatConsumer,
		@Nonnull Consumer<ResourceID> resourceManagerHeartbeatConsumer,
		@Nonnull Supplier<CompletableFuture<JobDetails>> requestJobDetailsSupplier,
		@Nonnull Supplier<CompletableFuture<ArchivedExecutionGraph>> requestJobSupplier,
		@Nonnull BiFunction<String, Boolean, CompletableFuture<String>> triggerSavepointFunction,
		@Nonnull Function<JobVertexID, CompletableFuture<OperatorBackPressureStatsResponse>> requestOperatorBackPressureStatsFunction,
		@Nonnull BiConsumer<AllocationID, Throwable> notifyAllocationFailureConsumer,
		@Nonnull Consumer<Tuple5<JobID, ExecutionAttemptID, Long, CheckpointMetrics, TaskStateSnapshot>> acknowledgeCheckpointConsumer,
		@Nonnull Consumer<DeclineCheckpoint> declineCheckpointConsumer,
		@Nonnull Supplier<JobMasterId> fencingTokenSupplier,
		@Nonnull BiFunction<JobID, String, CompletableFuture<KvStateLocation>> requestKvStateLocationFunction,
		@Nonnull Function<Tuple6<JobID, JobVertexID, KeyGroupRange, String, KvStateID, InetSocketAddress>, CompletableFuture<Acknowledge>> notifyKvStateRegisteredFunction,
		@Nonnull Function<Tuple4<JobID, JobVertexID, KeyGroupRange, String>, CompletableFuture<Acknowledge>> notifyKvStateUnregisteredFunction,
		@Nonnull TriFunction<String, Object, byte[], CompletableFuture<Object>> updateAggregateFunction) {
	this.address = address;
	this.hostname = hostname;
	this.cancelFunction = cancelFunction;
	this.stopFunction = stopFunction;
	this.rescalingJobFunction = rescalingJobFunction;
	this.rescalingOperatorsFunction = rescalingOperatorsFunction;
	this.updateTaskExecutionStateFunction = updateTaskExecutionStateFunction;
	this.requestNextInputSplitFunction = requestNextInputSplitFunction;
	this.requestPartitionStateFunction = requestPartitionStateFunction;
	this.scheduleOrUpdateConsumersFunction = scheduleOrUpdateConsumersFunction;
	this.disconnectTaskManagerFunction = disconnectTaskManagerFunction;
	this.disconnectResourceManagerConsumer = disconnectResourceManagerConsumer;
	this.classloadingPropsSupplier = classloadingPropsSupplier;
	this.offerSlotsFunction = offerSlotsFunction;
	this.failSlotConsumer = failSlotConsumer;
	this.registerTaskManagerFunction = registerTaskManagerFunction;
	this.taskManagerHeartbeatConsumer = taskManagerHeartbeatConsumer;
	this.resourceManagerHeartbeatConsumer = resourceManagerHeartbeatConsumer;
	this.requestJobDetailsSupplier = requestJobDetailsSupplier;
	this.requestJobSupplier = requestJobSupplier;
	this.triggerSavepointFunction = triggerSavepointFunction;
	this.requestOperatorBackPressureStatsFunction = requestOperatorBackPressureStatsFunction;
	this.notifyAllocationFailureConsumer = notifyAllocationFailureConsumer;
	this.acknowledgeCheckpointConsumer = acknowledgeCheckpointConsumer;
	this.declineCheckpointConsumer = declineCheckpointConsumer;
	this.fencingTokenSupplier = fencingTokenSupplier;
	this.requestKvStateLocationFunction = requestKvStateLocationFunction;
	this.notifyKvStateRegisteredFunction = notifyKvStateRegisteredFunction;
	this.notifyKvStateUnregisteredFunction = notifyKvStateUnregisteredFunction;
	this.updateAggregateFunction = updateAggregateFunction;
}
 

public TestingJobMasterGatewayBuilder setFailSlotConsumer(TriConsumer<ResourceID, AllocationID, Throwable> failSlotConsumer) {
	this.failSlotConsumer = failSlotConsumer;
	return this;
}
 

public TestingJobMasterGateway(
		@Nonnull String address,
		@Nonnull String hostname,
		@Nonnull Supplier<CompletableFuture<Acknowledge>> cancelFunction,
		@Nonnull Function<TaskExecutionState, CompletableFuture<Acknowledge>> updateTaskExecutionStateFunction,
		@Nonnull BiFunction<JobVertexID, ExecutionAttemptID, CompletableFuture<SerializedInputSplit>> requestNextInputSplitFunction,
		@Nonnull BiFunction<IntermediateDataSetID, ResultPartitionID, CompletableFuture<ExecutionState>> requestPartitionStateFunction,
		@Nonnull Function<ResultPartitionID, CompletableFuture<Acknowledge>> scheduleOrUpdateConsumersFunction,
		@Nonnull Function<ResourceID, CompletableFuture<Acknowledge>> disconnectTaskManagerFunction,
		@Nonnull Consumer<ResourceManagerId> disconnectResourceManagerConsumer,
		@Nonnull Supplier<CompletableFuture<ClassloadingProps>> classloadingPropsSupplier,
		@Nonnull BiFunction<ResourceID, Collection<SlotOffer>, CompletableFuture<Collection<SlotOffer>>> offerSlotsFunction,
		@Nonnull TriConsumer<ResourceID, AllocationID, Throwable> failSlotConsumer,
		@Nonnull BiFunction<String, TaskManagerLocation, CompletableFuture<RegistrationResponse>> registerTaskManagerFunction,
		@Nonnull BiConsumer<ResourceID, AccumulatorReport> taskManagerHeartbeatConsumer,
		@Nonnull Consumer<ResourceID> resourceManagerHeartbeatConsumer,
		@Nonnull Supplier<CompletableFuture<JobDetails>> requestJobDetailsSupplier,
		@Nonnull Supplier<CompletableFuture<ArchivedExecutionGraph>> requestJobSupplier,
		@Nonnull BiFunction<String, Boolean, CompletableFuture<String>> triggerSavepointFunction,
		@Nonnull BiFunction<String, Boolean, CompletableFuture<String>> stopWithSavepointFunction,
		@Nonnull Function<JobVertexID, CompletableFuture<OperatorBackPressureStatsResponse>> requestOperatorBackPressureStatsFunction,
		@Nonnull BiConsumer<AllocationID, Throwable> notifyAllocationFailureConsumer,
		@Nonnull Consumer<Tuple5<JobID, ExecutionAttemptID, Long, CheckpointMetrics, TaskStateSnapshot>> acknowledgeCheckpointConsumer,
		@Nonnull Consumer<DeclineCheckpoint> declineCheckpointConsumer,
		@Nonnull Supplier<JobMasterId> fencingTokenSupplier,
		@Nonnull BiFunction<JobID, String, CompletableFuture<KvStateLocation>> requestKvStateLocationFunction,
		@Nonnull Function<Tuple6<JobID, JobVertexID, KeyGroupRange, String, KvStateID, InetSocketAddress>, CompletableFuture<Acknowledge>> notifyKvStateRegisteredFunction,
		@Nonnull Function<Tuple4<JobID, JobVertexID, KeyGroupRange, String>, CompletableFuture<Acknowledge>> notifyKvStateUnregisteredFunction,
		@Nonnull TriFunction<String, Object, byte[], CompletableFuture<Object>> updateAggregateFunction) {
	this.address = address;
	this.hostname = hostname;
	this.cancelFunction = cancelFunction;
	this.updateTaskExecutionStateFunction = updateTaskExecutionStateFunction;
	this.requestNextInputSplitFunction = requestNextInputSplitFunction;
	this.requestPartitionStateFunction = requestPartitionStateFunction;
	this.scheduleOrUpdateConsumersFunction = scheduleOrUpdateConsumersFunction;
	this.disconnectTaskManagerFunction = disconnectTaskManagerFunction;
	this.disconnectResourceManagerConsumer = disconnectResourceManagerConsumer;
	this.classloadingPropsSupplier = classloadingPropsSupplier;
	this.offerSlotsFunction = offerSlotsFunction;
	this.failSlotConsumer = failSlotConsumer;
	this.registerTaskManagerFunction = registerTaskManagerFunction;
	this.taskManagerHeartbeatConsumer = taskManagerHeartbeatConsumer;
	this.resourceManagerHeartbeatConsumer = resourceManagerHeartbeatConsumer;
	this.requestJobDetailsSupplier = requestJobDetailsSupplier;
	this.requestJobSupplier = requestJobSupplier;
	this.triggerSavepointFunction = triggerSavepointFunction;
	this.stopWithSavepointFunction = stopWithSavepointFunction;
	this.requestOperatorBackPressureStatsFunction = requestOperatorBackPressureStatsFunction;
	this.notifyAllocationFailureConsumer = notifyAllocationFailureConsumer;
	this.acknowledgeCheckpointConsumer = acknowledgeCheckpointConsumer;
	this.declineCheckpointConsumer = declineCheckpointConsumer;
	this.fencingTokenSupplier = fencingTokenSupplier;
	this.requestKvStateLocationFunction = requestKvStateLocationFunction;
	this.notifyKvStateRegisteredFunction = notifyKvStateRegisteredFunction;
	this.notifyKvStateUnregisteredFunction = notifyKvStateUnregisteredFunction;
	this.updateAggregateFunction = updateAggregateFunction;
}
 

void setStartContainerAsyncConsumer(TriConsumer<Container, ContainerLaunchContext, CallbackHandler> startContainerAsyncConsumer) {
	this.startContainerAsyncConsumer = Preconditions.checkNotNull(startContainerAsyncConsumer);
}
 

void setStopContainerAsyncConsumer(TriConsumer<ContainerId, NodeId, CallbackHandler> stopContainerAsyncConsumer) {
	this.stopContainerAsyncConsumer = Preconditions.checkNotNull(stopContainerAsyncConsumer);
}
 

void setUnregisterApplicationMasterConsumer(
	TriConsumer<FinalApplicationStatus, String, String> unregisterApplicationMasterConsumer) {
	this.unregisterApplicationMasterConsumer = unregisterApplicationMasterConsumer;
}
 

public TestingJobMasterGatewayBuilder setFailSlotConsumer(TriConsumer<ResourceID, AllocationID, Throwable> failSlotConsumer) {
	this.failSlotConsumer = failSlotConsumer;
	return this;
}
 

public TestingJobMasterGateway(
		@Nonnull String address,
		@Nonnull String hostname,
		@Nonnull Supplier<CompletableFuture<Acknowledge>> cancelFunction,
		@Nonnull Function<TaskExecutionState, CompletableFuture<Acknowledge>> updateTaskExecutionStateFunction,
		@Nonnull BiFunction<JobVertexID, ExecutionAttemptID, CompletableFuture<SerializedInputSplit>> requestNextInputSplitFunction,
		@Nonnull BiFunction<IntermediateDataSetID, ResultPartitionID, CompletableFuture<ExecutionState>> requestPartitionStateFunction,
		@Nonnull Function<ResultPartitionID, CompletableFuture<Acknowledge>> scheduleOrUpdateConsumersFunction,
		@Nonnull Function<ResourceID, CompletableFuture<Acknowledge>> disconnectTaskManagerFunction,
		@Nonnull Consumer<ResourceManagerId> disconnectResourceManagerConsumer,
		@Nonnull BiFunction<ResourceID, Collection<SlotOffer>, CompletableFuture<Collection<SlotOffer>>> offerSlotsFunction,
		@Nonnull TriConsumer<ResourceID, AllocationID, Throwable> failSlotConsumer,
		@Nonnull BiFunction<String, UnresolvedTaskManagerLocation, CompletableFuture<RegistrationResponse>> registerTaskManagerFunction,
		@Nonnull BiConsumer<ResourceID, AccumulatorReport> taskManagerHeartbeatConsumer,
		@Nonnull Consumer<ResourceID> resourceManagerHeartbeatConsumer,
		@Nonnull Supplier<CompletableFuture<JobDetails>> requestJobDetailsSupplier,
		@Nonnull Supplier<CompletableFuture<ArchivedExecutionGraph>> requestJobSupplier,
		@Nonnull BiFunction<String, Boolean, CompletableFuture<String>> triggerSavepointFunction,
		@Nonnull BiFunction<String, Boolean, CompletableFuture<String>> stopWithSavepointFunction,
		@Nonnull Function<JobVertexID, CompletableFuture<OperatorBackPressureStatsResponse>> requestOperatorBackPressureStatsFunction,
		@Nonnull BiConsumer<AllocationID, Throwable> notifyAllocationFailureConsumer,
		@Nonnull Consumer<Tuple5<JobID, ExecutionAttemptID, Long, CheckpointMetrics, TaskStateSnapshot>> acknowledgeCheckpointConsumer,
		@Nonnull Consumer<DeclineCheckpoint> declineCheckpointConsumer,
		@Nonnull Supplier<JobMasterId> fencingTokenSupplier,
		@Nonnull BiFunction<JobID, String, CompletableFuture<KvStateLocation>> requestKvStateLocationFunction,
		@Nonnull Function<Tuple6<JobID, JobVertexID, KeyGroupRange, String, KvStateID, InetSocketAddress>, CompletableFuture<Acknowledge>> notifyKvStateRegisteredFunction,
		@Nonnull Function<Tuple4<JobID, JobVertexID, KeyGroupRange, String>, CompletableFuture<Acknowledge>> notifyKvStateUnregisteredFunction,
		@Nonnull TriFunction<String, Object, byte[], CompletableFuture<Object>> updateAggregateFunction,
		@Nonnull TriFunction<ExecutionAttemptID, OperatorID, SerializedValue<OperatorEvent>, CompletableFuture<Acknowledge>> operatorEventSender,
		@Nonnull BiFunction<OperatorID, SerializedValue<CoordinationRequest>, CompletableFuture<CoordinationResponse>> deliverCoordinationRequestFunction) {
	this.address = address;
	this.hostname = hostname;
	this.cancelFunction = cancelFunction;
	this.updateTaskExecutionStateFunction = updateTaskExecutionStateFunction;
	this.requestNextInputSplitFunction = requestNextInputSplitFunction;
	this.requestPartitionStateFunction = requestPartitionStateFunction;
	this.scheduleOrUpdateConsumersFunction = scheduleOrUpdateConsumersFunction;
	this.disconnectTaskManagerFunction = disconnectTaskManagerFunction;
	this.disconnectResourceManagerConsumer = disconnectResourceManagerConsumer;
	this.offerSlotsFunction = offerSlotsFunction;
	this.failSlotConsumer = failSlotConsumer;
	this.registerTaskManagerFunction = registerTaskManagerFunction;
	this.taskManagerHeartbeatConsumer = taskManagerHeartbeatConsumer;
	this.resourceManagerHeartbeatConsumer = resourceManagerHeartbeatConsumer;
	this.requestJobDetailsSupplier = requestJobDetailsSupplier;
	this.requestJobSupplier = requestJobSupplier;
	this.triggerSavepointFunction = triggerSavepointFunction;
	this.stopWithSavepointFunction = stopWithSavepointFunction;
	this.requestOperatorBackPressureStatsFunction = requestOperatorBackPressureStatsFunction;
	this.notifyAllocationFailureConsumer = notifyAllocationFailureConsumer;
	this.acknowledgeCheckpointConsumer = acknowledgeCheckpointConsumer;
	this.declineCheckpointConsumer = declineCheckpointConsumer;
	this.fencingTokenSupplier = fencingTokenSupplier;
	this.requestKvStateLocationFunction = requestKvStateLocationFunction;
	this.notifyKvStateRegisteredFunction = notifyKvStateRegisteredFunction;
	this.notifyKvStateUnregisteredFunction = notifyKvStateUnregisteredFunction;
	this.updateAggregateFunction = updateAggregateFunction;
	this.operatorEventSender = operatorEventSender;
	this.deliverCoordinationRequestFunction = deliverCoordinationRequestFunction;
}
 

private void testJobMasterConnectionTerminationAfterExternalReleaseOrPromotion(TriConsumer<TaskExecutorGateway, JobID, ResultPartitionID> releaseOrPromoteCall) throws Exception {
	final CompletableFuture<Void> disconnectFuture = new CompletableFuture<>();
	final JobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
		.setDisconnectTaskManagerFunction(resourceID -> {
			disconnectFuture.complete(null);
			return CompletableFuture.completedFuture(Acknowledge.get());
		}).build();

	final DefaultJobTable jobTable = DefaultJobTable.create();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setJobTable(jobTable)
		.setShuffleEnvironment(new NettyShuffleEnvironmentBuilder().build())
		.setTaskSlotTable(createTaskSlotTable())
		.build();

	final TestingTaskExecutorPartitionTracker partitionTracker = new TestingTaskExecutorPartitionTracker();

	final AtomicBoolean trackerIsTrackingPartitions = new AtomicBoolean(false);
	partitionTracker.setIsTrackingPartitionsForFunction(jobId -> trackerIsTrackingPartitions.get());

	final CompletableFuture<Collection<ResultPartitionID>> firstReleasePartitionsCallFuture = new CompletableFuture<>();
	partitionTracker.setStopTrackingAndReleasePartitionsConsumer(firstReleasePartitionsCallFuture::complete);

	final ResultPartitionDeploymentDescriptor resultPartitionDeploymentDescriptor = PartitionTestUtils.createPartitionDeploymentDescriptor(ResultPartitionType.BLOCKING);
	final ResultPartitionID resultPartitionId = resultPartitionDeploymentDescriptor.getShuffleDescriptor().getResultPartitionID();

	final TestingTaskExecutor taskExecutor = createTestingTaskExecutor(taskManagerServices, partitionTracker);

	try {
		taskExecutor.start();
		taskExecutor.waitUntilStarted();

		TaskSubmissionTestEnvironment.registerJobMasterConnection(
			jobTable,
			jobId,
			rpc,
			jobMasterGateway,
			new NoOpTaskManagerActions(),
			timeout,
			taskExecutor.getMainThreadExecutableForTesting());

		final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);

		trackerIsTrackingPartitions.set(true);
		assertThat(firstReleasePartitionsCallFuture.isDone(), is(false));

		taskExecutorGateway.releaseOrPromotePartitions(jobId, Collections.singleton(new ResultPartitionID()), Collections.emptySet());

		// at this point we only know that the TE has entered releasePartitions; we cannot be certain whether it
		// has already checked whether it should disconnect or not
		firstReleasePartitionsCallFuture.get();

		// connection should be kept alive since the table still contains partitions
		assertThat(disconnectFuture.isDone(), is(false));

		trackerIsTrackingPartitions.set(false);

		// the TM should check whether partitions are still stored, and afterwards terminate the connection
		releaseOrPromoteCall.accept(taskExecutorGateway, jobId, resultPartitionId);

		disconnectFuture.get();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

public TestingTaskExecutorGatewayBuilder setReleaseOrPromotePartitionsConsumer(TriConsumer<JobID, Set<ResultPartitionID>, Set<ResultPartitionID>> releasePartitionsConsumer) {
	this.releaseOrPromotePartitionsConsumer = releasePartitionsConsumer;
	return this;
}