KAFKA-19693: Added PersisterBatch record in Share Partition which includes updatedState and stateBatch (#20507)

The method rollbackOrProcessStateUpdates in SharePartition received 2
separate lists of updatedStates (InFlightState) and stateBatches
(PersisterStateBatch). This PR introduces a new subclass called
`PersisterBatch` which encompasses both these objects.

Reviewers: Apoorv Mittal <apoorvmittal10@gmail.com>
This commit is contained in:
Chirag Wadhwa 2025-09-09 15:51:42 +05:30 committed by GitHub
parent 620a01b74b
commit d5e624e918
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 78 additions and 82 deletions

View File

@ -884,8 +884,7 @@ public class SharePartition {
CompletableFuture<Void> future = new CompletableFuture<>();
Throwable throwable = null;
List<InFlightState> updatedStates = new ArrayList<>();
List<PersisterStateBatch> stateBatches = new ArrayList<>();
List<PersisterBatch> persisterBatches = new ArrayList<>();
lock.writeLock().lock();
try {
// Avoided using enhanced for loop as need to check if the last batch have offsets
@ -925,8 +924,7 @@ public class SharePartition {
batch,
recordStateMap,
subMap,
updatedStates,
stateBatches
persisterBatches
);
if (ackThrowable.isPresent()) {
@ -939,7 +937,7 @@ public class SharePartition {
}
// If the acknowledgement is successful then persist state, complete the state transition
// and update the cached state for start offset. Else rollback the state transition.
rollbackOrProcessStateUpdates(future, throwable, updatedStates, stateBatches);
rollbackOrProcessStateUpdates(future, throwable, persisterBatches);
return future;
}
@ -955,8 +953,7 @@ public class SharePartition {
CompletableFuture<Void> future = new CompletableFuture<>();
Throwable throwable = null;
List<InFlightState> updatedStates = new ArrayList<>();
List<PersisterStateBatch> stateBatches = new ArrayList<>();
List<PersisterBatch> persisterBatches = new ArrayList<>();
lock.writeLock().lock();
try {
@ -975,14 +972,14 @@ public class SharePartition {
}
if (inFlightBatch.offsetState() != null) {
Optional<Throwable> releaseAcquiredRecordsThrowable = releaseAcquiredRecordsForPerOffsetBatch(memberId, inFlightBatch, recordState, updatedStates, stateBatches);
Optional<Throwable> releaseAcquiredRecordsThrowable = releaseAcquiredRecordsForPerOffsetBatch(memberId, inFlightBatch, recordState, persisterBatches);
if (releaseAcquiredRecordsThrowable.isPresent()) {
throwable = releaseAcquiredRecordsThrowable.get();
break;
}
continue;
}
Optional<Throwable> releaseAcquiredRecordsThrowable = releaseAcquiredRecordsForCompleteBatch(memberId, inFlightBatch, recordState, updatedStates, stateBatches);
Optional<Throwable> releaseAcquiredRecordsThrowable = releaseAcquiredRecordsForCompleteBatch(memberId, inFlightBatch, recordState, persisterBatches);
if (releaseAcquiredRecordsThrowable.isPresent()) {
throwable = releaseAcquiredRecordsThrowable.get();
break;
@ -993,7 +990,7 @@ public class SharePartition {
}
// If the release acquired records is successful then persist state, complete the state transition
// and update the cached state for start offset. Else rollback the state transition.
rollbackOrProcessStateUpdates(future, throwable, updatedStates, stateBatches);
rollbackOrProcessStateUpdates(future, throwable, persisterBatches);
return future;
}
@ -1004,8 +1001,7 @@ public class SharePartition {
private Optional<Throwable> releaseAcquiredRecordsForPerOffsetBatch(String memberId,
InFlightBatch inFlightBatch,
RecordState recordState,
List<InFlightState> updatedStates,
List<PersisterStateBatch> stateBatches) {
List<PersisterBatch> persisterBatches) {
log.trace("Offset tracked batch record found, batch: {} for the share partition: {}-{}", inFlightBatch,
groupId, topicIdPartition);
@ -1032,10 +1028,9 @@ public class SharePartition {
return Optional.of(new InvalidRecordStateException("Unable to release acquired records for the offset"));
}
// Successfully updated the state of the offset.
updatedStates.add(updateResult);
stateBatches.add(new PersisterStateBatch(offsetState.getKey(), offsetState.getKey(),
updateResult.state().id(), (short) updateResult.deliveryCount()));
// Successfully updated the state of the offset and created a persister state batch for write to persister.
persisterBatches.add(new PersisterBatch(updateResult, new PersisterStateBatch(offsetState.getKey(),
offsetState.getKey(), updateResult.state().id(), (short) updateResult.deliveryCount())));
// Do not update the next fetch offset as the offset has not completed the transition yet.
}
}
@ -1045,8 +1040,7 @@ public class SharePartition {
private Optional<Throwable> releaseAcquiredRecordsForCompleteBatch(String memberId,
InFlightBatch inFlightBatch,
RecordState recordState,
List<InFlightState> updatedStates,
List<PersisterStateBatch> stateBatches) {
List<PersisterBatch> persisterBatches) {
// Check if member id is the owner of the batch.
if (!inFlightBatch.batchMemberId().equals(memberId) && !inFlightBatch.batchMemberId().equals(EMPTY_MEMBER_ID)) {
@ -1072,10 +1066,9 @@ public class SharePartition {
return Optional.of(new InvalidRecordStateException("Unable to release acquired records for the batch"));
}
// Successfully updated the state of the batch.
updatedStates.add(updateResult);
stateBatches.add(new PersisterStateBatch(inFlightBatch.firstOffset(), inFlightBatch.lastOffset(),
updateResult.state().id(), (short) updateResult.deliveryCount()));
// Successfully updated the state of the batch and created a persister state batch for write to persister.
persisterBatches.add(new PersisterBatch(updateResult, new PersisterStateBatch(inFlightBatch.firstOffset(),
inFlightBatch.lastOffset(), updateResult.state().id(), (short) updateResult.deliveryCount())));
// Do not update the next fetch offset as the batch has not completed the transition yet.
}
return Optional.empty();
@ -1826,8 +1819,7 @@ public class SharePartition {
ShareAcknowledgementBatch batch,
Map<Long, RecordState> recordStateMap,
NavigableMap<Long, InFlightBatch> subMap,
final List<InFlightState> updatedStates,
List<PersisterStateBatch> stateBatches
List<PersisterBatch> persisterBatches
) {
Optional<Throwable> throwable;
lock.writeLock().lock();
@ -1889,11 +1881,11 @@ public class SharePartition {
}
throwable = acknowledgePerOffsetBatchRecords(memberId, batch, inFlightBatch,
recordStateMap, updatedStates, stateBatches);
recordStateMap, persisterBatches);
} else {
// The in-flight batch is a full match hence change the state of the complete batch.
throwable = acknowledgeCompleteBatch(batch, inFlightBatch,
recordStateMap.get(batch.firstOffset()), updatedStates, stateBatches);
recordStateMap.get(batch.firstOffset()), persisterBatches);
}
if (throwable.isPresent()) {
@ -1930,8 +1922,7 @@ public class SharePartition {
ShareAcknowledgementBatch batch,
InFlightBatch inFlightBatch,
Map<Long, RecordState> recordStateMap,
List<InFlightState> updatedStates,
List<PersisterStateBatch> stateBatches
List<PersisterBatch> persisterBatches
) {
lock.writeLock().lock();
try {
@ -1995,10 +1986,9 @@ public class SharePartition {
return Optional.of(new InvalidRecordStateException(
"Unable to acknowledge records for the batch"));
}
// Successfully updated the state of the offset.
updatedStates.add(updateResult);
stateBatches.add(new PersisterStateBatch(offsetState.getKey(), offsetState.getKey(),
updateResult.state().id(), (short) updateResult.deliveryCount()));
// Successfully updated the state of the offset and created a persister state batch for write to persister.
persisterBatches.add(new PersisterBatch(updateResult, new PersisterStateBatch(offsetState.getKey(),
offsetState.getKey(), updateResult.state().id(), (short) updateResult.deliveryCount())));
// Do not update the nextFetchOffset as the offset has not completed the transition yet.
}
} finally {
@ -2011,8 +2001,7 @@ public class SharePartition {
ShareAcknowledgementBatch batch,
InFlightBatch inFlightBatch,
RecordState recordState,
List<InFlightState> updatedStates,
List<PersisterStateBatch> stateBatches
List<PersisterBatch> persisterBatches
) {
lock.writeLock().lock();
try {
@ -2044,11 +2033,9 @@ public class SharePartition {
new InvalidRecordStateException("Unable to acknowledge records for the batch"));
}
// Successfully updated the state of the batch.
updatedStates.add(updateResult);
stateBatches.add(
new PersisterStateBatch(inFlightBatch.firstOffset(), inFlightBatch.lastOffset(),
updateResult.state().id(), (short) updateResult.deliveryCount()));
// Successfully updated the state of the batch and created a persister state batch for write to persister.
persisterBatches.add(new PersisterBatch(updateResult, new PersisterStateBatch(inFlightBatch.firstOffset(),
inFlightBatch.lastOffset(), updateResult.state().id(), (short) updateResult.deliveryCount())));
// Do not update the next fetch offset as the batch has not completed the transition yet.
} finally {
lock.writeLock().unlock();
@ -2090,8 +2077,7 @@ public class SharePartition {
void rollbackOrProcessStateUpdates(
CompletableFuture<Void> future,
Throwable throwable,
List<InFlightState> updatedStates,
List<PersisterStateBatch> stateBatches
List<PersisterBatch> persisterBatches
) {
lock.writeLock().lock();
try {
@ -2099,9 +2085,9 @@ public class SharePartition {
// Log in DEBUG to avoid flooding of logs for a faulty client.
log.debug("Request failed for updating state, rollback any changed state"
+ " for the share partition: {}-{}", groupId, topicIdPartition);
updatedStates.forEach(state -> {
state.completeStateTransition(false);
if (state.state() == RecordState.AVAILABLE) {
persisterBatches.forEach(persisterBatch -> {
persisterBatch.updatedState.completeStateTransition(false);
if (persisterBatch.updatedState.state() == RecordState.AVAILABLE) {
updateFindNextFetchOffset(true);
}
});
@ -2109,7 +2095,7 @@ public class SharePartition {
return;
}
if (stateBatches.isEmpty() && updatedStates.isEmpty()) {
if (persisterBatches.isEmpty()) {
future.complete(null);
return;
}
@ -2117,7 +2103,8 @@ public class SharePartition {
lock.writeLock().unlock();
}
writeShareGroupState(stateBatches).whenComplete((result, exception) -> {
writeShareGroupState(persisterBatches.stream().map(PersisterBatch::stateBatch).toList())
.whenComplete((result, exception) -> {
// There can be a pending delayed share fetch requests for the share partition which are waiting
// on the startOffset to move ahead, hence track if the state is updated in the cache. If
// yes, then notify the delayed share fetch purgatory to complete the pending requests.
@ -2129,9 +2116,9 @@ public class SharePartition {
groupId, topicIdPartition, exception);
// In case of failure when transition state is rolled back then it should be rolled
// back to ACQUIRED state, unless acquisition lock for the state has expired.
updatedStates.forEach(state -> {
state.completeStateTransition(false);
if (state.state() == RecordState.AVAILABLE) {
persisterBatches.forEach(persisterBatch -> {
persisterBatch.updatedState.completeStateTransition(false);
if (persisterBatch.updatedState.state() == RecordState.AVAILABLE) {
updateFindNextFetchOffset(true);
}
});
@ -2141,9 +2128,9 @@ public class SharePartition {
log.trace("State change request successful for share partition: {}-{}",
groupId, topicIdPartition);
updatedStates.forEach(state -> {
state.completeStateTransition(true);
if (state.state() == RecordState.AVAILABLE) {
persisterBatches.forEach(persisterBatch -> {
persisterBatch.updatedState.completeStateTransition(true);
if (persisterBatch.updatedState.state() == RecordState.AVAILABLE) {
updateFindNextFetchOffset(true);
}
});
@ -2929,6 +2916,15 @@ public class SharePartition {
}
}
/**
* PersisterBatch class is used to record the state updates for a batch or an offset.
* It contains the updated in-flight state and the persister state batch to be sent to persister.
*/
private record PersisterBatch(
InFlightState updatedState,
PersisterStateBatch stateBatch
) { }
/**
* LastOffsetAndMaxRecords class is used to track the last offset to acquire and the maximum number
* of records that can be acquired in a fetch request.