Java Code Examples for org.apache.flink.runtime.concurrent.FutureUtils#ConjunctFuture

/**
 * Tests that concurrent release operations only trigger the failing of the payload and
 * the return of the slot once.
 */
@Test
public void testConcurrentReleaseOperations() throws Exception {
	final CountingSlotOwner countingSlotOwner = new CountingSlotOwner();
	final CountingFailPayload countingFailPayload = new CountingFailPayload();
	final SingleLogicalSlot singleLogicalSlot = createSingleLogicalSlot(countingSlotOwner);

	singleLogicalSlot.tryAssignPayload(countingFailPayload);

	final ExecutorService executorService = Executors.newFixedThreadPool(4);

	try {
		final int numberConcurrentOperations = 10;
		final Collection<CompletableFuture<?>> releaseOperationFutures = new ArrayList<>(numberConcurrentOperations);

		for (int i = 0; i < numberConcurrentOperations; i++) {
			final CompletableFuture<Void> releaseOperationFuture = CompletableFuture.runAsync(
				() -> {
					try {
						singleLogicalSlot.releaseSlot(new FlinkException("Test exception")).get();
					} catch (InterruptedException | ExecutionException e) {
						ExceptionUtils.checkInterrupted(e);
						throw new CompletionException(e);
					}
				});

			releaseOperationFutures.add(releaseOperationFuture);
		}

		final FutureUtils.ConjunctFuture<Void> releaseOperationsFuture = FutureUtils.waitForAll(releaseOperationFutures);

		releaseOperationsFuture.get();

		assertThat(countingSlotOwner.getReleaseCount(), is(1));
		assertThat(countingFailPayload.getFailCount(), is(1));
	} finally {
		executorService.shutdownNow();
	}
}
 

private FutureUtils.ConjunctFuture<Void> closeHandlersAsync() {
	return FutureUtils.waitForAll(handlers.stream()
		.map(tuple -> tuple.f1)
		.filter(handler -> handler instanceof AutoCloseableAsync)
		.map(handler -> ((AutoCloseableAsync) handler).closeAsync())
		.collect(Collectors.toList()));
}
 

private FutureUtils.ConjunctFuture<Void> closeHandlersAsync() {
	return FutureUtils.waitForAll(handlers.stream()
		.map(tuple -> tuple.f1)
		.filter(handler -> handler instanceof AutoCloseableAsync)
		.map(handler -> ((AutoCloseableAsync) handler).closeAsync())
		.collect(Collectors.toList()));
}
 

/**
 * Tests that concurrent release operations only trigger the failing of the payload and
 * the return of the slot once.
 */
@Test
public void testConcurrentReleaseOperations() throws Exception {
	final CountingSlotOwner countingSlotOwner = new CountingSlotOwner();
	final CountingFailPayload countingFailPayload = new CountingFailPayload();
	final SingleLogicalSlot singleLogicalSlot = createSingleLogicalSlot(countingSlotOwner);

	singleLogicalSlot.tryAssignPayload(countingFailPayload);

	final ExecutorService executorService = Executors.newFixedThreadPool(4);

	try {
		final int numberConcurrentOperations = 10;
		final Collection<CompletableFuture<?>> releaseOperationFutures = new ArrayList<>(numberConcurrentOperations);

		for (int i = 0; i < numberConcurrentOperations; i++) {
			final CompletableFuture<Void> releaseOperationFuture = CompletableFuture.runAsync(
				() -> {
					try {
						singleLogicalSlot.releaseSlot(new FlinkException("Test exception")).get();
					} catch (InterruptedException | ExecutionException e) {
						ExceptionUtils.checkInterrupted(e);
						throw new CompletionException(e);
					}
				});

			releaseOperationFutures.add(releaseOperationFuture);
		}

		final FutureUtils.ConjunctFuture<Void> releaseOperationsFuture = FutureUtils.waitForAll(releaseOperationFutures);

		releaseOperationsFuture.get();

		assertThat(countingSlotOwner.getReleaseCount(), is(1));
		assertThat(countingFailPayload.getFailCount(), is(1));
	} finally {
		executorService.shutdownNow();
	}
}
 

/**
 * Tests that concurrent release operations only trigger the failing of the payload and
 * the return of the slot once.
 */
@Test
public void testConcurrentReleaseOperations() throws Exception {
	final CountingSlotOwner countingSlotOwner = new CountingSlotOwner();
	final CountingFailPayload countingFailPayload = new CountingFailPayload();
	final SingleLogicalSlot singleLogicalSlot = createSingleLogicalSlot(countingSlotOwner);

	singleLogicalSlot.tryAssignPayload(countingFailPayload);

	final ExecutorService executorService = Executors.newFixedThreadPool(4);

	try {
		final int numberConcurrentOperations = 10;
		final Collection<CompletableFuture<?>> releaseOperationFutures = new ArrayList<>(numberConcurrentOperations);

		for (int i = 0; i < numberConcurrentOperations; i++) {
			final CompletableFuture<Void> releaseOperationFuture = CompletableFuture.runAsync(
				() -> {
					try {
						singleLogicalSlot.releaseSlot(new FlinkException("Test exception")).get();
					} catch (InterruptedException | ExecutionException e) {
						ExceptionUtils.checkInterrupted(e);
						throw new CompletionException(e);
					}
				});

			releaseOperationFutures.add(releaseOperationFuture);
		}

		final FutureUtils.ConjunctFuture<Void> releaseOperationsFuture = FutureUtils.waitForAll(releaseOperationFutures);

		releaseOperationsFuture.get();

		assertThat(countingSlotOwner.getReleaseCount(), is(1));
		assertThat(countingFailPayload.getFailCount(), is(1));
	} finally {
		executorService.shutdownNow();
	}
}
 

