Java Code Examples for org.apache.flink.annotation.VisibleForTesting

@VisibleForTesting
void startTaskExecutor() throws Exception {
	synchronized (lock) {
		final Configuration configuration = miniClusterConfiguration.getConfiguration();

		final TaskExecutor taskExecutor = TaskManagerRunner.startTaskManager(
			configuration,
			new ResourceID(UUID.randomUUID().toString()),
			taskManagerRpcServiceFactory.createRpcService(),
			haServices,
			heartbeatServices,
			metricRegistry,
			blobCacheService,
			useLocalCommunication(),
			taskManagerTerminatingFatalErrorHandlerFactory.create(taskManagers.size()));

		taskExecutor.start();
		taskManagers.add(taskExecutor);
	}
}
 

/**
 * Recovers all jobs persisted via the submitted job graph store.
 */
@VisibleForTesting
Collection<JobGraph> recoverJobs() throws Exception {
	log.info("Recovering all persisted jobs.");
	final Collection<JobID> jobIds = submittedJobGraphStore.getJobIds();

	try {
		return recoverJobGraphs(jobIds);
	} catch (Exception e) {
		// release all recovered job graphs
		for (JobID jobId : jobIds) {
			try {
				submittedJobGraphStore.releaseJobGraph(jobId);
			} catch (Exception ie) {
				e.addSuppressed(ie);
			}
		}
		throw e;
	}
}
 

@VisibleForTesting
protected Collection<? extends DispatcherResourceManagerComponent<?>> createDispatcherResourceManagerComponents(
		Configuration configuration,
		RpcServiceFactory rpcServiceFactory,
		HighAvailabilityServices haServices,
		BlobServer blobServer,
		HeartbeatServices heartbeatServices,
		MetricRegistry metricRegistry,
		MetricQueryServiceRetriever metricQueryServiceRetriever,
		FatalErrorHandler fatalErrorHandler) throws Exception {
	SessionDispatcherResourceManagerComponentFactory dispatcherResourceManagerComponentFactory = createDispatcherResourceManagerComponentFactory();
	return Collections.singleton(
		dispatcherResourceManagerComponentFactory.create(
			configuration,
			rpcServiceFactory.createRpcService(),
			haServices,
			blobServer,
			heartbeatServices,
			metricRegistry,
			new MemoryArchivedExecutionGraphStore(),
			metricQueryServiceRetriever,
			fatalErrorHandler));
}
 

/**
 * Convenience constructor for testing.
 */
@VisibleForTesting
ExecutionJobVertex(
		ExecutionGraph graph,
		JobVertex jobVertex,
		int defaultParallelism,
		Time timeout) throws JobException {

	this(
		graph,
		jobVertex,
		defaultParallelism,
		JobManagerOptions.MAX_ATTEMPTS_HISTORY_SIZE.defaultValue(),
		timeout,
		1L,
		System.currentTimeMillis());
}
 

@VisibleForTesting
ExecutionGraph(
		JobInformation jobInformation,
		ScheduledExecutorService futureExecutor,
		Executor ioExecutor,
		Time timeout,
		RestartStrategy restartStrategy,
		FailoverStrategy.Factory failoverStrategy,
		SlotProvider slotProvider) throws IOException {
	this(
		jobInformation,
		futureExecutor,
		ioExecutor,
		timeout,
		restartStrategy,
		failoverStrategy,
		slotProvider,
		ExecutionGraph.class.getClassLoader(),
		VoidBlobWriter.getInstance(),
		timeout);
}
 

@VisibleForTesting
void requestPartitions() throws IOException, InterruptedException {
	synchronized (requestLock) {
		if (!requestedPartitionsFlag) {
			if (closeFuture.isDone()) {
				throw new IllegalStateException("Already released.");
			}

			// Sanity checks
			if (numberOfInputChannels != inputChannels.size()) {
				throw new IllegalStateException(String.format(
					"Bug in input gate setup logic: mismatch between " +
					"number of total input channels [%s] and the currently set number of input " +
					"channels [%s].",
					inputChannels.size(),
					numberOfInputChannels));
			}

			for (InputChannel inputChannel : inputChannels.values()) {
				inputChannel.requestSubpartition(consumedSubpartitionIndex);
			}
		}

		requestedPartitionsFlag = true;
	}
}
 

