org.apache.flink.core.testutils.OneShotLatch Java Examples

public BlockingFSDataInputStream(
	@Nullable FSDataInputStream delegate,
	@Nullable OneShotLatch waitForBlock,
	@Nullable OneShotLatch triggerUnblock,
	long blockAtPosition) {

	this.delegate = delegate;
	this.triggerUnblock = triggerUnblock;
	this.waitUntilStreamBlocked = waitForBlock;
	this.blockAtPosition = blockAtPosition;
	if (delegate != null) {
		try {
			this.position = delegate.getPos();
		} catch (IOException e) {
			throw new RuntimeException(e);
		}
	} else {
		this.position = 0;
	}
	this.closed = new AtomicBoolean(false);
}
 

public BlockingCheckpointOutputStream(
	@Nullable FSDataOutputStream delegate,
	@Nullable OneShotLatch waitForBlocking,
	@Nullable OneShotLatch triggerUnblock,
	long blockAtPosition) {

	this.delegate = delegate;
	this.triggerUnblock = triggerUnblock;
	this.waitForBlocking = waitForBlocking;
	this.blockAtPosition = blockAtPosition;
	if (delegate != null) {
		try {
			this.position = delegate.getPos();
		} catch (IOException e) {
			throw new RuntimeException(e);
		}
	} else {
		this.position = 0;
	}
	this.closed = new AtomicBoolean(false);
}
 