private FutureUtils.ConjunctFuture<Void> closeHandlersAsync() {
	return FutureUtils.waitForAll(handlers.stream()
		.map(tuple -> tuple.f1)
		.filter(handler -> handler instanceof AutoCloseableAsync)
		.map(handler -> ((AutoCloseableAsync) handler).closeAsync())
		.collect(Collectors.toList()));
}
 

/**
 * FLINK-5667
 *
 * <p>Tests that a concurrent cancel operation discards the state handles of a not yet
 * acknowledged checkpoint and prevents sending an acknowledge message to the
 * CheckpointCoordinator. The situation can only happen if the cancel call is executed
 * before Environment.acknowledgeCheckpoint().
 */
@Test
public void testAsyncCheckpointingConcurrentCloseBeforeAcknowledge() throws Exception {

	final TestingKeyedStateHandle managedKeyedStateHandle = new TestingKeyedStateHandle();
	final TestingKeyedStateHandle rawKeyedStateHandle = new TestingKeyedStateHandle();
	final TestingOperatorStateHandle managedOperatorStateHandle = new TestingOperatorStateHandle();
	final TestingOperatorStateHandle rawOperatorStateHandle = new TestingOperatorStateHandle();

	final BlockingRunnableFuture<SnapshotResult<KeyedStateHandle>> rawKeyedStateHandleFuture = new BlockingRunnableFuture<>(2, SnapshotResult.of(rawKeyedStateHandle));
	OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(
		DoneFuture.of(SnapshotResult.of(managedKeyedStateHandle)),
		rawKeyedStateHandleFuture,
		DoneFuture.of(SnapshotResult.of(managedOperatorStateHandle)),
		DoneFuture.of(SnapshotResult.of(rawOperatorStateHandle)),
		DoneFuture.of(SnapshotResult.empty()),
		DoneFuture.of(SnapshotResult.empty()));

	final OneInputStreamOperator<String, String> streamOperator = streamOperatorWithSnapshot(operatorSnapshotResult);

	final AcknowledgeDummyEnvironment mockEnvironment = new AcknowledgeDummyEnvironment();

	RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(
		mockEnvironment,
		operatorChain(streamOperator)));

	waitTaskIsRunning(task.streamTask, task.invocationFuture);

	final long checkpointId = 42L;
	task.streamTask.triggerCheckpointAsync(
		new CheckpointMetaData(checkpointId, 1L),
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		false);

	rawKeyedStateHandleFuture.awaitRun();

	task.streamTask.cancel();

	final FutureUtils.ConjunctFuture<Void> discardFuture = FutureUtils.waitForAll(
		asList(
			managedKeyedStateHandle.getDiscardFuture(),
			rawKeyedStateHandle.getDiscardFuture(),
			managedOperatorStateHandle.getDiscardFuture(),
			rawOperatorStateHandle.getDiscardFuture()));

	// make sure that all state handles have been discarded
	discardFuture.get();

	try {
		mockEnvironment.getAcknowledgeCheckpointFuture().get(10L, TimeUnit.MILLISECONDS);
		fail("The checkpoint should not get acknowledged.");
	} catch (TimeoutException expected) {
		// future should not be completed
	}

	task.waitForTaskCompletion(true);
}
 

/**
 * [FLINK-6020]
 * Tests that concurrent put operations will only upload the file once to the {@link BlobStore}
 * and that the files are not corrupt at any time.
 *
 * @param jobId
 * 		job ID to use (or <tt>null</tt> if job-unrelated)
 * @param blobType
 * 		whether the BLOB should become permanent or transient
 */
private void testConcurrentPutOperations(
		@Nullable final JobID jobId, final BlobKey.BlobType blobType)
		throws IOException, InterruptedException, ExecutionException {
	final Configuration config = new Configuration();
	config.setString(BlobServerOptions.STORAGE_DIRECTORY, temporaryFolder.newFolder().getAbsolutePath());

	BlobStore blobStore = mock(BlobStore.class);
	int concurrentPutOperations = 2;
	int dataSize = 1024;

	final CountDownLatch countDownLatch = new CountDownLatch(concurrentPutOperations);
	final byte[] data = new byte[dataSize];

	ArrayList<CompletableFuture<BlobKey>> allFutures = new ArrayList<>(concurrentPutOperations);

	ExecutorService executor = Executors.newFixedThreadPool(concurrentPutOperations);

	try (final BlobServer server = new BlobServer(config, blobStore)) {

		server.start();

		for (int i = 0; i < concurrentPutOperations; i++) {
			CompletableFuture<BlobKey> putFuture = CompletableFuture
				.supplyAsync(
					() -> {
						try {
							BlockingInputStream inputStream =
								new BlockingInputStream(countDownLatch, data);
							BlobKey uploadedKey = put(server, jobId, inputStream, blobType);
							// check the uploaded file's contents (concurrently)
							verifyContents(server, jobId, uploadedKey, data);
							return uploadedKey;
						} catch (IOException e) {
							throw new CompletionException(new FlinkException(
								"Could not upload blob.", e));
						}
					},
					executor);

			allFutures.add(putFuture);
		}

		FutureUtils.ConjunctFuture<Collection<BlobKey>> conjunctFuture = FutureUtils.combineAll(allFutures);

		// wait until all operations have completed and check that no exception was thrown
		Collection<BlobKey> blobKeys = conjunctFuture.get();

		Iterator<BlobKey> blobKeyIterator = blobKeys.iterator();

		assertTrue(blobKeyIterator.hasNext());

		BlobKey blobKey = blobKeyIterator.next();

		// make sure that all blob keys are the same
		while (blobKeyIterator.hasNext()) {
			verifyKeyDifferentHashEquals(blobKey, blobKeyIterator.next());
		}

		// check the uploaded file's contents
		verifyContents(server, jobId, blobKey, data);

		// check that we only uploaded the file once to the blob store
		if (blobType == PERMANENT_BLOB) {
			verify(blobStore, times(1)).put(any(File.class), eq(jobId), eq(blobKey));
		} else {
			// can't really verify much in the other cases other than that the put operations should
			// work and not corrupt files
			verify(blobStore, times(0)).put(any(File.class), eq(jobId), eq(blobKey));
		}
	} finally {
		executor.shutdownNow();
	}
}
 

