Java Code Examples for org.apache.flink.core.testutils.OneShotLatch#trigger()

@Override
public void run(SourceContext<T> ctx) throws Exception {
	OneShotLatch latch = guards.get(guardId);
	try {
		source.run(ctx);
	} finally {
		latch.trigger();
	}

	while (running) {
		try {
			Thread.sleep(100);
		} catch (InterruptedException e) {
			// ignore
		}
	}
}
 

/**
 * Tests {@link NetworkBufferPool#requestMemorySegments(int)}, verifying it may be aborted and
 * remains in a defined state even if the waiting is interrupted.
 */
@Test
public void testRequestMemorySegmentsInterruptable2() throws Exception {
	final int numBuffers = 10;

	NetworkBufferPool globalPool = new NetworkBufferPool(numBuffers, 128);
	MemorySegment segment = globalPool.requestMemorySegment();
	assertNotNull(segment);

	final OneShotLatch isRunning = new OneShotLatch();
	CheckedThread asyncRequest = new CheckedThread() {
		@Override
		public void go() throws Exception {
			isRunning.trigger();
			globalPool.requestMemorySegments(10);
		}
	};
	asyncRequest.start();

	// We want the destroy call inside the blocking part of the globalPool.requestMemorySegments()
	// call above. We cannot guarantee this though but make it highly probable:
	isRunning.await();
	Thread.sleep(10);
	asyncRequest.interrupt();

	globalPool.recycle(segment);

	try {
		asyncRequest.sync();
	} catch (IOException e) {
		assertThat(e, hasProperty("cause", instanceOf(InterruptedException.class)));

		// test indirectly for NetworkBufferPool#numTotalRequiredBuffers being correct:
		// -> creating a new buffer pool should not fail
		globalPool.createBufferPool(10, 10);
	} finally {
		globalPool.destroy();
	}
}
 

/**
 * Tests that the JobLeaderService won't try to reconnect to JobMaster after it
 * has lost the leadership. See FLINK-16836.
 */
@Test
public void doesNotReconnectAfterTargetLostLeadership() throws Exception {
	final JobID jobId = new JobID();

	final SettableLeaderRetrievalService leaderRetrievalService = new SettableLeaderRetrievalService();
	final TestingHighAvailabilityServices haServices = new TestingHighAvailabilityServicesBuilder()
		.setJobMasterLeaderRetrieverFunction(ignored -> leaderRetrievalService)
		.build();
	final TestingJobMasterGateway jobMasterGateway = registerJobMaster();

	final OneShotLatch jobManagerGainedLeadership = new OneShotLatch();
	final TestingJobLeaderListener testingJobLeaderListener = new TestingJobLeaderListener(ignored -> jobManagerGainedLeadership.trigger());

	final JobLeaderService jobLeaderService = createAndStartJobLeaderService(haServices, testingJobLeaderListener);

	try {
		jobLeaderService.addJob(jobId, jobMasterGateway.getAddress());

		leaderRetrievalService.notifyListener(jobMasterGateway.getAddress(), UUID.randomUUID());

		jobManagerGainedLeadership.await();

		// revoke the leadership
		leaderRetrievalService.notifyListener(null, null);
		testingJobLeaderListener.waitUntilJobManagerLostLeadership();

		jobLeaderService.reconnect(jobId);
	} finally {
		jobLeaderService.stop();
	}
}
 

/**
 * Tests {@link NetworkBufferPool#requestMemorySegments()}, verifying it may be aborted and
 * remains in a defined state even if the waiting is interrupted.
 */
@Test
public void testRequestMemorySegmentsInterruptable2() throws Exception {
	final int numBuffers = 10;

	NetworkBufferPool globalPool = new NetworkBufferPool(numBuffers, 128, 10);
	MemorySegment segment = globalPool.requestMemorySegment();
	assertNotNull(segment);

	final OneShotLatch isRunning = new OneShotLatch();
	CheckedThread asyncRequest = new CheckedThread() {
		@Override
		public void go() throws Exception {
			isRunning.trigger();
			globalPool.requestMemorySegments();
		}
	};
	asyncRequest.start();

	// We want the destroy call inside the blocking part of the globalPool.requestMemorySegments()
	// call above. We cannot guarantee this though but make it highly probable:
	isRunning.await();
	Thread.sleep(10);
	asyncRequest.interrupt();

	globalPool.recycle(segment);

	try {
		asyncRequest.sync();
	} catch (IOException e) {
		assertThat(e, hasProperty("cause", instanceOf(InterruptedException.class)));

		// test indirectly for NetworkBufferPool#numTotalRequiredBuffers being correct:
		// -> creating a new buffer pool should not fail
		globalPool.createBufferPool(10, 10);
	} finally {
		globalPool.destroy();
	}
}
 