@VisibleForTesting
@SuppressWarnings("unchecked")
@Override
public int numKeyValueStateEntries() {
	int count = 0;

	for (RocksDbKvStateInfo metaInfo : kvStateInformation.values()) {
		//TODO maybe filterOrTransform only for k/v states
		try (RocksIteratorWrapper rocksIterator = RocksDBOperationUtils.getRocksIterator(db, metaInfo.columnFamilyHandle)) {
			rocksIterator.seekToFirst();

			while (rocksIterator.isValid()) {
				count++;
				rocksIterator.next();
			}
		}
	}

	return count;
}
 

@VisibleForTesting
public RefCountedBufferingFileStream(
		final RefCountedFile file,
		final int bufferSize) {

	checkArgument(bufferSize > 0L);

	this.currentTmpFile = checkNotNull(file);
	this.buffer = new byte[bufferSize];
	this.positionInBuffer = 0;
	this.closed = false;
}
 

@VisibleForTesting
S3RecoverableWriter(
		final S3AccessHelper s3AccessHelper,
		final S3RecoverableMultipartUploadFactory uploadFactory,
		final FunctionWithException<File, RefCountedFile, IOException> tempFileCreator,
		final long userDefinedMinPartSize) {

	this.s3AccessHelper = checkNotNull(s3AccessHelper);
	this.uploadFactory = checkNotNull(uploadFactory);
	this.tempFileCreator = checkNotNull(tempFileCreator);
	this.userDefinedMinPartSize = userDefinedMinPartSize;
}
 

/**
 * Build the {@link BulkProcessor}.
 *
 * <p>Note: this is exposed for testing purposes.
 */
@VisibleForTesting
protected BulkProcessor buildBulkProcessor(BulkProcessor.Listener listener) {
	checkNotNull(listener);

	BulkProcessor.Builder bulkProcessorBuilder = callBridge.createBulkProcessorBuilder(client, listener);

	// This makes flush() blocking
	bulkProcessorBuilder.setConcurrentRequests(0);

	if (bulkProcessorFlushMaxActions != null) {
		bulkProcessorBuilder.setBulkActions(bulkProcessorFlushMaxActions);
	}

	if (bulkProcessorFlushMaxSizeMb != null) {
		bulkProcessorBuilder.setBulkSize(new ByteSizeValue(bulkProcessorFlushMaxSizeMb, ByteSizeUnit.MB));
	}

	if (bulkProcessorFlushIntervalMillis != null) {
		bulkProcessorBuilder.setFlushInterval(TimeValue.timeValueMillis(bulkProcessorFlushIntervalMillis));
	}

	// if backoff retrying is disabled, bulkProcessorFlushBackoffPolicy will be null
	callBridge.configureBulkProcessorBackoff(bulkProcessorBuilder, bulkProcessorFlushBackoffPolicy);

	return bulkProcessorBuilder.build();
}
 

@VisibleForTesting
@Override
public boolean equals(Object o) {
	if (this == o) {
		return true;
	}
	if (o == null || getClass() != o.getClass()) {
		return false;
	}
	ListAggWsWithRetractAccumulator that = (ListAggWsWithRetractAccumulator) o;
	return Objects.equals(list, that.list) &&
		Objects.equals(retractList, that.retractList) &&
		Objects.equals(delimiter, that.delimiter);
}
 

/**
 * Computes and returns a set with the prior allocation ids from all execution vertices scheduled together.
 *
 * @param executionVertexSchedulingRequirements contains the execution vertices which are scheduled together
 */
@VisibleForTesting
static Set<AllocationID> computeAllPriorAllocationIds(
		Collection<ExecutionVertexSchedulingRequirements> executionVertexSchedulingRequirements) {
	return executionVertexSchedulingRequirements
		.stream()
		.map(ExecutionVertexSchedulingRequirements::getPreviousAllocationId)
		.filter(Objects::nonNull)
		.collect(Collectors.toSet());
}
 