/**
 * Tests that {@link TransientBlobCache} cleans up after a default TTL and keeps files which are
 * constantly accessed.
 */
private void testTransientBlobCleanup(@Nullable final JobID jobId)
		throws IOException, InterruptedException, ExecutionException {

	// 1s should be a safe-enough buffer to still check for existence after a BLOB's last access
	long cleanupInterval = 1L; // in seconds
	final int numberConcurrentGetOperations = 3;

	final List<CompletableFuture<Void>> getOperations = new ArrayList<>(numberConcurrentGetOperations);

	byte[] data = new byte[2000000];
	rnd.nextBytes(data);
	byte[] data2 = Arrays.copyOfRange(data, 10, 54);

	Configuration config = new Configuration();
	config.setString(BlobServerOptions.STORAGE_DIRECTORY,
		temporaryFolder.newFolder().getAbsolutePath());
	config.setLong(BlobServerOptions.CLEANUP_INTERVAL, cleanupInterval);

	long cleanupLowerBound;

	try (
		BlobServer server = new BlobServer(config, new VoidBlobStore());
		final BlobCacheService cache = new BlobCacheService(
			config, new VoidBlobStore(), new InetSocketAddress("localhost", server.getPort())
		)) {
		ConcurrentMap<Tuple2<JobID, TransientBlobKey>, Long> transientBlobExpiryTimes =
			cache.getTransientBlobService().getBlobExpiryTimes();

		server.start();

		final TransientBlobKey key1 =
			(TransientBlobKey) put(server, jobId, data, TRANSIENT_BLOB);
		final TransientBlobKey key2 =
			(TransientBlobKey) put(server, jobId, data2, TRANSIENT_BLOB);

		// access key1, verify expiry times
		cleanupLowerBound = System.currentTimeMillis() + cleanupInterval;
		verifyContents(cache, jobId, key1, data);
		final Long key1ExpiryFirstAccess = transientBlobExpiryTimes.get(Tuple2.of(jobId, key1));
		assertThat(key1ExpiryFirstAccess, greaterThanOrEqualTo(cleanupLowerBound));
		assertNull(transientBlobExpiryTimes.get(Tuple2.of(jobId, key2)));

		// access key2, verify expiry times (delay at least 1ms to also verify key1 expiry is unchanged)
		Thread.sleep(1);
		cleanupLowerBound = System.currentTimeMillis() + cleanupInterval;
		verifyContents(cache, jobId, key2, data2);
		assertEquals(key1ExpiryFirstAccess, transientBlobExpiryTimes.get(Tuple2.of(jobId, key1)));
		assertThat(transientBlobExpiryTimes.get(Tuple2.of(jobId, key2)),
			greaterThanOrEqualTo(cleanupLowerBound));

		// files are cached now for the given TTL - remove from server so that they are not re-downloaded
		if (jobId != null) {
			server.cleanupJob(jobId, true);
		} else {
			server.deleteFromCache(key1);
			server.deleteFromCache(key2);
		}
		checkFileCountForJob(0, jobId, server);

		// cleanup task is run every cleanupInterval seconds
		// => unaccessed file should remain at most 2*cleanupInterval seconds
		// (use 3*cleanupInterval to check that we can still access it)
		final long finishTime = System.currentTimeMillis() + 3 * cleanupInterval;

		final ExecutorService executor = Executors.newFixedThreadPool(numberConcurrentGetOperations);
		for (int i = 0; i < numberConcurrentGetOperations; i++) {
			CompletableFuture<Void> getOperation = CompletableFuture
				.supplyAsync(
					() -> {
						try {
							// constantly access key1 so this should not get deleted
							while (System.currentTimeMillis() < finishTime) {
								get(cache, jobId, key1);
							}

							return null;
						} catch (IOException e) {
							throw new CompletionException(new FlinkException(
								"Could not retrieve blob.", e));
						}
				}, executor);

			getOperations.add(getOperation);
		}

		FutureUtils.ConjunctFuture<Collection<Void>> filesFuture = FutureUtils.combineAll(getOperations);
		filesFuture.get();

		verifyDeletedEventually(server, jobId, key1, key2);
	}
}
 

/**
 * [FLINK-6020] Tests that concurrent get operations don't concurrently access the BlobStore to
 * download a blob.
 *
 * @param jobId
 * 		job ID to use (or <tt>null</tt> if job-unrelated)
 * @param blobType
 * 		whether the BLOB should become permanent or transient
 * @param cacheAccessesHAStore
 * 		whether the cache has access to the {@link BlobServer}'s HA store or not
 */
