Examples with VisibleForTesting org.apache.flink.annotation.VisibleForTesting used on opensource projects

Example 1 with VisibleForTesting

use of org.apache.flink.annotation.VisibleForTesting in project flink by apache.

the class SlotManager method addFreeSlot.

/**
	 * Add free slots directly to the free pool, this will not trigger pending requests allocation
	 *
	 * @param slot The resource slot
	 */
@VisibleForTesting
void addFreeSlot(final ResourceSlot slot) {
    final ResourceID resourceId = slot.getResourceID();
    final SlotID slotId = slot.getSlotId();
    if (!registeredSlots.containsKey(resourceId)) {
        registeredSlots.put(resourceId, new HashMap<SlotID, ResourceSlot>());
    }
    registeredSlots.get(resourceId).put(slot.getSlotId(), slot);
    freeSlots.put(slotId, slot);
}

Example 2 with VisibleForTesting

use of org.apache.flink.annotation.VisibleForTesting in project flink by apache.

the class FlinkKafkaConsumerBase method setSubscribedPartitions.

@VisibleForTesting
void setSubscribedPartitions(List<KafkaTopicPartition> allSubscribedPartitions) {
    checkNotNull(allSubscribedPartitions);
    this.subscribedPartitionsToStartOffsets = new HashMap<>();
    for (KafkaTopicPartition partition : allSubscribedPartitions) {
        this.subscribedPartitionsToStartOffsets.put(partition, null);
    }
}

Example 3 with VisibleForTesting

use of org.apache.flink.annotation.VisibleForTesting in project flink by apache.

the class CheckpointCoordinator method triggerCheckpoint.

