Skip to content

Commit 12ee299

Browse files
frediafredia
authored
[FLINK-35904][test] Test harness for async state processing operators (#25802)
1 parent 6248fd4 commit 12ee299

File tree

6 files changed

+1012
-169
lines changed

6 files changed

+1012
-169
lines changed

‎flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorTest.java

+122-87
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
import org.apache.flink.runtime.asyncprocessing.AsyncExecutionController;
2525
import org.apache.flink.runtime.asyncprocessing.StateRequestType;
2626
import org.apache.flink.runtime.jobgraph.OperatorID;
27-
import org.apache.flink.runtime.state.CheckpointStorageLocationReference;
2827
import org.apache.flink.runtime.state.VoidNamespace;
2928
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
3029
import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
@@ -33,19 +32,22 @@
3332
import org.apache.flink.streaming.api.operators.InternalTimerServiceAsyncImpl;
3433
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
3534
import org.apache.flink.streaming.api.operators.Triggerable;
35+
import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
3636
import org.apache.flink.streaming.api.watermark.Watermark;
3737
import org.apache.flink.streaming.runtime.io.RecordProcessorUtils;
3838
import org.apache.flink.streaming.runtime.operators.asyncprocessing.ElementOrder;
3939
import org.apache.flink.streaming.runtime.streamrecord.LatencyMarker;
4040
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
4141
import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
42-
import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
42+
import org.apache.flink.streaming.util.TestHarnessUtil;
43+
import org.apache.flink.streaming.util.asyncprocessing.AsyncKeyedOneInputStreamOperatorTestHarness;
44+
import org.apache.flink.streaming.util.asyncprocessing.AsyncKeyedTwoInputStreamOperatorTestHarness;
4345
import org.apache.flink.util.function.ThrowingConsumer;
4446

4547
import org.junit.jupiter.api.Test;
4648

47-
import java.util.concurrent.ExecutorService;
48-
import java.util.concurrent.Executors;
49+
import java.util.concurrent.CompletableFuture;
50+
import java.util.concurrent.ConcurrentLinkedQueue;
4951
import java.util.concurrent.atomic.AtomicInteger;
5052

5153
import static org.apache.flink.runtime.state.StateBackendTestUtils.buildAsyncStateBackend;
@@ -54,14 +56,14 @@
5456
/** Basic tests for {@link AbstractAsyncStateStreamOperator}. */
5557
public class AbstractAsyncStateStreamOperatorTest {
5658

57-
protected KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
59+
protected AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
5860
createTestHarness(
5961
int maxParalelism, int numSubtasks, int subtaskIndex, ElementOrder elementOrder)
6062
throws Exception {
6163
TestOperator testOperator = new TestOperator(elementOrder);
62-
KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
64+
AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
6365
testHarness =
64-
new KeyedOneInputStreamOperatorTestHarness<>(
66+
AsyncKeyedOneInputStreamOperatorTestHarness.create(
6567
testOperator,
6668
new TestKeySelector(),
6769
BasicTypeInfo.INT_TYPE_INFO,
@@ -74,7 +76,7 @@ public class AbstractAsyncStateStreamOperatorTest {
7476

7577
@Test
7678
void testCreateAsyncExecutionController() throws Exception {
77-
try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
79+
try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
7880
testHarness = createTestHarness(128, 1, 0, ElementOrder.RECORD_ORDER)) {
7981
testHarness.open();
8082
assertThat(testHarness.getOperator())
@@ -93,51 +95,32 @@ void testCreateAsyncExecutionController() throws Exception {
9395

9496
@Test
9597
void testRecordProcessorWithFirstStateOrder() throws Exception {
96-
try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
98+
try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
9799
testHarness = createTestHarness(128, 1, 0, ElementOrder.FIRST_STATE_ORDER)) {
98100
testHarness.open();
99101
TestOperator testOperator = (TestOperator) testHarness.getOperator();
100-
ThrowingConsumer<StreamRecord<Tuple2<Integer, String>>, Exception> processor =
101-
RecordProcessorUtils.getRecordProcessor(testOperator);
102-
ExecutorService anotherThread = Executors.newSingleThreadExecutor();
103-
// Trigger the processor
104-
anotherThread.execute(
105-
() -> {
106-
try {
107-
processor.accept(new StreamRecord<>(Tuple2.of(5, "5")));
108-
} catch (Exception e) {
109-
}
110-
});
102+
CompletableFuture<Void> future =
103+
testHarness.processElementInternal(new StreamRecord<>(Tuple2.of(5, "5")));
111104

112105
Thread.sleep(1000);
113106
assertThat(testOperator.getProcessed()).isEqualTo(1);
114107
assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(1);
115108

116109
// Proceed processing
117110
testOperator.proceed();
118-
anotherThread.shutdown();
119-
Thread.sleep(1000);
111+
future.get();
120112
assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(0);
121113
}
122114
}
123115

124116
@Test
125117
void testRecordProcessorWithRecordOrder() throws Exception {
126-
try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
118+
try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
127119
testHarness = createTestHarness(128, 1, 0, ElementOrder.RECORD_ORDER)) {
128120
testHarness.open();
129121
TestOperator testOperator = (TestOperator) testHarness.getOperator();
130-
ThrowingConsumer<StreamRecord<Tuple2<Integer, String>>, Exception> processor =
131-
RecordProcessorUtils.getRecordProcessor(testOperator);
132-
ExecutorService anotherThread = Executors.newSingleThreadExecutor();
133-
// Trigger the processor
134-
anotherThread.execute(
135-
() -> {
136-
try {
137-
processor.accept(new StreamRecord<>(Tuple2.of(5, "5")));
138-
} catch (Exception e) {
139-
}
140-
});
122+
CompletableFuture<Void> future =
123+
testHarness.processElementInternal(new StreamRecord<>(Tuple2.of(5, "5")));
141124

142125
Thread.sleep(1000);
143126
assertThat(testOperator.getProcessed()).isEqualTo(1);
@@ -147,19 +130,19 @@ void testRecordProcessorWithRecordOrder() throws Exception {
147130

148131
// Proceed processing
149132
testOperator.proceed();
150-
anotherThread.shutdown();
151-
Thread.sleep(1000);
133+
future.get();
152134
assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(0);
153135
}
154136
}
155137

156138
@Test
157139
void testAsyncProcessWithKey() throws Exception {
140+
158141
TestOperatorWithDirectAsyncProcess testOperator =
159142
new TestOperatorWithDirectAsyncProcess(ElementOrder.RECORD_ORDER);
160-
KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
143+
AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
161144
testHarness =
162-
new KeyedOneInputStreamOperatorTestHarness<>(
145+
AsyncKeyedOneInputStreamOperatorTestHarness.create(
163146
testOperator,
164147
new TestKeySelector(),
165148
BasicTypeInfo.INT_TYPE_INFO,
@@ -169,17 +152,8 @@ void testAsyncProcessWithKey() throws Exception {
169152
testHarness.setStateBackend(buildAsyncStateBackend(new HashMapStateBackend()));
170153
try {
171154
testHarness.open();
172-
ThrowingConsumer<StreamRecord<Tuple2<Integer, String>>, Exception> processor =
173-
RecordProcessorUtils.getRecordProcessor(testOperator);
174-
ExecutorService anotherThread = Executors.newSingleThreadExecutor();
175-
// Trigger the processor
176-
anotherThread.execute(
177-
() -> {
178-
try {
179-
processor.accept(new StreamRecord<>(Tuple2.of(5, "5")));
180-
} catch (Exception e) {
181-
}
182-
});
155+
CompletableFuture<Void> future =
156+
testHarness.processElementInternal(new StreamRecord<>(Tuple2.of(5, "5")));
183157

184158
Thread.sleep(1000);
185159
assertThat(testOperator.getProcessed()).isEqualTo(0);
@@ -189,8 +163,7 @@ void testAsyncProcessWithKey() throws Exception {
189163

190164
// Proceed processing
191165
testOperator.proceed();
192-
anotherThread.shutdown();
193-
Thread.sleep(1000);
166+
future.get();
194167
assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(0);
195168

196169
// We don't have the mailbox executor actually running, so the new context is blocked
@@ -203,11 +176,9 @@ void testAsyncProcessWithKey() throws Exception {
203176

204177
@Test
205178
void testCheckpointDrain() throws Exception {
206-
try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
179+
try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
207180
testHarness = createTestHarness(128, 1, 0, ElementOrder.RECORD_ORDER)) {
208181
testHarness.open();
209-
CheckpointStorageLocationReference locationReference =
210-
CheckpointStorageLocationReference.getDefault();
211182
AsyncExecutionController asyncExecutionController =
212183
((AbstractAsyncStateStreamOperator) testHarness.getOperator())
213184
.getAsyncExecutionController();
@@ -218,14 +189,14 @@ void testCheckpointDrain() throws Exception {
218189
((AbstractAsyncStateStreamOperator<String>) testHarness.getOperator())
219190
.postProcessElement();
220191
assertThat(asyncExecutionController.getInFlightRecordNum()).isEqualTo(1);
221-
testHarness.getOperator().prepareSnapshotPreBarrier(1);
192+
testHarness.drainStateRequests();
222193
assertThat(asyncExecutionController.getInFlightRecordNum()).isEqualTo(0);
223194
}
224195
}
225196

226197
@Test
227198
void testTimerServiceIsAsync() throws Exception {
228-
try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
199+
try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
229200
testHarness = createTestHarness(128, 1, 0, ElementOrder.RECORD_ORDER)) {
230201
testHarness.open();
231202
assertThat(testHarness.getOperator())
@@ -249,22 +220,14 @@ public void onProcessingTime(InternalTimer timer) throws Exception {}
249220

250221
@Test
251222
void testNonRecordProcess() throws Exception {
252-
try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
223+
try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
253224
testHarness = createTestHarness(128, 1, 0, ElementOrder.RECORD_ORDER)) {
254225
testHarness.open();
255226
TestOperator testOperator = (TestOperator) testHarness.getOperator();
256-
ThrowingConsumer<StreamRecord<Tuple2<Integer, String>>, Exception> processor =
257-
RecordProcessorUtils.getRecordProcessor(testOperator);
258-
ExecutorService anotherThread = Executors.newSingleThreadExecutor();
259-
anotherThread.execute(
260-
() -> {
261-
try {
262-
processor.accept(new StreamRecord<>(Tuple2.of(5, "5")));
263-
testOperator.processLatencyMarker(
264-
new LatencyMarker(1234, new OperatorID(), 0));
265-
} catch (Exception e) {
266-
}
267-
});
227+
testHarness.processElementInternal(new StreamRecord<>(Tuple2.of(5, "5")));
228+
CompletableFuture<Void> future =
229+
testHarness.processLatencyMarkerInternal(
230+
new LatencyMarker(1234, new OperatorID(), 0));
268231

269232
Thread.sleep(1000);
270233
assertThat(testOperator.getProcessed()).isEqualTo(1);
@@ -274,32 +237,24 @@ void testNonRecordProcess() throws Exception {
274237

275238
// Proceed processing
276239
testOperator.proceed();
277-
anotherThread.shutdown();
278-
Thread.sleep(1000);
240+
future.get();
279241
assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(0);
280242
assertThat(testOperator.getLatencyProcessed()).isEqualTo(1);
281243
}
282244
}
283245

284246
@Test
285247
void testWatermarkStatus() throws Exception {
286-
try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
248+
try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
287249
testHarness = createTestHarness(128, 1, 0, ElementOrder.RECORD_ORDER)) {
288250
testHarness.open();
289251
TestOperator testOperator = (TestOperator) testHarness.getOperator();
290252
ThrowingConsumer<StreamRecord<Tuple2<Integer, String>>, Exception> processor =
291253
RecordProcessorUtils.getRecordProcessor(testOperator);
292-
ExecutorService anotherThread = Executors.newSingleThreadExecutor();
293-
anotherThread.execute(
294-
() -> {
295-
try {
296-
processor.accept(new StreamRecord<>(Tuple2.of(5, "5")));
297-
testOperator.processWatermark1(new Watermark(205L));
298-
testOperator.processWatermark2(new Watermark(105L));
299-
testOperator.processWatermarkStatus(WatermarkStatus.IDLE, 1);
300-
} catch (Exception e) {
301-
}
302-
});
254+
testHarness.processElementInternal(new StreamRecord<>(Tuple2.of(5, "5")));
255+
testHarness.processWatermarkInternal(new Watermark(205L));
256+
CompletableFuture<Void> future =
257+
testHarness.processWatermarkStatusInternal(WatermarkStatus.IDLE);
303258

304259
Thread.sleep(1000);
305260
assertThat(testOperator.getProcessed()).isEqualTo(1);
@@ -310,15 +265,57 @@ void testWatermarkStatus() throws Exception {
310265

311266
// Proceed processing
312267
testOperator.proceed();
313-
anotherThread.shutdown();
314-
Thread.sleep(1000);
268+
future.get();
315269
assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(0);
316270
assertThat(testOperator.watermarkStatus.isActive()).isFalse();
317271
assertThat(testHarness.getOutput())
318272
.containsExactly(
319273
new StreamRecord<>("EventTimer-5-105"),
320-
new Watermark(105L),
321-
new Watermark(205L));
274+
new Watermark(205L),
275+
WatermarkStatus.IDLE);
276+
}
277+
}
278+
279+
@Test
280+
void testIdleWatermarkHandling() throws Exception {
281+
final WatermarkTestingOperator testOperator = new WatermarkTestingOperator();
282+
283+
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
284+
KeySelector<Long, Integer> dummyKeySelector = l -> 0;
285+
try (AsyncKeyedTwoInputStreamOperatorTestHarness<Integer, Long, Long, Long> testHarness =
286+
AsyncKeyedTwoInputStreamOperatorTestHarness.create(
287+
testOperator,
288+
dummyKeySelector,
289+
dummyKeySelector,
290+
BasicTypeInfo.INT_TYPE_INFO,
291+
1,
292+
1,
293+
0)) {
294+
testHarness.setup();
295+
testHarness.open();
296+
testHarness.processElement1(1L, 1L);
297+
testHarness.processElement1(3L, 3L);
298+
testHarness.processElement1(4L, 4L);
299+
testHarness.processWatermark1(new Watermark(1L));
300+
assertThat(testHarness.getOutput()).isEmpty();
301+
302+
testHarness.processWatermarkStatus2(WatermarkStatus.IDLE);
303+
expectedOutput.add(new StreamRecord<>(1L));
304+
expectedOutput.add(new Watermark(1L));
305+
TestHarnessUtil.assertOutputEquals(
306+
"Output was not correct", expectedOutput, testHarness.getOutput());
307+
308+
testHarness.processWatermark1(new Watermark(3L));
309+
expectedOutput.add(new StreamRecord<>(3L));
310+
expectedOutput.add(new Watermark(3L));
311+
TestHarnessUtil.assertOutputEquals(
312+
"Output was not correct", expectedOutput, testHarness.getOutput());
313+
314+
testHarness.processWatermarkStatus2(WatermarkStatus.ACTIVE);
315+
// the other input is active now, we should not emit the watermark
316+
testHarness.processWatermark1(new Watermark(4L));
317+
TestHarnessUtil.assertOutputEquals(
318+
"Output was not correct", expectedOutput, testHarness.getOutput());
322319
}
323320
}
324321

@@ -420,18 +417,56 @@ private static class TestOperatorWithDirectAsyncProcess extends TestOperator {
420417

421418
@Override
422419
public void processElement(StreamRecord<Tuple2<Integer, String>> element) throws Exception {
420+
System.out.println("processElement " + Thread.currentThread().getName());
423421
asyncProcessWithKey(
424422
element.getValue().f0,
425423
() -> {
424+
System.out.println(
425+
"asyncProcessWithKey " + Thread.currentThread().getName());
426426
processed.incrementAndGet();
427427
});
428428
synchronized (objectToWait) {
429429
objectToWait.wait();
430430
}
431+
System.out.println("processElement " + Thread.currentThread().getName());
431432
processed.incrementAndGet();
432433
}
433434
}
434435

436+
private static class WatermarkTestingOperator extends AbstractAsyncStateStreamOperator<Long>
437+
implements TwoInputStreamOperator<Long, Long, Long>,
438+
Triggerable<Integer, VoidNamespace> {
439+
440+
private transient InternalTimerService<VoidNamespace> timerService;
441+
442+
@Override
443+
public void open() throws Exception {
444+
super.open();
445+
446+
this.timerService =
447+
getInternalTimerService("test-timers", VoidNamespaceSerializer.INSTANCE, this);
448+
}
449+
450+
@Override
451+
public void onEventTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {
452+
output.collect(new StreamRecord<>(timer.getTimestamp()));
453+
}
454+
455+
@Override
456+
public void onProcessingTime(InternalTimer<Integer, VoidNamespace> timer)
457+
throws Exception {}
458+
459+
@Override
460+
public void processElement1(StreamRecord<Long> element) throws Exception {
461+
timerService.registerEventTimeTimer(VoidNamespace.INSTANCE, element.getValue());
462+
}
463+
464+
@Override
465+
public void processElement2(StreamRecord<Long> element) throws Exception {
466+
timerService.registerEventTimeTimer(VoidNamespace.INSTANCE, element.getValue());
467+
}
468+
}
469+
435470
/** {@link KeySelector} for tests. */
436471
public static class TestKeySelector implements KeySelector<Tuple2<Integer, String>, Integer> {
437472
private static final long serialVersionUID = 1L;

0 commit comments

Comments
 (0)