@VisibleForTesting
@Nonnull
protected Collection<DispatcherResourceManagerComponent<?>> getDispatcherResourceManagerComponents() {
	synchronized (lock) {
		return Collections.unmodifiableCollection(dispatcherResourceManagerComponents);
	}
}
 

@VisibleForTesting
protected KinesisDataFetcher(List<String> streams,
							SourceFunction.SourceContext<T> sourceContext,
							Object checkpointLock,
							RuntimeContext runtimeContext,
							Properties configProps,
							KinesisDeserializationSchema<T> deserializationSchema,
							KinesisShardAssigner shardAssigner,
							AssignerWithPeriodicWatermarks<T> periodicWatermarkAssigner,
							WatermarkTracker watermarkTracker,
							AtomicReference<Throwable> error,
							List<KinesisStreamShardState> subscribedShardsState,
							HashMap<String, String> subscribedStreamsToLastDiscoveredShardIds,
							FlinkKinesisProxyFactory kinesisProxyFactory) {
	this.streams = checkNotNull(streams);
	this.configProps = checkNotNull(configProps);
	this.sourceContext = checkNotNull(sourceContext);
	this.checkpointLock = checkNotNull(checkpointLock);
	this.runtimeContext = checkNotNull(runtimeContext);
	this.totalNumberOfConsumerSubtasks = runtimeContext.getNumberOfParallelSubtasks();
	this.indexOfThisConsumerSubtask = runtimeContext.getIndexOfThisSubtask();
	this.deserializationSchema = checkNotNull(deserializationSchema);
	this.shardAssigner = checkNotNull(shardAssigner);
	this.periodicWatermarkAssigner = periodicWatermarkAssigner;
	this.watermarkTracker = watermarkTracker;
	this.kinesisProxyFactory = checkNotNull(kinesisProxyFactory);
	this.kinesis = kinesisProxyFactory.create(configProps);

	this.consumerMetricGroup = runtimeContext.getMetricGroup()
		.addGroup(KinesisConsumerMetricConstants.KINESIS_CONSUMER_METRICS_GROUP);

	this.error = checkNotNull(error);
	this.subscribedShardsState = checkNotNull(subscribedShardsState);
	this.subscribedStreamsToLastDiscoveredShardIds = checkNotNull(subscribedStreamsToLastDiscoveredShardIds);

	this.shardConsumersExecutor =
		createShardConsumersThreadPool(runtimeContext.getTaskNameWithSubtasks());
	this.recordEmitter = createRecordEmitter(configProps);
}
 

@VisibleForTesting
public StateSnapshotContextSynchronousImpl(long checkpointId, long checkpointTimestamp) {
	this.checkpointId = checkpointId;
	this.checkpointTimestamp = checkpointTimestamp;
	this.streamFactory = null;
	this.keyGroupRange = KeyGroupRange.EMPTY_KEY_GROUP_RANGE;
	this.closableRegistry = new CloseableRegistry();
}
 

@VisibleForTesting
Map<N, Map<K, S>> getMapForKeyGroup(int keyGroupIndex) {
	final int pos = indexToOffset(keyGroupIndex);
	if (pos >= 0 && pos < state.length) {
		return state[pos];
	} else {
		return null;
	}
}
 

@VisibleForTesting
public int numEventTimeTimers() {
	int count = 0;
	for (InternalTimerServiceImpl<?, ?> timerService : timerServices.values()) {
		count += timerService.numEventTimeTimers();
	}
	return count;
}
 

