Skip to content

Fix flaky SharedTsBlockQueue concurrency test #18244

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

Open
Caideyipi wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Fix flaky SharedTsBlockQueue concurrency test #18244

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
170 changes: 49 additions & 121 deletions .../test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,15 +27,15 @@

import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import org.apache.tsfile.external.commons.lang3.Validate;
import org.apache.tsfile.read.common.block.TsBlock;
import org.junit.Assert;
import org.junit.Test;
import org.mockito.Mockito;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

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

Expand Down Expand Up @@ -117,18 +117,18 @@ public void testAsyncListenerAfterAbortDoesNotAddTsBlock() {
Assert.assertTrue(addFuture.isDone());
}

@Test(timeout = 15000L)
public void concurrencyTest() {
@Test
public void concurrencyTest() throws Exception {
final String queryId = "q0";
final long mockTsBlockSize = 1024L * 1024L;
final int numOfTsBlocks = 1000;

// Construct a mock LocalMemoryManager with capacity 5 * mockTsBlockSize per
// query.
LocalMemoryManager mockLocalMemoryManager = Mockito.mock(LocalMemoryManager.class);
MemoryManager memoryManager = Mockito.spy(new MemoryManager(10 * mockTsBlockSize));
MemoryPool spyMemoryPool =
Mockito.spy(new MemoryPool("test", memoryManager, 5 * mockTsBlockSize));
Mockito.when(mockLocalMemoryManager.getQueryPool()).thenReturn(spyMemoryPool);
MemoryManager memoryManager = new MemoryManager(10 * mockTsBlockSize);
MemoryPool memoryPool = new MemoryPool("test", memoryManager, 5 * mockTsBlockSize);
Mockito.when(mockLocalMemoryManager.getQueryPool()).thenReturn(memoryPool);
SharedTsBlockQueue queue =
new SharedTsBlockQueue(
new TFragmentInstanceId(queryId, 0, "0"),
Expand All @@ -139,119 +139,47 @@ public void concurrencyTest() {
queue.setMaxBytesCanReserve(Long.MAX_VALUE);

ExecutorService executor = Executors.newFixedThreadPool(2);
AtomicReference<Integer> numOfTimesSenderBlocked = new AtomicReference<>(0);
AtomicReference<Integer> numOfTimesReceiverBlocked = new AtomicReference<>(0);
AtomicReference<Integer> numOfTsBlocksToSend = new AtomicReference<>(1000);
AtomicReference<Integer> numOfTsBlocksToReceive = new AtomicReference<>(1000);
executor.submit(
new SendTask(
queue, mockTsBlockSize, numOfTsBlocksToSend, numOfTimesSenderBlocked, executor));
executor.submit(
new ReceiveTask(queue, numOfTsBlocksToReceive, numOfTimesReceiverBlocked, executor));

while (numOfTsBlocksToSend.get() != 0 && numOfTsBlocksToReceive.get() != 0) {
String message =
String.format(
"Sender %d: %d, Receiver %d: %d",
numOfTimesSenderBlocked.get(),
numOfTsBlocksToSend.get(),
numOfTimesReceiverBlocked.get(),
numOfTsBlocksToReceive.get());
System.out.println(message);
try {
Thread.sleep(10L);
} catch (InterruptedException e) {
Assert.fail(e.getMessage());
}
}
}

private static class SendTask implements Runnable {

private final SharedTsBlockQueue queue;
private final long mockTsBlockSize;
private final AtomicReference<Integer> numOfTsBlocksToSend;
private final AtomicReference<Integer> numOfTimesBlocked;
private final ExecutorService executor;

public SendTask(
SharedTsBlockQueue queue,
long mockTsBlockSize,
AtomicReference<Integer> numOfTsBlocksToSend,
AtomicReference<Integer> numOfTimesBlocked,
ExecutorService executor) {
this.queue = Validate.notNull(queue);
Validate.isTrue(mockTsBlockSize > 0L);
this.mockTsBlockSize = mockTsBlockSize;
this.numOfTsBlocksToSend = Validate.notNull(numOfTsBlocksToSend);
this.numOfTimesBlocked = Validate.notNull(numOfTimesBlocked);
this.executor = Validate.notNull(executor);
}

@Override
public void run() {
ListenableFuture<Void> blockedOnMemory = null;
while (numOfTsBlocksToSend.get() > 0) {
synchronized (queue) {
blockedOnMemory = queue.add(Utils.createMockTsBlock(mockTsBlockSize));
}
numOfTsBlocksToSend.updateAndGet(v -> v - 1);
if (!blockedOnMemory.isDone()) {
break;
}
}

if (blockedOnMemory != null) {
numOfTimesBlocked.updateAndGet(v -> v + 1);
blockedOnMemory.addListener(
new SendTask(queue, mockTsBlockSize, numOfTsBlocksToSend, numOfTimesBlocked, executor),
executor);
} else {
synchronized (queue) {
queue.setNoMoreTsBlocks(true);
}
}
}
}

private static class ReceiveTask implements Runnable {

private final SharedTsBlockQueue queue;
private final AtomicReference<Integer> numOfTsBlocksToReceive;
private final AtomicReference<Integer> numOfTimesBlocked;
private final ExecutorService executor;

public ReceiveTask(
SharedTsBlockQueue queue,
AtomicReference<Integer> numOfTsBlocksToReceive,
AtomicReference<Integer> numOfTimesBlocked,
ExecutorService executor) {
this.queue = Validate.notNull(queue);
this.numOfTsBlocksToReceive = Validate.notNull(numOfTsBlocksToReceive);
this.numOfTimesBlocked = Validate.notNull(numOfTimesBlocked);
this.executor = Validate.notNull(executor);
}

@Override
public void run() {
ListenableFuture<Void> blocked = null;
while (numOfTsBlocksToReceive.get() > 0) {
synchronized (queue) {
blocked = queue.isBlocked();
if (blocked.isDone()) {
queue.remove();
numOfTsBlocksToReceive.updateAndGet(v -> v - 1);
} else {
break;
}
}
}

if (blocked != null) {
numOfTimesBlocked.updateAndGet(v -> v + 1);
blocked.addListener(
new ReceiveTask(queue, numOfTsBlocksToReceive, numOfTimesBlocked, executor), executor);
}
try {
Future<Void> sender =
executor.submit(
() -> {
for (int i = 0; i < numOfTsBlocks; i++) {
ListenableFuture<Void> blockedOnMemory;
synchronized (queue) {
blockedOnMemory = queue.add(Utils.createMockTsBlock(mockTsBlockSize));
}
blockedOnMemory.get();
}
synchronized (queue) {
queue.setNoMoreTsBlocks(true);
}
return null;
});
Future<Void> receiver =
executor.submit(
() -> {
for (int i = 0; i < numOfTsBlocks; i++) {
ListenableFuture<Void> blocked;
synchronized (queue) {
blocked = queue.isBlocked();
}
blocked.get();
synchronized (queue) {
queue.remove();
}
}
return null;
});

sender.get(30, TimeUnit.SECONDS);

@JackieTien97 JackieTien97 Jul 20, 2026

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.

Could we wait for whichever worker completes first here? If the receiver fails before the sender finishes, it stops freeing queue memory; the sender can then block once the pool fills, so this Future.get call reports a 30-second TimeoutException and receiver.get is never reached. That hides the receiver's actual exception, even though this change aims to propagate worker failures. Using an ExecutorCompletionService with a shared deadline (and cancelling the sibling task on failure) would make this fail fast and preserve the root cause.

receiver.get(30, TimeUnit.SECONDS);

Assert.assertTrue(queue.hasNoMoreTsBlocks());
Assert.assertTrue(queue.isEmpty());
Assert.assertEquals(0L, memoryPool.getReservedBytes());
} finally {
executor.shutdownNow();
executor.awaitTermination(10, TimeUnit.SECONDS);
}
}
}
Loading