@Test
public void testConcurrentPartitionsDiscoveryAndLoopFetching() throws Exception {
    String tp = topicName("test", 2);
    TestSourceContext<Long> sourceContext = new TestSourceContext<Long>();
    Map<String, MessageId> offset = Collections.singletonMap(topicName(tp, 1), MessageId.latest);

    OneShotLatch fetchLoopWaitLatch = new OneShotLatch();
    OneShotLatch stateIterationBlockLatch = new OneShotLatch();

    TestFetcher fetcher = new TestFetcher(
            sourceContext,
            offset,
            null,
            null,
            new TestProcessingTimeService(),
            10,
            fetchLoopWaitLatch,
            stateIterationBlockLatch);

    // ----- run the fetcher -----

    final CheckedThread checkedThread = new CheckedThread() {
        @Override
        public void go() throws Exception {
            fetcher.runFetchLoop();
        }
    };
    checkedThread.start();

    fetchLoopWaitLatch.await();
    fetcher.addDiscoveredTopics(Sets.newSet(tp));

    stateIterationBlockLatch.trigger();
    checkedThread.sync();
}
 

/**
 * Test delayed registration of task executor where the delay is introduced during connection from resource manager
 * to the registering task executor.
 */
@Test
public void testDelayedRegisterTaskExecutor() throws Exception {
	final Time fastTimeout = Time.milliseconds(1L);
	try {
		final OneShotLatch startConnection = new OneShotLatch();
		final OneShotLatch finishConnection = new OneShotLatch();

		// first registration is with blocking connection
		rpcService.setRpcGatewayFutureFunction(rpcGateway ->
			CompletableFuture.supplyAsync(
				() -> {
					startConnection.trigger();
					try {
						finishConnection.await();
					} catch (InterruptedException ignored) {}
					return rpcGateway;
				},
				TestingUtils.defaultExecutor()));

		CompletableFuture<RegistrationResponse> firstFuture =
			rmGateway.registerTaskExecutor(taskExecutorGateway.getAddress(), taskExecutorResourceID, dataPort, hardwareDescription, fastTimeout);
		try {
			firstFuture.get();
			fail("Should have failed because connection to taskmanager is delayed beyond timeout");
		} catch (Exception e) {
			assertThat(ExceptionUtils.stripExecutionException(e), instanceOf(AskTimeoutException.class));
		}

		startConnection.await();

		// second registration after timeout is with no delay, expecting it to be succeeded
		rpcService.resetRpcGatewayFutureFunction();
		CompletableFuture<RegistrationResponse> secondFuture =
			rmGateway.registerTaskExecutor(taskExecutorGateway.getAddress(), taskExecutorResourceID, dataPort, hardwareDescription, TIMEOUT);
		RegistrationResponse response = secondFuture.get();
		assertTrue(response instanceof TaskExecutorRegistrationSuccess);

		// on success, send slot report for taskmanager registration
		final SlotReport slotReport = new SlotReport(new SlotStatus(new SlotID(taskExecutorResourceID, 0), ResourceProfile.UNKNOWN));
		rmGateway.sendSlotReport(taskExecutorResourceID,
			((TaskExecutorRegistrationSuccess) response).getRegistrationId(), slotReport, TIMEOUT).get();

		// let the remaining part of the first registration proceed
		finishConnection.trigger();
		Thread.sleep(1L);

		// verify that the latest registration is valid not being unregistered by the delayed one
		final TaskManagerInfo taskManagerInfo = rmGateway.requestTaskManagerInfo(
			taskExecutorResourceID,
			TIMEOUT).get();
		assertThat(taskManagerInfo.getResourceId(), equalTo(taskExecutorResourceID));
		assertThat(taskManagerInfo.getNumberSlots(), equalTo(1));
	} finally {
		rpcService.resetRpcGatewayFutureFunction();
	}
}
 

@Test
public void testSnapshotAsyncCancel() throws Exception {
	DefaultOperatorStateBackend operatorStateBackend =
		new DefaultOperatorStateBackendBuilder(
			OperatorStateBackendTest.class.getClassLoader(),
			new ExecutionConfig(),
			true,
			emptyStateHandles,
			new CloseableRegistry()).build();

	ListStateDescriptor<MutableType> stateDescriptor1 =
			new ListStateDescriptor<>("test1", new JavaSerializer<MutableType>());

	ListState<MutableType> listState1 = operatorStateBackend.getListState(stateDescriptor1);

	listState1.add(MutableType.of(42));
	listState1.add(MutableType.of(4711));

	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);

	OneShotLatch waiterLatch = new OneShotLatch();
	OneShotLatch blockerLatch = new OneShotLatch();

	streamFactory.setWaiterLatch(waiterLatch);
	streamFactory.setBlockerLatch(blockerLatch);

	RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture =
			operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

	ExecutorService executorService = Executors.newFixedThreadPool(1);

	executorService.submit(runnableFuture);

	// wait until the async checkpoint is in the stream's write code, then continue
	waiterLatch.await();

	// cancel the future, which should close the underlying stream
	runnableFuture.cancel(true);

	for (BlockingCheckpointOutputStream stream : streamFactory.getAllCreatedStreams()) {
		Assert.assertTrue(stream.isClosed());
	}

	// we allow the stream under test to proceed
	blockerLatch.trigger();

	try {
		runnableFuture.get(60, TimeUnit.SECONDS);
		Assert.fail();
	} catch (CancellationException ignore) {
	}
}
 

