Skip to content

KAFKA-8620: fix NPE due to race condition during shutdown while rebalancing #7021

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mjsax merged 10 commits into from
Jul 15, 2019
Merged

KAFKA-8620: fix NPE due to race condition during shutdown while rebalancing #7021

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
c6a5997
fix the state after poll request
Jul 3, 2019
1337616
add unit test
Jul 3, 2019
d5b7b5c
trying to find a way to pass assignment in
Jul 4, 2019
b2326b4
minor fix
Jul 7, 2019
58b7ce3
address comments
Jul 9, 2019
9c13d42
success reproduce
Jul 9, 2019
f1b1002
address minor suggestion
Jul 9, 2019
eb490a9
avoid info exposure of internal implementation
Jul 9, 2019
a99606b
address comments by initializing real task manager and one-time mock …
Jul 10, 2019
5c6ec62
address Bruno's comment
Jul 10, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 24 additions & 8 deletions streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,14 +190,20 @@ State setState(final State newState) {
oldState = state;

if (state == State.PENDING_SHUTDOWN && newState != State.DEAD) {
log.debug("Ignoring request to transit from PENDING_SHUTDOWN to {}: " +
"only DEAD state is a valid next state", newState);
Comment thread
abbccdda marked this conversation as resolved.
// when the state is already in PENDING_SHUTDOWN, all other transitions will be
// refused but we do not throw exception here
return null;
} else if (state == State.DEAD) {
log.debug("Ignoring request to transit from DEAD to {}: " +
"no valid next state after DEAD", newState);
Comment thread
abbccdda marked this conversation as resolved.
// when the state is already in NOT_RUNNING, all its transitions
// will be refused but we do not throw exception here
return null;
} else if (state == State.PARTITIONS_REVOKED && newState == State.PARTITIONS_REVOKED) {
log.debug("Ignoring request to transit from PARTITIONS_REVOKED to PARTITIONS_REVOKED: " +
"self transition is not allowed");

@mjsax mjsax Jul 9, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thinking about this, I am wondering if we should just change the FSM to allow this transition and simplify the code here? \cc @guozhangwang

@guozhangwang guozhangwang Jul 15, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought about this when tightening the FSM before but the unit tests reminds me of one thing: our current contract is that we only transit to PARTITIONS_REVOKED when calling onPartitionsRevoked, which is called only once at the beginning of the rebalance today, so keeping it strict is better just in case we have incorrect partial rebalance procedure.

With KIP-429 this may be violated so we need to revisit our FSM once Streams adopt cooperative protocols. cc @ableegoldman who's working on this.

// when the state is already in PARTITIONS_REVOKED, its transition to itself will be
// refused but we do not throw exception here
return null;
Expand Down Expand Up @@ -268,17 +274,19 @@ public void onPartitionsAssigned(final Collection<TopicPartition> assignment) {
final long start = time.milliseconds();
try {
if (streamThread.setState(State.PARTITIONS_ASSIGNED) == null) {
return;
Comment thread
abbccdda marked this conversation as resolved.
}
if (streamThread.assignmentErrorCode.get() == StreamsPartitionAssignor.Error.NONE.code()) {
log.debug("Skipping task creation in rebalance because we are already in {} phase.",
Comment thread
abbccdda marked this conversation as resolved.
Outdated
streamThread.state());
} else if (streamThread.assignmentErrorCode.get() != StreamsPartitionAssignor.Error.NONE.code()) {

@cadonna cadonna Jul 9, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't this a rather brittle fix? How does this code guarantee that after streamThread.assignmentErrorCode.get() != StreamsPartitionAssignor.Error.NONE.code() is false, the result of streamThread.assignmentErrorCode.get() does not change again before the tasks are created?

@abbccdda abbccdda Jul 9, 2019

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in PartitionsAssigned we should have already passed the assignment error placement, so I won't expect another flip for its value during this round rebalance.

@cadonna cadonna Jul 9, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense

log.debug("Encountered assignment error during partition assignment: {}. " +
"Will skip the task initialization", streamThread.assignmentErrorCode);
Comment thread
abbccdda marked this conversation as resolved.
Outdated

@cadonna cadonna Jul 10, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't the indentation be as follows?

log.debug(
    "Encountered assignment error during partition assignment: {}. Will skip the task initialization", 
    streamThread.assignmentErrorCode
);

The first parameter is one or two characters too long, though.

} else {
log.debug("Creating tasks based on assignment.");
taskManager.createTasks(assignment);
}
} catch (final Throwable t) {
log.error(
"Error caught during partition assignment, " +
"will abort the current process and re-throw at the end of rebalance: {}",
t
);
"will abort the current process and re-throw at the end of rebalance", t);

@guozhangwang guozhangwang Jul 15, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice catch.

streamThread.setRebalanceException(t);
} finally {
log.info("partition assignment took {} ms.\n" +
Expand Down Expand Up @@ -809,7 +817,6 @@ private void enforceRebalance() {
// Visible for testing
void runOnce() {
final ConsumerRecords<byte[], byte[]> records;

now = time.milliseconds();

if (state == State.PARTITIONS_ASSIGNED) {
Expand All @@ -830,6 +837,15 @@ void runOnce() {
throw new StreamsException(logPrefix + "Unexpected state " + state + " during normal iteration");
}

// Shutdown hook could potentially be triggered and transit the thread state to PENDING_SHUTDOWN during #pollRequests().
// The task manager internal states could be uninitialized if the state transition happens during #onPartitionAssigned().
// Should only proceed when the thread is still running after #pollRequests(), because no external state mutation
// could affect the task manager state beyond this point within #runOnce().
if (!isRunning()) {
log.debug("State already transits to {}, skipping the run once call after poll request", state);
return;
}

final long pollLatency = advanceNowAndComputeLatency();

if (records != null && !records.isEmpty()) {
Expand Down Expand Up @@ -1179,7 +1195,7 @@ private void completeShutdown(final boolean cleanRun) {
// intentionally do not check the returned flag
setState(State.PENDING_SHUTDOWN);

log.info("Shutting down");
log.info("Shutting down with clean flag to {}", cleanRun);
Comment thread
abbccdda marked this conversation as resolved.
Outdated

try {
taskManager.shutdown(cleanRun);
Expand Down
67 changes: 67 additions & 0 deletions streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,10 @@
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.InvalidOffsetException;
import org.apache.kafka.clients.consumer.MockConsumer;
import org.apache.kafka.clients.consumer.OffsetResetStrategy;
import org.apache.kafka.clients.producer.MockProducer;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.common.Cluster;
Expand Down Expand Up @@ -695,6 +697,45 @@ public void shouldShutdownTaskManagerOnCloseWithoutStart() {
EasyMock.verify(taskManager);
}

@Test
public void shouldNotAccessTaskManagerWhenPendingShutdownInRunOnce() {

@cadonna cadonna Jul 10, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would give this test a more meaningful name, e.g. shouldNotThrowWhenPendingShutdownInRunOnce.

final TaskManager taskManager = EasyMock.createNiceMock(TaskManager.class);
Comment thread
abbccdda marked this conversation as resolved.
Outdated
EasyMock.expect(taskManager.activeTask(t1p1)).andReturn(null);
EasyMock.replay(taskManager);

final MockStreamThreadConsumer<byte[], byte[]> mockStreamThreadConsumer =
new MockStreamThreadConsumer<byte[], byte[]>(OffsetResetStrategy.EARLIEST) {
@Override
public void onPoll() {
super.streamThread.shutdown();
}
};

final StreamsMetricsImpl streamsMetrics = new StreamsMetricsImpl(metrics, clientId);
final StreamThread thread = new StreamThread(
mockTime,
config,
null,
mockStreamThreadConsumer,
mockStreamThreadConsumer,
null,
taskManager,
streamsMetrics,
internalTopologyBuilder,
clientId,
new LogContext(""),
new AtomicInteger()
).updateThreadMetadata(getSharedAdminClientId(clientId));

mockStreamThreadConsumer.setStreamThread(thread);
mockStreamThreadConsumer.assign(Collections.singletonList(t1p1));
mockStreamThreadConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L));

addRecord(mockStreamThreadConsumer, 1L, 0L);
thread.setState(StreamThread.State.STARTING);
thread.runOnce();
}

@Test
public void shouldOnlyShutdownOnce() {
final Consumer<byte[], byte[]> consumer = EasyMock.createNiceMock(Consumer.class);
Expand Down Expand Up @@ -1570,4 +1611,30 @@ private void addRecord(final MockConsumer<byte[], byte[]> mockConsumer,
new byte[0],
new byte[0]));
}


/**
* Mock consumer that could alter stream thread state during #poll()
*/
private abstract class MockStreamThreadConsumer<K, V> extends MockConsumer<K, V> {

private StreamThread streamThread;

MockStreamThreadConsumer(final OffsetResetStrategy offsetResetStrategy) {
super(offsetResetStrategy);
}

@Override
public synchronized ConsumerRecords<K, V> poll(final Duration timeout) {
assertNotNull(streamThread);
onPoll();
return super.poll(timeout);
}

public abstract void onPoll();
Comment thread
abbccdda marked this conversation as resolved.
Outdated

void setStreamThread(final StreamThread streamThread) {
this.streamThread = streamThread;
}
}
}