private void testConcurrentGetOperations(final JobID jobId, final BlobKey.BlobType blobType,
		final boolean cacheAccessesHAStore)
		throws IOException, InterruptedException, ExecutionException {
	final Configuration config = new Configuration();
	config.setString(BlobServerOptions.STORAGE_DIRECTORY, temporaryFolder.newFolder().getAbsolutePath());

	final BlobStore blobStoreServer = mock(BlobStore.class);
	final BlobStore blobStoreCache = mock(BlobStore.class);

	final int numberConcurrentGetOperations = 3;
	final List<CompletableFuture<File>> getOperations = new ArrayList<>(numberConcurrentGetOperations);

	final byte[] data = {1, 2, 3, 4, 99, 42};

	final ExecutorService executor = Executors.newFixedThreadPool(numberConcurrentGetOperations);

	try (
		final BlobServer server = new BlobServer(config, blobStoreServer);
		final BlobCacheService cache = new BlobCacheService(config, cacheAccessesHAStore ? blobStoreServer : blobStoreCache, new InetSocketAddress("localhost", server.getPort())
		)) {

		server.start();

		// upload data first
		final BlobKey blobKey = put(server, jobId, data, blobType);

		// now try accessing it concurrently (only HA mode will be able to retrieve it from HA store!)
		for (int i = 0; i < numberConcurrentGetOperations; i++) {
			CompletableFuture<File> getOperation = CompletableFuture
				.supplyAsync(
					() -> {
						try {
							File file = get(cache, jobId, blobKey);
							// check that we have read the right data
							validateGetAndClose(new FileInputStream(file), data);
							return file;
						} catch (IOException e) {
							throw new CompletionException(new FlinkException(
								"Could not read blob for key " + blobKey + '.', e));
						}
					}, executor);

			getOperations.add(getOperation);
		}

		FutureUtils.ConjunctFuture<Collection<File>> filesFuture = FutureUtils.combineAll(getOperations);

		if (blobType == PERMANENT_BLOB) {
			// wait until all operations have completed and check that no exception was thrown
			filesFuture.get();
		} else {
			// wait for all futures to complete (do not abort on expected exceptions) and check
			// that at least one succeeded
			int completedSuccessfully = 0;
			for (CompletableFuture<File> op : getOperations) {
				try {
					op.get();
					++completedSuccessfully;
				} catch (Throwable t) {
					// transient BLOBs get deleted upon first access and only one request will be successful while all others will have an IOException caused by a FileNotFoundException
					if (!(ExceptionUtils.getRootCause(t) instanceof FileNotFoundException)) {
						// ignore
						org.apache.flink.util.ExceptionUtils.rethrowIOException(t);
					}
				}
			}
			// multiple clients may have accessed the BLOB successfully before it was
			// deleted, but always at least one:
			assertThat(completedSuccessfully, greaterThanOrEqualTo(1));
		}
	} finally {
		executor.shutdownNow();
	}
}
 

/**
 * Shuts down the mini cluster, failing all currently executing jobs.
 * The mini cluster can be started again by calling the {@link #start()} method again.
 *
 * <p>This method shuts down all started services and components,
 * even if an exception occurs in the process of shutting down some component.
 *
 * @return Future which is completed once the MiniCluster has been completely shut down
 */
@Override
public CompletableFuture<Void> closeAsync() {
	synchronized (lock) {
		if (running) {
			LOG.info("Shutting down Flink Mini Cluster");
			try {
				final long shutdownTimeoutMillis = miniClusterConfiguration.getConfiguration().getLong(ClusterOptions.CLUSTER_SERVICES_SHUTDOWN_TIMEOUT);
				final int numComponents = 2 + miniClusterConfiguration.getNumTaskManagers();
				final Collection<CompletableFuture<Void>> componentTerminationFutures = new ArrayList<>(numComponents);

				componentTerminationFutures.addAll(terminateTaskExecutors());

				componentTerminationFutures.add(shutDownResourceManagerComponents());

				final FutureUtils.ConjunctFuture<Void> componentsTerminationFuture = FutureUtils.completeAll(componentTerminationFutures);

				final CompletableFuture<Void> metricSystemTerminationFuture = FutureUtils.composeAfterwards(
					componentsTerminationFuture,
					this::closeMetricSystem);

				final CompletableFuture<Void> rpcServicesTerminationFuture = FutureUtils.composeAfterwards(
					metricSystemTerminationFuture,
					this::terminateRpcServices);

				final CompletableFuture<Void> remainingServicesTerminationFuture = FutureUtils.runAfterwards(
					rpcServicesTerminationFuture,
					this::terminateMiniClusterServices);

				final CompletableFuture<Void> executorsTerminationFuture = FutureUtils.composeAfterwards(
					remainingServicesTerminationFuture,
					() -> terminateExecutors(shutdownTimeoutMillis));

				executorsTerminationFuture.whenComplete(
						(Void ignored, Throwable throwable) -> {
							if (throwable != null) {
								terminationFuture.completeExceptionally(ExceptionUtils.stripCompletionException(throwable));
							} else {
								terminationFuture.complete(null);
							}
						});
			} finally {
				running = false;
			}
		}

		return terminationFuture;
	}
}
 

/**
 * FLINK-5667
 *
 * <p>Tests that a concurrent cancel operation discards the state handles of a not yet
 * acknowledged checkpoint and prevents sending an acknowledge message to the
 * CheckpointCoordinator. The situation can only happen if the cancel call is executed
 * before Environment.acknowledgeCheckpoint().
 */