@Test
public void testCancel() throws Exception {
	// test cancel properly
	JobID jid = new JobID();

	OneShotLatch cancelLatch = new OneShotLatch();

	String[] parameters = { jid.toString() };

	TestingClusterClient<String> clusterClient = new TestingClusterClient<>();

	clusterClient.setCancelFunction(jobID -> {
		cancelLatch.trigger();
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

	MockedCliFrontend testFrontend = new MockedCliFrontend(clusterClient);
	testFrontend.cancel(parameters);
	cancelLatch.await();
}
 

@Test
public void testRunAfterwards() throws Exception {
	final CompletableFuture<Void> inputFuture = new CompletableFuture<>();
	final OneShotLatch runnableLatch = new OneShotLatch();

	final CompletableFuture<Void> runFuture = FutureUtils.runAfterwards(
		inputFuture,
		runnableLatch::trigger);

	assertThat(runnableLatch.isTriggered(), is(false));
	assertThat(runFuture.isDone(), is(false));

	inputFuture.complete(null);

	assertThat(runnableLatch.isTriggered(), is(true));
	assertThat(runFuture.isDone(), is(true));

	// check that this future is not exceptionally completed
	runFuture.get();
}
 

@Test
public void testCancellationFailsWithBlockingLock() throws Exception {
	syncLatch = new OneShotLatch();

	StreamConfig cfg = new StreamConfig(new Configuration());
	Task task = createTask(CancelFailingTask.class, cfg, new Configuration());

	// start the task and wait until it runs
	// execution state RUNNING is not enough, we need to wait until the stream task's run() method
	// is entered
	task.startTaskThread();
	syncLatch.await();

	// cancel the execution - this should lead to smooth shutdown
	task.cancelExecution();
	task.getExecutingThread().join();

	assertEquals(ExecutionState.CANCELED, task.getExecutionState());
}
 

@Test
public void testStopWithMaxWMAndDefaultSavepointDir() throws Exception {
	JobID jid = new JobID();

	String[] parameters = { "-p", "-d", jid.toString() };
	OneShotLatch stopWithSavepointLatch = new OneShotLatch();
	TestingClusterClient<String> clusterClient = new TestingClusterClient<>();
	clusterClient.setStopWithSavepointFunction((jobID, advanceToEndOfEventTime, savepointDirectory) -> {
		assertThat(jobID, is(jid));
		assertThat(advanceToEndOfEventTime, is(true));
		assertNull(savepointDirectory);
		stopWithSavepointLatch.trigger();
		return CompletableFuture.completedFuture(savepointDirectory);
	});
	MockedCliFrontend testFrontend = new MockedCliFrontend(clusterClient);
	testFrontend.stop(parameters);

	stopWithSavepointLatch.await();
}
 

private FencedTestEndpoint(
		RpcService rpcService,
		ReentrantLock lock,
		AtomicBoolean concurrentAccess,
		UUID initialFencingToken,
		OneShotLatch enteringSetNewFencingToken,
		OneShotLatch triggerSetNewFencingToken) {
	super(rpcService);

	this.lock = lock;
	this.concurrentAccess = concurrentAccess;

	this.enteringSetNewFencingToken = enteringSetNewFencingToken;
	this.triggerSetNewFencingToken = triggerSetNewFencingToken;

	// make it look as if we are running in the main thread
	currentMainThread.set(Thread.currentThread());

	try {
		setFencingToken(initialFencingToken);
	} finally {
		currentMainThread.set(null);
	}
}
 

@Test
public void testComposeAfterwards() throws ExecutionException, InterruptedException {
	final CompletableFuture<Void> inputFuture = new CompletableFuture<>();
	final OneShotLatch composeLatch = new OneShotLatch();

	final CompletableFuture<Void> composeFuture = FutureUtils.composeAfterwards(
		inputFuture,
		() -> {
			composeLatch.trigger();
			return CompletableFuture.completedFuture(null);
		});

	assertThat(composeLatch.isTriggered(), is(false));
	assertThat(composeFuture.isDone(), is(false));

	inputFuture.complete(null);

	assertThat(composeLatch.isTriggered(), is(true));
	assertThat(composeFuture.isDone(), is(true));

	// check that tthis future is not exceptionally completed
	composeFuture.get();
}
 

public BlockingCheckpointOutputStream(
	@Nullable FSDataOutputStream delegate,
	@Nullable OneShotLatch waitForBlocking,
	@Nullable OneShotLatch triggerUnblock,
	long blockAtPosition) {

	this.delegate = delegate;
	this.triggerUnblock = triggerUnblock;
	this.waitForBlocking = waitForBlocking;
	this.blockAtPosition = blockAtPosition;
	if (delegate != null) {
		try {
			this.position = delegate.getPos();
		} catch (IOException e) {
			throw new RuntimeException(e);
		}
	} else {
		this.position = 0;
	}
	this.closed = new AtomicBoolean(false);
}
 

@Test
public void testLibraryCacheManagerRegistration() throws Exception {
	final OneShotLatch registerClassLoaderLatch = new OneShotLatch();
	final OneShotLatch closeClassLoaderLeaseLatch = new OneShotLatch();
	final TestingClassLoaderLease classLoaderLease = TestingClassLoaderLease.newBuilder()
		.setGetOrResolveClassLoaderFunction((permanentBlobKeys, urls) -> {
			registerClassLoaderLatch.trigger();
			return JobManagerRunnerImplTest.class.getClassLoader();
		})
		.setCloseRunnable(closeClassLoaderLeaseLatch::trigger)
		.build();
	final JobManagerRunner jobManagerRunner = createJobManagerRunner(classLoaderLease);

	try {
		jobManagerRunner.start();

		registerClassLoaderLatch.await();

		jobManagerRunner.close();

		closeClassLoaderLeaseLatch.await();
	} finally {
		jobManagerRunner.close();
	}
}
 

private FlinkKubeClient createTestingFlinkKubeClientAllocatingPodsAfter(
		int numberOfRetries,
		AtomicInteger retries,
		OneShotLatch podCreated) {
	ExecutorService kubeClientExecutorService = Executors.newDirectExecutorService();
	return new Fabric8FlinkKubeClient(flinkConfig, kubeClient, () -> kubeClientExecutorService) {
		@Override
		public CompletableFuture<Void> createTaskManagerPod(KubernetesPod kubernetesPod) {
			if (retries.getAndIncrement() < numberOfRetries) {
				return FutureUtils.completedExceptionally(new RuntimeException("Exception"));
			}
			podCreated.trigger();
			return super.createTaskManagerPod(kubernetesPod);
		}
	};
}
 

/**
 * 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();
	}
}
 

public TestableKinesisDataFetcher(
		List<String> fakeStreams,
		SourceFunction.SourceContext<T> sourceContext,
		Properties fakeConfiguration,
		KinesisDeserializationSchema<T> deserializationSchema,
		int fakeTotalCountOfSubtasks,
		int fakeIndexOfThisSubtask,
		AtomicReference<Throwable> thrownErrorUnderTest,
		LinkedList<KinesisStreamShardState> subscribedShardsStateUnderTest,
		HashMap<String, String> subscribedStreamsToLastDiscoveredShardIdsStateUnderTest,
		KinesisProxyInterface fakeKinesis) {
	super(
		fakeStreams,
		sourceContext,
		sourceContext.getCheckpointLock(),
		getMockedRuntimeContext(fakeTotalCountOfSubtasks, fakeIndexOfThisSubtask),
		fakeConfiguration,
		deserializationSchema,
		DEFAULT_SHARD_ASSIGNER,
		null,
		null,
		thrownErrorUnderTest,
		subscribedShardsStateUnderTest,
		subscribedStreamsToLastDiscoveredShardIdsStateUnderTest,
		(properties) -> fakeKinesis);

	this.runWaiter = new OneShotLatch();
	this.initialDiscoveryWaiter = new OneShotLatch();
	this.shutdownWaiter = new OneShotLatch();

	this.running = true;
}
 

@Test
public void closeAsync_duringJobRecovery_preventsDispatcherServiceCreation() throws Exception {
	final OneShotLatch jobRecoveryStartedLatch = new OneShotLatch();
	final OneShotLatch completeJobRecoveryLatch = new OneShotLatch();
	final OneShotLatch createDispatcherServiceLatch = new OneShotLatch();

	this.jobGraphStore = TestingJobGraphStore.newBuilder()
		.setJobIdsFunction(storedJobs -> {
			jobRecoveryStartedLatch.trigger();
			completeJobRecoveryLatch.await();
			return storedJobs;
		})
		.build();

	this.dispatcherServiceFactory = TestingDispatcherServiceFactory.newBuilder()
		.setCreateFunction(
			(ignoredA, ignoredB, ignoredC) -> {
				createDispatcherServiceLatch.trigger();
				return TestingDispatcherGatewayService.newBuilder().build();
			})
		.build();

	try (final SessionDispatcherLeaderProcess dispatcherLeaderProcess = createDispatcherLeaderProcess()) {
		dispatcherLeaderProcess.start();

		jobRecoveryStartedLatch.await();

		dispatcherLeaderProcess.closeAsync();

		completeJobRecoveryLatch.trigger();

		try {
			createDispatcherServiceLatch.await(10L, TimeUnit.MILLISECONDS);
			fail("No dispatcher service should be created after the process has been stopped.");
		} catch (TimeoutException expected) {}
	}
}
 

private static SystemProcessingTimeService createBlockingSystemProcessingTimeService(
		final OneShotLatch blockUntilTriggered,
		final AtomicBoolean check) {

	final OneShotLatch waitUntilTimerStarted = new OneShotLatch();

	Preconditions.checkState(!check.get());

	final SystemProcessingTimeService timeService = new SystemProcessingTimeService(exception -> {});

	timeService.scheduleAtFixedRate(
		timestamp -> {

			waitUntilTimerStarted.trigger();

			boolean unblocked = false;

			while (!unblocked) {
				try {
					blockUntilTriggered.await();
					unblocked = true;
				} catch (InterruptedException ignore) {
				}
			}

			check.set(true);
		},
		0L,
		10L);

	try {
		waitUntilTimerStarted.await();
	} catch (InterruptedException e) {
		Assert.fail("Problem while starting up service.");
	}

	return timeService;
}
 

@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();
}
 

/**
 * Tests {@link NetworkBufferPool#requestMemorySegments()}, verifying it may be aborted in
 * case of a concurrent {@link NetworkBufferPool#destroy()} call.
 */
@Test
public void testRequestMemorySegmentsInterruptable() 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);
	globalPool.destroy();

	segment.free();

	expectedException.expect(IllegalStateException.class);
	expectedException.expectMessage("destroyed");
	try {
		asyncRequest.sync();
	} finally {
		globalPool.destroy();
	}
}
 