@VisibleForTesting
BucketState<BucketID> deserializeV1(DataInputView in) throws IOException {
	final BucketID bucketId = SimpleVersionedSerialization.readVersionAndDeSerialize(bucketIdSerializer, in);
	final String bucketPathStr = in.readUTF();
	final long creationTime = in.readLong();

	// then get the current resumable stream
	RecoverableWriter.ResumeRecoverable current = null;
	if (in.readBoolean()) {
		current = SimpleVersionedSerialization.readVersionAndDeSerialize(resumableSerializer, in);
	}

	final int committableVersion = in.readInt();
	final int numCheckpoints = in.readInt();
	final HashMap<Long, List<RecoverableWriter.CommitRecoverable>> resumablesPerCheckpoint = new HashMap<>(numCheckpoints);

	for (int i = 0; i < numCheckpoints; i++) {
		final long checkpointId = in.readLong();
		final int noOfResumables = in.readInt();

		final List<RecoverableWriter.CommitRecoverable> resumables = new ArrayList<>(noOfResumables);
		for (int j = 0; j < noOfResumables; j++) {
			final byte[] bytes = new byte[in.readInt()];
			in.readFully(bytes);
			resumables.add(commitableSerializer.deserialize(committableVersion, bytes));
		}
		resumablesPerCheckpoint.put(checkpointId, resumables);
	}

	return new BucketState<>(
			bucketId,
			new Path(bucketPathStr),
			creationTime,
			current,
			resumablesPerCheckpoint);
}
 

@VisibleForTesting
static ShuffleDescriptor getConsumedPartitionShuffleDescriptor(
		ResultPartitionID consumedPartitionId,
		ResultPartitionType resultPartitionType,
		boolean isConsumable,
		ExecutionState producerState,
		boolean allowUnknownPartitions,
		@Nullable ResultPartitionDeploymentDescriptor consumedPartitionDescriptor) {
	// The producing task needs to be RUNNING or already FINISHED
	if ((resultPartitionType.isPipelined() || isConsumable) &&
		consumedPartitionDescriptor != null &&
		isProducerAvailable(producerState)) {
		// partition is already registered
		return consumedPartitionDescriptor.getShuffleDescriptor();
	}
	else if (allowUnknownPartitions) {
		// The producing task might not have registered the partition yet
		return new UnknownShuffleDescriptor(consumedPartitionId);
	}
	else {
		// throw respective exceptions
		handleConsumedPartitionShuffleDescriptorErrors(
			consumedPartitionId,
			resultPartitionType,
			isConsumable,
			producerState);
		return null; // should never happen
	}
}
 

@VisibleForTesting
protected void cancel() {
	closeSnapshotIO();
	if (resourceCleanupOwnershipTaken.compareAndSet(false, true)) {
		cleanup();
	}
}
 

@VisibleForTesting
public static void serializeOperatorStateHandle(
	OperatorStateHandle stateHandle, DataOutputStream dos) throws IOException {

	if (stateHandle != null) {
		dos.writeByte(PARTITIONABLE_OPERATOR_STATE_HANDLE);
		Map<String, OperatorStateHandle.StateMetaInfo> partitionOffsetsMap =
				stateHandle.getStateNameToPartitionOffsets();
		dos.writeInt(partitionOffsetsMap.size());
		for (Map.Entry<String, OperatorStateHandle.StateMetaInfo> entry : partitionOffsetsMap.entrySet()) {
			dos.writeUTF(entry.getKey());

			OperatorStateHandle.StateMetaInfo stateMetaInfo = entry.getValue();

			int mode = stateMetaInfo.getDistributionMode().ordinal();
			dos.writeByte(mode);

			long[] offsets = stateMetaInfo.getOffsets();
			dos.writeInt(offsets.length);
			for (long offset : offsets) {
				dos.writeLong(offset);
			}
		}
		serializeStreamStateHandle(stateHandle.getDelegateStateHandle(), dos);
	} else {
		dos.writeByte(NULL_HANDLE);
	}
}
 

@VisibleForTesting
public JobStatus waitUntilTerminal() throws InterruptedException {
	try {
		return terminationFuture.get();
	}
	catch (ExecutionException e) {
		// this should never happen
		// it would be a bug, so we  don't expect this to be handled and throw
		// an unchecked exception here
		throw new RuntimeException(e);
	}
}
 

@VisibleForTesting
int getNumTasksScheduled() {
	BlockingQueue<?> queue = timerService.getQueue();
	if (queue == null) {
		return 0;
	} else {
		return queue.size();
	}
}
 