@Test
public void testAsyncCheckpointingConcurrentCloseBeforeAcknowledge() throws Exception {

	final TestingKeyedStateHandle managedKeyedStateHandle = new TestingKeyedStateHandle();
	final TestingKeyedStateHandle rawKeyedStateHandle = new TestingKeyedStateHandle();
	final TestingOperatorStateHandle managedOperatorStateHandle = new TestingOperatorStateHandle();
	final TestingOperatorStateHandle rawOperatorStateHandle = new TestingOperatorStateHandle();

	final BlockingRunnableFuture<SnapshotResult<KeyedStateHandle>> rawKeyedStateHandleFuture = new BlockingRunnableFuture<>(2, SnapshotResult.of(rawKeyedStateHandle));
	OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(
		DoneFuture.of(SnapshotResult.of(managedKeyedStateHandle)),
		rawKeyedStateHandleFuture,
		DoneFuture.of(SnapshotResult.of(managedOperatorStateHandle)),
		DoneFuture.of(SnapshotResult.of(rawOperatorStateHandle)));

	final StreamOperator<?> streamOperator = streamOperatorWithSnapshot(operatorSnapshotResult);

	final AcknowledgeDummyEnvironment mockEnvironment = new AcknowledgeDummyEnvironment();

	RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(
		mockEnvironment,
		operatorChain(streamOperator)));

	waitTaskIsRunning(task.streamTask, task.invocationFuture);

	final long checkpointId = 42L;
	task.streamTask.triggerCheckpoint(
		new CheckpointMetaData(checkpointId, 1L),
		CheckpointOptions.forCheckpointWithDefaultLocation(),
		false);

	rawKeyedStateHandleFuture.awaitRun();

	task.streamTask.cancel();

	final FutureUtils.ConjunctFuture<Void> discardFuture = FutureUtils.waitForAll(
		Arrays.asList(
			managedKeyedStateHandle.getDiscardFuture(),
			rawKeyedStateHandle.getDiscardFuture(),
			managedOperatorStateHandle.getDiscardFuture(),
			rawOperatorStateHandle.getDiscardFuture()));

	// make sure that all state handles have been discarded
	discardFuture.get();

	try {
		mockEnvironment.getAcknowledgeCheckpointFuture().get(10L, TimeUnit.MILLISECONDS);
		fail("The checkpoint should not get acknowledged.");
	} catch (TimeoutException expected) {
		// future should not be completed
	}

	task.waitForTaskCompletion(true);
}
 

/**
 * [FLINK-6020]
 * Tests that concurrent put operations will only upload the file once to the {@link BlobStore}
 * and that the files are not corrupt at any time.
 *
 * @param jobId
 * 		job ID to use (or <tt>null</tt> if job-unrelated)
 * @param blobType
 * 		whether the BLOB should become permanent or transient
 */
private void testConcurrentPutOperations(
		@Nullable final JobID jobId, final BlobKey.BlobType blobType)
		throws IOException, InterruptedException, ExecutionException {
	final Configuration config = new Configuration();
	config.setString(BlobServerOptions.STORAGE_DIRECTORY, temporaryFolder.newFolder().getAbsolutePath());

	BlobStore blobStore = mock(BlobStore.class);
	int concurrentPutOperations = 2;
	int dataSize = 1024;

	final CountDownLatch countDownLatch = new CountDownLatch(concurrentPutOperations);
	final byte[] data = new byte[dataSize];

	ArrayList<CompletableFuture<BlobKey>> allFutures = new ArrayList<>(concurrentPutOperations);

	ExecutorService executor = Executors.newFixedThreadPool(concurrentPutOperations);

	try (final BlobServer server = new BlobServer(config, blobStore)) {

		server.start();

		for (int i = 0; i < concurrentPutOperations; i++) {
			CompletableFuture<BlobKey> putFuture = CompletableFuture
				.supplyAsync(
					() -> {
						try {
							BlockingInputStream inputStream =
								new BlockingInputStream(countDownLatch, data);
							BlobKey uploadedKey = put(server, jobId, inputStream, blobType);
							// check the uploaded file's contents (concurrently)
							verifyContents(server, jobId, uploadedKey, data);
							return uploadedKey;
						} catch (IOException e) {
							throw new CompletionException(new FlinkException(
								"Could not upload blob.", e));
						}
					},
					executor);

			allFutures.add(putFuture);
		}

		FutureUtils.ConjunctFuture<Collection<BlobKey>> conjunctFuture = FutureUtils.combineAll(allFutures);

		// wait until all operations have completed and check that no exception was thrown
		Collection<BlobKey> blobKeys = conjunctFuture.get();

		Iterator<BlobKey> blobKeyIterator = blobKeys.iterator();

		assertTrue(blobKeyIterator.hasNext());

		BlobKey blobKey = blobKeyIterator.next();

		// make sure that all blob keys are the same
		while (blobKeyIterator.hasNext()) {
			verifyKeyDifferentHashEquals(blobKey, blobKeyIterator.next());
		}

		// check the uploaded file's contents
		verifyContents(server, jobId, blobKey, data);

		// check that we only uploaded the file once to the blob store
		if (blobType == PERMANENT_BLOB) {
			verify(blobStore, times(1)).put(any(File.class), eq(jobId), eq(blobKey));
		} else {
			// can't really verify much in the other cases other than that the put operations should
			// work and not corrupt files
			verify(blobStore, times(0)).put(any(File.class), eq(jobId), eq(blobKey));
		}
	} finally {
		executor.shutdownNow();
	}
}
 

/**
 * Tests that {@link TransientBlobCache} cleans up after a default TTL and keeps files which are
 * constantly accessed.
 */