protected FencedTestEndpoint(
		RpcService rpcService,
		ReentrantLock lock,
		AtomicBoolean concurrentAccess) {
	this(
		rpcService,
		lock,
		concurrentAccess,
		UUID.randomUUID(),
		new OneShotLatch(),
		new OneShotLatch());
}
 

private FencedTestEndpoint(
		RpcService rpcService,
		ReentrantLock lock,
		AtomicBoolean concurrentAccess,
		UUID initialFencingToken,
		OneShotLatch enteringSetNewFencingToken,
		OneShotLatch triggerSetNewFencingToken) {
	super(rpcService, initialFencingToken);

	this.lock = lock;
	this.concurrentAccess = concurrentAccess;

	this.enteringSetNewFencingToken = enteringSetNewFencingToken;
	this.triggerSetNewFencingToken = triggerSetNewFencingToken;
}
 

protected FencedTestEndpoint(
		RpcService rpcService,
		ReentrantLock lock,
		AtomicBoolean concurrentAccess) {
	this(
		rpcService,
		lock,
		concurrentAccess,
		UUID.randomUUID(),
		new OneShotLatch(),
		new OneShotLatch());
}
 

private AllocateSlotNotifyingTaskSlotTable(
		Collection<ResourceProfile> resourceProfiles,
		TimerService<AllocationID> timerService,
		OneShotLatch allocateSlotLatch) {
	super(resourceProfiles, timerService);
	this.allocateSlotLatch = allocateSlotLatch;
}
 