@Test
public void testAsyncSnapshotCancellation() throws Exception {
	OneShotLatch blocker = new OneShotLatch();
	OneShotLatch waiter = new OneShotLatch();
	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);
	streamFactory.setWaiterLatch(waiter);
	streamFactory.setBlockerLatch(blocker);
	streamFactory.setAfterNumberInvocations(10);

	final AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	try {

		if (!backend.supportsAsynchronousSnapshots()) {
			return;
		}

		InternalValueState<Integer, VoidNamespace, Integer> valueState = backend.createInternalState(
				VoidNamespaceSerializer.INSTANCE,
				new ValueStateDescriptor<>("test", IntSerializer.INSTANCE));

		valueState.setCurrentNamespace(VoidNamespace.INSTANCE);

		for (int i = 0; i < 10; ++i) {
			backend.setCurrentKey(i);
			valueState.update(i);
		}

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot =
				backend.snapshot(0L, 0L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

		Thread runner = new Thread(snapshot);
		runner.start();

		// wait until the code reached some stream read
		waiter.await();

		// close the backend to see if the close is propagated to the stream
		IOUtils.closeQuietly(backend);

		//unblock the stream so that it can run into the IOException
		blocker.trigger();

		runner.join();

		try {
			snapshot.get();
			fail("Close was not propagated.");
		} catch (CancellationException ex) {
			//ignore
		}

	} finally {
		backend.dispose();
	}
}
 

@Test
public void testConcurrentPartitionsDiscoveryAndLoopFetching() throws Exception {
	// test data
	final KafkaTopicPartition testPartition = new KafkaTopicPartition("test", 42);

	// ----- create the test fetcher -----

	@SuppressWarnings("unchecked")
	SourceContext<String> sourceContext = new TestSourceContext<>();
	Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets =
		Collections.singletonMap(testPartition, KafkaTopicPartitionStateSentinel.GROUP_OFFSET);

	final OneShotLatch fetchLoopWaitLatch = new OneShotLatch();
	final OneShotLatch stateIterationBlockLatch = new OneShotLatch();

	final TestFetcher<String> fetcher = new TestFetcher<>(
		sourceContext,
		partitionsWithInitialOffsets,
		null, /* periodic assigner */
		null, /* punctuated assigner */
		new TestProcessingTimeService(),
		10,
		fetchLoopWaitLatch,
		stateIterationBlockLatch);

	// ----- run the fetcher -----

	final CheckedThread checkedThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			fetcher.runFetchLoop();
		}
	};
	checkedThread.start();

	// wait until state iteration begins before adding discovered partitions
	fetchLoopWaitLatch.await();
	fetcher.addDiscoveredPartitions(Collections.singletonList(testPartition));

	stateIterationBlockLatch.trigger();
	checkedThread.sync();
}
 

@Test
public void testSnapshotAsyncCancel() throws Exception {
	DefaultOperatorStateBackend operatorStateBackend =
		new DefaultOperatorStateBackendBuilder(
			OperatorStateBackendTest.class.getClassLoader(),
			new ExecutionConfig(),
			true,
			emptyStateHandles,
			new CloseableRegistry()).build();

	ListStateDescriptor<MutableType> stateDescriptor1 =
			new ListStateDescriptor<>("test1", new JavaSerializer<MutableType>());

	ListState<MutableType> listState1 = operatorStateBackend.getOperatorState(stateDescriptor1);

	listState1.add(MutableType.of(42));
	listState1.add(MutableType.of(4711));

	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);

	OneShotLatch waiterLatch = new OneShotLatch();
	OneShotLatch blockerLatch = new OneShotLatch();

	streamFactory.setWaiterLatch(waiterLatch);
	streamFactory.setBlockerLatch(blockerLatch);

	RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture =
			operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

	ExecutorService executorService = Executors.newFixedThreadPool(1);

	executorService.submit(runnableFuture);

	// wait until the async checkpoint is in the stream's write code, then continue
	waiterLatch.await();

	// cancel the future, which should close the underlying stream
	runnableFuture.cancel(true);

	for (BlockingCheckpointOutputStream stream : streamFactory.getAllCreatedStreams()) {
		Assert.assertTrue(stream.isClosed());
	}

	// we allow the stream under test to proceed
	blockerLatch.trigger();

	try {
		runnableFuture.get(60, TimeUnit.SECONDS);
		Assert.fail();
	} catch (CancellationException ignore) {
	}
}
 

@Test
public void testReaderScalingUp() throws Exception {
	// simulates the scenario of scaling up from 1 to 2 instances

	final OneShotLatch waitingLatch1 = new OneShotLatch();
	final OneShotLatch triggerLatch1 = new OneShotLatch();

	BlockingFileInputFormat format1 = new BlockingFileInputFormat(
		triggerLatch1, waitingLatch1, new Path("test"), 20, 5);
	FileInputSplit[] splits = format1.createInputSplits(2);

	OneInputStreamOperatorTestHarness<TimestampedFileInputSplit, String> testHarness1 = getTestHarness(format1, 1, 0);
	testHarness1.open();

	testHarness1.processElement(new StreamRecord<>(getTimestampedSplit(0, splits[0])));
	testHarness1.processElement(new StreamRecord<>(getTimestampedSplit(1, splits[1])));

	// wait until its arrives to element 5
	if (!triggerLatch1.isTriggered()) {
		triggerLatch1.await();
	}

	OperatorSubtaskState snapshot = testHarness1.snapshot(0, 0);

	// this will be the init state for new instance-0
	OperatorSubtaskState initState1 =
		AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot, maxParallelism, 1, 2, 0);

	// this will be the init state for new instance-1
	OperatorSubtaskState initState2 =
		AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot, maxParallelism, 1, 2, 1);

	// 1) clear the output of instance so that we can compare it with one created by the new instances, and
	// 2) let the operator process the rest of its state
	testHarness1.getOutput().clear();
	waitingLatch1.trigger();

	// create the second instance and let it process the second split till element 15
	final OneShotLatch triggerLatch2 = new OneShotLatch();
	final OneShotLatch waitingLatch2 = new OneShotLatch();

	BlockingFileInputFormat format2 = new BlockingFileInputFormat(
		triggerLatch2, waitingLatch2, new Path("test"), 20, 15);

	OneInputStreamOperatorTestHarness<TimestampedFileInputSplit, String> testHarness2 = getTestHarness(format2, 2, 0);
	testHarness2.setup();
	testHarness2.initializeState(initState1);
	testHarness2.open();

	BlockingFileInputFormat format3 = new BlockingFileInputFormat(
		triggerLatch2, waitingLatch2, new Path("test"), 20, 15);

	OneInputStreamOperatorTestHarness<TimestampedFileInputSplit, String> testHarness3 = getTestHarness(format3, 2, 1);
	testHarness3.setup();
	testHarness3.initializeState(initState2);
	testHarness3.open();

	triggerLatch2.trigger();
	waitingLatch2.trigger();

	// and wait for the processing to finish
	synchronized (testHarness1.getCheckpointLock()) {
		testHarness1.close();
	}
	synchronized (testHarness2.getCheckpointLock()) {
		testHarness2.close();
	}
	synchronized (testHarness3.getCheckpointLock()) {
		testHarness3.close();
	}

	Queue<Object> expectedResult = new ArrayDeque<>();
	putElementsInQ(expectedResult, testHarness1.getOutput());

	Queue<Object> actualResult = new ArrayDeque<>();
	putElementsInQ(actualResult, testHarness2.getOutput());
	putElementsInQ(actualResult, testHarness3.getOutput());

	Assert.assertEquals(35, actualResult.size());
	Assert.assertArrayEquals(expectedResult.toArray(), actualResult.toArray());
}
 

@Test
public void testReaderSnapshotRestore() throws Exception {
	String testBasePath = hdfsURI + "/" + UUID.randomUUID() + "/";

	TimestampedFileInputSplit split1 =
		new TimestampedFileInputSplit(0, 3, new Path("test/test1"), 0, 100, null);

	TimestampedFileInputSplit split2 =
		new TimestampedFileInputSplit(10, 2, new Path("test/test2"), 101, 200, null);

	TimestampedFileInputSplit split3 =
		new TimestampedFileInputSplit(10, 1, new Path("test/test2"), 0, 100, null);

	TimestampedFileInputSplit split4 =
		new TimestampedFileInputSplit(11, 0, new Path("test/test3"), 0, 100, null);

	final OneShotLatch latch = new OneShotLatch();

	BlockingFileInputFormat format = new BlockingFileInputFormat(latch, new Path(testBasePath));
	TypeInformation<FileInputSplit> typeInfo = TypeExtractor.getInputFormatTypes(format);

	OneInputStreamOperatorTestHarness<TimestampedFileInputSplit, FileInputSplit> initTestInstance = createHarness(format);
	initTestInstance.setTimeCharacteristic(TimeCharacteristic.EventTime);
	initTestInstance.open();

	// create some state in the reader
	initTestInstance.processElement(new StreamRecord<>(split1));
	initTestInstance.processElement(new StreamRecord<>(split2));
	initTestInstance.processElement(new StreamRecord<>(split3));
	initTestInstance.processElement(new StreamRecord<>(split4));

	// take a snapshot of the operator's state. This will be used
	// to initialize another reader and compare the results of the
	// two operators.

	final OperatorSubtaskState snapshot;
	synchronized (initTestInstance.getCheckpointLock()) {
		snapshot = initTestInstance.snapshot(0L, 0L);
	}

	OneInputStreamOperatorTestHarness<TimestampedFileInputSplit, FileInputSplit> restoredTestInstance  =
		createHarness(new BlockingFileInputFormat(latch, new Path(testBasePath)));
	restoredTestInstance.setTimeCharacteristic(TimeCharacteristic.EventTime);

	restoredTestInstance.initializeState(snapshot);
	restoredTestInstance.open();

	// now let computation start
	latch.trigger();

	// ... and wait for the operators to close gracefully

	synchronized (initTestInstance.getCheckpointLock()) {
		initTestInstance.close();
	}

	synchronized (restoredTestInstance.getCheckpointLock()) {
		restoredTestInstance.close();
	}

	FileInputSplit fsSplit1 = createSplitFromTimestampedSplit(split1);
	FileInputSplit fsSplit2 = createSplitFromTimestampedSplit(split2);
	FileInputSplit fsSplit3 = createSplitFromTimestampedSplit(split3);
	FileInputSplit fsSplit4 = createSplitFromTimestampedSplit(split4);

	// compare if the results contain what they should contain and also if
	// they are the same, as they should.

	Assert.assertTrue(initTestInstance.getOutput().contains(new StreamRecord<>(fsSplit1)));
	Assert.assertTrue(initTestInstance.getOutput().contains(new StreamRecord<>(fsSplit2)));
	Assert.assertTrue(initTestInstance.getOutput().contains(new StreamRecord<>(fsSplit3)));
	Assert.assertTrue(initTestInstance.getOutput().contains(new StreamRecord<>(fsSplit4)));

	Assert.assertArrayEquals(
		initTestInstance.getOutput().toArray(),
		restoredTestInstance.getOutput().toArray()
	);
}
 