@VisibleForTesting
public RefCountedBufferingFileStream(
		final RefCountedFile file,
		final int bufferSize) {

	checkArgument(bufferSize > 0L);

	this.currentTmpFile = checkNotNull(file);
	this.buffer = new byte[bufferSize];
	this.positionInBuffer = 0;
	this.closed = false;
}
 

/**
 * This constructor is for tests only, because it sets default values for many fields.
 */
@VisibleForTesting
ExecutionGraph(
		ScheduledExecutorService futureExecutor,
		Executor ioExecutor,
		JobID jobId,
		String jobName,
		Configuration jobConfig,
		SerializedValue<ExecutionConfig> serializedConfig,
		Time timeout,
		RestartStrategy restartStrategy,
		SlotProvider slotProvider) throws IOException {

	this(
		new JobInformation(
			jobId,
			jobName,
			serializedConfig,
			jobConfig,
			Collections.emptyList(),
			Collections.emptyList()),
		futureExecutor,
		ioExecutor,
		timeout,
		restartStrategy,
		slotProvider);
}
 

@VisibleForTesting
@Nullable
public Instance getInstance(ResourceID resourceId) {
	for (Instance instance : allInstances) {
		if (Objects.equals(resourceId, instance.getTaskManagerID())) {
			return instance;
		}
	}
	return null;
}
 

/**
 * Check whether the service monitors the given job.
 *
 * @param jobId identifying the job
 * @return True if the given job is monitored; otherwise false
 */
@VisibleForTesting
public boolean containsJob(JobID jobId) {
	Preconditions.checkState(JobLeaderService.State.STARTED == state, "The service is currently not running.");

	return jobLeaderServices.containsKey(jobId);
}
 

@VisibleForTesting
public PartitionDescriptor(
		IntermediateDataSetID resultId,
		IntermediateResultPartitionID partitionId,
		ResultPartitionType partitionType,
		int numberOfSubpartitions,
		int connectionIndex) {
	this.resultId = checkNotNull(resultId);
	this.partitionId = checkNotNull(partitionId);
	this.partitionType = checkNotNull(partitionType);
	checkArgument(numberOfSubpartitions >= 1);
	this.numberOfSubpartitions = numberOfSubpartitions;
	this.connectionIndex = connectionIndex;
}
 

@VisibleForTesting
public static OperatorStateHandle deserializeOperatorStateHandle(
		DataInputStream dis) throws IOException {

	final int type = dis.readByte();
	if (NULL_HANDLE == type) {
		return null;
	} else if (PARTITIONABLE_OPERATOR_STATE_HANDLE == type) {
		int mapSize = dis.readInt();
		Map<String, OperatorStateHandle.StateMetaInfo> offsetsMap = new HashMap<>(mapSize);
		for (int i = 0; i < mapSize; ++i) {
			String key = dis.readUTF();

			int modeOrdinal = dis.readByte();
			OperatorStateHandle.Mode mode = OperatorStateHandle.Mode.values()[modeOrdinal];

			long[] offsets = new long[dis.readInt()];
			for (int j = 0; j < offsets.length; ++j) {
				offsets[j] = dis.readLong();
			}

			OperatorStateHandle.StateMetaInfo metaInfo =
					new OperatorStateHandle.StateMetaInfo(offsets, mode);
			offsetsMap.put(key, metaInfo);
		}
		StreamStateHandle stateHandle = deserializeStreamStateHandle(dis);
		return new OperatorStreamStateHandle(offsetsMap, stateHandle);
	} else {
		throw new IllegalStateException("Reading invalid OperatorStateHandle, type: " + type);
	}
}
 

@VisibleForTesting
static void setDefaultExecutionModeIfNotConfigured(Configuration configuration) {
	if (isNoExecutionModeConfigured(configuration)) {
		// In contrast to other places, the default for standalone job clusters is ExecutionMode.DETACHED
		configuration.setString(ClusterEntrypoint.EXECUTION_MODE, ExecutionMode.DETACHED.toString());
	}
}