Jackie-Jiang commented on code in PR #19400:
URL: https://github.com/apache/pinot/pull/19400#discussion_r4430044048
##########
pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/PinotFSSegmentUploader.java:
##########
@@ -69,47 +80,116 @@ public URI uploadSegment(File segmentFile, LLCSegmentName
segmentName, int timeo
return null;
}
final String rawTableName =
TableNameBuilder.extractRawTableName(segmentName.getTableName());
- Callable<URI> uploadTask = () -> {
- URI destUri = new URI(StringUtil.join(File.separator,
_segmentStoreUriStr, segmentName.getTableName(),
+ final URI destUri;
+ try {
+ destUri = new URI(StringUtil.join(File.separator, _segmentStoreUriStr,
segmentName.getTableName(),
SegmentCompletionUtils.generateTmpSegmentFileName(segmentName.getSegmentName())));
- long startTime = System.currentTimeMillis();
- try {
- PinotFS pinotFS = PinotFSFactory.create(new
URI(_segmentStoreUriStr).getScheme());
- // Check and delete any existing segment file.
- if (pinotFS.exists(destUri)) {
- pinotFS.delete(destUri, true);
- }
- pinotFS.copyFromLocalFile(segmentFile, destUri);
- return destUri;
- } catch (Exception e) {
- LOGGER.warn("Failed copy segment tar file {} to segment store {}: {}",
segmentFile.getName(), destUri, e);
- } finally {
- long duration = System.currentTimeMillis() - startTime;
- _serverMetrics.addTimedTableValue(rawTableName,
ServerTimer.SEGMENT_UPLOAD_TIME_MS, duration,
- TimeUnit.MILLISECONDS);
- }
+ } catch (URISyntaxException e) {
+ LOGGER.warn("Failed copy segment tar file {} to segment store {}",
segmentFile.getName(), _segmentStoreUriStr,
+ e);
+ _serverMetrics.addMeteredTableValue(rawTableName,
ServerMeter.SEGMENT_UPLOAD_FAILURE, 1);
return null;
- };
- Future<URI> future = _executorService.submit(uploadTask);
+ }
+ AtomicReference<AttemptState> state = new
AtomicReference<>(AttemptState.IN_FLIGHT);
+ Future<URI> future = _executorService.submit(() ->
copySegment(segmentFile, destUri, rawTableName, state));
try {
URI segmentLocation = future.get(timeoutInMillis, TimeUnit.MILLISECONDS);
- LOGGER.info("Successfully upload segment {} to {}.", segmentName,
segmentLocation);
- _serverMetrics.addMeteredTableValue(rawTableName,
- segmentLocation == null ? ServerMeter.SEGMENT_UPLOAD_FAILURE :
ServerMeter.SEGMENT_UPLOAD_SUCCESS, 1);
- return segmentLocation;
+ if (segmentLocation != null) {
+ LOGGER.info("Successfully upload segment {} to {}.", segmentName,
segmentLocation);
+ _serverMetrics.addMeteredTableValue(rawTableName,
ServerMeter.SEGMENT_UPLOAD_SUCCESS, 1);
+ return segmentLocation;
+ }
} catch (InterruptedException e) {
LOGGER.info("Interrupted while waiting for segment upload of {} to {}.",
segmentName, _segmentStoreUriStr);
+ // future.get already cleared the interrupt flag. Leave it clear until a
published object is deleted.
+ // S3 and HDFS abort a delete when the calling thread is interrupted,
and the worker will not delete a URI
+ // it already published. Restore the flag after that delete.
+ if (abandonUnlessPublished(state)) {
+ deleteBestEffort(destUri);
+ }
Thread.currentThread().interrupt();
} catch (TimeoutException e) {
- // Emit a separate metric for timeout since this is relatively more
common than other errors.
+ if (abandonUnlessPublished(state)) {
+ LOGGER.info("Successfully upload segment {} to {}.", segmentName,
destUri);
+ _serverMetrics.addMeteredTableValue(rawTableName,
ServerMeter.SEGMENT_UPLOAD_SUCCESS, 1);
+ return destUri;
+ }
_serverMetrics.addMeteredTableValue(rawTableName,
ServerMeter.SEGMENT_UPLOAD_TIMEOUT, 1);
LOGGER.warn("Timed out waiting to upload segment: {} for table: {}",
segmentName.getSegmentName(), rawTableName);
} catch (Exception e) {
LOGGER.warn("Failed to upload file {} of segment {} for table {}",
- segmentFile.getAbsolutePath(), segmentName, rawTableName, e);
+ segmentFile.getAbsolutePath(), segmentName, rawTableName, e);
}
_serverMetrics.addMeteredTableValue(rawTableName,
ServerMeter.SEGMENT_UPLOAD_FAILURE, 1);
-
return null;
}
+
+ /// Settles a wait that ended without the worker's result (timeout or
interrupt) against the worker's publish.
+ /// Returns true when the worker already published, so the object belongs to
the caller. Otherwise the attempt is
+ /// abandoned (or had already failed) and the worker deletes whatever its
copy leaves behind, including a late
+ /// success. The caller must not delete in that case: an early delete can
race a copy that recreates the object.
+ /// Does not cancel the future. An interrupt mid-upload can leave multipart
parts this code cannot abort.
+ static boolean abandonUnlessPublished(AtomicReference<AttemptState> state) {
+ return !state.compareAndSet(AttemptState.IN_FLIGHT, AttemptState.ABANDONED)
Review Comment:
Nonblocking simplification: this can express the same ownership transition
with one atomic operation:
```java
return state.compareAndExchange(AttemptState.IN_FLIGHT,
AttemptState.ABANDONED)
== AttemptState.PUBLISHED;
```
The returned previous state tells us whether the worker already published,
removing the negated compareAndSet and subsequent get.
##########
pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/PinotFSSegmentUploader.java:
##########
@@ -69,47 +80,116 @@ public URI uploadSegment(File segmentFile, LLCSegmentName
segmentName, int timeo
return null;
}
final String rawTableName =
TableNameBuilder.extractRawTableName(segmentName.getTableName());
- Callable<URI> uploadTask = () -> {
- URI destUri = new URI(StringUtil.join(File.separator,
_segmentStoreUriStr, segmentName.getTableName(),
+ final URI destUri;
+ try {
+ destUri = new URI(StringUtil.join(File.separator, _segmentStoreUriStr,
segmentName.getTableName(),
SegmentCompletionUtils.generateTmpSegmentFileName(segmentName.getSegmentName())));
- long startTime = System.currentTimeMillis();
- try {
- PinotFS pinotFS = PinotFSFactory.create(new
URI(_segmentStoreUriStr).getScheme());
- // Check and delete any existing segment file.
- if (pinotFS.exists(destUri)) {
- pinotFS.delete(destUri, true);
- }
- pinotFS.copyFromLocalFile(segmentFile, destUri);
- return destUri;
- } catch (Exception e) {
- LOGGER.warn("Failed copy segment tar file {} to segment store {}: {}",
segmentFile.getName(), destUri, e);
- } finally {
- long duration = System.currentTimeMillis() - startTime;
- _serverMetrics.addTimedTableValue(rawTableName,
ServerTimer.SEGMENT_UPLOAD_TIME_MS, duration,
- TimeUnit.MILLISECONDS);
- }
+ } catch (URISyntaxException e) {
+ LOGGER.warn("Failed copy segment tar file {} to segment store {}",
segmentFile.getName(), _segmentStoreUriStr,
+ e);
+ _serverMetrics.addMeteredTableValue(rawTableName,
ServerMeter.SEGMENT_UPLOAD_FAILURE, 1);
return null;
- };
- Future<URI> future = _executorService.submit(uploadTask);
+ }
+ AtomicReference<AttemptState> state = new
AtomicReference<>(AttemptState.IN_FLIGHT);
+ Future<URI> future = _executorService.submit(() ->
copySegment(segmentFile, destUri, rawTableName, state));
Review Comment:
Nonblocking simplification: resolve the PinotFS once alongside destUri,
before submitting the task, using PinotFSFactory.create(destUri.getScheme()),
then pass that instance to the worker and cleanup. PinotFSFactory.create() only
looks up the registered instance, so this does not introduce remote I/O before
the timed wait. This would remove openSegmentStore(), the URI-only
deleteBestEffort() overload, and repeated URI parsing. Preserve the null return
and failure metric when URI construction or filesystem lookup fails.
##########
pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/PinotFSSegmentUploaderTest.java:
##########
@@ -24,72 +24,231 @@
import java.net.URI;
import java.net.URISyntaxException;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import java.util.UUID;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
import org.apache.commons.io.FileUtils;
import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.metrics.ServerMeter;
import org.apache.pinot.common.metrics.ServerMetrics;
import org.apache.pinot.common.utils.LLCSegmentName;
+import
org.apache.pinot.core.data.manager.realtime.PinotFSSegmentUploader.AttemptState;
import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.filesystem.BasePinotFS;
import org.apache.pinot.spi.filesystem.PinotFSFactory;
import org.apache.pinot.spi.utils.StringUtil;
-import org.mockito.Mockito;
-import org.testng.Assert;
import org.testng.annotations.BeforeClass;
+import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+
public class PinotFSSegmentUploaderTest {
private static final int TIMEOUT_IN_MS = 1000;
private File _file;
private LLCSegmentName _llcSegmentName;
- private ServerMetrics _serverMetrics = Mockito.mock(ServerMetrics.class);
+ private ServerMetrics _serverMetrics = mock(ServerMetrics.class);
@BeforeClass
public void setUp()
throws URISyntaxException, IOException, HttpErrorStatusException {
Map<String, Object> properties = new HashMap<>();
properties.put("class.hdfs",
"org.apache.pinot.core.data.manager.realtime.PinotFSSegmentUploaderTest$AlwaysSucceedPinotFS");
- properties.put("class.timeout",
-
"org.apache.pinot.core.data.manager.realtime.PinotFSSegmentUploaderTest$AlwaysTimeoutPinotFS");
- properties.put("class.existing",
-
"org.apache.pinot.core.data.manager.realtime.PinotFSSegmentUploaderTest$AlwaysExistPinotFS");
+ properties.put("class.record",
+
"org.apache.pinot.core.data.manager.realtime.PinotFSSegmentUploaderTest$RecordingPinotFS");
PinotFSFactory.init(new PinotConfiguration(properties));
_file = FileUtils.getFile(FileUtils.getTempDirectory(),
UUID.randomUUID().toString());
_file.deleteOnExit();
_llcSegmentName = new LLCSegmentName("test_REALTIME", 1, 0,
System.currentTimeMillis());
}
+ @BeforeMethod
+ public void resetRecordedUploads() {
+ RecordingPinotFS.reset();
+ clearInvocations(_serverMetrics);
+ }
+
@Test
public void testSuccessfulUpload() {
- SegmentUploader segmentUploader = new
PinotFSSegmentUploader("hdfs://root", TIMEOUT_IN_MS, _serverMetrics);
- URI segmentURI = segmentUploader.uploadSegment(_file, _llcSegmentName);
- Assert.assertTrue(segmentURI.toString().startsWith(StringUtil
- .join(File.separator, "hdfs://root", _llcSegmentName.getTableName(),
_llcSegmentName.getSegmentName())));
+ SegmentUploader segmentUploader = new
PinotFSSegmentUploader("record://root", TIMEOUT_IN_MS, _serverMetrics);
+ URI first = segmentUploader.uploadSegment(_file, _llcSegmentName);
+ URI second = segmentUploader.uploadSegment(_file, _llcSegmentName);
+ assertNotNull(first);
+ assertNotNull(second);
+ assertNotEquals(first, second);
+ assertUuidTempUri(first);
+ assertUuidTempUri(second);
+ assertEquals(RecordingPinotFS.COPIED_DEST_URIS.size(), 2);
+ assertTrue(RecordingPinotFS.DELETED_DEST_URIS.isEmpty());
}
@Test
- public void testSegmentAlreadyExist() {
- SegmentUploader segmentUploader = new
PinotFSSegmentUploader("existing://root", TIMEOUT_IN_MS, _serverMetrics);
+ public void testCopyFailureDeletesThatUri() {
+ RecordingPinotFS._throwOnCopy = true;
+ SegmentUploader segmentUploader = new
PinotFSSegmentUploader("record://root", TIMEOUT_IN_MS, _serverMetrics);
URI segmentURI = segmentUploader.uploadSegment(_file, _llcSegmentName);
- Assert.assertTrue(segmentURI.toString().startsWith(StringUtil
- .join(File.separator, "existing://root",
_llcSegmentName.getTableName(), _llcSegmentName.getSegmentName())));
+ assertNull(segmentURI);
+ assertEquals(RecordingPinotFS.COPIED_DEST_URIS.size(), 1);
+ assertEquals(RecordingPinotFS.DELETED_DEST_URIS,
RecordingPinotFS.COPIED_DEST_URIS);
}
@Test
- public void testUploadTimeOut() {
- SegmentUploader segmentUploader = new
PinotFSSegmentUploader("timeout://root", TIMEOUT_IN_MS, _serverMetrics);
- URI segmentURI = segmentUploader.uploadSegment(_file, _llcSegmentName);
- Assert.assertNull(segmentURI);
+ public void testCopyFailureSurvivesThrowingDelete() {
+ RecordingPinotFS._throwOnCopy = true;
+ RecordingPinotFS._throwOnDelete = true;
+ SegmentUploader segmentUploader = new
PinotFSSegmentUploader("record://root", TIMEOUT_IN_MS, _serverMetrics);
+ assertNull(segmentUploader.uploadSegment(_file, _llcSegmentName));
+ assertEquals(RecordingPinotFS.COPIED_DEST_URIS.size(), 1);
+ assertTrue(RecordingPinotFS.DELETED_DEST_URIS.isEmpty());
+ }
+
+ @Test
+ public void testTimeoutDeletesOnlyTheAbandonedAttempt()
+ throws Exception {
+ RecordingPinotFS._enteredCopy = new CountDownLatch(1);
+ RecordingPinotFS._releaseCopy = new CountDownLatch(1);
+ RecordingPinotFS._deleted = new CountDownLatch(1);
+ SegmentUploader segmentUploader = new
PinotFSSegmentUploader("record://root", TIMEOUT_IN_MS, _serverMetrics);
+ ExecutorService pool = Executors.newSingleThreadExecutor();
+ try {
+ Future<URI> first = pool.submit(() ->
segmentUploader.uploadSegment(_file, _llcSegmentName, 200));
+ assertTrue(RecordingPinotFS._enteredCopy.await(5, TimeUnit.SECONDS));
+ assertNull(first.get(5, TimeUnit.SECONDS));
+ assertTrue(RecordingPinotFS.DELETED_DEST_URIS.isEmpty());
+ URI abandoned = RecordingPinotFS.COPIED_DEST_URIS.get(0);
+ RecordingPinotFS._releaseCopy.countDown();
+ assertTrue(RecordingPinotFS._deleted.await(5, TimeUnit.SECONDS));
+ assertTrue(RecordingPinotFS.DELETED_DEST_URIS.contains(abandoned));
+
+ URI retry = segmentUploader.uploadSegment(_file, _llcSegmentName,
TIMEOUT_IN_MS);
+ assertNotNull(retry);
+ assertNotEquals(retry, abandoned);
+ assertFalse(RecordingPinotFS.DELETED_DEST_URIS.contains(retry));
+ } finally {
+ RecordingPinotFS.releaseCopy();
+ pool.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testOverlappingUploadsKeepDistinctUris()
+ throws Exception {
+ RecordingPinotFS._enteredCopy = new CountDownLatch(2);
+ RecordingPinotFS._releaseCopy = new CountDownLatch(1);
+ SegmentUploader segmentUploader = new
PinotFSSegmentUploader("record://root", 30_000, _serverMetrics);
+ ExecutorService pool = Executors.newFixedThreadPool(2);
+ try {
+ Future<URI> first = pool.submit(() ->
segmentUploader.uploadSegment(_file, _llcSegmentName));
+ Future<URI> second = pool.submit(() ->
segmentUploader.uploadSegment(_file, _llcSegmentName));
+ assertTrue(RecordingPinotFS._enteredCopy.await(5, TimeUnit.SECONDS));
+ RecordingPinotFS._releaseCopy.countDown();
+ URI uri1 = first.get(5, TimeUnit.SECONDS);
+ URI uri2 = second.get(5, TimeUnit.SECONDS);
+ assertNotNull(uri1);
+ assertNotNull(uri2);
+ assertNotEquals(uri1, uri2);
+ assertTrue(RecordingPinotFS.DELETED_DEST_URIS.isEmpty());
+ } finally {
+ RecordingPinotFS.releaseCopy();
+ pool.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testAbandonUnlessPublished() {
Review Comment:
Nonblocking testing suggestion: the helper test verifies the PUBLISHED
state, but the removed post-publication tests also covered the caller's
behavior. Please consider restoring deterministic cases for (1) timeout after
publication returning the URI without deleting it, and (2) interruption after
publication deleting the URI and restoring the interrupt flag. The mocked
addTimedTableValue() call in the worker's finally block can be paused with
latches after publication and before the Future completes, so these tests do
not require production TestHooks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