/**
 * Tests that the {@link SlotReport} sent to the RM does not contain
 * out dated/stale information as slots are being requested from the
 * TM.
 *
 * <p>This is a probabilistic test case and needs to be executed
 * several times to produce a failure without the fix for FLINK-12865.
 */
@Test
public void testSlotReportDoesNotContainStaleInformation() throws Exception {
	final OneShotLatch receivedSlotRequest = new OneShotLatch();
	final CompletableFuture<Void> verifySlotReportFuture = new CompletableFuture<>();
	final OneShotLatch terminateSlotReportVerification = new OneShotLatch();
	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	// Assertions for this test
	testingResourceManagerGateway.setTaskExecutorHeartbeatConsumer((ignored, heartbeatPayload) -> {
		try {
			final ArrayList<SlotStatus> slots = Lists.newArrayList(heartbeatPayload.getSlotReport());
			assertThat(slots, hasSize(1));
			final SlotStatus slotStatus = slots.get(0);

			log.info("Received SlotStatus: {}", slotStatus);

			if (receivedSlotRequest.isTriggered()) {
				assertThat(slotStatus.getAllocationID(), is(notNullValue()));
			} else {
				assertThat(slotStatus.getAllocationID(), is(nullValue()));
			}
		} catch (AssertionError e) {
			verifySlotReportFuture.completeExceptionally(e);
		}

		if (terminateSlotReportVerification.isTriggered()) {
			verifySlotReportFuture.complete(null);
		}
	});
	final CompletableFuture<ResourceID> taskExecutorRegistrationFuture = new CompletableFuture<>();

	testingResourceManagerGateway.setSendSlotReportFunction(ignored -> {
		taskExecutorRegistrationFuture.complete(null);
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

	rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
	resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskSlotTable(new AllocateSlotNotifyingTaskSlotTable(receivedSlotRequest))
		.build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);
	final ResourceID taskExecutorResourceId = taskManagerServices.getUnresolvedTaskManagerLocation().getResourceID();

	taskExecutor.start();

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

	final ScheduledExecutorService heartbeatExecutor = java.util.concurrent.Executors.newSingleThreadScheduledExecutor();

	try {
		taskExecutorRegistrationFuture.get();

		final OneShotLatch scheduleFirstHeartbeat = new OneShotLatch();
		final ResourceID resourceManagerResourceId = testingResourceManagerGateway.getOwnResourceId();
		final long heartbeatInterval = 5L;
		heartbeatExecutor.scheduleWithFixedDelay(
			() -> {
				scheduleFirstHeartbeat.trigger();
				taskExecutorGateway.heartbeatFromResourceManager(resourceManagerResourceId);
			},
			0L,
			heartbeatInterval,
			TimeUnit.MILLISECONDS);

		scheduleFirstHeartbeat.await();

		SlotID slotId = new SlotID(taskExecutorResourceId, 0);
		final CompletableFuture<Acknowledge> requestSlotFuture = taskExecutorGateway.requestSlot(
			slotId,
			jobId,
			new AllocationID(),
			ResourceProfile.ZERO,
			"foobar",
			testingResourceManagerGateway.getFencingToken(),
			timeout);

		requestSlotFuture.get();

		terminateSlotReportVerification.trigger();

		verifySlotReportFuture.get();
	} finally {
		ExecutorUtils.gracefulShutdown(timeout.toMilliseconds(), TimeUnit.MILLISECONDS, heartbeatExecutor);
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

private void runTest(final boolean useAwaits) throws Exception {
	final File tempFile = tempDir.newFile();
	final Path path1 = new Path(tempFile.getAbsolutePath(), "1");
	final Path path2 = new Path(tempFile.getAbsolutePath(), "2");

	final OneShotLatch deleteAwaitLatch1 = new OneShotLatch();
	final OneShotLatch deleteAwaitLatch2 = new OneShotLatch();
	final OneShotLatch mkdirsAwaitLatch1 = new OneShotLatch();
	final OneShotLatch mkdirsAwaitLatch2 = new OneShotLatch();

	final OneShotLatch deleteTriggerLatch1 = new OneShotLatch();
	final OneShotLatch deletetriggerLatch2 = new OneShotLatch();
	final OneShotLatch mkdirsTriggerLatch1 = new OneShotLatch();
	final OneShotLatch mkdirsTriggerLatch2 = new OneShotLatch();

	final OneShotLatch createAwaitLatch = new OneShotLatch();
	final OneShotLatch createTriggerLatch = new OneShotLatch();

	// this "new LocalDataOutputStream()" is in the end called by the async threads
	whenNew(LocalDataOutputStream.class).withAnyArguments().thenAnswer(new Answer<LocalDataOutputStream>() {

		@Override
		public LocalDataOutputStream answer(InvocationOnMock invocation) throws Throwable {
			createAwaitLatch.trigger();
			createTriggerLatch.await();

			final File file = (File) invocation.getArguments()[0];
			return new LocalDataOutputStream(file);
		}
	});

	final LocalFileSystem fs1 = new SyncedFileSystem(
			deleteAwaitLatch1, mkdirsAwaitLatch1,
			deleteTriggerLatch1, mkdirsTriggerLatch1);

	final LocalFileSystem fs2 = new SyncedFileSystem(
			deleteAwaitLatch2, mkdirsAwaitLatch2,
			deletetriggerLatch2, mkdirsTriggerLatch2);

	// start the concurrent file creators
	FileCreator thread1 = new FileCreator(fs1, path1);
	FileCreator thread2 = new FileCreator(fs2, path2);
	thread1.start();
	thread2.start();

	// wait until they both decide to delete the directory
	if (useAwaits) {
		deleteAwaitLatch1.await();
		deleteAwaitLatch2.await();
	} else {
		Thread.sleep(5);
	}

	// now send off #1 to delete the directory (it will pass the 'mkdirs' fast) and wait to create the file
	mkdirsTriggerLatch1.trigger();
	deleteTriggerLatch1.trigger();

	if (useAwaits) {
		createAwaitLatch.await();
	} else {
		// this needs a bit more sleep time, because here mockito is working
		Thread.sleep(100);
	}

	// now send off #2 to delete the directory - it waits at 'mkdirs'
	deletetriggerLatch2.trigger();
	if (useAwaits) {
		mkdirsAwaitLatch2.await();
	} else {
		Thread.sleep(5);
	}

	// let #1 try to create the file and see if it succeeded
	createTriggerLatch.trigger();
	if (useAwaits) {
		thread1.sync();
	} else {
		Thread.sleep(5);
	}

	// now let #1 finish up
	mkdirsTriggerLatch2.trigger();

	thread1.sync();
	thread2.sync();
}
 

/**
 * Tests that interleaved granting and revoking of the leadership won't interfere
 * with the job recovery and the resulting internal state of the Dispatcher.
 */
@Test
public void testGrantingRevokingLeadership() throws Exception {
	final TestingHighAvailabilityServices highAvailabilityServices = new TestingHighAvailabilityServices();
	final JobGraph nonEmptyJobGraph = createNonEmptyJobGraph();
	final SubmittedJobGraph submittedJobGraph = new SubmittedJobGraph(nonEmptyJobGraph);

	final OneShotLatch enterGetJobIdsLatch = new OneShotLatch();
	final OneShotLatch proceedGetJobIdsLatch = new OneShotLatch();
	highAvailabilityServices.setSubmittedJobGraphStore(new BlockingSubmittedJobGraphStore(submittedJobGraph, enterGetJobIdsLatch, proceedGetJobIdsLatch));
	final TestingLeaderElectionService dispatcherLeaderElectionService = new TestingLeaderElectionService();
	highAvailabilityServices.setDispatcherLeaderElectionService(dispatcherLeaderElectionService);

	final BlockingQueue<DispatcherId> fencingTokens = new ArrayBlockingQueue<>(2);

	final HATestingDispatcher dispatcher = createDispatcherWithObservableFencingTokens(highAvailabilityServices, fencingTokens);

	dispatcher.start();

	try {
		// wait until the election service has been started
		dispatcherLeaderElectionService.getStartFuture().get();

		final UUID leaderId = UUID.randomUUID();
		dispatcherLeaderElectionService.isLeader(leaderId);

		dispatcherLeaderElectionService.notLeader();

		final DispatcherId firstFencingToken = fencingTokens.take();

		assertThat(firstFencingToken, equalTo(NULL_FENCING_TOKEN));

		enterGetJobIdsLatch.await();
		proceedGetJobIdsLatch.trigger();

		assertThat(dispatcher.getNumberJobs(timeout).get(), is(0));

	} finally {
		RpcUtils.terminateRpcEndpoint(dispatcher, timeout);
	}
}
 

@Test
public void testFunctionRestore() throws Exception {
	String testBasePath = hdfsURI + "/" + UUID.randomUUID() + "/";

	org.apache.hadoop.fs.Path path = null;
	long fileModTime = Long.MIN_VALUE;
	for (int i = 0; i < 1; i++) {
		Tuple2<org.apache.hadoop.fs.Path, String> file = createFileAndFillWithData(testBasePath, "file", i, "This is test line.");
		path = file.f0;
		fileModTime = hdfs.getFileStatus(file.f0).getModificationTime();
	}

	TextInputFormat format = new TextInputFormat(new Path(testBasePath));

	final ContinuousFileMonitoringFunction<String> monitoringFunction =
		createTestContinuousFileMonitoringFunction(format, FileProcessingMode.PROCESS_CONTINUOUSLY);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> src =
		new StreamSource<>(monitoringFunction);

	final AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarness =
		new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
	testHarness.open();

	final Throwable[] error = new Throwable[1];

	final OneShotLatch latch = new OneShotLatch();

	final DummySourceContext sourceContext = new DummySourceContext() {
		@Override
		public void collect(TimestampedFileInputSplit element) {
			latch.trigger();
		}
	};

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				monitoringFunction.run(sourceContext);
			}
			catch (Throwable t) {
				t.printStackTrace();
				error[0] = t;
			}
		}
	};
	runner.start();

	// first condition for the source to have updated its state: emit at least one element
	if (!latch.isTriggered()) {
		latch.await();
	}

	// second condition for the source to have updated its state: it's not on the lock anymore,
	// this means it has processed all the splits and updated its state.
	synchronized (sourceContext.getCheckpointLock()) {}

	OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
	monitoringFunction.cancel();
	runner.join();

	testHarness.close();

	final ContinuousFileMonitoringFunction<String> monitoringFunctionCopy =
		createTestContinuousFileMonitoringFunction(format, FileProcessingMode.PROCESS_CONTINUOUSLY);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> srcCopy =
		new StreamSource<>(monitoringFunctionCopy);

	AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarnessCopy =
		new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);
	testHarnessCopy.initializeState(snapshot);
	testHarnessCopy.open();

	Assert.assertNull(error[0]);
	Assert.assertEquals(fileModTime, monitoringFunctionCopy.getGlobalModificationTime());

	hdfs.delete(path, false);
}
 