private void testTransientBlobCleanup(@Nullable final JobID jobId)
		throws IOException, InterruptedException, ExecutionException {

	// 1s should be a safe-enough buffer to still check for existence after a BLOB's last access
	long cleanupInterval = 1L; // in seconds
	final int numberConcurrentGetOperations = 3;

	final List<CompletableFuture<Void>> getOperations = new ArrayList<>(numberConcurrentGetOperations);

	byte[] data = new byte[2000000];
	rnd.nextBytes(data);
	byte[] data2 = Arrays.copyOfRange(data, 10, 54);

	Configuration config = new Configuration();
	config.setString(BlobServerOptions.STORAGE_DIRECTORY,
		temporaryFolder.newFolder().getAbsolutePath());
	config.setLong(BlobServerOptions.CLEANUP_INTERVAL, cleanupInterval);

	long cleanupLowerBound;

	try (
		BlobServer server = new BlobServer(config, new VoidBlobStore());
		final BlobCacheService cache = new BlobCacheService(
			config, new VoidBlobStore(), new InetSocketAddress("localhost", server.getPort())
		)) {
		ConcurrentMap<Tuple2<JobID, TransientBlobKey>, Long> transientBlobExpiryTimes =
			cache.getTransientBlobService().getBlobExpiryTimes();

		server.start();

		final TransientBlobKey key1 =
			(TransientBlobKey) put(server, jobId, data, TRANSIENT_BLOB);
		final TransientBlobKey key2 =
			(TransientBlobKey) put(server, jobId, data2, TRANSIENT_BLOB);

		// access key1, verify expiry times
		cleanupLowerBound = System.currentTimeMillis() + cleanupInterval;
		verifyContents(cache, jobId, key1, data);
		final Long key1ExpiryFirstAccess = transientBlobExpiryTimes.get(Tuple2.of(jobId, key1));
		assertThat(key1ExpiryFirstAccess, greaterThanOrEqualTo(cleanupLowerBound));
		assertNull(transientBlobExpiryTimes.get(Tuple2.of(jobId, key2)));

		// access key2, verify expiry times (delay at least 1ms to also verify key1 expiry is unchanged)
		Thread.sleep(1);
		cleanupLowerBound = System.currentTimeMillis() + cleanupInterval;
		verifyContents(cache, jobId, key2, data2);
		assertEquals(key1ExpiryFirstAccess, transientBlobExpiryTimes.get(Tuple2.of(jobId, key1)));
		assertThat(transientBlobExpiryTimes.get(Tuple2.of(jobId, key2)),
			greaterThanOrEqualTo(cleanupLowerBound));

		// files are cached now for the given TTL - remove from server so that they are not re-downloaded
		if (jobId != null) {
			server.cleanupJob(jobId, true);
		} else {
			server.deleteFromCache(key1);
			server.deleteFromCache(key2);
		}
		checkFileCountForJob(0, jobId, server);

		// cleanup task is run every cleanupInterval seconds
		// => unaccessed file should remain at most 2*cleanupInterval seconds
		// (use 3*cleanupInterval to check that we can still access it)
		final long finishTime = System.currentTimeMillis() + 3 * cleanupInterval;

		final ExecutorService executor = Executors.newFixedThreadPool(numberConcurrentGetOperations);
		for (int i = 0; i < numberConcurrentGetOperations; i++) {
			CompletableFuture<Void> getOperation = CompletableFuture
				.supplyAsync(
					() -> {
						try {
							// constantly access key1 so this should not get deleted
							while (System.currentTimeMillis() < finishTime) {
								get(cache, jobId, key1);
							}

							return null;
						} catch (IOException e) {
							throw new CompletionException(new FlinkException(
								"Could not retrieve blob.", e));
						}
				}, executor);

			getOperations.add(getOperation);
		}

		FutureUtils.ConjunctFuture<Collection<Void>> filesFuture = FutureUtils.combineAll(getOperations);
		filesFuture.get();

		verifyDeletedEventually(server, jobId, key1, key2);
	}
}
 

/**
 * [FLINK-6020] Tests that concurrent get operations don't concurrently access the BlobStore to
 * download a blob.
 *
 * @param jobId
 * 		job ID to use (or <tt>null</tt> if job-unrelated)
 * @param blobType
 * 		whether the BLOB should become permanent or transient
 * @param cacheAccessesHAStore
 * 		whether the cache has access to the {@link BlobServer}'s HA store or not
 */
private void testConcurrentGetOperations(final JobID jobId, final BlobKey.BlobType blobType,
		final boolean cacheAccessesHAStore)
		throws IOException, InterruptedException, ExecutionException {
	final Configuration config = new Configuration();
	config.setString(BlobServerOptions.STORAGE_DIRECTORY, temporaryFolder.newFolder().getAbsolutePath());

	final BlobStore blobStoreServer = mock(BlobStore.class);
	final BlobStore blobStoreCache = mock(BlobStore.class);

	final int numberConcurrentGetOperations = 3;
	final List<CompletableFuture<File>> getOperations = new ArrayList<>(numberConcurrentGetOperations);

	final byte[] data = {1, 2, 3, 4, 99, 42};

	final ExecutorService executor = Executors.newFixedThreadPool(numberConcurrentGetOperations);

	try (
		final BlobServer server = new BlobServer(config, blobStoreServer);
		final BlobCacheService cache = new BlobCacheService(config, cacheAccessesHAStore ? blobStoreServer : blobStoreCache, new InetSocketAddress("localhost", server.getPort())
		)) {

		server.start();

		// upload data first
		final BlobKey blobKey = put(server, jobId, data, blobType);

		// now try accessing it concurrently (only HA mode will be able to retrieve it from HA store!)
		for (int i = 0; i < numberConcurrentGetOperations; i++) {
			CompletableFuture<File> getOperation = CompletableFuture
				.supplyAsync(
					() -> {
						try {
							File file = get(cache, jobId, blobKey);
							// check that we have read the right data
							validateGetAndClose(new FileInputStream(file), data);
							return file;
						} catch (IOException e) {
							throw new CompletionException(new FlinkException(
								"Could not read blob for key " + blobKey + '.', e));
						}
					}, executor);

			getOperations.add(getOperation);
		}

		FutureUtils.ConjunctFuture<Collection<File>> filesFuture = FutureUtils.combineAll(getOperations);

		if (blobType == PERMANENT_BLOB) {
			// wait until all operations have completed and check that no exception was thrown
			filesFuture.get();
		} else {
			// wait for all futures to complete (do not abort on expected exceptions) and check
			// that at least one succeeded
			int completedSuccessfully = 0;
			for (CompletableFuture<File> op : getOperations) {
				try {
					op.get();
					++completedSuccessfully;
				} catch (Throwable t) {
					// transient BLOBs get deleted upon first access and only one request will be successful while all others will have an IOException caused by a FileNotFoundException
					if (!(ExceptionUtils.getRootCause(t) instanceof FileNotFoundException)) {
						// ignore
						org.apache.flink.util.ExceptionUtils.rethrowIOException(t);
					}
				}
			}
			// multiple clients may have accessed the BLOB successfully before it was
			// deleted, but always at least one:
			assertThat(completedSuccessfully, greaterThanOrEqualTo(1));
		}
	} finally {
		executor.shutdownNow();
	}
}
 

