Skip to content

Commit 24fd30e

Browse files
CaideyipiCaideyipi
authored
Stabilize query execution concurrency unit tests (#19079)
* Stabilize query execution concurrency unit tests * Stabilize blocked priority queue unit tests * Propagate worker failures and guarantee unit test cleanup
1 parent cebe663 commit 24fd30e

6 files changed

Lines changed: 152 additions & 206 deletions

File tree

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java‎

Lines changed: 52 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -29,15 +29,17 @@
2929

3030
import com.google.common.util.concurrent.ListenableFuture;
3131
import com.google.common.util.concurrent.SettableFuture;
32-
import org.apache.tsfile.external.commons.lang3.Validate;
3332
import org.apache.tsfile.read.common.block.TsBlock;
3433
import org.junit.Assert;
3534
import org.junit.Test;
3635
import org.mockito.Mockito;
3736

37+
import java.util.concurrent.CompletionService;
38+
import java.util.concurrent.ExecutorCompletionService;
3839
import java.util.concurrent.ExecutorService;
3940
import java.util.concurrent.Executors;
40-
import java.util.concurrent.atomic.AtomicReference;
41+
import java.util.concurrent.Future;
42+
import java.util.concurrent.TimeUnit;
4143

4244
import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
4345

@@ -119,8 +121,8 @@ public void testAsyncListenerAfterAbortDoesNotAddTsBlock() {
119121
Assert.assertTrue(addFuture.isDone());
120122
}
121123

122-
@Test(timeout = 15000L)
123-
public void concurrencyTest() {
124+
@Test
125+
public void concurrencyTest() throws Exception {
124126
final String queryId = "q0";
125127
final long mockTsBlockSize = 1024L * 1024L;
126128

@@ -140,120 +142,57 @@ public void concurrencyTest() {
140142
queue.getCanAddTsBlock().set(null);
141143
queue.setMaxBytesCanReserve(Long.MAX_VALUE);
142144

145+
// Reuse an immutable block so mock generation is outside the concurrency deadline.
146+
TsBlock block = Utils.createMockTsBlock(mockTsBlockSize);
147+
int blockCount = 1000;
143148
ExecutorService executor = Executors.newFixedThreadPool(2);
144-
AtomicReference<Integer> numOfTimesSenderBlocked = new AtomicReference<>(0);
145-
AtomicReference<Integer> numOfTimesReceiverBlocked = new AtomicReference<>(0);
146-
AtomicReference<Integer> numOfTsBlocksToSend = new AtomicReference<>(1000);
147-
AtomicReference<Integer> numOfTsBlocksToReceive = new AtomicReference<>(1000);
148-
executor.submit(
149-
new SendTask(
150-
queue, mockTsBlockSize, numOfTsBlocksToSend, numOfTimesSenderBlocked, executor));
151-
executor.submit(
152-
new ReceiveTask(queue, numOfTsBlocksToReceive, numOfTimesReceiverBlocked, executor));
153-
154-
while (numOfTsBlocksToSend.get() != 0 && numOfTsBlocksToReceive.get() != 0) {
155-
String message =
156-
String.format(
157-
"Sender %d: %d, Receiver %d: %d",
158-
numOfTimesSenderBlocked.get(),
159-
numOfTsBlocksToSend.get(),
160-
numOfTimesReceiverBlocked.get(),
161-
numOfTsBlocksToReceive.get());
162-
System.out.println(message);
163-
try {
164-
Thread.sleep(10L);
165-
} catch (InterruptedException e) {
166-
Assert.fail(e.getMessage());
167-
}
168-
}
169-
}
170-
171-
private static class SendTask implements Runnable {
172-
173-
private final SharedTsBlockQueue queue;
174-
private final long mockTsBlockSize;
175-
private final AtomicReference<Integer> numOfTsBlocksToSend;
176-
private final AtomicReference<Integer> numOfTimesBlocked;
177-
private final ExecutorService executor;
178-
179-
public SendTask(
180-
SharedTsBlockQueue queue,
181-
long mockTsBlockSize,
182-
AtomicReference<Integer> numOfTsBlocksToSend,
183-
AtomicReference<Integer> numOfTimesBlocked,
184-
ExecutorService executor) {
185-
this.queue = Validate.notNull(queue);
186-
Validate.isTrue(mockTsBlockSize > 0L);
187-
this.mockTsBlockSize = mockTsBlockSize;
188-
this.numOfTsBlocksToSend = Validate.notNull(numOfTsBlocksToSend);
189-
this.numOfTimesBlocked = Validate.notNull(numOfTimesBlocked);
190-
this.executor = Validate.notNull(executor);
191-
}
192-
193-
@Override
194-
public void run() {
195-
ListenableFuture<Void> blockedOnMemory = null;
196-
while (numOfTsBlocksToSend.get() > 0) {
197-
synchronized (queue) {
198-
blockedOnMemory = queue.add(Utils.createMockTsBlock(mockTsBlockSize));
199-
}
200-
numOfTsBlocksToSend.updateAndGet(v -> v - 1);
201-
if (!blockedOnMemory.isDone()) {
202-
break;
203-
}
204-
}
205-
206-
if (blockedOnMemory != null) {
207-
numOfTimesBlocked.updateAndGet(v -> v + 1);
208-
blockedOnMemory.addListener(
209-
new SendTask(queue, mockTsBlockSize, numOfTsBlocksToSend, numOfTimesBlocked, executor),
210-
executor);
211-
} else {
212-
synchronized (queue) {
213-
queue.setNoMoreTsBlocks(true);
214-
}
149+
CompletionService<Void> completedTasks = new ExecutorCompletionService<>(executor);
150+
try {
151+
completedTasks.submit(
152+
() -> {
153+
for (int i = 0; i < blockCount; i++) {
154+
ListenableFuture<Void> blocked;
155+
synchronized (queue) {
156+
blocked = queue.add(block);
157+
}
158+
blocked.get(10, TimeUnit.SECONDS);
159+
}
160+
synchronized (queue) {
161+
queue.setNoMoreTsBlocks(true);
162+
}
163+
return null;
164+
});
165+
completedTasks.submit(
166+
() -> {
167+
for (int i = 0; i < blockCount; i++) {
168+
ListenableFuture<Void> blocked;
169+
synchronized (queue) {
170+
blocked = queue.isBlocked();
171+
}
172+
blocked.get(10, TimeUnit.SECONDS);
173+
synchronized (queue) {
174+
Assert.assertSame(block, queue.remove());
175+
}
176+
}
177+
return null;
178+
});
179+
for (int i = 0; i < 2; i++) {
180+
Future<Void> completed = completedTasks.poll(10, TimeUnit.SECONDS);
181+
Assert.assertNotNull(completed);
182+
completed.get();
215183
}
216-
}
217-
}
218-
219-
private static class ReceiveTask implements Runnable {
220-
221-
private final SharedTsBlockQueue queue;
222-
private final AtomicReference<Integer> numOfTsBlocksToReceive;
223-
private final AtomicReference<Integer> numOfTimesBlocked;
224-
private final ExecutorService executor;
225-
226-
public ReceiveTask(
227-
SharedTsBlockQueue queue,
228-
AtomicReference<Integer> numOfTsBlocksToReceive,
229-
AtomicReference<Integer> numOfTimesBlocked,
230-
ExecutorService executor) {
231-
this.queue = Validate.notNull(queue);
232-
this.numOfTsBlocksToReceive = Validate.notNull(numOfTsBlocksToReceive);
233-
this.numOfTimesBlocked = Validate.notNull(numOfTimesBlocked);
234-
this.executor = Validate.notNull(executor);
235-
}
236-
237-
@Override
238-
public void run() {
239-
ListenableFuture<Void> blocked = null;
240-
while (numOfTsBlocksToReceive.get() > 0) {
184+
Assert.assertTrue(queue.hasNoMoreTsBlocks());
185+
Assert.assertTrue(queue.isEmpty());
186+
Assert.assertEquals(0, spyMemoryPool.getReservedBytes());
187+
} finally {
188+
executor.shutdownNow();
189+
try {
190+
Assert.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS));
191+
} finally {
241192
synchronized (queue) {
242-
blocked = queue.isBlocked();
243-
if (blocked.isDone()) {
244-
queue.remove();
245-
numOfTsBlocksToReceive.updateAndGet(v -> v - 1);
246-
} else {
247-
break;
248-
}
193+
queue.close();
249194
}
250195
}
251-
252-
if (blocked != null) {
253-
numOfTimesBlocked.updateAndGet(v -> v + 1);
254-
blocked.addListener(
255-
new ReceiveTask(queue, numOfTsBlocksToReceive, numOfTimesBlocked, executor), executor);
256-
}
257196
}
258197
}
259198

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ public void testFragmentInstanceExecution() throws InterruptedException {
134134
}
135135

136136
@Test
137-
public void testTVListOwnerTransfer() throws InterruptedException {
137+
public void testTVListOwnerTransfer() throws Exception {
138138
ExecutorService instanceNotificationExecutor =
139139
IoTDBThreadPoolFactory.newFixedThreadPool(1, "test-instance-notification");
140140

@@ -173,8 +173,9 @@ public void testTVListOwnerTransfer() throws InterruptedException {
173173
fragmentInstanceContext2.decrementNumOfUnClosedDriver();
174174

175175
fragmentInstanceContext1.getStateMachine().finished();
176-
Thread.sleep(100);
176+
instanceNotificationExecutor.submit(() -> {}).get(10, TimeUnit.SECONDS);
177177
fragmentInstanceContext2.getStateMachine().finished();
178+
instanceNotificationExecutor.submit(() -> {}).get(10, TimeUnit.SECONDS);
178179

179180
assertTrue(execution1.getInstanceState().isDone());
180181
assertTrue(execution2.getInstanceState().isDone());

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtilCancellationTest.java‎

Lines changed: 40 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -244,12 +244,12 @@ public void testCompactionContextWithoutStateMachine() throws IOException {
244244
assertEquals(2, loadedFiles);
245245
}
246246

247-
@Test(timeout = 15000)
247+
@Test
248248
public void testTimeoutUnblocksDriverResourceCleanup() throws Exception {
249249
assertFailureUnblocksDriverResourceCleanup(new QueryTimeoutException());
250250
}
251251

252-
@Test(timeout = 15000)
252+
@Test
253253
public void testSemanticFailureUnblocksDriverResourceCleanup() throws Exception {
254254
assertFailureUnblocksDriverResourceCleanup(
255255
new SemanticException("Scalar sub-query has returned multiple rows."));
@@ -260,6 +260,7 @@ private void assertFailureUnblocksDriverResourceCleanup(RuntimeException failure
260260
ExecutorService notifications = Executors.newSingleThreadExecutor();
261261
FragmentInstanceContext context = newContext(notifications);
262262
context.initializeNumOfDrivers(1);
263+
ExecutorService worker = Executors.newSingleThreadExecutor();
263264
try {
264265
CountDownLatch closeRequested = new CountDownLatch(1);
265266
SeriesScanUtil scanner =
@@ -314,10 +315,14 @@ public void close() {
314315
false,
315316
exchangeManager);
316317

317-
assertSame(
318-
failure,
319-
assertThrows(
320-
RuntimeException.class, () -> driver.processFor(new Duration(1, TimeUnit.SECONDS))));
318+
// Bound the driver operation rather than cold context and mock initialization in a CI fork.
319+
Future<RuntimeException> execution =
320+
worker.submit(
321+
() ->
322+
assertThrows(
323+
RuntimeException.class,
324+
() -> driver.processFor(new Duration(1, TimeUnit.SECONDS))));
325+
assertSame(failure, execution.get(5, TimeUnit.SECONDS));
321326
// This can complete only after the preceding cleanup callback gets past allDriversClosed.
322327
notifications.submit(() -> {}).get(5, TimeUnit.SECONDS);
323328
assertEquals(2, loadedFiles);
@@ -329,12 +334,17 @@ public void close() {
329334
} finally {
330335
// Ensure a failed assertion cannot leave the notification thread waiting on the latch.
331336
context.decrementNumOfUnClosedDriver();
332-
notifications.shutdownNow();
333-
assertTrue(notifications.awaitTermination(5, TimeUnit.SECONDS));
337+
worker.shutdownNow();
338+
notifications.shutdown();
339+
try {
340+
assertTrue(worker.awaitTermination(5, TimeUnit.SECONDS));
341+
} finally {
342+
assertTrue(notifications.awaitTermination(5, TimeUnit.SECONDS));
343+
}
334344
}
335345
}
336346

337-
@Test(timeout = 60000)
347+
@Test
338348
public void testFinishedScanUnblocksDriverResourceCleanup() throws Exception {
339349
for (boolean sequence : new boolean[] {true, false}) {
340350
for (boolean waitForClose : new boolean[] {false, true}) {
@@ -352,7 +362,6 @@ private void assertFinishedScanUnblocksDriverResourceCleanup(
352362
CountDownLatch resumeScan = new CountDownLatch(1);
353363
CountDownLatch closeRequested = new CountDownLatch(1);
354364
// Cover both orderings: the FI is FINISHED before driver.close(), and close is already pending.
355-
notifications.submit(() -> await(allowNotifications));
356365
FragmentInstanceContext context = newContext(notifications);
357366
context.initializeNumOfDrivers(1);
358367
Operator operator = mock(Operator.class);
@@ -401,6 +410,13 @@ public void close() {
401410
false,
402411
exchangeManager);
403412
try {
413+
// Release this gate in finally so a timeout cannot change the ordering under test.
414+
Future<?> notificationsAllowed =
415+
notifications.submit(
416+
() -> {
417+
allowNotifications.await();
418+
return null;
419+
});
404420
Future<?> execution =
405421
worker.submit(() -> driver.processFor(new Duration(1, TimeUnit.SECONDS)));
406422
await(scanEntered);
@@ -415,6 +431,7 @@ public void close() {
415431
assertEquals(FragmentInstanceState.FINISHED, context.getStateMachine().getState());
416432
assertTrue(context.getStateMachine().getFailureCauses().isEmpty());
417433
allowNotifications.countDown();
434+
notificationsAllowed.get(5, TimeUnit.SECONDS);
418435
notifications.submit(() -> {}).get(5, TimeUnit.SECONDS);
419436
verify(operator).close();
420437
verify(context.getMemoryReservationContext()).releaseAllReservedMemory();
@@ -426,12 +443,19 @@ public void close() {
426443
resumeScan.countDown();
427444
allowNotifications.countDown();
428445
worker.shutdownNow();
429-
assertTrue(worker.awaitTermination(5, TimeUnit.SECONDS));
430-
driver.close();
431-
context.decrementNumOfUnClosedDriver();
432-
// Let FI cleanup finish after opening its latch, without interrupting its driver-close wait.
433-
notifications.shutdown();
434-
assertTrue(notifications.awaitTermination(5, TimeUnit.SECONDS));
446+
try {
447+
assertTrue(worker.awaitTermination(5, TimeUnit.SECONDS));
448+
} finally {
449+
try {
450+
driver.close();
451+
} finally {
452+
context.decrementNumOfUnClosedDriver();
453+
// Let FI cleanup finish after opening its latch, without interrupting its driver-close
454+
// wait.
455+
notifications.shutdown();
456+
assertTrue(notifications.awaitTermination(5, TimeUnit.SECONDS));
457+
}
458+
}
435459
}
436460
}
437461

0 commit comments

Comments
 (0)