/**
 * Test delayed registration of task executor where the delay is introduced during connection from resource manager
 * to the registering task executor.
 */
@Test
public void testDelayedRegisterTaskExecutor() throws Exception {
	final Time fastTimeout = Time.milliseconds(1L);
	try {
		final OneShotLatch startConnection = new OneShotLatch();
		final OneShotLatch finishConnection = new OneShotLatch();

		// first registration is with blocking connection
		rpcService.setRpcGatewayFutureFunction(rpcGateway ->
			CompletableFuture.supplyAsync(
				() -> {
					startConnection.trigger();
					try {
						finishConnection.await();
					} catch (InterruptedException ignored) {}
					return rpcGateway;
				},
				TestingUtils.defaultExecutor()));

		TaskExecutorRegistration taskExecutorRegistration = new TaskExecutorRegistration(
			taskExecutorGateway.getAddress(),
			taskExecutorResourceID,
			dataPort,
			hardwareDescription,
			ResourceProfile.ZERO,
			ResourceProfile.ZERO);

		CompletableFuture<RegistrationResponse> firstFuture =
			rmGateway.registerTaskExecutor(taskExecutorRegistration, fastTimeout);
		try {
			firstFuture.get();
			fail("Should have failed because connection to taskmanager is delayed beyond timeout");
		} catch (Exception e) {
			final Throwable cause = ExceptionUtils.stripExecutionException(e);
			assertThat(cause, instanceOf(TimeoutException.class));
			assertThat(cause.getMessage(), containsString("ResourceManagerGateway.registerTaskExecutor"));
		}

		startConnection.await();

		// second registration after timeout is with no delay, expecting it to be succeeded
		rpcService.resetRpcGatewayFutureFunction();
		CompletableFuture<RegistrationResponse> secondFuture =
			rmGateway.registerTaskExecutor(taskExecutorRegistration, TIMEOUT);
		RegistrationResponse response = secondFuture.get();
		assertTrue(response instanceof TaskExecutorRegistrationSuccess);

		// on success, send slot report for taskmanager registration
		final SlotReport slotReport = new SlotReport(new SlotStatus(new SlotID(taskExecutorResourceID, 0), ResourceProfile.ANY));
		rmGateway.sendSlotReport(taskExecutorResourceID,
			((TaskExecutorRegistrationSuccess) response).getRegistrationId(), slotReport, TIMEOUT).get();

		// let the remaining part of the first registration proceed
		finishConnection.trigger();
		Thread.sleep(1L);

		// verify that the latest registration is valid not being unregistered by the delayed one
		final TaskManagerInfo taskManagerInfo = rmGateway.requestTaskManagerInfo(
			taskExecutorResourceID,
			TIMEOUT).get();
		assertThat(taskManagerInfo.getResourceId(), equalTo(taskExecutorResourceID));
		assertThat(taskManagerInfo.getNumberSlots(), equalTo(1));
	} finally {
		rpcService.resetRpcGatewayFutureFunction();
	}
}
 

