Skip to content

Commit

Permalink
[FLINK-35025][Runtime/State] Split AsyncStateProcessing and introdu…
Browse files Browse the repository at this point in the history 
…ce element order for state processing
  • Loading branch information
@Zakelly
Zakelly committed Apr 16, 2024
1 parent e99640a commit af6b88d
Show file tree
Hide file tree
Showing 14 changed files with 475 additions and 113 deletions.
21 changes: 21 additions & 0 deletions ...time/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,9 @@
import org.apache.flink.api.common.operators.MailboxExecutor;
import org.apache.flink.api.common.state.v2.State;
import org.apache.flink.core.state.InternalStateFuture;
import org.apache.flink.util.FlinkRuntimeException;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.function.ThrowingRunnable;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand Down Expand Up @@ -248,4 +250,23 @@ private void seizeCapacity() {
setCurrentContext(storedContext);
inFlightRecordNum.incrementAndGet();
}

/**
* A helper to request a {@link StateRequestType#SYNC_POINT} and run a callback if it finishes
* (once the record is not blocked).
*
* @param callback the callback to run if it finishes (once the record is not blocked).
*/
public void syncPointRequestWithCallback(ThrowingRunnable<Exception> callback) {
handleRequest(null, StateRequestType.SYNC_POINT, null)
.thenAccept(
v -> {
try {
callback.run();
} catch (Exception e) {
// TODO: Properly handle the exception and fail the entire job.
throw new FlinkRuntimeException("Unexpected runtime exception", e);
}
});
}
}
8 changes: 6 additions & 2 deletions flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestBuffer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,11 @@ public StateRequestBuffer() {
}

void enqueueToActive(StateRequest<K, ?, ?> request) {
activeQueue.add(request);
if (request.getRequestType() == StateRequestType.SYNC_POINT) {
request.getFuture().complete(null);
} else {
activeQueue.add(request);
}
}

void enqueueToBlocking(StateRequest<K, ?, ?> request) {
Expand All @@ -83,7 +87,7 @@ RecordContext<K> tryActivateOneByKey(K key) {
}

StateRequest<K, ?, ?> stateRequest = blockingQueue.get(key).removeFirst();
activeQueue.add(stateRequest);
enqueueToActive(stateRequest);
if (blockingQueue.get(key).isEmpty()) {
blockingQueue.remove(key);
}
Expand Down
44 changes: 44 additions & 0 deletions .../src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,50 @@ void testInFlightRecordControl() {
}
}

@Test
public void testSyncPoint() {
AtomicInteger counter = new AtomicInteger(0);

// Test the sync point processing without a key occupied.
RecordContext<String> recordContext = aec.buildContext("record", "key");
aec.setCurrentContext(recordContext);
recordContext.retain();
aec.syncPointRequestWithCallback(counter::incrementAndGet);
assertThat(counter.get()).isEqualTo(1);
assertThat(recordContext.getReferenceCount()).isEqualTo(1);
assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(0);
assertThat(aec.stateRequestsBuffer.blockingQueueSize()).isEqualTo(0);
assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(1);
recordContext.release();
assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(0);

counter.set(0);
// Test the sync point processing with a key occupied.
RecordContext<String> recordContext1 = aec.buildContext("record1", "occupied");
aec.setCurrentContext(recordContext1);
userCode.run();

RecordContext<String> recordContext2 = aec.buildContext("record2", "occupied");
aec.setCurrentContext(recordContext2);
aec.syncPointRequestWithCallback(counter::incrementAndGet);
recordContext2.retain();
assertThat(counter.get()).isEqualTo(0);
assertThat(recordContext2.getReferenceCount()).isGreaterThan(1);
assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(1);
assertThat(aec.stateRequestsBuffer.blockingQueueSize()).isEqualTo(1);
aec.triggerIfNeeded(true);
assertThat(counter.get()).isEqualTo(0);
assertThat(recordContext2.getReferenceCount()).isGreaterThan(1);
assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(1);
assertThat(aec.stateRequestsBuffer.blockingQueueSize()).isEqualTo(1);
aec.triggerIfNeeded(true);
assertThat(counter.get()).isEqualTo(1);
assertThat(recordContext2.getReferenceCount()).isEqualTo(1);
assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(0);
assertThat(aec.stateRequestsBuffer.blockingQueueSize()).isEqualTo(0);
recordContext2.release();
}