/**
 * [FLINK-6020]
 * Tests that concurrent put operations will only upload the file once to the {@link BlobStore}
 * and that the files are not corrupt at any time.
 *
 * @param jobId
 * 		job ID to use (or <tt>null</tt> if job-unrelated)
 * @param blobType
 * 		whether the BLOB should become permanent or transient
 */
private void testConcurrentPutOperations(
		@Nullable final JobID jobId, final BlobKey.BlobType blobType)
		throws IOException, InterruptedException, ExecutionException {
	final Configuration config = new Configuration();
	config.setString(BlobServerOptions.STORAGE_DIRECTORY, temporaryFolder.newFolder().getAbsolutePath());

	BlobStore blobStore = mock(BlobStore.class);
	int concurrentPutOperations = 2;
	int dataSize = 1024;

	final CountDownLatch countDownLatch = new CountDownLatch(concurrentPutOperations);
	final byte[] data = new byte[dataSize];

	ArrayList<CompletableFuture<BlobKey>> allFutures = new ArrayList<>(concurrentPutOperations);

	ExecutorService executor = Executors.newFixedThreadPool(concurrentPutOperations);

	try (final BlobServer server = new BlobServer(config, blobStore)) {

		server.start();

		for (int i = 0; i < concurrentPutOperations; i++) {
			CompletableFuture<BlobKey> putFuture = CompletableFuture
				.supplyAsync(
					() -> {
						try {
							BlockingInputStream inputStream =
								new BlockingInputStream(countDownLatch, data);
							BlobKey uploadedKey = put(server, jobId, inputStream, blobType);
							// check the uploaded file's contents (concurrently)
							verifyContents(server, jobId, uploadedKey, data);
							return uploadedKey;
						} catch (IOException e) {
							throw new CompletionException(new FlinkException(
								"Could not upload blob.", e));
						}
					},
					executor);

			allFutures.add(putFuture);
		}

		FutureUtils.ConjunctFuture<Collection<BlobKey>> conjunctFuture = FutureUtils.combineAll(allFutures);

		// wait until all operations have completed and check that no exception was thrown
		Collection<BlobKey> blobKeys = conjunctFuture.get();

		Iterator<BlobKey> blobKeyIterator = blobKeys.iterator();

		assertTrue(blobKeyIterator.hasNext());

		BlobKey blobKey = blobKeyIterator.next();

		// make sure that all blob keys are the same
		while (blobKeyIterator.hasNext()) {
			verifyKeyDifferentHashEquals(blobKey, blobKeyIterator.next());
		}

		// check the uploaded file's contents
		verifyContents(server, jobId, blobKey, data);

		// check that we only uploaded the file once to the blob store
		if (blobType == PERMANENT_BLOB) {
			verify(blobStore, times(1)).put(any(File.class), eq(jobId), eq(blobKey));
		} else {
			// can't really verify much in the other cases other than that the put operations should
			// work and not corrupt files
			verify(blobStore, times(0)).put(any(File.class), eq(jobId), eq(blobKey));
		}
	} finally {
		executor.shutdownNow();
	}
}
 

/**
 * Tests that {@link TransientBlobCache} cleans up after a default TTL and keeps files which are
 * constantly accessed.
 */