/**
 * The purpose of this test is to check that parallel snapshots are possible, and work even if a previous snapshot
 * is still running and blocking.
 */
@Test
public void testParallelAsyncSnapshots() throws Exception {
	OneShotLatch blocker = new OneShotLatch();
	OneShotLatch waiter = new OneShotLatch();
	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);
	streamFactory.setWaiterLatch(waiter);
	streamFactory.setBlockerLatch(blocker);
	streamFactory.setAfterNumberInvocations(10);

	final AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	try {

		if (!backend.supportsAsynchronousSnapshots()) {
			return;
		}

		// insert some data to the backend.
		InternalValueState<Integer, VoidNamespace, Integer> valueState = backend.createInternalState(
			VoidNamespaceSerializer.INSTANCE,
			new ValueStateDescriptor<>("test", IntSerializer.INSTANCE));

		valueState.setCurrentNamespace(VoidNamespace.INSTANCE);

		for (int i = 0; i < 10; ++i) {
			backend.setCurrentKey(i);
			valueState.update(i);
		}

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot1 =
			backend.snapshot(0L, 0L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

		Thread runner1 = new Thread(snapshot1, "snapshot-1-runner");
		runner1.start();
		// after this call returns, we have a running snapshot-1 that is blocked in IO.
		waiter.await();

		// do some updates in between the snapshots.
		for (int i = 5; i < 15; ++i) {
			backend.setCurrentKey(i);
			valueState.update(i + 1);
		}

		// we don't want to block the second snapshot.
		streamFactory.setWaiterLatch(null);
		streamFactory.setBlockerLatch(null);

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot2 =
			backend.snapshot(1L, 1L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

		Thread runner2 = new Thread(snapshot2,"snapshot-2-runner");
		runner2.start();
		// snapshot-2 should run and succeed, while snapshot-1 is still running and blocked in IO.
		snapshot2.get();

		// we release the blocking IO so that snapshot-1 can also finish and succeed.
		blocker.trigger();
		snapshot1.get();

	} finally {
		backend.dispose();
	}
}
 

/**
 * Tests that we can submit a task to the TaskManager given that we've allocated a slot there.
 */
@Test(timeout = 10000L)
public void testTaskSubmission() throws Exception {
	final AllocationID allocationId = new AllocationID();
	final JobMasterId jobMasterId = JobMasterId.generate();
	final JobVertexID jobVertexId = new JobVertexID();

	JobInformation jobInformation = new JobInformation(
			jobId,
			testName.getMethodName(),
			new SerializedValue<>(new ExecutionConfig()),
			new Configuration(),
			Collections.emptyList(),
			Collections.emptyList());

	TaskInformation taskInformation = new TaskInformation(
			jobVertexId,
			"test task",
			1,
			1,
			TestInvokable.class.getName(),
			new Configuration());

	SerializedValue<JobInformation> serializedJobInformation = new SerializedValue<>(jobInformation);
	SerializedValue<TaskInformation> serializedJobVertexInformation = new SerializedValue<>(taskInformation);

	final TaskDeploymentDescriptor tdd = new TaskDeploymentDescriptor(
			jobId,
			new TaskDeploymentDescriptor.NonOffloaded<>(serializedJobInformation),
			new TaskDeploymentDescriptor.NonOffloaded<>(serializedJobVertexInformation),
			new ExecutionAttemptID(),
			allocationId,
			0,
			0,
			0,
			null,
			Collections.emptyList(),
			Collections.emptyList());

	final LibraryCacheManager libraryCacheManager = mock(LibraryCacheManager.class);
	when(libraryCacheManager.getClassLoader(any(JobID.class))).thenReturn(ClassLoader.getSystemClassLoader());

	final JobMasterGateway jobMasterGateway = mock(JobMasterGateway.class);
	when(jobMasterGateway.getFencingToken()).thenReturn(jobMasterId);

	final OneShotLatch taskInTerminalState = new OneShotLatch();
	final TaskManagerActions taskManagerActions = new NoOpTaskManagerActions() {
		@Override
		public void updateTaskExecutionState(TaskExecutionState taskExecutionState) {
			if (taskExecutionState.getExecutionState().isTerminal()) {
				taskInTerminalState.trigger();
			}
		}
	};
	final JobManagerConnection jobManagerConnection = new JobManagerConnection(
		jobId,
		ResourceID.generate(),
		jobMasterGateway,
		taskManagerActions,
		mock(CheckpointResponder.class),
		new TestGlobalAggregateManager(),
		libraryCacheManager,
		new NoOpResultPartitionConsumableNotifier(),
		mock(PartitionProducerStateChecker.class));

	final JobManagerTable jobManagerTable = new JobManagerTable();
	jobManagerTable.put(jobId, jobManagerConnection);

	final TaskSlotTable taskSlotTable = mock(TaskSlotTable.class);
	when(taskSlotTable.tryMarkSlotActive(eq(jobId), eq(allocationId))).thenReturn(true);
	when(taskSlotTable.addTask(any(Task.class))).thenReturn(true);

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setShuffleEnvironment(nettyShuffleEnvironment)
		.setTaskSlotTable(taskSlotTable)
		.setJobManagerTable(jobManagerTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();

		final TaskExecutorGateway tmGateway = taskManager.getSelfGateway(TaskExecutorGateway.class);

		tmGateway.submitTask(tdd, jobMasterId, timeout);

		CompletableFuture<Boolean> completionFuture = TestInvokable.COMPLETABLE_FUTURE;

		completionFuture.get();

		taskInTerminalState.await();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

/**
 * Tests that a blocking {@link JobManagerRunner} creation, e.g. due to blocking FileSystem access,
 * does not block the {@link Dispatcher}.
 *
 * <p>See FLINK-10314
 */
@Test
public void testBlockingJobManagerRunner() throws Exception {
	final OneShotLatch jobManagerRunnerCreationLatch = new OneShotLatch();
	dispatcher = createAndStartDispatcher(
		heartbeatServices,
		haServices,
		new BlockingJobManagerRunnerFactory(jobManagerRunnerCreationLatch::await));

	dispatcherLeaderElectionService.isLeader(UUID.randomUUID()).get();

	final DispatcherGateway dispatcherGateway = dispatcher.getSelfGateway(DispatcherGateway.class);

	final CompletableFuture<Acknowledge> submissionFuture = dispatcherGateway.submitJob(jobGraph, TIMEOUT);

	assertThat(submissionFuture.isDone(), is(false));

	final CompletableFuture<Collection<String>> metricQueryServiceAddressesFuture = dispatcherGateway.requestMetricQueryServiceAddresses(Time.seconds(5L));

	assertThat(metricQueryServiceAddressesFuture.get(), is(empty()));

	assertThat(submissionFuture.isDone(), is(false));

	jobManagerRunnerCreationLatch.trigger();

	submissionFuture.get();
}