/** Simulate the underlying state that is actually used to execute the request. */
static class TestUnderlyingState {

Expand Down
31 changes: 8 additions & 23 deletions ...-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractInput.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.operators.asyncprocessing.AsyncStateProcessing;
import org.apache.flink.streaming.runtime.operators.asyncprocessing.AsyncStateProcessingOperator;
import org.apache.flink.streaming.runtime.streamrecord.LatencyMarker;
import org.apache.flink.streaming.runtime.streamrecord.RecordAttributes;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
Expand Down Expand Up @@ -94,32 +95,16 @@ public boolean hasKeyContext() {
@Internal
@Override
public final boolean isAsyncStateProcessingEnabled() {
return (owner instanceof AsyncStateProcessing)
&& ((AsyncStateProcessing) owner).isAsyncStateProcessingEnabled();
return (owner instanceof AsyncStateProcessingOperator)
&& ((AsyncStateProcessingOperator) owner).isAsyncStateProcessingEnabled();
}

@Internal
@Override
@SuppressWarnings("unchecked")
public final <T> void setAsyncKeyedContextElement(
StreamRecord<T> record, KeySelector<T, ?> keySelector) throws Exception {
((AsyncStateProcessing) owner).setAsyncKeyedContextElement(record, keySelector);
}

@Internal
@Override
public final void postProcessElement() {
((AsyncStateProcessing) owner).postProcessElement();
}

@Internal
@Override
@SuppressWarnings({"unchecked", "rawtypes"})
public final ThrowingConsumer<StreamRecord, Exception> getRecordProcessor(int inputId) {
return record -> {
setAsyncKeyedContextElement(record, stateKeySelector);
processElement(record);
postProcessElement();
};
public final ThrowingConsumer<StreamRecord<IN>, Exception> getRecordProcessor(int inputId) {
return AsyncStateProcessing.makeRecordProcessor(
(AsyncStateProcessingOperator) owner,
(KeySelector) stateKeySelector,
this::processElement);
}
}
44 changes: 28 additions & 16 deletions ...e/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
import org.apache.flink.util.function.ThrowingConsumer;
import org.apache.flink.util.function.ThrowingRunnable;

/**
* This operator is an abstract class that give the {@link AbstractStreamOperator} the ability to
Expand All @@ -41,7 +42,7 @@
@Internal
@SuppressWarnings("rawtypes")
public abstract class AbstractAsyncStateStreamOperator<OUT> extends AbstractStreamOperator<OUT>
implements AsyncStateProcessing {
implements AsyncStateProcessingOperator {

private AsyncExecutionController asyncExecutionController;

Expand All @@ -65,6 +66,11 @@ public boolean isAsyncStateProcessingEnabled() {
return true;
}

@Override
public ElementOrder getElementOrder() {
return ElementOrder.RECORD_ORDER;
}

@Override
@SuppressWarnings("unchecked")
public final <T> void setAsyncKeyedContextElement(
Expand All @@ -91,6 +97,12 @@ public final void postProcessElement() {
currentProcessingContext.release();
}

@Override
@SuppressWarnings("unchecked")
public final void preserveRecordOrderAndProcess(ThrowingRunnable<Exception> processing) {
asyncExecutionController.syncPointRequestWithCallback(processing);
}

@Override
@SuppressWarnings("unchecked")
public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor(int inputId) {
Expand All @@ -99,26 +111,21 @@ public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor
if (this instanceof TwoInputStreamOperator) {
switch (inputId) {
case 1:
return record -> {
setAsyncKeyedContextElement(record, (KeySelector<T, ?>) stateKeySelector1);
((TwoInputStreamOperator) this).processElement1(record);
postProcessElement();
};
return AsyncStateProcessing.<T>makeRecordProcessor(
this,
(KeySelector) stateKeySelector1,
((TwoInputStreamOperator) this)::processElement1);
case 2:
return record -> {
setAsyncKeyedContextElement(record, (KeySelector<T, ?>) stateKeySelector2);
((TwoInputStreamOperator) this).processElement2(record);
postProcessElement();
};
return AsyncStateProcessing.<T>makeRecordProcessor(
this,
(KeySelector) stateKeySelector2,
((TwoInputStreamOperator) this)::processElement2);
default:
break;
}
} else if (this instanceof Input && inputId == 1) {
return record -> {
setAsyncKeyedContextElement(record, (KeySelector<T, ?>) stateKeySelector1);
((Input) this).processElement(record);
postProcessElement();
};
return AsyncStateProcessing.<T>makeRecordProcessor(
this, (KeySelector) stateKeySelector1, ((Input) this)::processElement);
}
throw new IllegalArgumentException(
String.format(
Expand All @@ -130,4 +137,9 @@ public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor
AsyncExecutionController<?> getAsyncExecutionController() {
return asyncExecutionController;
}

@VisibleForTesting
RecordContext getCurrentProcessingContext() {
return currentProcessingContext;
}
}
19 changes: 18 additions & 1 deletion ...flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import org.apache.flink.streaming.api.operators.StreamTaskStateInitializer;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.util.function.ThrowingConsumer;
import org.apache.flink.util.function.ThrowingRunnable;

/**
* This operator is an abstract class that give the {@link AbstractStreamOperatorV2} the ability to
Expand All @@ -38,7 +39,7 @@
@Internal
@SuppressWarnings("rawtypes")
public abstract class AbstractAsyncStateStreamOperatorV2<OUT> extends AbstractStreamOperatorV2<OUT>
implements AsyncStateProcessing {
implements AsyncStateProcessingOperator {

private AsyncExecutionController asyncExecutionController;

Expand All @@ -65,6 +66,11 @@ public boolean isAsyncStateProcessingEnabled() {
return true;
}

@Override
public ElementOrder getElementOrder() {
return ElementOrder.RECORD_ORDER;
}

@Override
@SuppressWarnings("unchecked")
public final <T> void setAsyncKeyedContextElement(
Expand All @@ -91,6 +97,12 @@ public final void postProcessElement() {
currentProcessingContext.release();
}

@Override
@SuppressWarnings("unchecked")
public final void preserveRecordOrderAndProcess(ThrowingRunnable<Exception> processing) {
asyncExecutionController.syncPointRequestWithCallback(processing);
}

@Override
public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor(int inputId) {
// The real logic should be in First/SecondInputOfTwoInput#getRecordProcessor.
Expand All @@ -103,4 +115,9 @@ public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor
AsyncExecutionController<?> getAsyncExecutionController() {
return asyncExecutionController;
}

@VisibleForTesting
RecordContext getCurrentProcessingContext() {
return currentProcessingContext;
}
}
45 changes: 32 additions & 13 deletions ...va/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,19 +34,6 @@ public interface AsyncStateProcessing {
*/
boolean isAsyncStateProcessingEnabled();

/**
* Set key context for async state processing.
*
* @param record the record.
* @param keySelector the key selector to select a key from record.
* @param <T> the type of the record.
*/
<T> void setAsyncKeyedContextElement(StreamRecord<T> record, KeySelector<T, ?> keySelector)
throws Exception;

/** A callback that will be triggered after an element finishes {@code processElement}. */
void postProcessElement();

/**
* Get the record processor that could process record from input, which is the only entry for
* async processing.
Expand All @@ -56,4 +43,36 @@ <T> void setAsyncKeyedContextElement(StreamRecord<T> record, KeySelector<T, ?> k
* there is multiple inputs for the instance.
*/
<T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor(int inputId);

/**
* Static method helper to make a record processor with given infos.
*
* @param asyncOperator the operator that can process state asynchronously.
* @param keySelector the key selector.
* @param processor the record processing logic.
* @return the built record processor that can returned by {@link #getRecordProcessor(int)}.
*/
static <T> ThrowingConsumer<StreamRecord<T>, Exception> makeRecordProcessor(
AsyncStateProcessingOperator asyncOperator,
KeySelector<T, ?> keySelector,
ThrowingConsumer<StreamRecord<T>, Exception> processor) {
switch (asyncOperator.getElementOrder()) {
case RECORD_ORDER:
return (record) -> {
asyncOperator.setAsyncKeyedContextElement(record, keySelector);
asyncOperator.preserveRecordOrderAndProcess(() -> processor.accept(record));
asyncOperator.postProcessElement();
};
case FIRST_STATE_ORDER:
return (record) -> {
asyncOperator.setAsyncKeyedContextElement(record, keySelector);
processor.accept(record);
asyncOperator.postProcessElement();
};
default:
throw new UnsupportedOperationException(
"Unknown element order for async processing:"
+ asyncOperator.getElementOrder());
}
}
}
Loading

0 comments on commit af6b88d

Please sign in to comment.