software.amazon.kinesis.lifecycle.ShardConsumer Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of amazon-kinesis-client Show documentation
Show all versions of amazon-kinesis-client Show documentation
The Amazon Kinesis Client Library for Java enables Java developers to easily consume and process data
from Amazon Kinesis.
/*
* Copyright 2019 Amazon.com, Inc. or its affiliates.
* Licensed under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package software.amazon.kinesis.lifecycle;
import java.time.Duration;
import java.time.Instant;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.RejectedExecutionException;
import com.google.common.annotations.VisibleForTesting;
import lombok.AccessLevel;
import lombok.Getter;
import lombok.NonNull;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Subscription;
import software.amazon.kinesis.annotations.KinesisClientInternalApi;
import software.amazon.kinesis.exceptions.internal.BlockedOnParentShardException;
import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.lifecycle.events.ProcessRecordsInput;
import software.amazon.kinesis.lifecycle.events.TaskExecutionListenerInput;
import software.amazon.kinesis.retrieval.RecordsPublisher;
/**
* Responsible for consuming data records of a (specified) shard.
* The instance should be shutdown when we lose the primary responsibility for a shard.
* A new instance should be created if the primary responsibility is reassigned back to this process.
*/
@Getter(AccessLevel.PACKAGE)
@Accessors(fluent = true)
@Slf4j
@KinesisClientInternalApi
public class ShardConsumer {
public static final int MAX_TIME_BETWEEN_REQUEST_RESPONSE = 60 * 1000;
private final RecordsPublisher recordsPublisher;
private final ExecutorService executorService;
private final ShardInfo shardInfo;
private final ShardConsumerArgument shardConsumerArgument;
@NonNull
private final Optional logWarningForTaskAfterMillis;
private final int bufferSize;
private final TaskExecutionListener taskExecutionListener;
private final String streamIdentifier;
private ConsumerTask currentTask;
private TaskOutcome taskOutcome;
private CompletableFuture stateChangeFuture;
private boolean needsInitialization = true;
private volatile Instant taskDispatchedAt;
private volatile boolean taskIsRunning = false;
/**
* Tracks current state. It is only updated via the consumeStream/shutdown APIs. Therefore we don't do
* much coordination/synchronization to handle concurrent reads/updates.
*/
private ConsumerState currentState;
private final Object shutdownLock = new Object();
@Getter(AccessLevel.PUBLIC)
private volatile ShutdownReason shutdownReason;
private volatile ShutdownNotification shutdownNotification;
private final ShardConsumerSubscriber subscriber;
private ProcessRecordsInput shardEndProcessRecordsInput;
private final ConsumerTaskFactory taskFactory;
public ShardConsumer(
RecordsPublisher recordsPublisher,
ExecutorService executorService,
ShardInfo shardInfo,
Optional logWarningForTaskAfterMillis,
ShardConsumerArgument shardConsumerArgument,
TaskExecutionListener taskExecutionListener,
int readTimeoutsToIgnoreBeforeWarning) {
this(
recordsPublisher,
executorService,
shardInfo,
logWarningForTaskAfterMillis,
shardConsumerArgument,
ConsumerStates.INITIAL_STATE,
8,
taskExecutionListener,
readTimeoutsToIgnoreBeforeWarning,
new KinesisConsumerTaskFactory());
}
//
// TODO: Make bufferSize configurable
//
public ShardConsumer(
RecordsPublisher recordsPublisher,
ExecutorService executorService,
ShardInfo shardInfo,
Optional logWarningForTaskAfterMillis,
ShardConsumerArgument shardConsumerArgument,
ConsumerState initialState,
int bufferSize,
TaskExecutionListener taskExecutionListener,
int readTimeoutsToIgnoreBeforeWarning) {
this(
recordsPublisher,
executorService,
shardInfo,
logWarningForTaskAfterMillis,
shardConsumerArgument,
initialState,
bufferSize,
taskExecutionListener,
readTimeoutsToIgnoreBeforeWarning,
new KinesisConsumerTaskFactory());
}
public ShardConsumer(
RecordsPublisher recordsPublisher,
ExecutorService executorService,
ShardInfo shardInfo,
Optional logWarningForTaskAfterMillis,
ShardConsumerArgument shardConsumerArgument,
ConsumerState initialState,
int bufferSize,
TaskExecutionListener taskExecutionListener,
int readTimeoutsToIgnoreBeforeWarning,
ConsumerTaskFactory taskFactory) {
this.recordsPublisher = recordsPublisher;
this.executorService = executorService;
this.shardInfo = shardInfo;
this.streamIdentifier = shardInfo.streamIdentifierSerOpt().orElse("single_stream_mode");
this.shardConsumerArgument = shardConsumerArgument;
this.logWarningForTaskAfterMillis = logWarningForTaskAfterMillis;
this.taskExecutionListener = taskExecutionListener;
this.currentState = initialState;
subscriber = new ShardConsumerSubscriber(
recordsPublisher, executorService, bufferSize, this, readTimeoutsToIgnoreBeforeWarning);
this.bufferSize = bufferSize;
if (this.shardInfo.isCompleted()) {
markForShutdown(ShutdownReason.SHARD_END);
}
this.taskFactory = taskFactory;
}
synchronized void handleInput(ProcessRecordsInput input, Subscription subscription) {
if (isShutdownRequested()) {
subscription.cancel();
return;
}
processData(input);
if (taskOutcome == TaskOutcome.END_OF_SHARD) {
markForShutdown(ShutdownReason.SHARD_END);
shardEndProcessRecordsInput = input;
subscription.cancel();
return;
}
subscription.request(1);
}
public void executeLifecycle() {
if (isShutdown()) {
return;
}
if (stateChangeFuture != null && !stateChangeFuture.isDone()) {
return;
}
try {
if (isShutdownRequested()) {
stateChangeFuture = shutdownComplete();
} else if (needsInitialization) {
if (stateChangeFuture != null) {
if (stateChangeFuture.get()) {
// Task rejection during the subscribe() call will not be propagated back as it not executed
// in the context of the Scheduler thread. Hence we should not assume the subscription will
// always be successful.
// But if subscription was not successful, then it will recover
// during healthCheck which will restart subscription.
// From Shardconsumer point of view, initialization after the below subscribe call
// is complete
subscribe();
needsInitialization = false;
}
}
stateChangeFuture = initializeComplete();
}
} catch (InterruptedException e) {
//
// Ignored should be handled by scheduler
//
} catch (ExecutionException e) {
throw new RuntimeException(e);
} catch (RejectedExecutionException e) {
// It is possible the tasks submitted to the executor service by the Scheduler thread might get rejected
// due to various reasons. Such failed executions must be captured and marked as failure to prevent
// the state transitions.
taskOutcome = TaskOutcome.FAILURE;
}
if (ConsumerStates.ShardConsumerState.PROCESSING.equals(currentState.state())) {
Throwable t = healthCheck();
if (t instanceof Error) {
throw (Error) t;
}
}
}
@VisibleForTesting
Throwable healthCheck() {
logNoDataRetrievedAfterTime();
logLongRunningTask();
Throwable failure = subscriber.healthCheck(MAX_TIME_BETWEEN_REQUEST_RESPONSE);
if (failure != null) {
return failure;
}
Throwable dispatchFailure = subscriber.getAndResetDispatchFailure();
if (dispatchFailure != null) {
log.warn(
"{} : Exception occurred while dispatching incoming data. The incoming data has been skipped",
streamIdentifier,
dispatchFailure);
return dispatchFailure;
}
return null;
}
Duration taskRunningTime() {
if (taskDispatchedAt != null && taskIsRunning) {
return Duration.between(taskDispatchedAt, Instant.now());
}
return null;
}
String longRunningTaskMessage(Duration taken) {
if (taken != null) {
return String.format(
"Previous %s task still pending for shard %s since %s ago. ",
currentTask.taskType(), shardInfo.shardId(), taken);
}
return null;
}
private void logNoDataRetrievedAfterTime() {
logWarningForTaskAfterMillis.ifPresent(value -> {
Instant lastDataArrival = subscriber.lastDataArrival();
if (lastDataArrival != null) {
Instant now = Instant.now();
Duration timeSince = Duration.between(subscriber.lastDataArrival(), now);
if (timeSince.toMillis() > value) {
log.warn("{} : Last time data arrived: {} ({})", streamIdentifier, lastDataArrival, timeSince);
}
}
});
}
private void logLongRunningTask() {
Duration taken = taskRunningTime();
if (taken != null) {
String message = longRunningTaskMessage(taken);
if (log.isDebugEnabled()) {
log.debug("{} : {} Not submitting new task.", streamIdentifier, message);
}
logWarningForTaskAfterMillis.ifPresent(value -> {
if (taken.toMillis() > value) {
log.warn("{} : {}", streamIdentifier, message);
}
});
}
}
@VisibleForTesting
void subscribe() {
subscriber.startSubscriptions();
}
@VisibleForTesting
synchronized CompletableFuture initializeComplete() {
if (!needsInitialization) {
// initialization already complete, this must be a no-op.
// ShardConsumer must be in ProcessingState and
// any further activity will be driven by publisher pushing data to subscriber
// which invokes handleInput and that triggers ProcessTask.
// Scheduler is only meant to do health-checks to ensure the consumer
// is not stuck for any reason and to do shutdown handling.
return CompletableFuture.completedFuture(true);
}
if (taskOutcome != null) {
updateState(taskOutcome);
}
if (currentState.state() == ConsumerStates.ShardConsumerState.PROCESSING) {
return CompletableFuture.completedFuture(true);
}
return CompletableFuture.supplyAsync(
() -> {
if (isShutdownRequested()) {
throw new IllegalStateException("Shutdown requested while initializing");
}
executeTask(null);
if (isShutdownRequested()) {
throw new IllegalStateException("Shutdown requested while initializing");
}
return false;
},
executorService);
}
@VisibleForTesting
CompletableFuture shutdownComplete() {
return CompletableFuture.supplyAsync(
() -> {
synchronized (this) {
if (taskOutcome != null) {
updateState(taskOutcome);
} else {
//
// ShardConsumer has been asked to shutdown before the first task even had a chance to run.
// In this case generate a successful task outcome, and allow the shutdown to continue.
// This should only happen if the lease was lost before the initial state had a chance to
// run.
//
updateState(TaskOutcome.SUCCESSFUL);
}
if (isShutdown()) {
return true;
}
executeTask(shardEndProcessRecordsInput);
// call shutdownNotification.shutdownComplete() if shutting down as part of gracefulShutdown
if (currentState.state() == ConsumerStates.ShardConsumerState.SHUTTING_DOWN
&& taskOutcome == TaskOutcome.SUCCESSFUL
&& shutdownNotification != null) {
shutdownNotification.shutdownComplete();
}
return false;
}
},
executorService);
}
private synchronized void processData(ProcessRecordsInput input) {
executeTask(input);
}
private synchronized void executeTask(ProcessRecordsInput input) {
TaskExecutionListenerInput taskExecutionListenerInput = TaskExecutionListenerInput.builder()
.shardInfo(shardInfo)
.taskType(currentState.taskType())
.build();
taskExecutionListener.beforeTaskExecution(taskExecutionListenerInput);
ConsumerTask task = currentState.createTask(shardConsumerArgument, ShardConsumer.this, input, taskFactory);
if (task != null) {
taskDispatchedAt = Instant.now();
currentTask = task;
taskIsRunning = true;
TaskResult result;
try {
result = task.call();
} finally {
taskIsRunning = false;
}
taskOutcome = resultToOutcome(result);
taskExecutionListenerInput = taskExecutionListenerInput.toBuilder()
.taskOutcome(taskOutcome)
.build();
}
taskExecutionListener.afterTaskExecution(taskExecutionListenerInput);
}
private TaskOutcome resultToOutcome(TaskResult result) {
if (result.getException() == null) {
if (result.isShardEndReached()) {
return TaskOutcome.END_OF_SHARD;
}
return TaskOutcome.SUCCESSFUL;
}
logTaskException(result);
return TaskOutcome.FAILURE;
}
private synchronized void updateState(TaskOutcome outcome) {
ConsumerState nextState = currentState;
switch (outcome) {
case SUCCESSFUL:
nextState = currentState.successTransition();
break;
case END_OF_SHARD:
markForShutdown(ShutdownReason.SHARD_END);
break;
case FAILURE:
nextState = currentState.failureTransition();
break;
default:
log.error("{} : No handler for outcome of {}", streamIdentifier, outcome.name());
nextState = currentState.failureTransition();
break;
}
nextState = handleShutdownTransition(outcome, nextState);
currentState = nextState;
}
private ConsumerState handleShutdownTransition(TaskOutcome outcome, ConsumerState nextState) {
synchronized (shutdownLock) {
if (isShutdownRequested() && outcome != TaskOutcome.FAILURE) {
return currentState.shutdownTransition(shutdownReason);
}
return nextState;
}
}
private void logTaskException(TaskResult taskResult) {
if (log.isDebugEnabled()) {
Exception taskException = taskResult.getException();
if (taskException instanceof BlockedOnParentShardException) {
// No need to log the stack trace for this exception (it is very specific).
log.debug(
"{} : Shard {} is blocked on completion of parent shard.",
streamIdentifier,
shardInfo.shardId());
} else {
log.debug(
"{} : Caught exception running {} task: ",
streamIdentifier,
currentTask.taskType(),
taskResult.getException());
}
}
}
/**
* Requests the shutdown of the ShardConsumer. This should give the record processor a chance to checkpoint
* before being shutdown.
*
* @param shutdownNotification used to signal that the record processor has been given the chance to shut down.
*/
public void gracefulShutdown(ShutdownNotification shutdownNotification) {
if (subscriber != null) {
subscriber.cancel();
}
if (shutdownNotification != null) {
this.shutdownNotification = shutdownNotification;
}
markForShutdown(ShutdownReason.REQUESTED);
}
/**
* Shutdown this ShardConsumer (including invoking the ShardRecordProcessor shutdown API).
* This is called by Worker when it loses responsibility for a shard.
*
* @return true if shutdown is complete (false if shutdown is still in progress)
*/
public boolean leaseLost() {
log.debug("{} : Shutdown({}): Lease lost triggered.", streamIdentifier, shardInfo.shardId());
if (subscriber != null) {
subscriber.cancel();
log.debug("{} : Shutdown({}): Subscriber cancelled.", streamIdentifier, shardInfo.shardId());
}
markForShutdown(ShutdownReason.LEASE_LOST);
return isShutdown();
}
void markForShutdown(ShutdownReason reason) {
synchronized (shutdownLock) {
//
// ShutdownReason.LEASE_LOST takes precedence over SHARD_END
// (we won't be able to save checkpoint at end of shard)
//
if (shutdownReason == null || shutdownReason.canTransitionTo(reason)) {
shutdownReason = reason;
}
}
}
/**
* Used (by Worker) to check if this ShardConsumer instance has been shutdown
* ShardRecordProcessor shutdown() has been invoked, as appropriate.
*
* @return true if shutdown is complete
*/
public boolean isShutdown() {
return currentState.isTerminal();
}
@VisibleForTesting
public boolean isShutdownRequested() {
synchronized (shutdownLock) {
return shutdownReason != null;
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy