This is an automated email from the ASF dual-hosted git repository.
leborchuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/cloudberry-pxf.git
The following commit(s) were added to refs/heads/main by this push:
new 3f1a6048 Fix cancelled reads being reported as completed (#143)
3f1a6048 is described below
commit 3f1a6048740778bd0ff574ec02733cf7f810cc71
Author: liuxiaoyu <[email protected]>
AuthorDate: Thu Sep 3 14:58:03 2026 +0800
Fix cancelled reads being reported as completed (#143)
---
.../pxf/service/controller/ReadServiceImpl.java | 12 ++++-
.../service/controller/ReadServiceImplTest.java | 53 +++++++++++++++++++++-
2 files changed, 62 insertions(+), 3 deletions(-)
diff --git
a/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/controller/ReadServiceImpl.java
b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/controller/ReadServiceImpl.java
index d03a53e0..17e49b56 100644
---
a/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/controller/ReadServiceImpl.java
+++
b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/controller/ReadServiceImpl.java
@@ -161,18 +161,20 @@ public class ReadServiceImpl extends
BaseServiceImpl<OperationStats> implements
bridge = getBridge(context);
// expose the bridge so pxf_cancel_backend can end it mid-read
attachBridge(bridge);
+ throwIfCancelled(context);
if (!bridge.beginIteration()) {
log.debug("Skipping streaming fragment {} of resource {}",
context.getFragmentIndex(), context.getDataSource());
} else {
log.debug("Starting streaming fragment {} of resource {}",
context.getFragmentIndex(), context.getDataSource());
- while ((record = bridge.getNext()) != null) {
+ while (!isCancelled() && (record = bridge.getNext()) != null) {
record.write(dos);
// fragment's current byte count is relative to the
previous stream's byte count
fragmentStats.reportCompletedRecord(countingOutputStream.getCount() -
previousStreamByteCount);
}
}
+ throwIfCancelled(context);
success = true;
} finally {
if (bridge != null) {
@@ -203,6 +205,14 @@ public class ReadServiceImpl extends
BaseServiceImpl<OperationStats> implements
}
}
+ private void throwIfCancelled(RequestContext context) {
+ if (isCancelled()) {
+ throw new PxfRuntimeException(String.format(
+ "Read of resource %s cancelled by pxf_cancel_backend",
+ context.getDataSource()));
+ }
+ }
+
private void updateProfile(RequestContext context, String profile) {
context.setProfile(profile);
PluginConf pluginConf = context.getPluginConf();
diff --git
a/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/controller/ReadServiceImplTest.java
b/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/controller/ReadServiceImplTest.java
index 6f3699ff..733ffd57 100644
---
a/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/controller/ReadServiceImplTest.java
+++
b/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/controller/ReadServiceImplTest.java
@@ -26,8 +26,16 @@ import java.nio.charset.StandardCharsets;
import java.security.PrivilegedAction;
import java.time.Duration;
import java.util.List;
-
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.ArgumentMatchers.same;
@@ -65,6 +73,7 @@ public class ReadServiceImplTest {
private RequestContext mockContext;
private ReadServiceImpl readService;
+ private ActiveRequestRegistry activeRequestRegistry;
@BeforeEach
public void setup() throws Exception {
@@ -76,7 +85,8 @@ public class ReadServiceImplTest {
return result;
});
- readService = new ReadServiceImpl(mockConfigurationFactory,
mockBridgeFactory, mockSecurityService, mockFragmenterService,
mockMetricReporter, new ActiveRequestRegistry());
+ activeRequestRegistry = new ActiveRequestRegistry();
+ readService = new ReadServiceImpl(mockConfigurationFactory,
mockBridgeFactory, mockSecurityService, mockFragmenterService,
mockMetricReporter, activeRequestRegistry);
}
@Test
@@ -256,6 +266,45 @@ public class ReadServiceImplTest {
inOrder.verifyNoMoreInteractions();
}
+ @Test
+ public void testCancelledLastFragmentIsReportedAsFailure() throws
Exception {
+ when(mockMetricReporter.getReportFrequency()).thenReturn(1L);
+ when(mockFragmentList.size()).thenReturn(1);
+ when(mockFragmentList.get(0)).thenReturn(mockFragment1);
+ when(mockBridgeFactory.getBridge(mockContext)).thenReturn(mockBridge1);
+ when(mockBridge1.beginIteration()).thenReturn(true);
+ when(mockContext.getSegmentId()).thenReturn(0);
+ when(mockContext.getGpSessionId()).thenReturn(42);
+ when(mockContext.getDataSource()).thenReturn("test-resource");
+
+ CountDownLatch getNextStarted = new CountDownLatch(1);
+ CountDownLatch cancelBridge = new CountDownLatch(1);
+ when(mockBridge1.getNext()).thenAnswer(invocation -> {
+ getNextStarted.countDown();
+ assertTrue(cancelBridge.await(5, TimeUnit.SECONDS), "cancel did
not end the bridge");
+ return null;
+ });
+ doAnswer(invocation -> {
+ cancelBridge.countDown();
+ return null;
+ }).when(mockBridge1).endIteration();
+
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> read = executor.submit(() ->
readService.readData(mockContext, mockOutputStream));
+ assertTrue(getNextStarted.await(5, TimeUnit.SECONDS), "read did
not reach getNext");
+ assertEquals(1, activeRequestRegistry.cancel(0, 42));
+
+ ExecutionException exception = assertThrows(
+ ExecutionException.class,
+ () -> read.get(5, TimeUnit.SECONDS));
+ assertTrue(exception.getCause() instanceof PxfRuntimeException);
+ assertTrue(exception.getCause().getMessage().contains("cancelled
by pxf_cancel_backend"));
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
// helper for writing mock record to a mock output stream
// mockOutputStream -> CountingOutputStream -> DataOutputStream
// in order for the us to see the side-effect of CountingOutputStream,
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