@VisibleForTesting
CheckpointTriggerResult triggerCheckpoint(long timestamp, CheckpointProperties props, String targetDirectory, boolean isPeriodic) {
    // Sanity check
    if (props.externalizeCheckpoint() && targetDirectory == null) {
        throw new IllegalStateException("No target directory specified to persist checkpoint to.");
    }
    // make some eager pre-checks
    synchronized (lock) {
        // abort if the coordinator has been shutdown in the meantime
        if (shutdown) {
            return new CheckpointTriggerResult(CheckpointDeclineReason.COORDINATOR_SHUTDOWN);
        }
        // Don't allow periodic checkpoint if scheduling has been disabled
        if (isPeriodic && !periodicScheduling) {
            return new CheckpointTriggerResult(CheckpointDeclineReason.PERIODIC_SCHEDULER_SHUTDOWN);
        }
        // these checks are not relevant for savepoints
        if (!props.forceCheckpoint()) {
            // sanity check: there should never be more than one trigger request queued
            if (triggerRequestQueued) {
                LOG.warn("Trying to trigger another checkpoint while one was queued already");
                return new CheckpointTriggerResult(CheckpointDeclineReason.ALREADY_QUEUED);
            }
            // if too many checkpoints are currently in progress, we need to mark that a request is queued
            if (pendingCheckpoints.size() >= maxConcurrentCheckpointAttempts) {
                triggerRequestQueued = true;
                if (currentPeriodicTrigger != null) {
                    currentPeriodicTrigger.cancel(false);
                    currentPeriodicTrigger = null;
                }
                return new CheckpointTriggerResult(CheckpointDeclineReason.TOO_MANY_CONCURRENT_CHECKPOINTS);
            }
            // make sure the minimum interval between checkpoints has passed
            final long earliestNext = lastCheckpointCompletionNanos + minPauseBetweenCheckpointsNanos;
            final long durationTillNextMillis = (earliestNext - System.nanoTime()) / 1_000_000;
            if (durationTillNextMillis > 0) {
                if (currentPeriodicTrigger != null) {
                    currentPeriodicTrigger.cancel(false);
                    currentPeriodicTrigger = null;
                }
                // Reassign the new trigger to the currentPeriodicTrigger
                currentPeriodicTrigger = timer.scheduleAtFixedRate(new ScheduledTrigger(), durationTillNextMillis, baseInterval, TimeUnit.MILLISECONDS);
                return new CheckpointTriggerResult(CheckpointDeclineReason.MINIMUM_TIME_BETWEEN_CHECKPOINTS);
            }
        }
    }
    // check if all tasks that we need to trigger are running.
    // if not, abort the checkpoint
    Execution[] executions = new Execution[tasksToTrigger.length];
    for (int i = 0; i < tasksToTrigger.length; i++) {
        Execution ee = tasksToTrigger[i].getCurrentExecutionAttempt();
        if (ee != null && ee.getState() == ExecutionState.RUNNING) {
            executions[i] = ee;
        } else {
            LOG.info("Checkpoint triggering task {} is not being executed at the moment. Aborting checkpoint.", tasksToTrigger[i].getSimpleName());
            return new CheckpointTriggerResult(CheckpointDeclineReason.NOT_ALL_REQUIRED_TASKS_RUNNING);
        }
    }
    // next, check if all tasks that need to acknowledge the checkpoint are running.
    // if not, abort the checkpoint
    Map<ExecutionAttemptID, ExecutionVertex> ackTasks = new HashMap<>(tasksToWaitFor.length);
    for (ExecutionVertex ev : tasksToWaitFor) {
        Execution ee = ev.getCurrentExecutionAttempt();
        if (ee != null) {
            ackTasks.put(ee.getAttemptId(), ev);
        } else {
            LOG.info("Checkpoint acknowledging task {} is not being executed at the moment. Aborting checkpoint.", ev.getSimpleName());
            return new CheckpointTriggerResult(CheckpointDeclineReason.NOT_ALL_REQUIRED_TASKS_RUNNING);
        }
    }
    // we avoid blocking the processing of 'acknowledge/decline' messages during that time.
    synchronized (triggerLock) {
        final long checkpointID;
        try {
            // this must happen outside the coordinator-wide lock, because it communicates
            // with external services (in HA mode) and may block for a while.
            checkpointID = checkpointIdCounter.getAndIncrement();
        } catch (Throwable t) {
            int numUnsuccessful = numUnsuccessfulCheckpointsTriggers.incrementAndGet();
            LOG.warn("Failed to trigger checkpoint (" + numUnsuccessful + " consecutive failed attempts so far)", t);
            return new CheckpointTriggerResult(CheckpointDeclineReason.EXCEPTION);
        }
        final PendingCheckpoint checkpoint = new PendingCheckpoint(job, checkpointID, timestamp, ackTasks, props, targetDirectory, executor);
        if (statsTracker != null) {
            PendingCheckpointStats callback = statsTracker.reportPendingCheckpoint(checkpointID, timestamp, props);
            checkpoint.setStatsCallback(callback);
        }
        // schedule the timer that will clean up the expired checkpoints
        final Runnable canceller = new Runnable() {

            @Override
            public void run() {
                synchronized (lock) {
                    // note that checkpoint completion discards the pending checkpoint object
                    if (!checkpoint.isDiscarded()) {
                        LOG.info("Checkpoint " + checkpointID + " expired before completing.");
                        checkpoint.abortExpired();
                        pendingCheckpoints.remove(checkpointID);
                        rememberRecentCheckpointId(checkpointID);
                        triggerQueuedRequests();
                    }
                }
            }
        };
        try {
            // re-acquire the coordinator-wide lock
            synchronized (lock) {
                // that the conditions still hold.
                if (shutdown) {
                    return new CheckpointTriggerResult(CheckpointDeclineReason.COORDINATOR_SHUTDOWN);
                } else if (!props.forceCheckpoint()) {
                    if (triggerRequestQueued) {
                        LOG.warn("Trying to trigger another checkpoint while one was queued already");
                        return new CheckpointTriggerResult(CheckpointDeclineReason.ALREADY_QUEUED);
                    }
                    if (pendingCheckpoints.size() >= maxConcurrentCheckpointAttempts) {
                        triggerRequestQueued = true;
                        if (currentPeriodicTrigger != null) {
                            currentPeriodicTrigger.cancel(false);
                            currentPeriodicTrigger = null;
                        }
                        return new CheckpointTriggerResult(CheckpointDeclineReason.TOO_MANY_CONCURRENT_CHECKPOINTS);
                    }
                    // make sure the minimum interval between checkpoints has passed
                    final long earliestNext = lastCheckpointCompletionNanos + minPauseBetweenCheckpointsNanos;
                    final long durationTillNextMillis = (earliestNext - System.nanoTime()) / 1_000_000;
                    if (durationTillNextMillis > 0) {
                        if (currentPeriodicTrigger != null) {
                            currentPeriodicTrigger.cancel(false);
                            currentPeriodicTrigger = null;
                        }
                        // Reassign the new trigger to the currentPeriodicTrigger
                        currentPeriodicTrigger = timer.scheduleAtFixedRate(new ScheduledTrigger(), durationTillNextMillis, baseInterval, TimeUnit.MILLISECONDS);
                        return new CheckpointTriggerResult(CheckpointDeclineReason.MINIMUM_TIME_BETWEEN_CHECKPOINTS);
                    }
                }
                LOG.info("Triggering checkpoint " + checkpointID + " @ " + timestamp);
                pendingCheckpoints.put(checkpointID, checkpoint);
                ScheduledFuture<?> cancellerHandle = timer.schedule(canceller, checkpointTimeout, TimeUnit.MILLISECONDS);
                if (!checkpoint.setCancellerHandle(cancellerHandle)) {
                    // checkpoint is already disposed!
                    cancellerHandle.cancel(false);
                }
            }
            // end of lock scope
            CheckpointOptions checkpointOptions;
            if (!props.isSavepoint()) {
                checkpointOptions = CheckpointOptions.forFullCheckpoint();
            } else {
                checkpointOptions = CheckpointOptions.forSavepoint(targetDirectory);
            }
            // send the messages to the tasks that trigger their checkpoint
            for (Execution execution : executions) {
                execution.triggerCheckpoint(checkpointID, timestamp, checkpointOptions);
            }
            numUnsuccessfulCheckpointsTriggers.set(0);
            return new CheckpointTriggerResult(checkpoint);
        } catch (Throwable t) {
            // guard the map against concurrent modifications
            synchronized (lock) {
                pendingCheckpoints.remove(checkpointID);
            }
            int numUnsuccessful = numUnsuccessfulCheckpointsTriggers.incrementAndGet();
            LOG.warn("Failed to trigger checkpoint (" + numUnsuccessful + " consecutive failed attempts so far)", t);
            if (!checkpoint.isDiscarded()) {
                checkpoint.abortError(new Exception("Failed to trigger checkpoint"));
            }
            return new CheckpointTriggerResult(CheckpointDeclineReason.EXCEPTION);
        }
    }
// end trigger lock
}