@Override
protected void run() throws Exception {
	final OneShotLatch latch = new OneShotLatch();
	final Object lock = new Object();

	LockHolder holder = new LockHolder(lock, latch);
	holder.start();
	try {
		// cancellation should try and cancel this
		getCancelables().registerCloseable(holder);

		// wait till the lock holder has the lock
		latch.await();

		// we are at the point where cancelling can happen
		syncLatch.trigger();

		// try to acquire the lock - this is not possible as long as the lock holder
		// thread lives
		//noinspection SynchronizationOnLocalVariableOrMethodParameter
		synchronized (lock) {
			// nothing
		}
	}
	finally {
		holder.close();
	}

}
 

@Override
protected void performDefaultAction(DefaultActionContext context) throws Exception {
	final OneShotLatch latch = new OneShotLatch();
	final Object lock = new Object();

	LockHolder holder = new LockHolder(lock, latch);
	holder.start();
	try {
		// cancellation should try and cancel this
		getCancelables().registerCloseable(holder);

		// wait till the lock holder has the lock
		latch.await();

		// we are at the point where cancelling can happen
		syncLatch.trigger();

		// try to acquire the lock - this is not possible as long as the lock holder
		// thread lives
		//noinspection SynchronizationOnLocalVariableOrMethodParameter
		synchronized (lock) {
			// nothing
		}
	}
	finally {
		holder.close();
	}
	context.allActionsCompleted();
}
 

@Test
public void testScheduleRunnable() throws Exception {
	final OneShotLatch latch = new OneShotLatch();
	final long delay = 100L;
	final long start = System.nanoTime();

	ScheduledFuture<?> scheduledFuture = akkaRpcService.scheduleRunnable(latch::trigger, delay, TimeUnit.MILLISECONDS);

	scheduledFuture.get();

	assertTrue(latch.isTriggered());
	final long stop = System.nanoTime();

	assertTrue("call was not properly delayed", ((stop - start) / 1000000) >= delay);
}
 

public BlockingSourceContext(String name, OneShotLatch latchToTrigger, OneShotLatch latchToWait,
							ConcurrentHashMap<String, List<T>> output, int elemToFire) {
	this.name = name;
	this.lock = new Object();
	this.latchToTrigger = latchToTrigger;
	this.latchToWait = latchToWait;
	this.collector = output;
	this.threshold = elemToFire;

	this.localOutput = new ArrayList<>();
	List<T> prev = collector.put(name, localOutput);
	if (prev != null) {
		Assert.fail();
	}
}
 

public WaitingSource(SourceFunction<T> source, TypeInformation<T> returnType) {
	this.source = source;
	this.returnType = returnType;
	this.guardId = UUID.randomUUID().toString();

	guards.put(guardId, new OneShotLatch());
	this.running = true;
}
 

@Before
public void setup() {
	awaitLatch = new OneShotLatch();
	triggerLatch = new OneShotLatch();

	shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build();
}
 

@Test
public void testComposeAfterwardsFirstExceptional() throws InterruptedException {
	final CompletableFuture<Void> inputFuture = new CompletableFuture<>();
	final OneShotLatch composeLatch = new OneShotLatch();
	final FlinkException testException = new FlinkException("Test exception");

	final CompletableFuture<Void> composeFuture = FutureUtils.composeAfterwards(
		inputFuture,
		() -> {
			composeLatch.trigger();
			return CompletableFuture.completedFuture(null);
		});

	assertThat(composeLatch.isTriggered(), is(false));
	assertThat(composeFuture.isDone(), is(false));

	inputFuture.completeExceptionally(testException);

	assertThat(composeLatch.isTriggered(), is(true));
	assertThat(composeFuture.isDone(), is(true));

	// check that this future is not exceptionally completed
	try {
		composeFuture.get();
		fail("Expected an exceptional completion");
	} catch (ExecutionException ee) {
		assertThat(ExceptionUtils.stripExecutionException(ee), is(testException));
	}
}
 

SyncedFileSystem(
		OneShotLatch deleteTriggerLatch,
		OneShotLatch mkdirsTriggerLatch,
		OneShotLatch deleteAwaitLatch,
		OneShotLatch mkdirsAwaitLatch) {

	this.deleteTriggerLatch = deleteTriggerLatch;
	this.mkdirsTriggerLatch = mkdirsTriggerLatch;
	this.deleteAwaitLatch = deleteAwaitLatch;
	this.mkdirsAwaitLatch = mkdirsAwaitLatch;
}
 

public TestAsyncWaitOperator(
		AsyncFunction<IN, OUT> asyncFunction,
		long timeout,
		int capacity,
		AsyncDataStream.OutputMode outputMode,
		OneShotLatch closingLatch) {
	super(asyncFunction, timeout, capacity, outputMode);

	this.closingLatch = Preconditions.checkNotNull(closingLatch);
}