private void testTransientBlobCleanup(@Nullable final JobID jobId)
		throws IOException, InterruptedException, ExecutionException {

	// 1s should be a safe-enough buffer to still check for existence after a BLOB's last access
	long cleanupInterval = 1L; // in seconds
	final int numberConcurrentGetOperations = 3;

	final List<CompletableFuture<Void>> getOperations = new ArrayList<>(numberConcurrentGetOperations);

	byte[] data = new byte[2000000];
	rnd.nextBytes(data);
	byte[] data2 = Arrays.copyOfRange(data, 10, 54);

	Configuration config = new Configuration();
	config.setString(BlobServerOptions.STORAGE_DIRECTORY,
		temporaryFolder.newFolder().getAbsolutePath());
	config.setLong(BlobServerOptions.CLEANUP_INTERVAL, cleanupInterval);

	long cleanupLowerBound;

	try (
		BlobServer server = new BlobServer(config, new VoidBlobStore());
		final BlobCacheService cache = new BlobCacheService(
			config, new VoidBlobStore(), new InetSocketAddress("localhost", server.getPort())
		)) {
		ConcurrentMap<Tuple2<JobID, TransientBlobKey>, Long> transientBlobExpiryTimes =
			cache.getTransientBlobService().getBlobExpiryTimes();

		server.start();

		final TransientBlobKey key1 =
			(TransientBlobKey) put(server, jobId, data, TRANSIENT_BLOB);
		final TransientBlobKey key2 =
			(TransientBlobKey) put(server, jobId, data2, TRANSIENT_BLOB);

		// access key1, verify expiry times
		cleanupLowerBound = System.currentTimeMillis() + cleanupInterval;
		verifyContents(cache, jobId, key1, data);
		final Long key1ExpiryFirstAccess = transientBlobExpiryTimes.get(Tuple2.of(jobId, key1));
		assertThat(key1ExpiryFirstAccess, greaterThanOrEqualTo(cleanupLowerBound));
		assertNull(transientBlobExpiryTimes.get(Tuple2.of(jobId, key2)));

		// access key2, verify expiry times (delay at least 1ms to also verify key1 expiry is unchanged)
		Thread.sleep(1);
		cleanupLowerBound = System.currentTimeMillis() + cleanupInterval;
		verifyContents(cache, jobId, key2, data2);
		assertEquals(key1ExpiryFirstAccess, transientBlobExpiryTimes.get(Tuple2.of(jobId, key1)));
		assertThat(transientBlobExpiryTimes.get(Tuple2.of(jobId, key2)),
			greaterThanOrEqualTo(cleanupLowerBound));

		// files are cached now for the given TTL - remove from server so that they are not re-downloaded
		if (jobId != null) {
			server.cleanupJob(jobId, true);
		} else {
			server.deleteFromCache(key1);
			server.deleteFromCache(key2);
		}
		checkFileCountForJob(0, jobId, server);

		// cleanup task is run every cleanupInterval seconds
		// => unaccessed file should remain at most 2*cleanupInterval seconds
		// (use 3*cleanupInterval to check that we can still access it)
		final long finishTime = System.currentTimeMillis() + 3 * cleanupInterval;

		final ExecutorService executor = Executors.newFixedThreadPool(numberConcurrentGetOperations);
		for (int i = 0; i < numberConcurrentGetOperations; i++) {
			CompletableFuture<Void> getOperation = CompletableFuture
				.supplyAsync(
					() -> {
						try {
							// constantly access key1 so this should not get deleted
							while (System.currentTimeMillis() < finishTime) {
								get(cache, jobId, key1);
							}

							return null;
						} catch (IOException e) {
							throw new CompletionException(new FlinkException(
								"Could not retrieve blob.", e));
						}
				}, executor);

			getOperations.add(getOperation);
		}

		FutureUtils.ConjunctFuture<Collection<Void>> filesFuture = FutureUtils.combineAll(getOperations);
		filesFuture.get();

		verifyDeletedEventually(server, jobId, key1, key2);
	}
}
 

/**
 * [FLINK-6020] Tests that concurrent get operations don't concurrently access the BlobStore to
 * download a blob.
 *
 * @param jobId
 * 		job ID to use (or <tt>null</tt> if job-unrelated)
 * @param blobType
 * 		whether the BLOB should become permanent or transient
 * @param cacheAccessesHAStore
 * 		whether the cache has access to the {@link BlobServer}'s HA store or not
 */
private void testConcurrentGetOperations(final JobID jobId, final BlobKey.BlobType blobType,
		final boolean cacheAccessesHAStore)
		throws IOException, InterruptedException, ExecutionException {
	final Configuration config = new Configuration();
	config.setString(BlobServerOptions.STORAGE_DIRECTORY, temporaryFolder.newFolder().getAbsolutePath());

	final BlobStore blobStoreServer = mock(BlobStore.class);
	final BlobStore blobStoreCache = mock(BlobStore.class);

	final int numberConcurrentGetOperations = 3;
	final List<CompletableFuture<File>> getOperations = new ArrayList<>(numberConcurrentGetOperations);

	final byte[] data = {1, 2, 3, 4, 99, 42};

	final ExecutorService executor = Executors.newFixedThreadPool(numberConcurrentGetOperations);

	try (
		final BlobServer server = new BlobServer(config, blobStoreServer);
		final BlobCacheService cache = new BlobCacheService(config, cacheAccessesHAStore ? blobStoreServer : blobStoreCache, new InetSocketAddress("localhost", server.getPort())
		)) {

		server.start();

		// upload data first
		final BlobKey blobKey = put(server, jobId, data, blobType);

		// now try accessing it concurrently (only HA mode will be able to retrieve it from HA store!)
		for (int i = 0; i < numberConcurrentGetOperations; i++) {
			CompletableFuture<File> getOperation = CompletableFuture
				.supplyAsync(
					() -> {
						try {
							File file = get(cache, jobId, blobKey);
							// check that we have read the right data
							validateGetAndClose(new FileInputStream(file), data);
							return file;
						} catch (IOException e) {
							throw new CompletionException(new FlinkException(
								"Could not read blob for key " + blobKey + '.', e));
						}
					}, executor);

			getOperations.add(getOperation);
		}

		FutureUtils.ConjunctFuture<Collection<File>> filesFuture = FutureUtils.combineAll(getOperations);

		if (blobType == PERMANENT_BLOB) {
			// wait until all operations have completed and check that no exception was thrown
			filesFuture.get();
		} else {
			// wait for all futures to complete (do not abort on expected exceptions) and check
			// that at least one succeeded
			int completedSuccessfully = 0;
			for (CompletableFuture<File> op : getOperations) {
				try {
					op.get();
					++completedSuccessfully;
				} catch (Throwable t) {
					// transient BLOBs get deleted upon first access and only one request will be successful while all others will have an IOException caused by a FileNotFoundException
					if (!(ExceptionUtils.getRootCause(t) instanceof FileNotFoundException)) {
						// ignore
						org.apache.flink.util.ExceptionUtils.rethrowIOException(t);
					}
				}
			}
			// multiple clients may have accessed the BLOB successfully before it was
			// deleted, but always at least one:
			assertThat(completedSuccessfully, greaterThanOrEqualTo(1));
		}
	} finally {
		executor.shutdownNow();
	}
}