Example 4 with VisibleForTesting

use of org.apache.flink.annotation.VisibleForTesting in project flink by apache.

the class KinesisFirehoseSinkBuilder method getClientPropertiesWithDefaultHttpProtocol.

@VisibleForTesting
Properties getClientPropertiesWithDefaultHttpProtocol() {
    Properties clientProperties = Optional.ofNullable(firehoseClientProperties).orElse(new Properties());
    clientProperties.putIfAbsent(HTTP_PROTOCOL_VERSION, DEFAULT_HTTP_PROTOCOL.toString());
    return clientProperties;
}

Example 5 with VisibleForTesting

use of org.apache.flink.annotation.VisibleForTesting in project flink by apache.

the class JarUploadHandler method handleRequest.

@Override
@VisibleForTesting
public CompletableFuture<JarUploadResponseBody> handleRequest(@Nonnull final HandlerRequest<EmptyRequestBody> request, @Nonnull final RestfulGateway gateway) throws RestHandlerException {
    Collection<File> uploadedFiles = request.getUploadedFiles();
    if (uploadedFiles.size() != 1) {
        throw new RestHandlerException("Exactly 1 file must be sent, received " + uploadedFiles.size() + '.', HttpResponseStatus.BAD_REQUEST);
    }
    final Path fileUpload = uploadedFiles.iterator().next().toPath();
    return CompletableFuture.supplyAsync(() -> {
        if (!fileUpload.getFileName().toString().endsWith(".jar")) {
            throw new CompletionException(new RestHandlerException("Only Jar files are allowed.", HttpResponseStatus.BAD_REQUEST));
        } else {
            final Path destination = jarDir.resolve(UUID.randomUUID() + "_" + fileUpload.getFileName());
            try {
                Files.move(fileUpload, destination);
            } catch (IOException e) {
                throw new CompletionException(new RestHandlerException(String.format("Could not move uploaded jar file [%s] to [%s].", fileUpload, destination), HttpResponseStatus.INTERNAL_SERVER_ERROR, e));
            }
            return new JarUploadResponseBody(destination.normalize().toString());
        }
    }, executor);
}