You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hugegraph.apache.org by ji...@apache.org on 2026/10/16 14:46:23 UTC
(hugegraph-toolchain) branch master updated: feat(hubble): close loaders on task exit (#814)
This is an automated email from the ASF dual-hosted git repository.
imbajin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hugegraph-toolchain.git
The following commit(s) were added to refs/heads/master by this push:
new f993f2db3 feat(hubble): close loaders on task exit (#814)
f993f2db3 is described below
commit f993f2db30c3683b9de57fc0a34233510f7b6d52
Author: imbajin <ji...@apache.org>
AuthorDate: Fri Oct 16 22:46:17 2026 +0800
feat(hubble): close loaders on task exit (#814)
Hubble could leave a Loader open after task exit or failed preparation/dispatch. Closing it also exposed an unsafe checkpoint: parsed rows could be skipped by Resume even when they had never been inserted.
---
hugegraph-hubble/README.md | 58 +
.../controller/ingest/IngestController.java | 2 +-
.../controller/load/FileMappingController.java | 22 +-
.../controller/load/FileUploadController.java | 358 ++--
.../org/apache/hugegraph/entity/load/LoadTask.java | 95 +-
.../hugegraph/service/load/FileMappingService.java | 64 +-
.../hugegraph/service/load/JobManagerService.java | 161 +-
.../hugegraph/service/load/LoadTaskService.java | 905 +++++++++-
.../load/IngestTransactionIntegrationTest.java | 699 +++++++-
.../service/load/JobAdmissionIntegrationTest.java | 1822 ++++++++++++++++++++
.../testutil/NativeCleanupTestSupport.java | 171 ++
.../hugegraph/unit/FileUploadControllerTest.java | 127 +-
.../hugegraph/unit/JobManagerServiceTest.java | 205 ++-
.../hugegraph/unit/LoadTaskCheckpointTest.java | 259 +++
.../apache/hugegraph/unit/LoadTaskServiceTest.java | 431 +++++
.../hugegraph/unit/LoaderScopeControllerTest.java | 4 +
.../org/apache/hugegraph/unit/UnitTestSuite.java | 2 +
hugegraph-loader/README.md | 4 +
.../apache/hugegraph/loader/HugeGraphLoader.java | 77 +-
.../hugegraph/loader/executor/LoadContext.java | 6 +-
.../loader/test/unit/LoadContextTest.java | 360 ++++
21 files changed, 5502 insertions(+), 330 deletions(-)
diff --git a/hugegraph-hubble/README.md b/hugegraph-hubble/README.md
index 9338594eb..75ffe2c1d 100644
--- a/hugegraph-hubble/README.md
+++ b/hugegraph-hubble/README.md
@@ -49,6 +49,63 @@ The compatibility boundary is deliberately small:
Use matching Server, PD, and Store minor versions in a distributed cluster. **Server/PD 1.8 or later is strongly recommended for the best Hubble experience.** Support for 1.5 and 1.7 is intentionally limited to minimum usability and does not imply feature parity with the current release. Version checks stay in the client adapter/resolver rather than being scattered through controllers or pages. See [`AGENTS.md`](AGENTS.md) for verification rules.
+## Import resource cleanup
+
+Import tasks close their Loader before publishing completion. The built-in example import also closes its Loader.
+A failed save/update or dispatch closes the unused Loader, including resumed and retried tasks.
+Rollback of the outer creation transaction also closes the new Loader.
+Dispatch rejection marks the task and its job as failed. An independent shutdown failure marks the task as failed;
+check the logs and graph loading mode before retrying. A shutdown that only rethrows the existing load failure preserves
+a user-requested paused or stopped status.
+Hubble closes its Loaders without saving parsed file offsets as checkpoints. Pause/resume and FAILED-task resume
+therefore reread the source from its existing checkpoint (or the beginning), including rows that were parsed but never
+inserted. Already inserted rows may be replayed, as before this cleanup change. The standalone Loader retains its
+existing checkpoint behavior. Checkpoints written by older builds are not invalidated automatically. If their progress
+may be unsafe, stop a paused task and use Retry to reread the entire source instead of resuming that checkpoint.
+Deleting a job signals all accepted tasks to stop before waiting for their native cleanup and final status callbacks.
+Single-task, job and source-file deletion wait for those callbacks before starting the short metadata transaction.
+An internal caller that already owns a database transaction is refused before changing data if callbacks remain pending;
+finish that transaction and retry. Completed tasks can still be deleted within the caller's existing transaction.
+A final callback committed before deletion remains committed if that deletion later rolls back. A task can complete cleanup after the
+task-list snapshot. Existing task IDs stay reserved
+against retry/resume until the deletion transaction commits or rolls back. An independent cleanup failure preserves
+the task metadata and runtime Loader, including when unused-Loader shutdown fails. Single-task deletion uses the same
+stop and cleanup checks. Runtime handles remain until deletion commits. After rollback, an already stopped task stays
+in memory while its persisted status may still be RUNNING or PAUSED. Ordinary stop can persist STOPPED using that
+completed-cleanup certificate without closing the Loader again; another deletion can also reuse it.
+Task creation and job deletion are serialized within the same Hubble process. A new task checks that its parent job
+and graph still exist before constructing a Loader or inserting metadata; deletion includes tasks admitted before its
+snapshot. Runtime ownership is published before asynchronous execution can begin. Multi-file requests keep their
+existing per-file behavior; these guards do not coordinate separate Hubble processes.
+The task-page deletion route uses the same job cleanup. Upload requests and source-file detach share the job
+reservation, so deletion includes an admitted upload and a detached source cannot start from a stale mapping.
+Different files of the same job can upload concurrently. Chunks for the same job and filename are serialized so a merge
+only sees completed chunks. This limits concurrency within a single file, while retaining parallel I/O across files.
+Tokens are opaque and bound to the job and filename. Committed job deletion invalidates even unused tokens;
+committed mapping deletion invalidates its pending upload tokens. Rollback keeps them available. Background
+cleanup also invalidates the affected upload before deleting its files, without cancelling same-name files in another job.
+Clear-all takes the job admission gate before reading its mapping list, so it includes uploads already in progress.
+For partial uploads, different names that the configured filesystem resolves to the same directory are rejected
+before writing data. Native filesystem identity is used; distinct case-sensitive filenames remain allowed.
+Token issuance waits for same-file failed-upload cleanup before returning a replacement token.
+Token issuance does not create upload directories; a failed mapping reservation removes only its new empty directory.
+Mappings needed by running, paused, failed or stopped tasks remain attached. Cleanup failures before a task INSERT
+or after transaction rollback keep their runtime handle and prevent parent-job deletion even if no task row remains.
+Retry/resume refuse to replace unfinished or cleanup-failed runtimes. After native cleanup, they wait for the previous
+status callback; its failure is surfaced without starting a replacement Loader. Stop and drain all imports, have an administrator
+inspect and recover the graph mode, then restart Hubble before retrying.
+
+Database failures during progress polling or after rejected dispatch require operator recovery; this change does not
+add automatic per-task persistence repair. A failed final write can leave a task RUNNING or its job LOADING in metadata.
+While a runtime remains in this Hubble process, unresolved callback/cleanup evidence blocks unsafe replacement or deletion.
+After an abnormal process exit, the old runtime is gone: a leftover RUNNING task can be deleted directly. This removes
+Hubble metadata only; it does not restore Server graph mode or release any leftover Server loading owner.
+Server recovery is a separate operation. Restore the metadata database; before starting new imports, stop related imports
+and have an administrator inspect and recover the Server graph mode/owner. If keeping the task for another attempt,
+restart Hubble with the same metadata and source files, use the existing Pause action for a persisted RUNNING row,
+then Resume (or Stop followed by Retry when old checkpoints are unsafe) with current authentication. These actions
+do not repair a database that is still unavailable.
+
## Local development feedback loop
Run the frontend with third-party source-map noise disabled:
@@ -250,6 +307,7 @@ Hubble 1.8 requires Java 17 and uses Spring Boot 3 with H2 2.x for metadata. MyS
The startup script includes `--add-opens=java.base/java.net=ALL-UNNAMED` for the embedded Loader's Hive ORC reader. Include this option when using a custom Java launch command. Packaged archives are generated under `target/` and exclude local runtime data.
+
## Graph switching
Query execution waits until the graph context matches the current route, including keyboard shortcuts during a graph switch. Standalone Server 1.5 supports core graph, schema and data operations; GraphSpace management and PD mode require a newer server.
diff --git a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/ingest/IngestController.java b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/ingest/IngestController.java
index 3fadd9af0..dc1bd7791 100644
--- a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/ingest/IngestController.java
+++ b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/ingest/IngestController.java
@@ -378,7 +378,7 @@ public class IngestController extends BaseController {
.message("Task not found: " + id).build();
}
this.requireJobWrite(job);
- jobManagerService.remove(id);
+ jobManagerService.deleteJob(id);
return Response.builder().status(Constant.STATUS_OK).build();
}
diff --git a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/load/FileMappingController.java b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/load/FileMappingController.java
index 43381c508..c4ed8f67f 100644
--- a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/load/FileMappingController.java
+++ b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/load/FileMappingController.java
@@ -127,16 +127,20 @@ public class FileMappingController extends BaseController {
@PathVariable("graph") String graph,
@PathVariable("jobId") int jobId) {
this.requireGraphSpaceWrite(graphSpace);
- List<FileMapping> mappings = this.service.listByJob(graphSpace, graph,
- jobId);
- Set<Integer> fileIds = new HashSet<>();
- for (FileMapping mapping : mappings) {
- fileIds.add(mapping.getId());
- }
- // Check them all upfront, don't delete part of the mappings then fail
- this.checkNoRunningTask(jobId, fileIds);
+ try (var reservation = this.jobService.reserveJob(jobId)) {
+ Ex.check(this.jobService.get(graphSpace, graph, jobId) != null,
+ "job-manager.not-exist.id", jobId);
+ List<FileMapping> mappings = this.service.listByJob(graphSpace, graph,
+ jobId);
+ Set<Integer> fileIds = new HashSet<>();
+ for (FileMapping mapping : mappings) {
+ fileIds.add(mapping.getId());
+ }
+ // Check them all upfront, don't delete part of the mappings then fail
+ this.checkNoRunningTask(jobId, fileIds);
- this.jobService.deleteMappings(jobId, mappings);
+ this.jobService.deleteMappings(jobId, mappings);
+ }
}
@PostMapping("{id}/file-setting")
diff --git a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/load/FileUploadController.java b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/load/FileUploadController.java
index cbbe7cbc6..540ec31dd 100644
--- a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/load/FileUploadController.java
+++ b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/controller/load/FileUploadController.java
@@ -23,6 +23,8 @@ import static org.apache.hugegraph.service.load.FileMappingService.JOB_PREIFX;
import java.io.File;
import java.io.IOException;
+import java.nio.file.FileAlreadyExistsException;
+import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Collections;
@@ -88,30 +90,35 @@ public class FileUploadController extends BaseController {
@RequestParam("names")
List<String> fileNames) {
this.requireGraphSpaceWrite(graphSpace);
- Ex.check(this.jobService.get(graphSpace, graph, jobId) != null,
- "job-manager.not-exist.id", jobId);
- Ex.check(CollectionUtil.allUnique(fileNames),
- "load.upload.file.duplicate-name");
- Map<String, String> tokens = new HashMap<>();
- Map<String, ReadWriteLock> reservations = new HashMap<>();
- try {
- for (String fileName : fileNames) {
- this.checkFileNameValid(fileName);
- String token = this.service.generateFileToken(fileName);
- ReadWriteLock lock = new ReentrantReadWriteLock();
- // Atomic reservation to avoid a containsKey-then-put race
- ReadWriteLock previous = this.uploadingTokenLocks()
- .putIfAbsent(token, lock);
- Ex.check(previous == null, "load.upload.file.token.existed");
- reservations.put(token, lock);
- tokens.put(fileName, token);
+ try (var reservation = this.jobService.reserveUpload(jobId)) {
+ Ex.check(this.jobService.get(graphSpace, graph, jobId) != null,
+ "job-manager.not-exist.id", jobId);
+ Ex.check(CollectionUtil.allUnique(fileNames),
+ "load.upload.file.duplicate-name");
+ Map<String, String> tokens = new HashMap<>();
+ Map<String, ReadWriteLock> reservations = new HashMap<>();
+ try {
+ for (String fileName : fileNames) {
+ this.checkFileNameValid(fileName);
+ synchronized (this.service) {
+ String token = this.service.generateFileToken(jobId, fileName);
+ // Aliases for one physical file share its existing native lock.
+ ReadWriteLock lock = this.uploadingTokenLocks().entrySet().stream()
+ .filter(entry -> this.service.matchesFileToken(jobId, fileName, entry.getKey()))
+ .map(Map.Entry::getValue).findFirst().orElseGet(ReentrantReadWriteLock::new);
+ ReadWriteLock previous = this.uploadingTokenLocks().putIfAbsent(token, lock);
+ Ex.check(previous == null, "load.upload.file.token.existed");
+ reservations.put(token, lock);
+ tokens.put(fileName, token);
+ }
+ }
+ return tokens;
+ } catch (RuntimeException e) {
+ reservations.forEach((token, lock) -> {
+ this.uploadingTokenLocks().remove(token, lock);
+ });
+ throw e;
}
- return tokens;
- } catch (RuntimeException e) {
- reservations.forEach((token, lock) -> {
- this.uploadingTokenLocks().remove(token, lock);
- });
- throw e;
}
}
@@ -128,100 +135,108 @@ public class FileUploadController extends BaseController {
@RequestParam("total") int total,
@RequestParam("index") int index) {
this.requireGraphSpaceWrite(graphSpace);
- this.checkTotalAndIndexValid(total, index);
- this.checkFileNameValid(fileName);
- this.checkFileNameMatchToken(fileName, token);
- JobManager jobEntity = this.jobService.get(graphSpace, graph, jobId);
- Long sourceFileSize = this.resolveSourceFileSize(file, fileSize,
- total, index);
- this.checkFileValid(graphSpace, graph, jobId, jobEntity, file, fileName);
- if (jobEntity.getJobStatus() == JobStatus.DEFAULT) {
- jobEntity.setJobStatus(JobStatus.UPLOADING);
- this.jobService.update(jobEntity);
- }
- // Ensure location exist and generate file path
- String filePath = this.generateFilePath(graphSpace, graph, jobId,
- fileName);
- // Check this file deleted before
- ReadWriteLock lock = this.uploadingTokenLocks().get(token);
- FileUploadResult result;
- if (lock == null) {
- result = new FileUploadResult();
- result.setName(file.getOriginalFilename());
- result.setSize(file.getSize());
- result.setStatus(FileUploadResult.Status.FAILURE);
- result.setCause("File has been deleted");
- return result;
- }
-
- lock.readLock().lock();
- try {
- FileMapping reservedMapping;
- synchronized (this.service) {
- reservedMapping = this.reserveUploadQuota(graphSpace, graph, jobId,
- fileName, filePath,
- sourceFileSize);
+ try (var reservation = this.jobService.reserveUpload(jobId)) {
+ this.checkTotalAndIndexValid(total, index);
+ this.checkFileNameValid(fileName);
+ this.checkFileNameMatchToken(jobId, fileName, token);
+ JobManager jobEntity = this.jobService.get(graphSpace, graph, jobId);
+ Long sourceFileSize = this.resolveSourceFileSize(file, fileSize,
+ total, index);
+ this.checkFileValid(graphSpace, graph, jobId, jobEntity, file, fileName);
+ // Check this file deleted before
+ ReadWriteLock lock = this.uploadingTokenLocks().get(token);
+ FileUploadResult result;
+ if (lock == null) {
+ return this.deletedUploadResult(file);
}
- result = this.service.uploadFile(file, fileName, index, filePath);
- if (result.getStatus() == FileUploadResult.Status.FAILURE) {
- return result;
+
+ if (jobEntity.getJobStatus() == JobStatus.DEFAULT) {
+ jobEntity.setJobStatus(JobStatus.UPLOADING);
+ this.jobService.update(jobEntity);
}
- synchronized (this.service) {
- // Verify the existence of fragmented files
- FileMapping mapping = this.service.get(graphSpace, graph, jobId,
- fileName);
- if (mapping == null) {
- mapping = reservedMapping;
+ lock.writeLock().lock();
+ try {
+ if (this.uploadingTokenLocks().get(token) != lock) {
+ return this.deletedUploadResult(file);
}
- Ex.check(mapping != null, "load.file-mapping.not-exist.name",
- fileName);
- if (mapping.getFileStatus() == FileMappingStatus.COMPLETED) {
- result.setId(mapping.getId());
- this.uploadingTokenLocks().remove(token);
- return result;
+ // Do not recreate directories for a revoked token waiting on its lock.
+ String filePath = this.generateFilePath(graphSpace, graph, jobId, fileName);
+ FileMapping reservedMapping;
+ synchronized (this.service) {
+ reservedMapping = this.reserveUploadQuota(graphSpace, graph, jobId,
+ fileName, filePath,
+ sourceFileSize);
}
- mapping.setUpdateTime(HubbleUtil.nowDate());
- // Determine whether all the parts have been uploaded, then merge them
- boolean merged = this.service.tryMergePartFiles(filePath, total);
- if (!merged) {
- this.service.update(mapping);
+ result = this.service.uploadFile(file, fileName, index, filePath);
+ if (result.getStatus() == FileUploadResult.Status.FAILURE) {
return result;
}
- JobManager currentJob = this.jobService.get(graphSpace, graph,
- jobId);
- long actualFileSize;
- try {
- actualFileSize = this.resolveUploadedFileSize(
- mapping.getPath());
- Ex.check(currentJob != null, "job-manager.not-exist.id",
- jobId);
- long reservedUploadingSize =
- this.sumReservedUploadingSize(jobId,
- mapping.getId());
- this.checkFileSizeLimit(actualFileSize,
- currentJob.getJobSize(),
- reservedUploadingSize);
- } catch (RuntimeException e) {
- this.cleanupFailedUpload(mapping, token);
- throw e;
+ synchronized (this.service) {
+ // Verify the existence of fragmented files
+ FileMapping mapping = this.service.get(graphSpace, graph, jobId,
+ fileName);
+ if (mapping == null) {
+ mapping = reservedMapping;
+ }
+ Ex.check(mapping != null, "load.file-mapping.not-exist.name",
+ fileName);
+ if (mapping.getFileStatus() == FileMappingStatus.COMPLETED) {
+ result.setId(mapping.getId());
+ this.service.removeUploadingTokens(jobId, fileName);
+ return result;
+ }
+ mapping.setUpdateTime(HubbleUtil.nowDate());
+ // Determine whether all the parts have been uploaded, then merge them
+ boolean merged = this.service.tryMergePartFiles(filePath, total);
+ if (!merged) {
+ this.service.update(mapping);
+ return result;
+ }
+ JobManager currentJob = this.jobService.get(graphSpace, graph,
+ jobId);
+ long actualFileSize;
+ try {
+ actualFileSize = this.resolveUploadedFileSize(
+ mapping.getPath());
+ Ex.check(currentJob != null, "job-manager.not-exist.id",
+ jobId);
+ long reservedUploadingSize =
+ this.sumReservedUploadingSize(jobId,
+ mapping.getId());
+ this.checkFileSizeLimit(actualFileSize,
+ currentJob.getJobSize(),
+ reservedUploadingSize);
+ } catch (RuntimeException e) {
+ this.cleanupFailedUpload(mapping, token);
+ throw e;
+ }
+ // Read column names and values then fill it
+ this.service.extractColumns(mapping);
+ mapping.setFileStatus(FileMappingStatus.COMPLETED);
+ mapping.setTotalLines(FileUtil.countLines(mapping.getPath()));
+ mapping.setTotalSize(actualFileSize);
+
+ this.finalizeUpload(mapping);
+ result.setId(mapping.getId());
+ // Remove uploading file token
+ this.service.removeUploadingTokens(jobId, fileName);
}
- // Read column names and values then fill it
- this.service.extractColumns(mapping);
- mapping.setFileStatus(FileMappingStatus.COMPLETED);
- mapping.setTotalLines(FileUtil.countLines(mapping.getPath()));
- mapping.setTotalSize(actualFileSize);
-
- this.finalizeUpload(mapping);
- result.setId(mapping.getId());
- // Remove uploading file token
- this.uploadingTokenLocks().remove(token);
+ return result;
+ } finally {
+ lock.writeLock().unlock();
}
- return result;
- } finally {
- lock.readLock().unlock();
}
}
+ private FileUploadResult deletedUploadResult(MultipartFile file) {
+ FileUploadResult result = new FileUploadResult();
+ result.setName(file.getOriginalFilename());
+ result.setSize(file.getSize());
+ result.setStatus(FileUploadResult.Status.FAILURE);
+ result.setCause("File has been deleted");
+ return result;
+ }
+
private void finalizeUpload(FileMapping mapping) {
String originalPath = mapping.getPath();
String movedPath = this.service.moveToNextLevelDir(mapping);
@@ -256,31 +271,31 @@ public class FileUploadController extends BaseController {
@RequestParam("name") String fileName,
@RequestParam("token") String token) {
this.requireGraphSpaceWrite(graphSpace);
- this.checkFileNameValid(fileName);
- JobManager jobEntity = this.jobService.get(graphSpace, graph, jobId);
- Ex.check(jobEntity != null, "job-manager.not-exist.id", jobId);
- Ex.check(jobEntity.getJobStatus() == JobStatus.UPLOADING ||
- jobEntity.getJobStatus() == JobStatus.MAPPING ||
- jobEntity.getJobStatus() == JobStatus.SETTING,
- "deleted.file.no-permission");
- FileMapping mapping = this.service.get(graphSpace, graph, jobId,
- fileName);
- Ex.check(mapping != null, "load.file-mapping.not-exist.name", fileName);
-
- ReadWriteLock lock = this.uploadingTokenLocks().get(token);
- if (lock != null) {
- lock.writeLock().lock();
- }
- try {
- this.jobService.deleteMappings(jobId,
- Collections.singletonList(mapping));
+ try (var reservation = this.jobService.reserveJob(jobId)) {
+ this.checkFileNameValid(fileName);
+ this.checkFileNameMatchToken(jobId, fileName, token);
+ JobManager jobEntity = this.jobService.get(graphSpace, graph, jobId);
+ Ex.check(jobEntity != null, "job-manager.not-exist.id", jobId);
+ Ex.check(jobEntity.getJobStatus() == JobStatus.UPLOADING ||
+ jobEntity.getJobStatus() == JobStatus.MAPPING ||
+ jobEntity.getJobStatus() == JobStatus.SETTING,
+ "deleted.file.no-permission");
+ FileMapping mapping = this.service.get(graphSpace, graph, jobId,
+ fileName);
+ Ex.check(mapping != null, "load.file-mapping.not-exist.name", fileName);
+
+ ReadWriteLock lock = this.uploadingTokenLocks().get(token);
if (lock != null) {
- this.uploadingTokenLocks().remove(token);
+ lock.writeLock().lock();
}
- return true;
- } finally {
- if (lock != null) {
- lock.writeLock().unlock();
+ try {
+ this.jobService.deleteMappings(jobId,
+ Collections.singletonList(mapping));
+ return true;
+ } finally {
+ if (lock != null) {
+ lock.writeLock().unlock();
+ }
}
}
}
@@ -317,9 +332,8 @@ public class FileUploadController extends BaseController {
}
}
- private void checkFileNameMatchToken(String fileName, String token) {
- String md5Prefix = HubbleUtil.md5(fileName);
- Ex.check(StringUtils.isNotEmpty(token) && token.startsWith(md5Prefix),
+ private void checkFileNameMatchToken(int jobId, String fileName, String token) {
+ Ex.check(this.service.matchesFileToken(jobId, fileName, token),
"load.upload.file.name-token.unmatch");
}
@@ -382,18 +396,80 @@ public class FileUploadController extends BaseController {
this.checkFileSizeLimit(reservedFileSize, currentJob.getJobSize(),
reservedUploadingSize);
- if (mapping == null) {
- mapping = new FileMapping(graphSpace, graph, fileName, filePath);
- mapping.setJobId(jobId);
+ Path directory = this.service.requirePathUnderUploadRoot(filePath).toPath();
+ // Later chunks reuse the directory already reserved by this UPLOADING mapping.
+ boolean needsReservation = mapping == null || !filePath.equals(mapping.getPath()) ||
+ !Files.isDirectory(directory);
+ boolean createdDirectory = needsReservation &&
+ this.reservePhysicalUploadDirectory(jobId, fileName, filePath);
+ try {
+ if (mapping == null) {
+ mapping = new FileMapping(graphSpace, graph, fileName, filePath);
+ mapping.setJobId(jobId);
+ this.fillUploadingReservation(mapping, reservedFileSize);
+ this.service.save(mapping);
+ return mapping;
+ }
+
+ mapping.setPath(filePath);
this.fillUploadingReservation(mapping, reservedFileSize);
- this.service.save(mapping);
+ this.service.update(mapping);
return mapping;
+ } catch (RuntimeException | Error failure) {
+ if (createdDirectory) {
+ this.removeUnusedUploadDirectory(Path.of(filePath), failure);
+ }
+ throw failure;
+ }
+ }
+
+ private boolean reservePhysicalUploadDirectory(int jobId, String fileName, String filePath) {
+ Path directory = this.service.requirePathUnderUploadRoot(filePath).toPath();
+ boolean created = false;
+ try {
+ try {
+ Files.createDirectory(directory);
+ created = true;
+ } catch (FileAlreadyExistsException existing) {
+ if (!Files.isDirectory(directory)) {
+ throw existing;
+ }
+ }
+ // The leaf now exists, so use native filesystem identity rather than spelling rules.
+ directory = this.service.requirePathUnderUploadRoot(directory.toString()).toPath();
+ for (FileMapping current : this.service.listByJob(jobId)) {
+ if (fileName.equals(current.getName()) ||
+ (current.getFileStatus() != FileMappingStatus.UPLOADING &&
+ current.getFileStatus() != FileMappingStatus.FAILURE)) {
+ continue;
+ }
+ Path existing = this.service.requirePathUnderUploadRoot(current.getPath()).toPath();
+ Ex.check(!Files.exists(existing) || !Files.isSameFile(directory, existing),
+ "load.upload.file.existed", fileName);
+ }
+ return created;
+ } catch (IOException error) {
+ InternalException failure = new InternalException("Failed to reserve upload directory '%s'", error,
+ filePath);
+ if (created) {
+ this.removeUnusedUploadDirectory(directory, failure);
+ }
+ throw failure;
+ } catch (RuntimeException | Error failure) {
+ if (created) {
+ this.removeUnusedUploadDirectory(directory, failure);
+ }
+ throw failure;
}
+ }
- mapping.setPath(filePath);
- this.fillUploadingReservation(mapping, reservedFileSize);
- this.service.update(mapping);
- return mapping;
+ private void removeUnusedUploadDirectory(Path directory, Throwable failure) {
+ try {
+ // Only a directory created by this request is eligible, and only if it remains empty.
+ Files.deleteIfExists(directory);
+ } catch (IOException cleanupFailure) {
+ failure.addSuppressed(cleanupFailure);
+ }
}
private Long resolveSourceFileSize(MultipartFile file, Long fileSize,
@@ -446,7 +522,7 @@ public class FileUploadController extends BaseController {
}
private void cleanupFailedUpload(FileMapping mapping, String token) {
- this.uploadingTokenLocks().remove(token);
+ this.service.removeUploadingTokens(mapping.getJobId(), mapping.getName());
this.service.cleanupMappings(Collections.singletonList(mapping));
}
diff --git a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/entity/load/LoadTask.java b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/entity/load/LoadTask.java
index ce6fd04af..6f2a57212 100644
--- a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/entity/load/LoadTask.java
+++ b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/entity/load/LoadTask.java
@@ -28,6 +28,7 @@ import org.apache.hugegraph.entity.GraphConnection;
import org.apache.hugegraph.entity.enums.LoadStatus;
import org.apache.hugegraph.exception.InternalException;
import org.apache.hugegraph.loader.HugeGraphLoader;
+import org.apache.hugegraph.loader.exception.LoadException;
import org.apache.hugegraph.loader.executor.LoadContext;
import org.apache.hugegraph.loader.executor.LoadOptions;
import org.apache.hugegraph.util.Ex;
@@ -68,6 +69,14 @@ public class LoadTask implements Runnable {
@JsonIgnore
private transient volatile boolean finished;
+ @TableField(exist = false)
+ @JsonIgnore
+ private transient volatile Throwable cleanupFailure;
+
+ @TableField(exist = false)
+ @JsonIgnore
+ private transient volatile boolean cleanupFailed;
+
@TableId(type = IdType.AUTO)
@MergeProperty(useNew = false)
@JsonProperty("id")
@@ -169,7 +178,7 @@ public class LoadTask implements Runnable {
this.currDuration = 0L;
this.createTime = HubbleUtil.nowDate();
- this.loader = new HugeGraphLoader(this.options);
+ this.loader = new HugeGraphLoader(this.options, false);
}
@Override
@@ -177,18 +186,61 @@ public class LoadTask implements Runnable {
Ex.check(this.options != null, "The load options shouldn't be null");
log.info("LoadTask is start running : {}", this.id);
- boolean noError;
+ this.cleanupFailure = null;
+ boolean noError = false;
+ this.cleanupFailed = false;
+ Throwable failure = null;
try {
- noError = this.loader.load();
+ // A queued task can be paused before its worker starts.
+ noError = this.context().stopped() ? this.context().noError() : this.loader.load();
} catch (Throwable e) {
- noError = false;
- log.error("Run task {} failed", this.id, e);
+ failure = e;
+ } finally {
+ Throwable recordedFailure = failure;
+ if (failure != null) {
+ // Only read the recorded error: shutdown can wrap a checked
+ // failure differently from load(), or rethrow it unchanged.
+ try {
+ this.context().throwIfFailed();
+ } catch (Throwable recorded) {
+ recordedFailure = recorded;
+ if (recorded != failure) {
+ // Retain an asynchronously recorded error even when it
+ // differs from the error propagated by load().
+ log.error("Recorded load failure for task {}", this.id, recorded);
+ }
+ }
+ }
+ Throwable recordedCause = recordedFailure == null ? null : recordedFailure.getCause();
+ int suppressed = recordedFailure == null ? 0 : recordedFailure.getSuppressed().length;
+ int causeSuppressed = recordedCause == null ? 0 : recordedCause.getSuppressed().length;
+ try {
+ this.loader.shutdown();
+ } catch (Throwable e) {
+ noError = false;
+ this.cleanupFailed = !sameLoadFailure(recordedFailure, e) ||
+ recordedFailure.getSuppressed().length != suppressed ||
+ (recordedCause != null &&
+ recordedCause.getSuppressed().length != causeSuppressed);
+ if (this.cleanupFailed) {
+ this.cleanupFailure = e;
+ log.error("Failed to close load task {}; check graph loading mode before retrying", this.id, e);
+ }
+ if (failure == null) {
+ failure = e;
+ } else if (this.cleanupFailed && failure != e && !referencesFailure(e, failure)) {
+ failure.addSuppressed(e);
+ }
+ }
+ }
+ if (failure != null) {
+ log.error("Run task {} failed", this.id, failure);
}
this.lock.lock();
try {
// Pay attention to whether the user stops actively or
// the program stops by itself
- if (this.status.inRunning()) {
+ if (this.cleanupFailed || this.status.inRunning()) {
if (noError) {
this.status = LoadStatus.SUCCEED;
} else {
@@ -209,6 +261,31 @@ public class LoadTask implements Runnable {
}
}
+ private static boolean sameLoadFailure(Throwable recorded, Throwable cleanup) {
+ if (recorded == null) {
+ return false;
+ }
+ if (recorded == cleanup) {
+ return true;
+ }
+ // LoadContext.throwIfFailed() wraps only non-runtime failures this way.
+ return recorded.getClass() == LoadException.class && cleanup.getClass() == LoadException.class &&
+ recorded.getCause() != null && !(recorded.getCause() instanceof RuntimeException) &&
+ recorded.getCause() == cleanup.getCause() && cleanup.getSuppressed().length == 0;
+ }
+
+ private static boolean referencesFailure(Throwable cleanup, Throwable failure) {
+ if (cleanup.getCause() == failure) {
+ return true;
+ }
+ for (Throwable suppressed : cleanup.getSuppressed()) {
+ if (suppressed == failure) {
+ return true;
+ }
+ }
+ return false;
+ }
+
public void lock() {
this.lock.lock();
}
@@ -226,7 +303,9 @@ public class LoadTask implements Runnable {
// pass
}
}
- this.loader = null;
+ if (this.context().closed() && !this.cleanupFailed) {
+ this.loader = null;
+ }
log.info("LoadTask {} stopped", this.id);
}
@@ -240,7 +319,7 @@ public class LoadTask implements Runnable {
runtimeOptions.token = token;
runtimeOptions.pdToken = null;
runtimeOptions.trustStoreToken = null;
- this.loader = new HugeGraphLoader(runtimeOptions);
+ this.loader = new HugeGraphLoader(runtimeOptions, false);
this.finished = false;
} catch (CloneNotSupportedException e) {
throw new InternalException("Failed to prepare load task options",
diff --git a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/FileMappingService.java b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/FileMappingService.java
index 0e531df83..5ae13458d 100644
--- a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/FileMappingService.java
+++ b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/FileMappingService.java
@@ -36,7 +36,6 @@ import java.util.Arrays;
import java.util.Collections;
import java.util.Date;
import java.util.HashSet;
-import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -96,6 +95,8 @@ public class FileMappingService {
private FileMappingMapper mapper;
@Autowired
private JobManagerMapper jobManagerMapper;
+ @Autowired
+ private LoadTaskService taskService;
private final Map<String, ReadWriteLock> uploadingTokenLocks;
@@ -190,18 +191,17 @@ public class FileMappingService {
}
}
- public String generateFileToken(String fileName) {
- /*
- * Append a short random suffix so that tokens cannot be guessed
- * from the file name and upload time alone. Note that
- * checkFileNameMatchToken() only checks the md5 prefix, so the
- * token still passes that check.
- */
- return this.fileTokenPrefix(fileName) +
+ public String generateFileToken(int jobId, String fileName) {
+ // Tokens are opaque to clients and bound to this job and filename.
+ return this.fileTokenPrefix(jobId, fileName) +
HubbleUtil.nowTime().getEpochSecond() + "-" +
randomTokenSuffix();
}
+ public boolean matchesFileToken(int jobId, String fileName, String token) {
+ return token != null && token.startsWith(this.fileTokenPrefix(jobId, fileName));
+ }
+
private static String randomTokenSuffix() {
StringBuilder suffix = new StringBuilder(TOKEN_SUFFIX_LENGTH);
for (int i = 0; i < TOKEN_SUFFIX_LENGTH; i++) {
@@ -529,6 +529,9 @@ public class FileMappingService {
}
private void tryCleanupMapping(FileMapping mapping) {
+ if (mapping.getJobId() != null) {
+ this.removeUploadingTokens(mapping.getJobId(), mapping.getName());
+ }
try {
this.deleteDiskFile(mapping);
this.removeCleanupRecord(mapping.getId());
@@ -585,6 +588,26 @@ public class FileMappingService {
}
private void tryDeleteUnfinishedMapping(FileMapping mapping) {
+ if (mapping.getJobId() == null) {
+ this.tryCleanupMapping(mapping);
+ return;
+ }
+ try (LoadTaskService.TaskReservation reservation = this.taskService.reserveJob(mapping.getJobId())) {
+ FileMapping current = this.get(mapping.getId());
+ if (current == null || (current.getFileStatus() != FileMappingStatus.UPLOADING &&
+ current.getFileStatus() != FileMappingStatus.FAILURE)) {
+ return;
+ }
+ long threshold = this.config.get(HubbleOptions.UPLOAD_FILE_MAX_TIME_CONSUMING) * 1000;
+ if (HubbleUtil.nowDate().getTime() - current.getUpdateTime().getTime() <= threshold) {
+ return;
+ }
+ this.removeUploadingTokens(mapping.getJobId(), mapping.getName());
+ this.deleteUnfinishedMapping(current);
+ }
+ }
+
+ private void deleteUnfinishedMapping(FileMapping mapping) {
String filePath = mapping.getPath();
try {
File file = this.requirePathUnderUploadRoot(filePath);
@@ -601,22 +624,19 @@ public class FileMappingService {
return;
}
this.remove(mapping.getId());
- this.removeUploadingTokens(mapping.getName());
}
- private void removeUploadingTokens(String fileName) {
- String tokenPrefix = this.fileTokenPrefix(fileName);
- Iterator<Map.Entry<String, ReadWriteLock>> iter;
- iter = this.uploadingTokenLocks.entrySet().iterator();
- while (iter.hasNext()) {
- Map.Entry<String, ReadWriteLock> entry = iter.next();
- if (entry.getKey().startsWith(tokenPrefix)) {
- iter.remove();
- }
- }
+ public void removeUploadingTokens(int jobId, String fileName) {
+ String prefix = this.fileTokenPrefix(jobId, fileName);
+ this.uploadingTokenLocks.keySet().removeIf(token -> token.startsWith(prefix));
+ }
+
+ public void removeJobUploadingTokens(int jobId) {
+ String prefix = jobId + "-";
+ this.uploadingTokenLocks.keySet().removeIf(token -> token.startsWith(prefix));
}
- private String fileTokenPrefix(String fileName) {
- return HubbleUtil.md5(fileName) + "-";
+ private String fileTokenPrefix(int jobId, String fileName) {
+ return jobId + "-" + HubbleUtil.md5(fileName) + "-";
}
}
diff --git a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/JobManagerService.java b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/JobManagerService.java
index deb05a8f7..12af1707a 100644
--- a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/JobManagerService.java
+++ b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/JobManagerService.java
@@ -42,12 +42,16 @@ import org.apache.hugegraph.entity.load.LoadTask;
import org.apache.hugegraph.exception.ExternalException;
import org.apache.hugegraph.exception.InternalException;
import org.apache.hugegraph.mapper.load.JobManagerMapper;
+import org.apache.hugegraph.util.Ex;
import org.apache.hugegraph.util.HubbleUtil;
import org.apache.hugegraph.util.PageUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
+import org.springframework.transaction.PlatformTransactionManager;
+import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.annotation.Isolation;
import org.springframework.transaction.annotation.Transactional;
+import org.springframework.transaction.support.TransactionTemplate;
import org.springframework.transaction.support.TransactionSynchronization;
import org.springframework.transaction.support.TransactionSynchronizationManager;
@@ -61,6 +65,8 @@ public class JobManagerService {
private LoadTaskService taskService;
@Autowired
private FileMappingService fileMappingService;
+ @Autowired
+ private PlatformTransactionManager transactionManager;
public int count() {
return Math.toIntExact(this.mapper.selectCount(null));
@@ -228,27 +234,58 @@ public class JobManagerService {
}
}
- @Transactional(isolation = Isolation.READ_COMMITTED)
public void deleteMappings(int jobId, List<FileMapping> mappings) {
- if (mappings.isEmpty()) {
- return;
- }
- long releasedSize = 0L;
- for (FileMapping mapping : mappings) {
- if (mapping.getJobId() == null || mapping.getJobId() != jobId) {
- throw new InternalException("File mapping does not belong to " +
- "job %s", jobId);
+ try (LoadTaskService.TaskReservation reservation = this.reserveJob(jobId)) {
+ if (mappings.isEmpty()) {
+ return;
}
- if (mapping.getFileStatus() == FileMappingStatus.COMPLETED) {
- releasedSize = Math.addExact(releasedSize,
- mapping.getTotalSize());
+ Ex.check(this.get(jobId) != null, "job-manager.not-exist.id", jobId);
+ List<LoadTask> tasks = this.taskService.taskListByJob(jobId);
+ List<FileMapping> currentMappings = new ArrayList<>();
+ long releasedSize = 0L;
+ for (FileMapping mapping : mappings) {
+ if (mapping.getJobId() == null || mapping.getJobId() != jobId) {
+ throw new InternalException("File mapping does not belong to " +
+ "job %s", jobId);
+ }
+ FileMapping current = this.fileMappingService.get(mapping.getId());
+ if (current == null || current.getJobId() == null || current.getJobId() != jobId) {
+ throw new InternalException("File mapping does not belong to job %s", jobId);
+ }
+ for (LoadTask task : tasks) {
+ LoadStatus status = task.getStatus();
+ boolean needsSource = status == LoadStatus.RUNNING || status == LoadStatus.PAUSED ||
+ status == LoadStatus.FAILED || status == LoadStatus.STOPPED;
+ Ex.check(!needsSource || !java.util.Objects.equals(task.getFileId(), current.getId()),
+ "load.file-mapping.delete.task-in-running", current.getId());
+ if (java.util.Objects.equals(task.getFileId(), current.getId())) {
+ this.taskService.checkCallbackWaitAllowed(task.getId());
+ }
+ }
+ if (current.getFileStatus() == FileMappingStatus.COMPLETED) {
+ releasedSize = Math.addExact(releasedSize,
+ current.getTotalSize());
+ }
+ currentMappings.add(current);
}
- this.fileMappingService.detachFromJob(mapping.getId(), jobId);
- }
- if (this.mapper.decreaseJobSize(jobId, releasedSize) != 1) {
- throw new InternalException("entity.update.failed", jobId);
+ for (LoadTask task : tasks) {
+ if (currentMappings.stream().anyMatch(mapping ->
+ java.util.Objects.equals(task.getFileId(), mapping.getId()))) {
+ this.taskService.awaitCallbacksForDelete(task.getId());
+ }
+ }
+ long size = releasedSize;
+ this.inDeleteTransaction(() -> {
+ for (FileMapping current : currentMappings) {
+ this.fileMappingService.detachFromJob(current.getId(), jobId);
+ }
+ if (this.mapper.decreaseJobSize(jobId, size) != 1) {
+ throw new InternalException("entity.update.failed", jobId);
+ }
+ this.deleteDiskFilesAfterCommit(jobId, currentMappings, false);
+ });
+
}
- this.deleteDiskFilesAfterCommit(mappings);
}
@Transactional(isolation = Isolation.READ_COMMITTED)
@@ -264,28 +301,69 @@ public class JobManagerService {
"graph", graph));
}
- @Transactional(isolation = Isolation.READ_COMMITTED)
+ public LoadTaskService.TaskReservation reserveJob(int id) {
+ return this.taskService.reserveJob(id);
+ }
+
+ public LoadTaskService.TaskReservation reserveUpload(int id) {
+ return this.taskService.reserveUpload(id);
+ }
+
public void deleteJob(int id) {
- JobManager job = this.get(id);
- if (job == null) {
- throw new ExternalException("job.manager.not-exist.id", id);
- }
+ try (LoadTaskService.TaskReservation jobReservation = this.taskService.reserveJob(id)) {
+ JobManager job = this.get(id);
+ if (job == null) {
+ throw new ExternalException("job.manager.not-exist.id", id);
+ }
- List<LoadTask> loadTasks = this.taskService.taskListByJob(id);
- for (LoadTask loadTask : loadTasks) {
- if (loadTask.getStatus().inRunning() ||
- loadTask.getStatus() == LoadStatus.PAUSED) {
- this.taskService.stop(loadTask.getId());
+ this.taskService.ensureJobCleanupComplete(id);
+ List<LoadTask> loadTasks = this.taskService.taskListByJob(id);
+ try (LoadTaskService.TaskReservation reservation = this.taskService.reserveTasks(
+ loadTasks.stream().map(LoadTask::getId).toList())) {
+ // The job gate keeps these IDs stable. Re-read after all task
+ // reservations, using selectById rather than the cached list query.
+ List<LoadTask> currentTasks = loadTasks.stream()
+ .map(task -> this.taskService.get(task.getId()))
+ .toList();
+ for (LoadTask task : loadTasks) {
+ this.taskService.checkCallbackWaitAllowed(task.getId());
+ }
+ // Publish every stop intent before waiting for a queued owner.
+ this.taskService.requestJobStopForDelete(id);
+ for (LoadTask task : loadTasks) {
+ this.taskService.awaitCallbacksForDelete(task.getId());
+ }
+ this.inDeleteTransaction(() -> {
+ for (int index = 0; index < loadTasks.size(); index++) {
+ LoadTask loadTask = loadTasks.get(index);
+ LoadTask current = currentTasks.get(index);
+ if (current == null) {
+ continue;
+ }
+ // An initially live row still needs its native certificate;
+ // a refreshed terminal status alone cannot provide one.
+ if (loadTask.getStatus().inRunning() || loadTask.getStatus() == LoadStatus.PAUSED ||
+ current.getStatus().inRunning() || current.getStatus() == LoadStatus.PAUSED) {
+ this.taskService.stopForDelete(loadTask.getId());
+ }
+ this.taskService.ensureRuntimeCleanupComplete(loadTask.getId());
+ this.taskService.remove(loadTask.getId());
+ }
+
+ List<FileMapping> mappings = this.fileMappingService.listByJob(id);
+ this.remove(id);
+ this.deleteDiskFilesAfterCommit(id, mappings, true);
+ });
}
- this.taskService.remove(loadTask.getId());
}
+ }
- List<FileMapping> mappings = this.fileMappingService.listByJob(id);
- this.remove(id);
- this.deleteDiskFilesAfterCommit(mappings);
+ private void inDeleteTransaction(Runnable action) {
+ TransactionTemplate transaction = new TransactionTemplate(this.transactionManager);
+ transaction.setIsolationLevel(TransactionDefinition.ISOLATION_READ_COMMITTED);
+ transaction.executeWithoutResult(status -> action.run());
}
- @Transactional(isolation = Isolation.READ_COMMITTED)
public void deleteJob(String graphSpace, String graph, int id) {
JobManager job = this.get(graphSpace, graph, id);
if (job == null) {
@@ -294,14 +372,10 @@ public class JobManagerService {
this.deleteJob(id);
}
- private void deleteDiskFilesAfterCommit(List<FileMapping> mappings) {
- if (mappings.isEmpty()) {
- return;
- }
-
+ private void deleteDiskFilesAfterCommit(int jobId, List<FileMapping> mappings, boolean wholeJob) {
List<FileMapping> copiedMappings = new ArrayList<>(mappings);
if (!TransactionSynchronizationManager.isSynchronizationActive()) {
- this.deleteDiskFiles(copiedMappings);
+ this.deleteDiskFiles(jobId, copiedMappings, wholeJob);
return;
}
@@ -309,12 +383,19 @@ public class JobManagerService {
new TransactionSynchronization() {
@Override
public void afterCommit() {
- deleteDiskFiles(copiedMappings);
+ deleteDiskFiles(jobId, copiedMappings, wholeJob);
}
});
}
- private void deleteDiskFiles(List<FileMapping> mappings) {
+ private void deleteDiskFiles(int jobId, List<FileMapping> mappings, boolean wholeJob) {
+ if (wholeJob) {
+ this.fileMappingService.removeJobUploadingTokens(jobId);
+ } else {
+ for (FileMapping mapping : mappings) {
+ this.fileMappingService.removeUploadingTokens(jobId, mapping.getName());
+ }
+ }
this.fileMappingService.cleanupMappings(mappings);
}
}
diff --git a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/LoadTaskService.java b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/LoadTaskService.java
index 23b3656c4..bc2a1796d 100644
--- a/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/LoadTaskService.java
+++ b/hugegraph-hubble/hubble-be/src/main/java/org/apache/hugegraph/service/load/LoadTaskService.java
@@ -54,6 +54,7 @@ import org.apache.hugegraph.loader.source.file.FileSource;
import org.apache.hugegraph.loader.util.MappingUtil;
import org.apache.hugegraph.loader.util.Printer;
import org.apache.hugegraph.mapper.load.JobManagerMapper;
+import org.apache.hugegraph.mapper.load.FileMappingMapper;
import org.apache.hugegraph.mapper.load.LoadTaskMapper;
import org.apache.hugegraph.service.schema.EdgeLabelService;
import org.apache.hugegraph.service.schema.VertexLabelService;
@@ -79,11 +80,21 @@ import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
+import java.util.Collections;
+import java.util.Objects;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CompletionException;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
+
+import com.google.common.util.concurrent.Striped;
+import java.util.concurrent.RejectedExecutionException;
import org.apache.commons.lang3.StringUtils;
@@ -98,6 +109,8 @@ public class LoadTaskService {
@Autowired
private LoadTaskMapper mapper;
@Autowired
+ private FileMappingMapper fileMappingMapper;
+ @Autowired
private JobManagerMapper jobManagerMapper;
@Autowired
private VertexLabelService vlService;
@@ -112,6 +125,137 @@ public class LoadTaskService {
private Map<Integer, LoadTask> runningTaskContainer;
+ private final Object admissionLock = new Object();
+ private final Set<AcceptedRun> acceptedRuns = new HashSet<>();
+ private boolean closing;
+ private final Striped<Lock> taskControls = Striped.lock(Runtime.getRuntime().availableProcessors() << 2);
+
+ private final Striped<ReadWriteLock> jobControls =
+ Striped.readWriteLock(Runtime.getRuntime().availableProcessors() << 2);
+
+ TaskReservation reserveJob(int id) {
+ return this.reserve(Collections.singletonList(this.jobControls.get(id).writeLock()));
+ }
+
+ TaskReservation reserveUpload(int id) {
+ return this.reserve(Collections.singletonList(this.jobControls.get(id).readLock()));
+ }
+
+ TaskReservation reserveTasks(List<Integer> ids) {
+ return this.reserve(this.taskControls, ids);
+ }
+
+ private TaskReservation reserve(Striped<Lock> controls, List<Integer> ids) {
+ List<Lock> locks = new ArrayList<>();
+ // Guava orders by stripe identity, including colliding task IDs.
+ controls.bulkGet(ids).forEach(locks::add);
+ return this.reserve(locks);
+ }
+
+ private TaskReservation reserve(List<Lock> locks) {
+ int acquired = 0;
+ try {
+ for (Lock lock : locks) {
+ lock.lock();
+ acquired++;
+ }
+ return new TaskReservation(locks);
+ } catch (RuntimeException | Error failure) {
+ for (int i = acquired - 1; i >= 0; i--) {
+ locks.get(i).unlock();
+ }
+ throw failure;
+ }
+ }
+
+ public final class TaskReservation implements AutoCloseable {
+
+ private final List<Lock> locks;
+ private boolean closed;
+
+ private TaskReservation(List<Lock> locks) {
+ this.locks = locks;
+ }
+
+ private void unlock() {
+ for (int i = this.locks.size() - 1; i >= 0; i--) {
+ this.locks.get(i).unlock();
+ }
+ }
+
+ @Override
+ public void close() {
+ if (this.closed) {
+ return;
+ }
+ this.closed = true;
+ if (TransactionSynchronizationManager.isActualTransactionActive() &&
+ TransactionSynchronizationManager.isSynchronizationActive()) {
+ try {
+ TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronization() {
+ @Override
+ public void afterCompletion(int status) {
+ unlock();
+ }
+ });
+ return;
+ } catch (RuntimeException | Error failure) {
+ this.unlock();
+ throw failure;
+ }
+ }
+ this.unlock();
+ }
+ }
+
+ private void forgetRuntimeAfterCommit(int taskId, LoadTask retained) {
+ Runnable forget = () -> {
+ if (retained != null) {
+ this.forgetRuntime(taskId, retained);
+ }
+ };
+ if (TransactionSynchronizationManager.isActualTransactionActive() &&
+ TransactionSynchronizationManager.isSynchronizationActive()) {
+ TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronization() {
+ @Override
+ public void afterCommit() {
+ forget.run();
+ }
+ });
+ } else {
+ forget.run();
+ }
+ }
+
+ void ensureJobCleanupComplete(int jobId) {
+ List<LoadTask> failures = new ArrayList<>();
+ synchronized (this.admissionLock) {
+ for (AcceptedRun run : this.acceptedRuns) {
+ failures.add(run.task);
+ }
+ }
+ failures.addAll(this.runningTaskContainer.values());
+ for (LoadTask task : failures) {
+ if (Objects.equals(task.getJobId(), jobId)) {
+ Ex.check(!task.isCleanupFailed(),
+ "Task cleanup is incomplete; inspect graph mode and recover before deleting the job");
+ }
+ }
+ }
+
+ void ensureRuntimeCleanupComplete(int taskId) {
+ LoadTask retained = this.runningTaskContainer.get(taskId);
+ if (retained == null) {
+ return;
+ }
+ retained.lock();
+ try {
+ Ex.check(retained.isFinished() && !retained.isCleanupFailed(),
+ "Task cleanup is incomplete; stop imports, inspect graph mode and recover before retrying");
+ } finally {
+ retained.unlock();
+ }
+ }
public LoadTaskService() {
this.runningTaskContainer = new ConcurrentHashMap<>();
@@ -174,6 +318,21 @@ public class LoadTaskService {
@Transactional(isolation = Isolation.READ_COMMITTED)
public void save(LoadTask entity) {
+ Integer jobId = entity.getJobId();
+ Ex.check(jobId != null, "Load task has no job id");
+ try (TaskReservation reservation = this.reserveJob(jobId)) {
+ this.ensureParentExists(jobId, entity.getGraphSpace(), entity.getGraph());
+ this.insertTask(entity);
+ }
+ }
+
+ private void ensureParentExists(int jobId, String graphSpace, String graph) {
+ JobManager parent = this.jobManagerMapper.selectById(jobId);
+ Ex.check(parent != null && Objects.equals(parent.getGraphSpace(), graphSpace) &&
+ Objects.equals(parent.getGraph(), graph), "job-manager.not-exist.id", jobId);
+ }
+
+ private void insertTask(LoadTask entity) {
LoadOptions runtimeOptions = entity.getOptions();
entity.setOptions(withoutCredentials(runtimeOptions));
try {
@@ -215,14 +374,34 @@ public class LoadTaskService {
}
}
- @Transactional(isolation = Isolation.READ_COMMITTED)
public void remove(int id) {
- this.runningTaskContainer.remove(id);
- if (this.mapper.deleteById(id) != 1) {
- throw new InternalException("entity.delete.failed", id);
+ try (TaskReservation reservation = this.reserveTasks(Collections.singletonList(id))) {
+ this.checkCallbackWaitAllowed(id);
+ LoadTask retained = this.runningTaskContainer.get(id);
+ if (retained != null) {
+ this.requestStopForDelete(retained);
+ }
+ this.awaitCallbacksForDelete(id);
+ this.ensureRuntimeCleanupComplete(id);
+ this.inDeleteTransaction(() -> {
+ LoadTask task = this.get(id);
+ if (task != null && (task.getStatus().inRunning() || task.getStatus() == LoadStatus.PAUSED)) {
+ this.stopForDelete(id);
+ }
+ if (this.mapper.deleteById(id) != 1) {
+ throw new InternalException("entity.delete.failed", id);
+ }
+ this.forgetRuntimeAfterCommit(id, retained);
+ });
}
}
+ private void inDeleteTransaction(Runnable action) {
+ TransactionTemplate transaction = new TransactionTemplate(this.transactionManager);
+ transaction.setIsolationLevel(TransactionDefinition.ISOLATION_READ_COMMITTED);
+ transaction.executeWithoutResult(status -> action.run());
+ }
+
public List<LoadTask> batchTasks(int jobId) {
QueryWrapper<LoadTask> query = Wrappers.query();
query.eq("job_id", jobId);
@@ -231,10 +410,34 @@ public class LoadTaskService {
public LoadTask start(GraphConnection connection, FileMapping fileMapping,
HugeClient client) {
- LoadTask task = this.buildLoadTask(connection, fileMapping, client);
- this.save(task);
- this.executeAfterCommit(task);
- return task;
+ synchronized (this.admissionLock) {
+ this.checkAcceptingLoads();
+ }
+ Integer jobId = fileMapping.getJobId();
+ Ex.check(jobId != null, "File mapping has no job id");
+ try (TaskReservation reservation = this.reserveJob(jobId)) {
+ this.ensureParentExists(jobId, connection.getGraphSpace(), connection.getGraph());
+ Ex.check(Objects.equals(fileMapping.getGraphSpace(), connection.getGraphSpace()) &&
+ Objects.equals(fileMapping.getGraph(), connection.getGraph()),
+ "File mapping does not belong to the load task graph");
+ FileMapping current = this.fileMappingMapper.selectById(fileMapping.getId());
+ Ex.check(current != null && Objects.equals(current.getJobId(), jobId) &&
+ Objects.equals(current.getGraphSpace(), connection.getGraphSpace()) &&
+ Objects.equals(current.getGraph(), connection.getGraph()),
+ "File mapping does not belong to the load task job");
+ synchronized (this.admissionLock) {
+ this.checkAcceptingLoads();
+ LoadTask task = this.buildLoadTask(connection, current, client);
+ try {
+ this.insertTask(task);
+ this.executeAfterCommit(task);
+ return task;
+ } catch (RuntimeException | Error failure) {
+ this.closeUnscheduled(task, failure);
+ throw failure;
+ }
+ }
+ }
}
protected void executeAfterCommit(LoadTask task) {
@@ -242,28 +445,58 @@ public class LoadTaskService {
this.executeSafely(task);
return;
}
- TransactionSynchronizationManager.registerSynchronization(
- new TransactionSynchronization() {
- @Override
- public void afterCommit() {
- executeSafely(task);
- }
- });
+ synchronized (this.admissionLock) {
+ AcceptedRun run = this.accept(task);
+ try {
+ TransactionSynchronizationManager.registerSynchronization(
+ new TransactionSynchronization() {
+ @Override
+ public void afterCommit() {
+ dispatchCommitted(run);
+ }
+
+ @Override
+ public void afterCompletion(int status) {
+ if (status != STATUS_COMMITTED) {
+ discardRolledBack(run);
+ }
+ }
+ });
+ } catch (RuntimeException | Error failure) {
+ this.release(run);
+ this.forgetRuntime(task.getId(), task);
+ throw failure;
+ }
+ }
}
private void executeSafely(LoadTask task) {
try {
this.execute(task);
} catch (RuntimeException e) {
- this.markDispatchFailed(task, e);
+ // execute() already completed native cleanup and its failure marker.
+ log.error("Failed to dispatch load task {}", task.getId(), e);
}
}
- private void markDispatchFailed(LoadTask task, RuntimeException cause) {
+ private boolean markDispatchFailed(LoadTask task, Throwable cause) {
log.error("Failed to dispatch load task {}, marking task and job " +
"as FAILED", task.getId(), cause);
+ try {
+ this.cleanupUnstarted(task);
+ } catch (RuntimeException | Error cleanupFailure) {
+ if (cause != cleanupFailure) {
+ cause.addSuppressed(cleanupFailure);
+ }
+ log.warn("Failed to close undispatched load task {}", task.getId(), cleanupFailure);
+ }
+ task.setFinished(true);
task.setStatus(LoadStatus.FAILED);
- this.runningTaskContainer.remove(task.getId());
+ if (task.isCleanupFailed()) {
+ this.runningTaskContainer.putIfAbsent(task.getId(), task);
+ } else {
+ this.forgetRuntime(task.getId(), task);
+ }
try {
TransactionTemplate transaction =
@@ -294,83 +527,374 @@ public class LoadTaskService {
}
return null;
});
- } catch (RuntimeException e) {
+ // TransactionTemplate returns only after REQUIRES_NEW actually commits.
+ return !task.isCleanupFailed();
+ } catch (RuntimeException | Error e) {
// The ingest transaction is already committed. Never turn a
// dispatch failure into a retryable request failure.
+ if (cause != e) {
+ cause.addSuppressed(e);
+ }
log.error("Failed to persist dispatch failure for load task {}",
task.getId(), e);
+ return false;
}
}
private void execute(LoadTask task) {
- // Executed in other threads
- this.taskExecutor.execute(task, () -> this.update(task));
- // Save current load task
- this.runningTaskContainer.put(task.getId(), task);
+ synchronized (this.admissionLock) {
+ AcceptedRun run = this.accept(task);
+ this.runningTaskContainer.put(task.getId(), task);
+ try {
+ this.taskExecutor.execute(task, () -> this.complete(run));
+ } catch (RuntimeException | Error failure) {
+ this.failDispatch(run, failure);
+ throw failure;
+ }
+ }
}
- public LoadTask pause(int taskId) {
- /*
- * The in-memory container is not repopulated on startup, so a task
- * that was RUNNING before a restart is only visible in the database
- */
- LoadTask task = this.runningTaskContainer.get(taskId);
- boolean loading = task != null;
- if (!loading) {
- task = this.get(taskId);
- }
- Ex.check(task != null, "load.task.not-exist.id", taskId);
- Ex.check(task.getStatus() == LoadStatus.RUNNING,
- "Can only pause the RUNNING task");
- // Mark status as paused, should set before context.stopLoading()
- task.setStatus(LoadStatus.PAUSED);
- if (loading) {
- // Let HugeGraphLoader stop
- task.stop();
- }
+ private AcceptedRun accept(LoadTask task) {
+ this.checkAcceptingLoads();
+ AcceptedRun run = new AcceptedRun(task);
+ this.acceptedRuns.add(run);
+ return run;
+ }
+ private void cleanupUnstarted(LoadTask task) {
task.lock();
try {
- this.update(task);
- this.runningTaskContainer.remove(taskId);
+ if (task.getLoader() != null) {
+ try {
+ task.getLoader().shutdown();
+ } catch (RuntimeException | Error failure) {
+ task.setCleanupFailed(true);
+ if (task.getCleanupFailure() == null) {
+ task.setCleanupFailure(failure);
+ }
+ throw failure;
+ }
+ // A repeated shutdown can return after context.close even when
+ // an earlier owner release failed. Keep that recovery handle.
+ if (!task.isCleanupFailed()) {
+ task.setLoader(null);
+ }
+ }
} finally {
task.unlock();
}
- return task;
+ }
+
+ private void updateAndExecute(LoadTask task) {
+ try {
+ this.update(task);
+ } catch (RuntimeException | Error failure) {
+ this.closeUnscheduled(task, failure);
+ throw failure;
+ }
+ // Frontend resume/retry must preserve a synchronous dispatch failure.
+ this.execute(task);
+ }
+
+ private void closeUnscheduled(LoadTask task, Throwable failure) {
+ try {
+ this.cleanupUnstarted(task);
+ } catch (RuntimeException | Error cleanupFailure) {
+ if (failure != cleanupFailure) {
+ failure.addSuppressed(cleanupFailure);
+ }
+ } finally {
+ task.setFinished(true);
+ }
+ if (!task.isCleanupFailed()) {
+ // The persistence operation never admitted this run. Once its
+ // native resources are closed, no shutdown responsibility remains.
+ return;
+ }
+ task.setStatus(LoadStatus.FAILED);
+ if (task.getId() != null) {
+ this.runningTaskContainer.compute(task.getId(), (id, current) ->
+ current == null || (current.isFinished() && !current.isCleanupFailed()) ? task : current);
+ }
+ // A failed insert can have no id. Keep its unresolved native cleanup
+ // in the same private ticket set without publishing a null map key.
+ for (AcceptedRun accepted : this.acceptedRuns) {
+ if (accepted.task == task) {
+ return;
+ }
+ }
+ AcceptedRun run = new AcceptedRun(task);
+ run.completed.completeExceptionally(failure);
+ this.acceptedRuns.add(run);
+ }
+
+ private void forgetRuntime(Integer taskId, LoadTask expected) {
+ if (taskId != null) {
+ this.runningTaskContainer.computeIfPresent(taskId, (id, current) -> current == expected ? null : current);
+ }
+ }
+
+ private void discardRolledBack(AcceptedRun run) {
+ try {
+ this.cleanupUnstarted(run.task);
+ run.task.setFinished(true);
+ run.completed.complete(null);
+ } catch (RuntimeException | Error failure) {
+ run.task.setFinished(true);
+ run.completed.completeExceptionally(failure);
+ throw failure;
+ } finally {
+ this.forgetRuntime(run.task.getId(), run.task);
+ this.release(run);
+ }
+ }
+
+ private void dispatchCommitted(AcceptedRun run) {
+ try {
+ // This run was admitted before close. The pause hook keeps the
+ // executor and metadata dependencies alive until its ticket ends.
+ this.runningTaskContainer.put(run.task.getId(), run.task);
+ this.taskExecutor.execute(run.task, () -> this.complete(run));
+ } catch (RuntimeException | Error failure) {
+ this.failDispatch(run, failure);
+ if (failure instanceof Error) {
+ throw (Error) failure;
+ }
+ }
+ }
+
+ private void failDispatch(AcceptedRun run, Throwable failure) {
+ boolean resolved = false;
+ try {
+ resolved = this.markDispatchFailed(run.task, failure);
+ } finally {
+ if (resolved) {
+ // The request still receives its dispatch error, but native
+ // cleanup and its committed FAILED marker need no recovery.
+ run.completed.complete(null);
+ } else {
+ run.completed.completeExceptionally(failure);
+ }
+ this.release(run);
+ }
+ }
+
+ public LoadTask pause(int taskId) {
+ try (TaskReservation reservation = this.reserveTasks(Collections.singletonList(taskId))) {
+ /*
+ * The in-memory container is not repopulated on startup, so a task
+ * that was RUNNING before a restart is only visible in the database
+ */
+ LoadTask task = this.runningTaskContainer.get(taskId);
+ boolean loading = task != null;
+ if (!loading) {
+ task = this.get(taskId);
+ }
+ Ex.check(task != null, "load.task.not-exist.id", taskId);
+ Ex.check(task.getStatus() == LoadStatus.RUNNING,
+ "Can only pause the RUNNING task");
+ // Mark status as paused, should set before context.stopLoading()
+ task.setStatus(LoadStatus.PAUSED);
+ if (loading) {
+ // Let HugeGraphLoader stop
+ task.stop();
+ }
+
+ task.lock();
+ try {
+ this.update(task);
+ if (!task.isCleanupFailed()) {
+ this.forgetRuntime(taskId, task);
+ }
+ } finally {
+ task.unlock();
+ }
+ return task;
+ }
}
public LoadTask resume(int taskId, String token) {
- LoadTask task = this.get(taskId);
- Ex.check(task.getStatus() == LoadStatus.PAUSED ||
- task.getStatus() == LoadStatus.FAILED,
- "Can only resume the PAUSED or FAILED task");
+ try (TaskReservation reservation = this.reserveTasks(Collections.singletonList(taskId))) {
+ synchronized (this.admissionLock) {
+ this.checkAcceptingLoads();
+ this.ensureRuntimeCleanupComplete(taskId);
+ }
+ // A stopped native runner may still own its final metadata write.
+ // Its callback does not take the task reservation; release admission
+ // and the entity lock before joining, then read fresh state below.
+ this.awaitTaskCallbacks(taskId);
+ synchronized (this.admissionLock) {
+ this.checkAcceptingLoads();
+ this.ensureRuntimeCleanupComplete(taskId);
+ LoadTask task = this.get(taskId);
+ Ex.check(task != null, "load.task.not-exist.id", taskId);
+ Ex.check(task.getStatus() == LoadStatus.PAUSED ||
+ task.getStatus() == LoadStatus.FAILED,
+ "Can only resume the PAUSED or FAILED task");
+ task.lock();
+ try {
+ // Set work mode in incrental mode, load from last breakpoint
+ task.getOptions().incrementalMode = true;
+ task.reconnect(token);
+ task.setStatus(LoadStatus.RUNNING);
+ this.updateAndExecute(task);
+ } finally {
+ task.unlock();
+ }
+ return task;
+ }
+ }
+ }
+
+ public LoadTask stop(int taskId) {
+ try (TaskReservation reservation = this.reserveTasks(Collections.singletonList(taskId))) {
+ return this.stop(taskId, false);
+ }
+ }
+
+ LoadTask stopForDelete(int taskId) {
+ LoadTask retained = this.runningTaskContainer.get(taskId);
+ if (retained != null) {
+ this.requestStopForDelete(retained);
+ }
+ // Do not write the row while its callback still needs that same row.
+ // Neither admission nor the task entity lock is held during this join.
+ this.awaitCallbacksForDelete(taskId);
+ if (retained != null && !retained.isFinished()) {
+ // Retained work from callers predating ticket registration still
+ // needs the native completion certificate before deletion.
+ retained.stop();
+ }
+ return this.stop(taskId, true);
+ }
+
+ void requestJobStopForDelete(int jobId) {
+ Set<LoadTask> tasks = Collections.newSetFromMap(new java.util.IdentityHashMap<>());
+ synchronized (this.admissionLock) {
+ for (AcceptedRun run : this.acceptedRuns) {
+ if (Objects.equals(run.task.getJobId(), jobId)) {
+ tasks.add(run.task);
+ }
+ }
+ for (LoadTask task : this.runningTaskContainer.values()) {
+ if (Objects.equals(task.getJobId(), jobId)) {
+ tasks.add(task);
+ }
+ }
+ }
+ Throwable failure = null;
+ for (LoadTask task : tasks) {
+ try {
+ this.requestStopForDelete(task);
+ } catch (RuntimeException | Error stopFailure) {
+ failure = collectFailure(failure, stopFailure);
+ }
+ }
+ throwFailure(failure);
+ }
+
+ private void requestStopForDelete(LoadTask task) {
task.lock();
try {
- // Set work mode in incrental mode, load from last breakpoint
- task.getOptions().incrementalMode = true;
- task.reconnect(token);
- task.setStatus(LoadStatus.RUNNING);
- this.update(task);
- this.taskExecutor.execute(task, () -> this.update(task));
- this.runningTaskContainer.put(taskId, task);
+ if (task.getStatus().inRunning() || task.getStatus() == LoadStatus.PAUSED) {
+ task.setStatus(LoadStatus.STOPPED);
+ if (task.getLoader() != null) {
+ task.context().stopLoading();
+ }
+ }
} finally {
task.unlock();
}
- return task;
}
- public LoadTask stop(int taskId) {
+ void awaitCallbacksForDelete(int taskId) {
+ this.awaitTaskCallbacks(taskId);
+ }
+
+ void checkCallbackWaitAllowed(int taskId) {
+ if (!TransactionSynchronizationManager.isActualTransactionActive()) {
+ return;
+ }
+ synchronized (this.admissionLock) {
+ for (AcceptedRun run : this.acceptedRuns) {
+ Ex.check(!Objects.equals(run.task.getId(), taskId) || run.completed.isDone(),
+ "Task callbacks are still running; retry after the current transaction completes");
+ }
+ }
+ }
+
+ private void awaitTaskCallbacks(int taskId) {
+ this.checkCallbackWaitAllowed(taskId);
+ List<AcceptedRun> runs = new ArrayList<>();
+ synchronized (this.admissionLock) {
+ for (AcceptedRun run : this.acceptedRuns) {
+ if (Objects.equals(run.task.getId(), taskId)) {
+ runs.add(run);
+ }
+ }
+ }
+ Throwable failure = null;
+ for (AcceptedRun run : runs) {
+ try {
+ run.completed.join();
+ } catch (CancellationException cancellation) {
+ failure = collectFailure(failure, cancellation);
+ } catch (CompletionException completion) {
+ failure = collectFailure(failure, completion.getCause() == null ? completion : completion.getCause());
+ }
+ }
+ throwFailure(failure);
+ }
+
+ private static void throwFailure(Throwable failure) {
+ if (failure instanceof RuntimeException) {
+ throw (RuntimeException) failure;
+ }
+ if (failure instanceof Error) {
+ throw (Error) failure;
+ }
+ if (failure != null) {
+ throw new CompletionException(failure);
+ }
+ }
+
+ private LoadTask stop(int taskId, boolean deleting) {
LoadTask task = this.runningTaskContainer.get(taskId);
boolean loading = task != null;
if (!loading) {
task = this.get(taskId);
}
Ex.check(task != null, "load.task.not-exist.id", taskId);
- LoadStatus status = task.getStatus();
- Ex.check(status == LoadStatus.RUNNING || status == LoadStatus.PAUSED,
- "Can only stop the RUNNING or PAUSED task");
- // Mark status as stopped
- task.setStatus(LoadStatus.STOPPED);
+ LoadStatus status;
+ task.lock();
+ try {
+ status = task.getStatus();
+ if (deleting && (status == LoadStatus.SUCCEED || status == LoadStatus.FAILED ||
+ status == LoadStatus.STOPPED)) {
+ // Only a retained runtime task can certify completed cleanup.
+ Ex.check(loading && task.isFinished() && !task.isCleanupFailed(),
+ "Task cleanup is not confirmed; check graph loading mode before deleting");
+ return task;
+ }
+ if (!deleting && loading && status == LoadStatus.STOPPED &&
+ task.isFinished() && !task.isCleanupFailed()) {
+ // A failed deletion can roll metadata back after native stop.
+ // Reuse that safe runtime; do not stop or close its Loader twice.
+ LoadTask persisted = this.get(taskId);
+ Ex.check(persisted != null &&
+ (persisted.getStatus() == LoadStatus.RUNNING ||
+ persisted.getStatus() == LoadStatus.PAUSED),
+ "Can only stop the RUNNING or PAUSED task");
+ this.update(task);
+ this.forgetRuntimeAfterCommit(taskId, task);
+ return task;
+ }
+ Ex.check(status == LoadStatus.RUNNING || status == LoadStatus.PAUSED,
+ "Can only stop the RUNNING or PAUSED task");
+ task.setStatus(LoadStatus.STOPPED);
+ } finally {
+ task.unlock();
+ }
/*
* Only a task taken from the in-memory container owns a loader, one
* restored from the database has nothing running to interrupt, so
@@ -382,8 +906,18 @@ public class LoadTaskService {
task.lock();
try {
+ if (deleting && loading) {
+ // stop() waits outside the lock used by the runner's final
+ // publication; independent cleanup errors must retain evidence.
+ Ex.check(task.isFinished() && !task.isCleanupFailed(),
+ "Task cleanup failed; check graph loading mode before deleting");
+ }
this.update(task);
- this.runningTaskContainer.remove(taskId);
+ // Deletion may still roll back after native stop. Keep the real
+ // terminal runtime until remove() observes transaction commit.
+ if (!deleting && !task.isCleanupFailed()) {
+ this.forgetRuntime(taskId, task);
+ }
} finally {
task.unlock();
}
@@ -391,25 +925,38 @@ public class LoadTaskService {
}
public LoadTask retry(int taskId, String token) {
- LoadTask task = this.get(taskId);
- Ex.check(task.getStatus() == LoadStatus.FAILED ||
- task.getStatus() == LoadStatus.STOPPED,
- "Can only retry the FAILED or STOPPED task");
- task.lock();
- try {
- // Set work mode in normal mode, load from begin
- task.getOptions().incrementalMode = false;
- task.reconnect(token);
- task.setStatus(LoadStatus.RUNNING);
- task.setLastDuration(0L);
- task.setCurrDuration(0L);
- this.update(task);
- this.taskExecutor.execute(task, () -> this.update(task));
- this.runningTaskContainer.put(taskId, task);
- } finally {
- task.unlock();
+ try (TaskReservation reservation = this.reserveTasks(Collections.singletonList(taskId))) {
+ synchronized (this.admissionLock) {
+ this.checkAcceptingLoads();
+ this.ensureRuntimeCleanupComplete(taskId);
+ }
+ // A stopped native runner may still own its final metadata write.
+ // Its callback does not take the task reservation; release admission
+ // and the entity lock before joining, then read fresh state below.
+ this.awaitTaskCallbacks(taskId);
+ synchronized (this.admissionLock) {
+ this.checkAcceptingLoads();
+ this.ensureRuntimeCleanupComplete(taskId);
+ LoadTask task = this.get(taskId);
+ Ex.check(task != null, "load.task.not-exist.id", taskId);
+ Ex.check(task.getStatus() == LoadStatus.FAILED ||
+ task.getStatus() == LoadStatus.STOPPED,
+ "Can only retry the FAILED or STOPPED task");
+ task.lock();
+ try {
+ // Set work mode in normal mode, load from begin
+ task.getOptions().incrementalMode = false;
+ task.reconnect(token);
+ task.setStatus(LoadStatus.RUNNING);
+ task.setLastDuration(0L);
+ task.setCurrDuration(0L);
+ this.updateAndExecute(task);
+ } finally {
+ task.unlock();
+ }
+ return task;
+ }
}
- return task;
}
public String readLoadFailedReason(FileMapping mapping) {
@@ -449,11 +996,178 @@ public class LoadTaskService {
}
public void pauseAllTasks() {
- List<LoadTask> tasks = this.listAll();
- for (LoadTask task : tasks) {
+ List<AcceptedRun> runs;
+ List<LoadTask> runtimeTasks;
+ synchronized (this.admissionLock) {
+ this.closing = true;
+ runs = new ArrayList<>(this.acceptedRuns);
+ runtimeTasks = new ArrayList<>(this.runningTaskContainer.values());
+ for (AcceptedRun run : runs) {
+ if (runtimeTasks.stream().noneMatch(task -> task == run.task)) {
+ runtimeTasks.add(run.task);
+ }
+ }
+ }
+ Set<Integer> runtimeIds = new HashSet<>();
+ for (LoadTask task : runtimeTasks) {
+ runtimeIds.add(task.getId());
+ }
+ Throwable failure = null;
+ List<LoadTask> failedTasks = new ArrayList<>();
+ // Signal every accepted runner before waiting: a queued runner needs
+ // the active workers to finish before it can observe its stop intent.
+ for (AcceptedRun run : runs) {
+ try {
+ this.requestPause(run.task);
+ } catch (RuntimeException | Error e) {
+ failure = collectFailure(failure, e);
+ }
+ }
+ for (AcceptedRun run : runs) {
+ try {
+ // finished only covers LoadTask.run(); this ticket also covers
+ // its async callback and the final metadata update.
+ run.completed.join();
+ } catch (CancellationException e) {
+ failure = collectFailure(failure, e);
+ failedTasks.add(run.task);
+ } catch (CompletionException e) {
+ Throwable cause = e.getCause() == null ? e : e.getCause();
+ failure = collectFailure(failure, cause);
+ failedTasks.add(run.task);
+ }
+ }
+ for (LoadTask task : runtimeTasks) {
+ boolean handled = failedTasks.stream().noneMatch(failed -> failed == task);
+ try {
+ // Always attempt native cleanup before dropping a handle,
+ // including when a callback or earlier cleanup failed.
+ this.cleanupUnstarted(task);
+ } catch (RuntimeException | Error e) {
+ failure = collectFailure(failure, e);
+ handled = false;
+ }
+ if (task.isCleanupFailed() && handled) {
+ Throwable cleanup = task.getCleanupFailure();
+ failure = collectFailure(failure, cleanup);
+ handled = false;
+ }
+ if (handled) {
+ this.forgetRuntime(task.getId(), task);
+ } else {
+ // Keep the unpersisted runtime and exceptional ticket visible.
+ if (task.getId() != null) {
+ this.runningTaskContainer.putIfAbsent(task.getId(), task);
+ }
+ }
+ }
+ try {
+ // Rows surviving a restart have no accepted runtime or Loader.
+ for (LoadTask task : this.listAll()) {
+ if (task.getStatus().inRunning() &&
+ !runtimeIds.contains(task.getId())) {
+ try {
+ task.setStatus(LoadStatus.PAUSED);
+ this.update(task);
+ } catch (RuntimeException | Error e) {
+ failure = collectFailure(failure, e);
+ }
+ }
+ }
+ } catch (RuntimeException | Error e) {
+ failure = collectFailure(failure, e);
+ }
+ if (failure instanceof RuntimeException) {
+ throw (RuntimeException) failure;
+ }
+ if (failure instanceof Error) {
+ throw (Error) failure;
+ }
+ if (failure != null) {
+ throw new CompletionException(failure);
+ }
+ }
+
+ private void checkAcceptingLoads() {
+ if (this.closing) {
+ throw new RejectedExecutionException("Hubble is closing");
+ }
+ }
+
+ private void complete(AcceptedRun run) {
+ Throwable failure = run.task.getCleanupFailure();
+ run.task.lock();
+ try {
+ // Native cleanup has already run in LoadTask.run(). Still attempt
+ // metadata persistence when cleanup failed, retaining both causes.
+ this.update(run.task);
+ } catch (RuntimeException | Error metadataFailure) {
+ if (failure == null) {
+ failure = metadataFailure;
+ } else if (failure != metadataFailure) {
+ failure.addSuppressed(metadataFailure);
+ }
+ } finally {
+ run.task.unlock();
+ }
+ if (failure == null) {
+ run.completed.complete(null);
+ this.release(run);
+ return;
+ }
+ run.completed.completeExceptionally(failure);
+ // Exceptional tickets remain registered even when the callback ended
+ // before close, until their unresolved metadata/cleanup is reported.
+ if (failure instanceof RuntimeException) {
+ throw (RuntimeException) failure;
+ }
+ if (failure instanceof Error) {
+ throw (Error) failure;
+ }
+ throw new CompletionException(failure);
+ }
+
+ private void release(AcceptedRun run) {
+ synchronized (this.admissionLock) {
+ if (!run.completed.isCompletedExceptionally()) {
+ this.acceptedRuns.remove(run);
+ }
+ }
+ }
+
+ private void requestPause(LoadTask task) {
+ task.lock();
+ try {
if (task.getStatus().inRunning()) {
- this.pause(task.getId());
+ task.setStatus(LoadStatus.PAUSED);
+ // Native stopLoading() only publishes a stop reason. It does
+ // not wait for the queued/active runner or its callback.
+ if (task.getLoader() != null) {
+ task.context().stopLoading();
+ }
}
+ } finally {
+ task.unlock();
+ }
+ }
+
+ private static Throwable collectFailure(Throwable primary, Throwable failure) {
+ if (primary == null) {
+ return failure;
+ }
+ if (primary != failure) {
+ primary.addSuppressed(failure);
+ }
+ return primary;
+ }
+
+ private static final class AcceptedRun {
+
+ private final LoadTask task;
+ private final CompletableFuture<Void> completed = new CompletableFuture<>();
+
+ private AcceptedRun(LoadTask task) {
+ this.task = task;
}
}
@@ -771,13 +1485,28 @@ public class LoadTaskService {
}
public void loader(LoadOptions options) {
- HugeGraphLoader loader;
+ HugeGraphLoader loader = null;
+ Throwable failure = null;
try {
loader = new HugeGraphLoader(options);
loader.load();
} catch (Throwable e) {
- Printer.printError("Failed to start loading", e);
- return;
+ failure = e;
+ } finally {
+ if (loader != null) {
+ try {
+ loader.shutdown();
+ } catch (Throwable e) {
+ if (failure == null) {
+ failure = e;
+ } else if (failure != e) {
+ failure.addSuppressed(e);
+ }
+ }
+ }
+ }
+ if (failure != null) {
+ Printer.printError("Failed to start loading", failure);
}
}
}
diff --git a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/service/load/IngestTransactionIntegrationTest.java b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/service/load/IngestTransactionIntegrationTest.java
index 38f997593..16986aa35 100644
--- a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/service/load/IngestTransactionIntegrationTest.java
+++ b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/service/load/IngestTransactionIntegrationTest.java
@@ -20,20 +20,27 @@ package org.apache.hugegraph.service.load;
import java.util.Collections;
import java.util.Date;
+import java.util.List;
+import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
+import java.util.concurrent.TimeoutException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
import javax.sql.DataSource;
import org.junit.After;
+import org.junit.Rule;
+import org.junit.rules.TemporaryFolder;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mockito;
+import org.mockito.MockedConstruction;
import org.mybatis.spring.annotation.MapperScan;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
@@ -45,8 +52,15 @@ import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.test.context.junit4.SpringRunner;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.annotation.EnableTransactionManagement;
+import org.springframework.transaction.support.TransactionTemplate;
import org.apache.hugegraph.config.HugeConfig;
+import org.apache.hugegraph.controller.ingest.IngestController;
+import org.springframework.test.web.servlet.MockMvc;
+import org.springframework.test.web.servlet.setup.MockMvcBuilders;
+import org.springframework.test.util.ReflectionTestUtils;
+import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.delete;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
import org.apache.hugegraph.driver.HugeClient;
import org.apache.hugegraph.entity.GraphConnection;
import org.apache.hugegraph.entity.enums.FileMappingStatus;
@@ -60,10 +74,13 @@ import org.apache.hugegraph.entity.load.LoadParameter;
import org.apache.hugegraph.entity.load.LoadTask;
import org.apache.hugegraph.handler.LoadTaskExecutor;
import org.apache.hugegraph.loader.executor.LoadOptions;
+import org.apache.hugegraph.loader.HugeGraphLoader;
+import org.apache.hugegraph.loader.executor.LoadContext;
import org.apache.hugegraph.options.HubbleOptions;
import org.apache.hugegraph.service.schema.EdgeLabelService;
import org.apache.hugegraph.service.schema.VertexLabelService;
import org.apache.hugegraph.testutil.Assert;
+import org.apache.hugegraph.testutil.NativeCleanupTestSupport;
@RunWith(SpringRunner.class)
@SpringBootTest(classes = IngestTransactionIntegrationTest.TestConfiguration.class,
@@ -89,6 +106,9 @@ import org.apache.hugegraph.testutil.Assert;
})
public class IngestTransactionIntegrationTest {
+ @Rule
+ public TemporaryFolder nativeFolder = new TemporaryFolder();
+
private static final String UPLOAD_ROOT =
"/tmp/hubble-ingest-transaction-test";
@@ -104,6 +124,8 @@ public class IngestTransactionIntegrationTest {
private DataSource dataSource;
@Autowired
private HugeConfig config;
+ @Autowired
+ private PlatformTransactionManager transactionManager;
@MockBean
private VertexLabelService vertexLabelService;
@MockBean
@@ -117,6 +139,10 @@ public class IngestTransactionIntegrationTest {
this.cleanup();
this.executor.reset();
this.taskService.failTaskInsert(false);
+ this.taskService.afterTaskSnapshot(null);
+ this.taskService.afterStopCertificate(null);
+ this.taskService.stubReconnect(false);
+ this.taskService.beforeRestartRead(null);
Mockito.when(this.config.get(HubbleOptions.UPLOAD_FILE_LOCATION))
.thenReturn(UPLOAD_ROOT);
}
@@ -144,10 +170,17 @@ public class IngestTransactionIntegrationTest {
public void testTaskInsertFailureRollsBackAllRowsAndDoesNotDispatch() {
this.taskService.failTaskInsert(true);
- Assert.assertThrows(RuntimeException.class, () -> {
+ Throwable failure = Assert.assertThrows(RuntimeException.class, () -> {
this.jobService.createIngestTask(this.job(), this.mapping(),
this.connection(), this.client());
});
+ Throwable cause = failure;
+ while (cause.getCause() != null) {
+ cause = cause.getCause();
+ }
+ Assert.assertTrue("The task INSERT must reach the real H2 constraint", cause instanceof java.sql.SQLException);
+ Assert.assertEquals("23502", ((java.sql.SQLException) cause).getSQLState());
+ Assert.assertTrue(cause.getMessage().toLowerCase(java.util.Locale.ROOT).contains("file_id"));
this.assertTableCounts(0, 0, 0);
Assert.assertEquals(0, this.executor.executions());
@@ -307,6 +340,565 @@ public class IngestTransactionIntegrationTest {
this.jobService.get(task.getJobId()).getJobStatus());
}
+ @Test
+ public void testUnscheduledLoadersCloseAfterSaveFailureRollbackAndRejection() {
+ for (int outcome = 0; outcome < 3; outcome++) {
+ this.cleanup();
+ this.executor.reset();
+ this.executor.reject(outcome == 2);
+ JobManager job = this.job();
+ this.jobService.save(job);
+ FileMapping mapping = this.mapping();
+ mapping.setJobId(job.getId());
+ mapping.setPath(UPLOAD_ROOT + "/cleanup/people.csv");
+ this.mappingService.save(mapping);
+ TransactionTemplate transaction = new TransactionTemplate(this.transactionManager);
+ if (outcome == 0) {
+ this.jdbc.execute("ALTER TABLE load_task ADD CONSTRAINT reject_task CHECK (id < 0)");
+ }
+ final int result = outcome;
+ try (MockedConstruction<HugeGraphLoader> construction = Mockito.mockConstruction(HugeGraphLoader.class)) {
+ if (outcome == 0) {
+ Assert.assertThrows(RuntimeException.class, () -> transaction.execute(status ->
+ this.taskService.startLoader(this.connection(), mapping, this.client())));
+ } else {
+ transaction.execute(status -> {
+ this.taskService.startLoader(this.connection(), mapping, this.client());
+ if (result == 1) {
+ status.setRollbackOnly();
+ }
+ return null;
+ });
+ }
+ Assert.assertEquals(1, construction.constructed().size());
+ Mockito.verify(construction.constructed().get(0)).shutdown();
+ if (outcome != 0) {
+ Assert.assertTrue(this.taskService.createdTask().isFinished());
+ }
+ Assert.assertEquals(outcome == 2 ? 1 : 0, this.count("load_task"));
+ Assert.assertEquals(outcome == 2 ? 1 : 0, this.executor.executions());
+ } finally {
+ if (outcome == 0) {
+ this.jdbc.execute("ALTER TABLE load_task DROP CONSTRAINT reject_task");
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testDeleteJobAcceptsCompletionAfterActualDatabaseSnapshot() throws Exception {
+ this.deleteAfterSnapshot(false);
+ }
+
+ @Test
+ public void testDeleteJobRetainsRowsAfterIndependentCleanupFailure() throws Exception {
+ this.deleteAfterSnapshot(true);
+ }
+
+ private void deleteAfterSnapshot(boolean cleanupFailure) throws Exception {
+ FileMapping mapping = this.mapping();
+ mapping.setPath(UPLOAD_ROOT + "/delete-race-" + UUID.randomUUID() + "/people.csv");
+ LoadTask task = this.jobService.createIngestTask(this.job(), mapping,
+ this.connection(), this.client());
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ LoadContext context = Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS);
+ Mockito.when(loader.context()).thenReturn(context);
+ Mockito.when(loader.load()).thenReturn(!cleanupFailure);
+ if (cleanupFailure) {
+ Mockito.doThrow(new IllegalStateException("release failed")).when(loader).shutdown();
+ }
+ task.setLoader(loader);
+ ExecutorService runner = Executors.newSingleThreadExecutor();
+ this.taskService.afterTaskSnapshot(() -> {
+ try {
+ // Independent thread uses the real Spring/MyBatis callback path
+ // after the deletion transaction reads its RUNNING snapshot.
+ runner.submit(() -> {
+ task.run();
+ this.executor.complete(task);
+ }).get(5, TimeUnit.SECONDS);
+ } catch (Exception e) {
+ throw new AssertionError(e);
+ }
+ });
+ try {
+ if (cleanupFailure) {
+ Assert.assertThrows(RuntimeException.class, () -> this.jobService.deleteJob(task.getJobId()));
+ this.assertTableCounts(1, 1, 1);
+ Assert.assertEquals(LoadStatus.FAILED, this.taskService.get(task.getId()).getStatus());
+ } else {
+ this.jobService.deleteJob(task.getJobId());
+ Assert.assertEquals(0, this.count("job_manager"));
+ Assert.assertEquals(0, this.count("load_task"));
+ }
+ Mockito.verify(loader).shutdown();
+ Assert.assertTrue(task.isFinished());
+ } finally {
+ runner.shutdownNow();
+ Assert.assertTrue(runner.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test(timeout = 15000)
+ public void testDeleteCertificateReservesIdentityAgainstFreshRetry() throws Exception {
+ FileMapping mapping = this.mapping();
+ mapping.setPath(UPLOAD_ROOT + "/reserved-delete-" + UUID.randomUUID() + "/people.csv");
+ LoadTask task = this.jobService.createIngestTask(this.job(), mapping,
+ this.connection(), this.client());
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ Mockito.when(loader.context()).thenReturn(Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS));
+ task.setLoader(loader);
+ this.taskService.stubReconnect(true);
+ ExecutorService callbacks = Executors.newSingleThreadExecutor();
+ this.taskService.afterTaskSnapshot(() -> {
+ try {
+ callbacks.submit(() -> {
+ task.run();
+ this.executor.complete(task);
+ }).get(5, TimeUnit.SECONDS);
+ } catch (Exception e) {
+ throw new AssertionError(e);
+ }
+ });
+ CountDownLatch certified = new CountDownLatch(1);
+ CountDownLatch delete = new CountDownLatch(1);
+ this.taskService.afterStopCertificate(() -> {
+ certified.countDown();
+ try {
+ Assert.assertTrue(delete.await(5, TimeUnit.SECONDS));
+ } catch (InterruptedException e) {
+ throw new AssertionError(e);
+ }
+ });
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ Future<?> deleting = requests.submit(() -> this.jobService.deleteJob(task.getJobId()));
+ Assert.assertTrue(certified.await(5, TimeUnit.SECONDS));
+ CountDownLatch retryEntered = new CountDownLatch(1);
+ Future<LoadTask> retry = requests.submit(() -> {
+ retryEntered.countDown();
+ return this.taskService.retry(task.getId(), "current-session-token");
+ });
+ Assert.assertTrue(retryEntered.await(5, TimeUnit.SECONDS));
+ try {
+ retry.get(300, TimeUnit.MILLISECONDS);
+ Assert.fail("retry crossed the completed-cleanup deletion reservation");
+ } catch (TimeoutException expected) {
+ // Deletion owns the task identity through its transaction.
+ }
+ delete.countDown();
+ deleting.get(5, TimeUnit.SECONDS);
+ Throwable rejected = Assert.assertThrows(
+ java.util.concurrent.ExecutionException.class, () -> retry.get(5, TimeUnit.SECONDS));
+ Assert.assertTrue(rejected.getCause() instanceof RuntimeException);
+ Assert.assertEquals(1, this.executor.executions());
+ Assert.assertEquals(0, this.count("job_manager"));
+ Assert.assertEquals(0, this.count("load_task"));
+ } finally {
+ delete.countDown();
+ requests.shutdownNow();
+ callbacks.shutdownNow();
+ Assert.assertTrue(requests.awaitTermination(5, TimeUnit.SECONDS));
+ Assert.assertTrue(callbacks.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ public void testNativeResetFailureCannotCertifyDeletionAfterLoadFailure() throws Exception {
+ FileMapping mapping = this.mapping();
+ mapping.setPath(UPLOAD_ROOT + "/native-delete-" + UUID.randomUUID() + "/people.csv");
+ LoadTask task = this.jobService.createIngestTask(this.job(), mapping,
+ this.connection(), this.client());
+ ExecutorService callback = Executors.newSingleThreadExecutor();
+ try (NativeCleanupTestSupport nativeCleanup = new NativeCleanupTestSupport(
+ task, this.nativeFolder.newFile("mode-reset.json").toPath(), true, LoadStatus.RUNNING)) {
+ this.taskService.afterTaskSnapshot(() -> {
+ try {
+ callback.submit(() -> {
+ task.run();
+ this.executor.complete(task);
+ }).get(5, TimeUnit.SECONDS);
+ } catch (Exception e) {
+ throw new AssertionError(e);
+ }
+ });
+
+ Assert.assertThrows(RuntimeException.class, () -> this.jobService.deleteJob(task.getJobId()));
+
+ this.assertTableCounts(1, 1, 1);
+ Assert.assertEquals(LoadStatus.FAILED, this.taskService.get(task.getId()).getStatus());
+ Assert.assertTrue(task.isCleanupFailed());
+ Assert.assertNotSame(nativeCleanup.original, nativeCleanup.shutdownFailure.get());
+ Assert.assertEquals("LOADING", nativeCleanup.mode());
+ nativeCleanup.assertResourcesClosed();
+ try (MockedConstruction<HugeGraphLoader> replacements = Mockito.mockConstruction(HugeGraphLoader.class)) {
+ Assert.assertThrows(RuntimeException.class,
+ () -> this.taskService.retry(task.getId(), "current-session-token"));
+ Assert.assertThrows(RuntimeException.class,
+ () -> this.taskService.resume(task.getId(), "current-session-token"));
+ Assert.assertTrue(replacements.constructed().isEmpty());
+ Assert.assertSame(nativeCleanup.loader, task.getLoader());
+ Assert.assertEquals(1, this.executor.executions());
+ }
+ } finally {
+ callback.shutdownNow();
+ Assert.assertTrue(callback.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test(timeout = 15000)
+ public void testDeletionReservationSurvivesOuterTransactionUntilCommit() throws Exception {
+ this.reserveUntilOuterCompletion(false);
+ }
+
+ @Test(timeout = 15000)
+ public void testDeletionRollbackRetainsRuntimeAndReleasesReservation() throws Exception {
+ this.reserveUntilOuterCompletion(true);
+ }
+
+ @Test(timeout = 15000)
+ public void testIngestDeleteReservationBlocksFreshRetryUntilOuterCommit() throws Exception {
+ this.reserveUntilOuterCompletion(false, true);
+ }
+
+ private void reserveUntilOuterCompletion(boolean rollback) throws Exception {
+ this.reserveUntilOuterCompletion(rollback, false);
+ }
+
+ private void reserveUntilOuterCompletion(boolean rollback, boolean endpoint) throws Exception {
+ FileMapping mapping = this.mapping();
+ mapping.setPath(UPLOAD_ROOT + "/outer-delete-" + UUID.randomUUID() + "/people.csv");
+ LoadTask task = this.jobService.createIngestTask(this.job(), mapping,
+ this.connection(), this.client());
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ Mockito.when(loader.context()).thenReturn(Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS));
+ task.setLoader(loader);
+ task.run();
+ this.executor.complete(task);
+ this.taskService.stubReconnect(true);
+ CountDownLatch deleted = new CountDownLatch(1);
+ CountDownLatch complete = new CountDownLatch(1);
+ CountDownLatch restartRead = new CountDownLatch(1);
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ Future<?> deleting = requests.submit(() -> new TransactionTemplate(this.transactionManager)
+ .execute(status -> {
+ if (endpoint) {
+ try {
+ this.ingestEndpoint().perform(delete(org.apache.hugegraph.common.Constant.API_VERSION +
+ "ingest/jobs/" + task.getId()))
+ .andExpect(status().isOk());
+ } catch (Exception e) {
+ throw new AssertionError(e);
+ }
+ } else {
+ this.jobService.deleteJob(task.getJobId());
+ }
+ deleted.countDown();
+ try {
+ Assert.assertTrue(complete.await(5, TimeUnit.SECONDS));
+ } catch (InterruptedException e) {
+ throw new AssertionError(e);
+ }
+ if (rollback) {
+ status.setRollbackOnly();
+ }
+ return null;
+ }));
+ Assert.assertTrue(deleted.await(5, TimeUnit.SECONDS));
+ // remove() now legitimately reads the task as well. Observe only
+ // the fresh restart request after deletion has returned.
+ this.taskService.beforeRestartRead(restartRead::countDown);
+ Assert.assertSame(task, this.taskService.retainedTask(task.getId()));
+ CountDownLatch retryEntered = new CountDownLatch(1);
+ Future<LoadTask> retry = requests.submit(() -> {
+ retryEntered.countDown();
+ return this.taskService.retry(task.getId(), "current-session-token");
+ });
+ Assert.assertTrue(retryEntered.await(5, TimeUnit.SECONDS));
+ Assert.assertFalse("restart read bypassed transaction reservation",
+ restartRead.await(300, TimeUnit.MILLISECONDS));
+ complete.countDown();
+ deleting.get(5, TimeUnit.SECONDS);
+ if (rollback) {
+ Assert.assertEquals(LoadStatus.RUNNING, retry.get(5, TimeUnit.SECONDS).getStatus());
+ this.assertTableCounts(1, 1, 1);
+ Assert.assertEquals(2, this.executor.executions());
+ } else {
+ Assert.assertThrows(java.util.concurrent.ExecutionException.class,
+ () -> retry.get(5, TimeUnit.SECONDS));
+ Assert.assertEquals(0, this.count("load_task"));
+ Assert.assertEquals(endpoint ? 1 : 0, this.count("job_manager"));
+ Assert.assertEquals(1, this.executor.executions());
+ }
+ Assert.assertTrue(restartRead.await(5, TimeUnit.SECONDS));
+ } finally {
+ complete.countDown();
+ requests.shutdownNow();
+ Assert.assertTrue(requests.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ public void testIngestDeleteEndpointStopsActiveTaskBeforeRemovingRow() throws Exception {
+ LoadTask task = this.jobService.createIngestTask(this.job(), this.mapping(),
+ this.connection(), this.client());
+ HugeGraphLoader loader = this.stopOnNativeSignal(task);
+
+ this.ingestEndpoint().perform(delete(org.apache.hugegraph.common.Constant.API_VERSION +
+ "ingest/jobs/" + task.getId()))
+ .andExpect(status().isOk());
+
+ Mockito.verify(loader).shutdown();
+ Assert.assertTrue(task.isFinished());
+ Assert.assertEquals(LoadStatus.STOPPED, task.getStatus());
+ Assert.assertEquals(0, this.count("load_task"));
+ Assert.assertNull(this.taskService.retainedTask(task.getId()));
+ }
+
+ @Test
+ public void testTaskPageDeletionStopsActiveChildAndCleansMetadata() throws Exception {
+ FileMapping mapping = this.mapping();
+ mapping.setPath(UPLOAD_ROOT + "/legacy-delete-" + UUID.randomUUID() + "/people.csv");
+ LoadTask task = this.jobService.createIngestTask(this.job(), mapping,
+ this.connection(), this.client());
+ HugeGraphLoader loader = this.stopOnNativeSignal(task);
+ this.ingestEndpoint().perform(delete(org.apache.hugegraph.common.Constant.API_VERSION +
+ "ingest/tasks/" + task.getJobId()))
+ .andExpect(status().isOk());
+ this.assertTableCounts(0, 0, 0);
+ Assert.assertNull(this.taskService.retainedTask(task.getId()));
+ Mockito.verify(loader).shutdown();
+ }
+
+ @Test
+ public void testIngestDeleteEndpointRetainsNativeCleanupFailure() throws Exception {
+ LoadTask task = this.jobService.createIngestTask(this.job(), this.mapping(),
+ this.connection(), this.client());
+ try (NativeCleanupTestSupport nativeCleanup = new NativeCleanupTestSupport(
+ task, this.nativeFolder.newFile("endpoint-reset.json").toPath(), true, LoadStatus.RUNNING)) {
+ task.run();
+ this.executor.complete(task);
+
+ Assert.assertThrows(Exception.class, () -> this.ingestEndpoint().perform(
+ delete(org.apache.hugegraph.common.Constant.API_VERSION + "ingest/jobs/" + task.getId())));
+
+ this.assertTableCounts(1, 1, 1);
+ Assert.assertSame(task, this.taskService.retainedTask(task.getId()));
+ Assert.assertTrue(task.isCleanupFailed());
+ Assert.assertEquals("LOADING", nativeCleanup.mode());
+ nativeCleanup.assertResourcesClosed();
+ }
+ }
+
+ @Test
+ public void testActiveStopDeletionFailureRollsBackAndKeepsRuntimeForDeletionRetry() {
+ LoadTask task = this.jobService.createIngestTask(this.job(), this.mapping(),
+ this.connection(), this.client());
+ HugeGraphLoader loader = this.stopOnNativeSignal(task);
+ // Fail the later job-row DELETE after native stop and task removal,
+ // using real H2 referential integrity rather than a mocked mapper.
+ this.jdbc.execute("CREATE TABLE delete_guard (job_id INT REFERENCES job_manager(id))");
+ try {
+ this.jdbc.update("INSERT INTO delete_guard VALUES (?)", task.getJobId());
+ Assert.assertThrows(RuntimeException.class, () -> this.jobService.deleteJob(task.getJobId()));
+ } finally {
+ this.jdbc.execute("DROP TABLE delete_guard");
+ }
+
+ this.assertTableCounts(1, 1, 1);
+ // The final callback commits before the short deletion transaction.
+ // A failed DELETE rolls back metadata removal, not that completed callback.
+ Assert.assertEquals(LoadStatus.STOPPED, this.taskService.get(task.getId()).getStatus());
+ Assert.assertSame(task, this.taskService.retainedTask(task.getId()));
+ Assert.assertEquals(LoadStatus.STOPPED, task.getStatus());
+ Assert.assertTrue(task.isFinished());
+ Mockito.verify(loader).shutdown();
+ // Periodic progress skips terminal runtimes; do not claim automatic
+ // DB reconciliation. A later deletion can reuse the real certificate.
+ this.jobService.deleteJob(task.getJobId());
+ Assert.assertEquals(0, this.count("load_task"));
+ Assert.assertEquals(0, this.count("job_manager"));
+ Assert.assertNull(this.taskService.retainedTask(task.getId()));
+ Mockito.verify(loader).shutdown();
+ }
+
+ @Test
+ public void testPublicStopRecoversRunningRowAfterDeletionRollback() throws Exception {
+ this.assertPublicStopAfterDeletionRollback(LoadStatus.RUNNING);
+ }
+
+ @Test
+ public void testPublicStopRecoversPausedRowAfterDeletionRollback() throws Exception {
+ this.assertPublicStopAfterDeletionRollback(LoadStatus.PAUSED);
+ }
+
+ @Test
+ public void testRecoveredStopRetainsRuntimeOnOuterRollback() throws Exception {
+ this.assertPublicStopAfterDeletionRollback(LoadStatus.RUNNING, true);
+ }
+
+ private void assertPublicStopAfterDeletionRollback(LoadStatus persisted) throws Exception {
+ this.assertPublicStopAfterDeletionRollback(persisted, false);
+ }
+
+ private void assertPublicStopAfterDeletionRollback(LoadStatus persisted, boolean rollback) throws Exception {
+ LoadTask task = this.jobService.createIngestTask(this.job(), this.mapping(),
+ this.connection(), this.client());
+ HugeGraphLoader loader = this.stopOnNativeSignal(task);
+ this.jdbc.execute("CREATE TABLE stop_recovery_guard (job_id INT REFERENCES job_manager(id))");
+ try {
+ this.jdbc.update("INSERT INTO stop_recovery_guard VALUES (?)", task.getJobId());
+ Assert.assertThrows(RuntimeException.class, () -> this.jobService.deleteJob(task.getJobId()));
+ } finally {
+ this.jdbc.execute("DROP TABLE stop_recovery_guard");
+ }
+ this.jdbc.update("UPDATE load_task SET load_status = ? WHERE id = ?", persisted.getValue(), task.getId());
+ Assert.assertEquals(persisted, this.taskService.get(task.getId()).getStatus());
+ Assert.assertSame(task, this.taskService.retainedTask(task.getId()));
+ Assert.assertTrue(task.isFinished());
+ Assert.assertEquals(LoadStatus.STOPPED, task.getStatus());
+
+ task.setCleanupFailed(true);
+ Assert.assertThrows(RuntimeException.class, () -> this.taskService.stop(task.getId()));
+ Assert.assertEquals(persisted, this.taskService.get(task.getId()).getStatus());
+ Assert.assertSame(task, this.taskService.retainedTask(task.getId()));
+ task.setCleanupFailed(false);
+ task.setFinished(false);
+ Assert.assertThrows(RuntimeException.class, () -> this.taskService.stop(task.getId()));
+ Assert.assertEquals(persisted, this.taskService.get(task.getId()).getStatus());
+ task.setFinished(true);
+
+ if (rollback) {
+ new TransactionTemplate(this.transactionManager).execute(transaction -> {
+ this.taskService.stop(task.getId());
+ Assert.assertEquals(LoadStatus.STOPPED, this.taskService.get(task.getId()).getStatus());
+ Assert.assertSame(task, this.taskService.retainedTask(task.getId()));
+ transaction.setRollbackOnly();
+ return null;
+ });
+ Assert.assertEquals(persisted, this.taskService.get(task.getId()).getStatus());
+ Assert.assertSame(task, this.taskService.retainedTask(task.getId()));
+ }
+
+ String endpoint = org.apache.hugegraph.common.Constant.API_VERSION +
+ "graphspaces/DEFAULT/graphs/hugegraph/job-manager/" + task.getJobId() + "/load-tasks/stop";
+ this.loadTaskEndpoint().perform(org.springframework.test.web.servlet.request.MockMvcRequestBuilders
+ .post(endpoint)
+ .param("task_id", task.getId().toString()))
+ .andExpect(status().isOk());
+ Assert.assertEquals(LoadStatus.STOPPED, this.taskService.get(task.getId()).getStatus());
+ Assert.assertNull(this.taskService.retainedTask(task.getId()));
+ Mockito.verify(loader).shutdown();
+ Assert.assertThrows(Exception.class, () -> this.loadTaskEndpoint().perform(
+ org.springframework.test.web.servlet.request.MockMvcRequestBuilders.post(endpoint)
+ .param("task_id", task.getId().toString())));
+ this.taskService.stubReconnect(true);
+ int dispatched = this.executor.executions();
+ this.taskService.retry(task.getId(), "fixture-current-token");
+ Assert.assertEquals(dispatched + 1, this.executor.executions());
+ Assert.assertEquals(LoadStatus.RUNNING, this.taskService.get(task.getId()).getStatus());
+ Mockito.verify(loader).shutdown();
+ }
+
+ private MockMvc loadTaskEndpoint() {
+ org.apache.hugegraph.controller.load.LoadTaskController controller =
+ new org.apache.hugegraph.controller.load.LoadTaskController(this.taskService) {
+ @Override
+ protected HugeClient requireGraphSpaceWrite(String graphSpace) {
+ return null;
+ }
+ };
+ ReflectionTestUtils.setField(controller, "jobService", this.jobService);
+ return MockMvcBuilders.standaloneSetup(controller).build();
+ }
+
+ @Test
+ public void testRejectedDispatchNativeCleanupFailureRetainsRowsAndBlocksReplacement() throws Exception {
+ this.executor.reject(true);
+ AtomicReference<LoadTask> created = new AtomicReference<>();
+ AtomicReference<NativeCleanupTestSupport> boundary =
+ new AtomicReference<>();
+ try {
+ new TransactionTemplate(this.transactionManager).execute(status -> {
+ LoadTask task = this.jobService.createIngestTask(this.job(), this.mapping(),
+ this.connection(), this.client());
+ created.set(task);
+ try {
+ NativeCleanupTestSupport nativeCleanup = new NativeCleanupTestSupport(
+ task, this.nativeFolder.newFile("dispatch-reset.json").toPath(), true, LoadStatus.RUNNING);
+ boundary.set(nativeCleanup);
+ nativeCleanup.loader.context().failLoading(nativeCleanup.original);
+ } catch (Exception e) {
+ throw new AssertionError(e);
+ }
+ return null;
+ });
+ LoadTask task = created.get();
+ NativeCleanupTestSupport nativeCleanup = boundary.get();
+ Assert.assertTrue(task.isFinished());
+ Assert.assertTrue(task.isCleanupFailed());
+ Assert.assertEquals(LoadStatus.FAILED, task.getStatus());
+ Assert.assertSame(task, this.taskService.retainedTask(task.getId()));
+ this.assertTableCounts(1, 1, 1);
+ Assert.assertEquals(LoadStatus.FAILED, this.taskService.get(task.getId()).getStatus());
+ Assert.assertEquals("LOADING", nativeCleanup.mode());
+ nativeCleanup.assertResourcesClosed();
+ try (MockedConstruction<HugeGraphLoader> replacements = Mockito.mockConstruction(HugeGraphLoader.class)) {
+ Assert.assertThrows(RuntimeException.class, () -> this.taskService.retry(task.getId(),
+ "current-token"));
+ Assert.assertThrows(RuntimeException.class, () -> this.taskService.resume(task.getId(),
+ "current-token"));
+ Assert.assertThrows(RuntimeException.class, () -> this.taskService.remove(task.getId()));
+ Assert.assertTrue(replacements.constructed().isEmpty());
+ this.assertTableCounts(1, 1, 1);
+ Assert.assertEquals(1, this.executor.executions());
+ }
+ } finally {
+ if (boundary.get() != null) {
+ boundary.get().close();
+ }
+ }
+ }
+
+ @Test
+ public void testMissingTaskDeletionKeepsExistingContracts() throws Exception {
+ this.ingestEndpoint().perform(delete(org.apache.hugegraph.common.Constant.API_VERSION + "ingest/jobs/-1"))
+ .andExpect(status().isOk())
+ .andExpect(org.springframework.test.web.servlet.result.MockMvcResultMatchers
+ .jsonPath("$.status").value(404));
+ Assert.assertThrows(org.apache.hugegraph.exception.InternalException.class,
+ () -> this.taskService.remove(-1));
+ this.assertTableCounts(0, 0, 0);
+ }
+
+ private HugeGraphLoader stopOnNativeSignal(LoadTask task) {
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ LoadContext context = Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS);
+ Mockito.when(loader.context()).thenReturn(context);
+ Mockito.doAnswer(invocation -> {
+ task.run();
+ this.executor.complete(task);
+ return null;
+ }).when(context).stopLoading();
+ task.setLoader(loader);
+ return loader;
+ }
+
+ private MockMvc ingestEndpoint() {
+ IngestController controller = new IngestController() {
+ @Override
+ protected HugeClient requireGraphSpaceWrite(String graphSpace) {
+ // Only authentication is isolated: routing, controller, Spring
+ // service proxy, transaction and MyBatis persistence are real.
+ return null;
+ }
+ };
+ ReflectionTestUtils.setField(controller, "loadTaskService", this.taskService);
+ ReflectionTestUtils.setField(controller, "jobManagerService", this.jobService);
+ return MockMvcBuilders.standaloneSetup(controller).build();
+ }
+
private JobManager job() {
Date now = new Date();
return JobManager.builder()
@@ -454,22 +1046,104 @@ public class IngestTransactionIntegrationTest {
public static class TestLoadTaskService extends LoadTaskService {
private boolean failTaskInsert;
+ private LoadTask createdTask;
+ private Runnable afterTaskSnapshot;
+ private Runnable afterStopCertificate;
+ private boolean stubReconnect;
+ private Runnable beforeRestartRead;
+
+ public LoadTask retainedTask(int id) {
+ try {
+ java.lang.reflect.Field field = LoadTaskService.class.getDeclaredField("runningTaskContainer");
+ field.setAccessible(true);
+ @SuppressWarnings("unchecked")
+ java.util.Map<Integer, LoadTask> running = (java.util.Map<Integer, LoadTask>) field.get(this);
+ return running.get(id);
+ } catch (ReflectiveOperationException e) {
+ throw new AssertionError(e);
+ }
+ }
+
+ public void beforeRestartRead(Runnable action) {
+ this.beforeRestartRead = action;
+ }
+
+ public void afterStopCertificate(Runnable action) {
+ this.afterStopCertificate = action;
+ }
+
+ public void stubReconnect(boolean enabled) {
+ this.stubReconnect = enabled;
+ }
+
+ @Override
+ public LoadTask get(int id) {
+ if (this.stubReconnect && this.beforeRestartRead != null) {
+ this.beforeRestartRead.run();
+ }
+ LoadTask task = super.get(id);
+ if (this.stubReconnect && task != null) {
+ task = Mockito.spy(task);
+ Mockito.doNothing().when(task).reconnect(Mockito.anyString());
+ }
+ return task;
+ }
+
+ @Override
+ LoadTask stopForDelete(int id) {
+ LoadTask task = super.stopForDelete(id);
+ Runnable action = this.afterStopCertificate;
+ this.afterStopCertificate = null;
+ if (action != null) {
+ action.run();
+ }
+ return task;
+ }
+
+ public void afterTaskSnapshot(Runnable action) {
+ this.afterTaskSnapshot = action;
+ }
+
+ @Override
+ public List<LoadTask> taskListByJob(int jobId) {
+ List<LoadTask> snapshot = super.taskListByJob(jobId);
+ Runnable action = this.afterTaskSnapshot;
+ this.afterTaskSnapshot = null;
+ if (action != null) {
+ action.run();
+ }
+ return snapshot;
+ }
public void failTaskInsert(boolean fail) {
this.failTaskInsert = fail;
}
+ public LoadTask startLoader(GraphConnection connection, FileMapping mapping, HugeClient client) {
+ this.createdTask = super.start(connection, mapping, client);
+ return this.createdTask;
+ }
+
+ @Override
+ public void save(LoadTask task) {
+ this.createdTask = task;
+ super.save(task);
+ }
+
+ public LoadTask createdTask() {
+ return this.createdTask;
+ }
+
@Override
public LoadTask start(GraphConnection connection, FileMapping mapping,
HugeClient client) {
- String graph = this.failTaskInsert ? null : connection.getGraph();
LoadOptions options = new LoadOptions();
options.token = connection.getToken();
LoadTask task = LoadTask.builder()
.graphSpace(connection.getGraphSpace())
- .graph(graph)
+ .graph(connection.getGraph())
.jobId(mapping.getJobId())
- .fileId(mapping.getId())
+ .fileId(this.failTaskInsert ? null : mapping.getId())
.fileName(mapping.getName())
.options(options)
.vertices(Collections.emptySet())
@@ -490,6 +1164,8 @@ public class IngestTransactionIntegrationTest {
public static class RecordingLoadTaskExecutor extends LoadTaskExecutor {
private final AtomicInteger executions = new AtomicInteger();
+ private final java.util.Map<LoadTask, Runnable> callbacks =
+ Collections.synchronizedMap(new java.util.IdentityHashMap<>());
private volatile boolean reject;
@Override
@@ -499,6 +1175,20 @@ public class IngestTransactionIntegrationTest {
throw new org.springframework.core.task.TaskRejectedException(
"test executor queue is full");
}
+ this.callbacks.put(task, callback);
+ }
+
+ public void complete(LoadTask task) {
+ Runnable callback = this.callbacks.remove(task);
+ Assert.assertNotNull("A recorded dispatch must run its actual service callback", callback);
+ try {
+ callback.run();
+ } catch (RuntimeException | Error failure) {
+ if (!task.isCleanupFailed()) {
+ throw failure;
+ }
+ Assert.assertSame(task.getCleanupFailure(), failure);
+ }
}
public int executions() {
@@ -511,6 +1201,7 @@ public class IngestTransactionIntegrationTest {
public void reset() {
this.executions.set(0);
+ this.callbacks.clear();
this.reject = false;
}
}
diff --git a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/service/load/JobAdmissionIntegrationTest.java b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/service/load/JobAdmissionIntegrationTest.java
new file mode 100644
index 000000000..5d9e13343
--- /dev/null
+++ b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/service/load/JobAdmissionIntegrationTest.java
@@ -0,0 +1,1822 @@
+/*
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with this
+ * work for additional information regarding copyright ownership. The ASF
+ * licenses this file to You under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance with the
+ * License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hugegraph.service.load;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Collections;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.Date;
+import java.util.List;
+import java.util.UUID;
+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.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import javax.sql.DataSource;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.mockito.MockedConstruction;
+import org.mockito.Mockito;
+import org.mybatis.spring.annotation.MapperScan;
+import org.springframework.aop.support.AopUtils;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.SpringBootConfiguration;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.test.mock.mockito.MockBean;
+import org.springframework.context.annotation.Bean;
+import org.springframework.jdbc.core.JdbcTemplate;
+import org.springframework.mock.web.MockMultipartFile;
+import org.springframework.scheduling.annotation.EnableAsync;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
+import org.springframework.test.context.junit4.SpringRunner;
+import org.springframework.test.util.ReflectionTestUtils;
+import org.springframework.transaction.PlatformTransactionManager;
+import org.springframework.transaction.annotation.EnableTransactionManagement;
+import org.springframework.transaction.support.TransactionTemplate;
+
+import org.apache.hugegraph.config.HugeConfig;
+import org.apache.hugegraph.controller.load.LoadTaskController;
+import org.apache.hugegraph.controller.load.FileUploadController;
+import org.apache.hugegraph.controller.load.FileMappingController;
+import org.apache.hugegraph.driver.HugeClient;
+import org.apache.hugegraph.entity.GraphConnection;
+import org.apache.hugegraph.entity.enums.FileMappingStatus;
+import org.apache.hugegraph.entity.enums.JobStatus;
+import org.apache.hugegraph.entity.enums.LoadStatus;
+import org.apache.hugegraph.entity.load.FileMapping;
+import org.apache.hugegraph.entity.load.FileSetting;
+import org.apache.hugegraph.entity.load.JobManager;
+import org.apache.hugegraph.entity.load.ListFormat;
+import org.apache.hugegraph.entity.load.LoadParameter;
+import org.apache.hugegraph.entity.load.LoadTask;
+import org.apache.hugegraph.handler.LoadTaskExecutor;
+import org.apache.hugegraph.loader.HugeGraphLoader;
+import org.apache.hugegraph.loader.executor.LoadContext;
+import org.apache.hugegraph.loader.executor.LoadOptions;
+import org.apache.hugegraph.loader.builder.SchemaCache;
+import org.apache.hugegraph.loader.task.TaskManager;
+import org.springframework.scheduling.annotation.Async;
+import org.apache.hugegraph.options.HubbleOptions;
+import org.apache.hugegraph.service.schema.EdgeLabelService;
+import org.apache.hugegraph.service.schema.VertexLabelService;
+import org.apache.hugegraph.testutil.Assert;
+
+@RunWith(SpringRunner.class)
+@SpringBootTest(classes = JobAdmissionIntegrationTest.Configuration.class,
+ webEnvironment = SpringBootTest.WebEnvironment.NONE,
+ properties = {
+ "spring.datasource.url=jdbc:h2:mem:job-admission;DB_CLOSE_DELAY=-1",
+ "spring.datasource.driver-class-name=org.h2.Driver",
+ "spring.datasource.username=sa", "spring.datasource.password=",
+ "spring.sql.init.schema-locations=classpath:database/ingest-transaction-schema.sql",
+ "spring.sql.init.mode=always", "spring.datasource.hikari.maximum-pool-size=3",
+ "spring.autoconfigure.exclude=" +
+ "org.mybatis.spring.boot.autoconfigure.MybatisAutoConfiguration," +
+ "org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchClientAutoConfiguration," +
+ "org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration",
+ "mybatis-plus.configuration.map-underscore-to-camel-case=true",
+ "mybatis-plus.configuration.use-generated-keys=true", "mybatis-plus.global-config.banner=false"
+ })
+public class JobAdmissionIntegrationTest {
+
+ @Rule
+ public TemporaryFolder folder = new TemporaryFolder();
+ @Autowired
+ private GateTaskService tasks;
+ @Autowired
+ private JobManagerService jobs;
+ @Autowired
+ private GateMappingService mappings;
+ @Autowired
+ private HugeConfig config;
+ @Autowired
+ private DataSource dataSource;
+ @Autowired
+ private PlatformTransactionManager transactionManager;
+ @Autowired
+ private ThreadPoolTaskExecutor executor;
+ @Autowired
+ private GateLoadTaskExecutor taskExecutor;
+ @MockBean
+ private VertexLabelService vertexLabels;
+ @MockBean
+ private EdgeLabelService edgeLabels;
+
+ private JdbcTemplate jdbc;
+ private Path uploads;
+
+ @Before
+ public void setup() throws Exception {
+ this.jdbc = new JdbcTemplate(this.dataSource);
+ this.clearRows();
+ this.mappings.getUploadingTokenLocks().clear();
+ this.tasks.afterSnapshot(null);
+ this.tasks.queuedFirst(null);
+ this.tasks.prepareRestart(null);
+ this.tasks.afterReservation(null);
+ CallbackWriteFailure.thread = null;
+ this.taskExecutor.callbackGate(null);
+ this.taskExecutor.callbackFinished(null);
+ this.mappings.afterSnapshot(null);
+ this.uploads = this.folder.newFolder("uploads").toPath();
+ Mockito.when(this.config.get(HubbleOptions.PD_ENABLED)).thenReturn(false);
+ Mockito.when(this.config.get(HubbleOptions.SERVER_URL)).thenReturn("http://127.0.0.1:65530");
+ Mockito.when(this.config.get(HubbleOptions.UPLOAD_FILE_LOCATION)).thenReturn(this.uploads.toString());
+ Mockito.when(this.config.get(HubbleOptions.UPLOAD_FILE_FORMAT_LIST)).thenReturn(List.of("csv"));
+ Mockito.when(this.config.get(HubbleOptions.UPLOAD_SINGLE_FILE_SIZE_LIMIT)).thenReturn(1024L);
+ Mockito.when(this.config.get(HubbleOptions.UPLOAD_TOTAL_FILE_SIZE_LIMIT)).thenReturn(4096L);
+ }
+
+ @After
+ public void teardown() throws Exception {
+ this.executor.getThreadPoolExecutor().submit(() -> { }).get(5, TimeUnit.SECONDS);
+ this.clearRows();
+ }
+
+ @Test(timeout = 15000)
+ public void testDeleteSnapshotRejectsLaterControllerStartBeforeLoaderConstruction() throws Exception {
+ FileMapping mapping = this.seed();
+ CountDownLatch snapshot = new CountDownLatch(1);
+ CountDownLatch delete = new CountDownLatch(1);
+ this.tasks.afterSnapshot(() -> {
+ snapshot.countDown();
+ await(delete);
+ });
+ AtomicInteger constructions = new AtomicInteger();
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ Future<?> deleting = requests.submit(() -> this.jobs.deleteJob(mapping.getJobId()));
+ Assert.assertTrue(snapshot.await(5, TimeUnit.SECONDS));
+ CountDownLatch entering = new CountDownLatch(1);
+ Future<?> starting = requests.submit(() -> {
+ try (MockedConstruction<HugeGraphLoader> ignored = Mockito.mockConstruction(
+ HugeGraphLoader.class, (loader, context) -> constructions.incrementAndGet())) {
+ entering.countDown();
+ return this.controller().start("DEFAULT", "hugegraph", mapping.getJobId(),
+ List.of(mapping.getId()));
+ }
+ });
+ Assert.assertTrue(entering.await(5, TimeUnit.SECONDS));
+ assertBlocked(starting);
+ Assert.assertEquals(0, constructions.get());
+ Assert.assertEquals(0, this.count("load_task"));
+ delete.countDown();
+ deleting.get(5, TimeUnit.SECONDS);
+ Assert.assertThrows(java.util.concurrent.ExecutionException.class, () -> starting.get(5, TimeUnit.SECONDS));
+ Assert.assertEquals(0, constructions.get());
+ Assert.assertEquals(0, this.count("job_manager"));
+ Assert.assertEquals(0, this.count("load_task"));
+ } finally {
+ delete.countDown();
+ stop(requests);
+ }
+ }
+
+ @Test(timeout = 15000)
+ public void testControllerStartBeforeDeleteMustFinishNativeCleanupBeforeCommit() throws Exception {
+ FileMapping mapping = this.seed();
+ CountDownLatch loading = new CountDownLatch(1);
+ CountDownLatch load = new CountDownLatch(1);
+ CountDownLatch cleaning = new CountDownLatch(1);
+ CountDownLatch cleanup = new CountDownLatch(1);
+ AtomicReference<HugeGraphLoader> constructed = new AtomicReference<>();
+ ExecutorService requests = Executors.newSingleThreadExecutor();
+ try (MockedConstruction<HugeGraphLoader> ignored = Mockito.mockConstruction(HugeGraphLoader.class,
+ (loader, construction) -> {
+ constructed.set(loader);
+ LoadContext context = Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS);
+ Mockito.when(loader.context()).thenReturn(context);
+ Mockito.when(loader.load()).thenAnswer(call -> {
+ loading.countDown();
+ await(load);
+ return false;
+ });
+ Mockito.doAnswer(call -> {
+ load.countDown();
+ return null;
+ }).when(context).stopLoading();
+ Mockito.doAnswer(call -> {
+ cleaning.countDown();
+ await(cleanup);
+ return null;
+ }).when(loader).shutdown();
+ })) {
+ LoadTask task = this.controller().start("DEFAULT", "hugegraph", mapping.getJobId(),
+ List.of(mapping.getId())).get(0);
+ Assert.assertTrue(AopUtils.isAopProxy(this.tasks));
+ Assert.assertTrue(AopUtils.isAopProxy(this.taskExecutor));
+ Assert.assertTrue(loading.await(5, TimeUnit.SECONDS));
+ Assert.assertSame(task, this.tasks.retained(task.getId()));
+ Future<?> deleting = requests.submit(() -> this.jobs.deleteJob(mapping.getJobId()));
+ Assert.assertTrue(cleaning.await(5, TimeUnit.SECONDS));
+ assertBlocked(deleting);
+ Assert.assertFalse(task.isFinished());
+ Assert.assertEquals(1, this.count("job_manager"));
+ Assert.assertEquals(1, this.count("load_task"));
+ cleanup.countDown();
+ deleting.get(5, TimeUnit.SECONDS);
+ Assert.assertEquals(0, this.count("job_manager"));
+ Assert.assertEquals(0, this.count("load_task"));
+ Assert.assertNull(this.tasks.retained(task.getId()));
+ Mockito.verify(constructed.get()).shutdown();
+ } finally {
+ load.countDown();
+ cleanup.countDown();
+ stop(requests);
+ }
+ }
+
+ @Test(timeout = 15000)
+ public void testPublicSaveCannotInsertAfterDeleteSnapshot() throws Exception {
+ this.publicSaveAfterDeletion(false);
+ }
+
+ @Test(timeout = 15000)
+ public void testJobReservationSurvivesDeleteMethodAndReleasesOnOuterRollback() throws Exception {
+ this.publicSaveAfterDeletion(true);
+ }
+
+ private void publicSaveAfterDeletion(boolean rollback) throws Exception {
+ FileMapping mapping = this.seed();
+ LoadTask task = this.metadata(mapping);
+ CountDownLatch deleted = new CountDownLatch(1);
+ CountDownLatch complete = new CountDownLatch(1);
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ TransactionTemplate transaction = new TransactionTemplate(this.transactionManager);
+ Future<?> deleting = requests.submit(() -> transaction.execute(status -> {
+ this.jobs.deleteJob(mapping.getJobId());
+ deleted.countDown();
+ await(complete);
+ if (rollback) {
+ status.setRollbackOnly();
+ }
+ return null;
+ }));
+ Assert.assertTrue(deleted.await(5, TimeUnit.SECONDS));
+ CountDownLatch entering = new CountDownLatch(1);
+ Future<?> saving = requests.submit(() -> {
+ entering.countDown();
+ this.tasks.save(task);
+ });
+ Assert.assertTrue(entering.await(5, TimeUnit.SECONDS));
+ assertBlocked(saving);
+ Assert.assertEquals(0, this.count("load_task"));
+ complete.countDown();
+ deleting.get(5, TimeUnit.SECONDS);
+ if (rollback) {
+ saving.get(5, TimeUnit.SECONDS);
+ Assert.assertEquals(1, this.count("job_manager"));
+ Assert.assertEquals(1, this.count("load_task"));
+ } else {
+ Assert.assertThrows(java.util.concurrent.ExecutionException.class, () -> saving.get(5,
+ TimeUnit.SECONDS));
+ Assert.assertEquals(0, this.count("job_manager"));
+ Assert.assertEquals(0, this.count("load_task"));
+ }
+ } finally {
+ complete.countDown();
+ stop(requests);
+ }
+ }
+
+ @Test
+ public void testPublicStartChecksFreshParentIdentityBeforeConstructingLoader() throws Exception {
+ FileMapping mapping = this.seed();
+ GraphConnection wrongGraph = this.connection();
+ wrongGraph.setGraph("another-graph");
+ try (MockedConstruction<HugeGraphLoader> constructed = Mockito.mockConstruction(HugeGraphLoader.class)) {
+ Assert.assertThrows(RuntimeException.class, () -> this.tasks.start(wrongGraph, mapping, this.client()));
+ LoadTask wrongTask = this.metadata(mapping);
+ wrongTask.setGraphSpace("another-space");
+ Assert.assertThrows(RuntimeException.class, () -> this.tasks.save(wrongTask));
+ this.jobs.deleteJob(mapping.getJobId());
+ Assert.assertThrows(RuntimeException.class, () -> this.tasks.start(this.connection(), mapping,
+ this.client()));
+ Assert.assertTrue(constructed.constructed().isEmpty());
+ Assert.assertEquals(0, this.count("load_task"));
+ }
+ }
+
+ @Test
+ public void testPublicStartAcceptsIngestLoadingParentAndClosesAfterOuterRollback() throws Exception {
+ FileMapping mapping = this.seed();
+ JobManager job = this.jobs.get(mapping.getJobId());
+ job.setJobStatus(JobStatus.LOADING);
+ this.jobs.update(job);
+ AtomicReference<LoadTask> created = new AtomicReference<>();
+ try (MockedConstruction<HugeGraphLoader> constructed = Mockito.mockConstruction(HugeGraphLoader.class)) {
+ new TransactionTemplate(this.transactionManager).execute(status -> {
+ created.set(this.tasks.start(this.connection(), mapping, this.client()));
+ status.setRollbackOnly();
+ return null;
+ });
+ Assert.assertEquals(1, constructed.constructed().size());
+ Mockito.verify(constructed.constructed().get(0)).shutdown();
+ Assert.assertTrue(created.get().isFinished());
+ Assert.assertEquals(0, this.count("load_task"));
+ Assert.assertEquals(1, this.count("job_manager"));
+ }
+ // The outer rollback must have released the job stripe.
+ this.tasks.save(this.metadata(mapping));
+ Assert.assertEquals(1, this.count("load_task"));
+ }
+
+ @Test(timeout = 15000)
+ public void testRuntimePublishedBeforeExecutorSubmissionCanBecomeObservable() throws Exception {
+ FileMapping mapping = this.seed();
+ LoadTaskExecutor original = this.taskExecutor;
+ AtomicInteger submits = new AtomicInteger();
+ LoadTaskExecutor observing = new LoadTaskExecutor() {
+ @Override
+ public void execute(LoadTask task, Runnable callback) {
+ submits.incrementAndGet();
+ Assert.assertSame("runtime must be registered before submission", task,
+ JobAdmissionIntegrationTest.this.tasks.retained(task.getId()));
+ super.execute(task, callback);
+ }
+ };
+ ReflectionTestUtils.setField(this.tasks, "taskExecutor", observing);
+ try (MockedConstruction<HugeGraphLoader> ignored = Mockito.mockConstruction(HugeGraphLoader.class,
+ (loader, construction) -> Mockito.when(loader.context()).thenReturn(
+ Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS)))) {
+ LoadTask task = this.tasks.start(this.connection(), mapping, this.client());
+ Assert.assertEquals(1, submits.get());
+ Assert.assertTrue(task.isFinished());
+ Assert.assertSame(task, this.tasks.retained(task.getId()));
+ } finally {
+ ReflectionTestUtils.setField(this.tasks, "taskExecutor", original);
+ }
+ }
+
+ @Test(timeout = 15000)
+ public void testClearEmptySnapshotIncludesAnAlreadyAdmittedUpload() throws Exception {
+ this.assertClearIncludesAdmittedUpload(false);
+ }
+
+ @Test(timeout = 15000)
+ public void testClearNonemptySnapshotIncludesAnAlreadyAdmittedUpload() throws Exception {
+ this.assertClearIncludesAdmittedUpload(true);
+ }
+
+ private void assertClearIncludesAdmittedUpload(boolean nonempty) throws Exception {
+ JobManager job = nonempty ? this.jobs.get(this.seed().getJobId()) : this.seedJob();
+ FileUploadController upload = this.uploadController();
+ FileMappingController clear = this.mappingController();
+ CountDownLatch admitted = new CountDownLatch(1);
+ CountDownLatch createMapping = new CountDownLatch(1);
+ CountDownLatch snapshot = new CountDownLatch(1);
+ FileMappingService observed = Mockito.mock(FileMappingService.class, call -> {
+ Object result = call.getMethod().invoke(this.mappings, call.getArguments());
+ if (call.getMethod().getName().equals("listByJob") && call.getArguments().length == 3) {
+ snapshot.countDown();
+ }
+ return result;
+ });
+ ReflectionTestUtils.setField(clear, "service", observed);
+ java.util.concurrent.atomic.AtomicBoolean firstSize = new java.util.concurrent.atomic.AtomicBoolean(true);
+ MockMultipartFile part = new MockMultipartFile("file", "admitted.csv", "text/csv", "name\nalice\n".getBytes()) {
+ @Override
+ public long getSize() {
+ if (firstSize.compareAndSet(true, false)) {
+ admitted.countDown();
+ await(createMapping);
+ }
+ return super.getSize();
+ }
+ };
+ String token = upload.fileToken("DEFAULT", "hugegraph", job.getId(), List.of("admitted.csv"))
+ .get("admitted.csv");
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ Future<?> uploading = requests.submit(() -> upload.upload("DEFAULT", "hugegraph", job.getId(),
+ part, "admitted.csv", 20L, token, 2, 0));
+ Assert.assertTrue(admitted.await(5, TimeUnit.SECONDS));
+ Future<?> clearing = requests.submit(() -> clear.clear("DEFAULT", "hugegraph", job.getId()));
+ boolean prematureSnapshot = snapshot.await(500, TimeUnit.MILLISECONDS);
+ assertBlocked(clearing);
+ createMapping.countDown();
+ uploading.get(5, TimeUnit.SECONDS);
+ clearing.get(5, TimeUnit.SECONDS);
+ Assert.assertTrue(this.mappings.listByJob(job.getId()).isEmpty());
+ Assert.assertFalse(prematureSnapshot);
+ } finally {
+ createMapping.countDown();
+ stop(requests);
+ }
+ }
+
+ @Test
+ public void testNativeCaseAliasCannotOverwriteAnExistingPartialUpload() throws Exception {
+ this.assertNativeAliasProtection("name.csv", "NAME.csv");
+ }
+
+ @Test
+ public void testNativeUnicodeAliasCannotOverwriteAnExistingPartialUpload() throws Exception {
+ this.assertNativeAliasProtection("café.csv", "café.csv");
+ }
+
+ private void assertNativeAliasProtection(String name, String alias) throws Exception {
+ JobManager job = this.seedJob();
+ FileUploadController controller = this.uploadController();
+ var tokens = controller.fileToken("DEFAULT", "hugegraph", job.getId(), List.of(name, alias));
+ MockMultipartFile first = new MockMultipartFile("file", name, "text/csv", "name\nalpha\n".getBytes());
+ controller.upload("DEFAULT", "hugegraph", job.getId(), first, name, 20L, tokens.get(name), 2, 0);
+ FileMapping mapping = this.mappings.get("DEFAULT", "hugegraph", job.getId(), name);
+ Path directory = Path.of(mapping.getPath());
+ Path aliasDirectory = directory.resolveSibling(alias);
+ boolean equivalent = Files.exists(aliasDirectory) && Files.isSameFile(directory, aliasDirectory);
+ AtomicInteger aliasWrites = new AtomicInteger();
+ MockMultipartFile second = new MockMultipartFile("file", alias, "text/csv", "other\nbeta\n".getBytes()) {
+ @Override
+ public void transferTo(File destination) throws IOException {
+ aliasWrites.incrementAndGet();
+ super.transferTo(destination);
+ }
+ };
+ if (equivalent) {
+ Assert.assertThrows(RuntimeException.class, () -> controller.upload("DEFAULT", "hugegraph", job.getId(),
+ second, alias, 20L, tokens.get(alias), 2, 0));
+ Assert.assertEquals(0, aliasWrites.get());
+ Assert.assertTrue(this.mappings.get("DEFAULT", "hugegraph", job.getId(), alias) == null);
+ Assert.assertEquals("name\nalpha\n", Files.readString(directory.resolve(name + "-0")));
+ } else {
+ var accepted = controller.upload("DEFAULT", "hugegraph", job.getId(), second, alias, 20L,
+ tokens.get(alias), 2, 0);
+ Assert.assertFalse(accepted.getStatus() ==
+ org.apache.hugegraph.entity.load.FileUploadResult.Status.FAILURE);
+ Assert.assertFalse(Files.isSameFile(directory, aliasDirectory));
+ Assert.assertEquals("name\nalpha\n", Files.readString(directory.resolve(name + "-0")));
+ Assert.assertEquals("other\nbeta\n", Files.readString(aliasDirectory.resolve(alias + "-0")));
+ }
+ }
+
+ @Test
+ public void testFailedMappingInsertRemovesOnlyItsNewEmptyPartDirectory() throws Exception {
+ JobManager job = this.seedJob();
+ FileUploadController controller = this.uploadController();
+ String token = controller.fileToken("DEFAULT", "hugegraph", job.getId(), List.of("rejected.csv"))
+ .get("rejected.csv");
+ try (var paths = Files.walk(this.uploads)) {
+ Assert.assertEquals(1L, paths.count());
+ }
+ String constraint = "reject_upload_fixture";
+ this.jdbc.execute("ALTER TABLE file_mapping ADD CONSTRAINT " + constraint +
+ " CHECK (job_id <> " + job.getId() + ")");
+ try {
+ Assert.assertThrows(RuntimeException.class, () -> controller.upload("DEFAULT", "hugegraph", job.getId(),
+ this.file("rejected.csv"), "rejected.csv", 20L, token, 2, 0));
+ Assert.assertTrue(this.mappings.listByJob(job.getId()).isEmpty());
+ Path directory = this.uploads.resolve("graph-connection-DEFAULT-hugegraph")
+ .resolve("job-" + job.getId()).resolve("rejected.csv");
+ Assert.assertFalse(Files.exists(directory));
+ } finally {
+ this.jdbc.execute("ALTER TABLE file_mapping DROP CONSTRAINT " + constraint);
+ }
+ }
+
+ @Test
+ public void testClearRejectsWrongGraphRouteBeforeTakingAMappingSnapshot() throws Exception {
+ FileMapping mapping = this.seed();
+ Assert.assertThrows(RuntimeException.class,
+ () -> this.mappingController().clear("DEFAULT", "other", mapping.getJobId()));
+ Assert.assertEquals((Object) mapping.getJobId(), this.mappings.get(mapping.getId()).getJobId());
+ Assert.assertTrue(Files.isRegularFile(Path.of(mapping.getPath())));
+ }
+
+ @Test(timeout = 15000)
+ public void testSameJobFilesWriteConcurrentlyAndDeleteWaitsForBoth() throws Exception {
+ JobManager job = this.seedJob();
+ FileUploadController controller = this.uploadController();
+ var tokens = controller.fileToken("DEFAULT", "hugegraph", job.getId(), List.of("a.csv", "b.csv"));
+ CountDownLatch writing = new CountDownLatch(2);
+ CountDownLatch finish = new CountDownLatch(1);
+ ExecutorService requests = Executors.newFixedThreadPool(3);
+ try {
+ Future<?> first = requests.submit(() -> controller.upload("DEFAULT", "hugegraph", job.getId(),
+ this.blockedFile("a.csv", writing, finish), "a.csv", 20L, tokens.get("a.csv"), 2, 0));
+ Future<?> second = requests.submit(() -> controller.upload("DEFAULT", "hugegraph", job.getId(),
+ this.blockedFile("b.csv", writing, finish), "b.csv", 20L, tokens.get("b.csv"), 2, 0));
+ Assert.assertTrue(writing.await(2, TimeUnit.SECONDS));
+ Future<?> deleting = requests.submit(() -> this.jobs.deleteJob(job.getId()));
+ assertBlocked(deleting);
+ finish.countDown();
+ first.get(5, TimeUnit.SECONDS);
+ second.get(5, TimeUnit.SECONDS);
+ deleting.get(5, TimeUnit.SECONDS);
+ Assert.assertEquals(0, this.count("file_mapping"));
+ Assert.assertTrue(this.mappings.getUploadingTokenLocks().isEmpty());
+ } finally {
+ finish.countDown();
+ stop(requests);
+ }
+ }
+
+ @Test(timeout = 15000)
+ public void testSameFileTokenAliasesCompleteBeforeMergeAndPreserveContent() throws Exception {
+ JobManager job = this.seedJob();
+ FileUploadController controller = this.uploadController();
+ String token = controller.fileToken("DEFAULT", "hugegraph", job.getId(), List.of("chunks.csv"))
+ .get("chunks.csv");
+ String alias = controller.fileToken("DEFAULT", "hugegraph", job.getId(), List.of("chunks.csv"))
+ .get("chunks.csv");
+ JobManager other = this.seedJob();
+ String otherToken = controller.fileToken("DEFAULT", "hugegraph", other.getId(), List.of("chunks.csv"))
+ .get("chunks.csv");
+ Assert.assertSame(this.mappings.getUploadingTokenLocks().get(token),
+ this.mappings.getUploadingTokenLocks().get(alias));
+ CountDownLatch writing = new CountDownLatch(1);
+ CountDownLatch finish = new CountDownLatch(1);
+ AtomicInteger secondWrites = new AtomicInteger();
+ MockMultipartFile firstPart = new MockMultipartFile("file", "chunks.csv", "text/csv", "name\n".getBytes()) {
+ @Override
+ public void transferTo(File destination) throws IOException {
+ super.transferTo(destination);
+ writing.countDown();
+ await(finish);
+ }
+ };
+ MockMultipartFile secondPart = new MockMultipartFile("file", "chunks.csv", "text/csv", "alice\n".getBytes()) {
+ @Override
+ public void transferTo(File destination) throws IOException {
+ secondWrites.incrementAndGet();
+ super.transferTo(destination);
+ }
+ };
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ Future<?> first = requests.submit(() -> controller.upload("DEFAULT", "hugegraph", job.getId(),
+ firstPart, "chunks.csv", 11L, token, 2, 0));
+ Assert.assertTrue(writing.await(5, TimeUnit.SECONDS));
+ Future<?> second = requests.submit(() -> controller.upload("DEFAULT", "hugegraph", job.getId(),
+ secondPart, "chunks.csv", 11L, alias, 2, 1));
+ assertBlocked(second);
+ Assert.assertEquals(0, secondWrites.get());
+ finish.countDown();
+ first.get(5, TimeUnit.SECONDS);
+ second.get(5, TimeUnit.SECONDS);
+ FileMapping mapping = this.mappings.get("DEFAULT", "hugegraph", job.getId(), "chunks.csv");
+ Assert.assertEquals(FileMappingStatus.COMPLETED, mapping.getFileStatus());
+ Assert.assertEquals("name\nalice\n", Files.readString(Path.of(mapping.getPath())));
+ Assert.assertEquals(11L, (long) mapping.getTotalSize());
+ Assert.assertEquals(List.of("col-1"), mapping.getFileSetting().getColumnNames());
+ Assert.assertEquals(List.of("name"), mapping.getFileSetting().getColumnValues());
+ Assert.assertEquals(11L, (long) this.jobs.get(job.getId()).getJobSize());
+ Assert.assertFalse(this.mappings.getUploadingTokenLocks().containsKey(token));
+ Assert.assertFalse(this.mappings.getUploadingTokenLocks().containsKey(alias));
+ Assert.assertTrue(this.mappings.getUploadingTokenLocks().containsKey(otherToken));
+ } finally {
+ finish.countDown();
+ stop(requests);
+ }
+ }
+
+ @Test
+ public void testMappingIdDeleteRevokesOldMultipartTokenOnlyInItsJob() throws Exception {
+ this.assertMappingDeleteRevokesToken(false);
+ }
+
+ @Test
+ public void testMappingClearRevokesOldMultipartTokenOnlyInItsJob() throws Exception {
+ this.assertMappingDeleteRevokesToken(true);
+ }
+
+ private void assertMappingDeleteRevokesToken(boolean clear) throws Exception {
+ JobManager job = this.seedJob();
+ JobManager other = this.seedJob();
+ FileUploadController controller = this.uploadController();
+ String token = controller.fileToken("DEFAULT", "hugegraph", job.getId(), List.of("partial.csv"))
+ .get("partial.csv");
+ String otherToken = controller.fileToken("DEFAULT", "hugegraph", other.getId(), List.of("partial.csv"))
+ .get("partial.csv");
+ controller.upload("DEFAULT", "hugegraph", job.getId(), this.file("partial.csv"),
+ "partial.csv", 20L, token, 2, 0);
+ FileMapping partial = this.mappings.get("DEFAULT", "hugegraph", job.getId(), "partial.csv");
+ FileMappingController mappingController = this.mappingController();
+ if (clear) {
+ mappingController.clear("DEFAULT", "hugegraph", job.getId());
+ } else {
+ mappingController.delete("DEFAULT", "hugegraph", job.getId(), partial.getId());
+ }
+ var late = controller.upload("DEFAULT", "hugegraph", job.getId(), this.file("partial.csv"),
+ "partial.csv", 20L, token, 2, 1);
+ Assert.assertEquals(org.apache.hugegraph.entity.load.FileUploadResult.Status.FAILURE, late.getStatus());
+ Assert.assertTrue(this.mappings.get("DEFAULT", "hugegraph", job.getId(), "partial.csv") == null);
+ Assert.assertFalse(Files.exists(Path.of(partial.getPath())));
+ Assert.assertTrue(this.mappings.getUploadingTokenLocks().containsKey(otherToken));
+ var accepted = controller.upload("DEFAULT", "hugegraph", other.getId(), this.file("partial.csv"),
+ "partial.csv", 20L, otherToken, 2, 0);
+ Assert.assertFalse(accepted.getStatus() ==
+ org.apache.hugegraph.entity.load.FileUploadResult.Status.FAILURE);
+ }
+
+ @Test
+ public void testTokenOnlyJobDeleteInvalidatesTokensAfterCommitButNotRollback() {
+ JobManager job = this.seedJob();
+ FileUploadController controller = this.uploadController();
+ String token = controller.fileToken("DEFAULT", "hugegraph", job.getId(), List.of("pending.csv"))
+ .get("pending.csv");
+ new TransactionTemplate(this.transactionManager).execute(status -> {
+ this.jobs.deleteJob(job.getId());
+ Assert.assertTrue(this.mappings.getUploadingTokenLocks().containsKey(token));
+ status.setRollbackOnly();
+ return null;
+ });
+ Assert.assertTrue(this.mappings.getUploadingTokenLocks().containsKey(token));
+ Assert.assertTrue(this.jobs.get(job.getId()) != null);
+ this.jobs.deleteJob(job.getId());
+ Assert.assertFalse(this.mappings.getUploadingTokenLocks().containsKey(token));
+ Assert.assertTrue(this.jobs.get(job.getId()) == null);
+ }
+
+ @Test
+ public void testTokenDeleteRollbackKeepsPartialUploadAndItsToken() {
+ JobManager job = this.seedJob();
+ FileUploadController controller = this.uploadController();
+ String token = controller.fileToken("DEFAULT", "hugegraph", job.getId(), List.of("rollback.csv"))
+ .get("rollback.csv");
+ controller.upload("DEFAULT", "hugegraph", job.getId(), this.file("rollback.csv"),
+ "rollback.csv", 20L, token, 2, 0);
+ FileMapping partial = this.mappings.get("DEFAULT", "hugegraph", job.getId(), "rollback.csv");
+ new TransactionTemplate(this.transactionManager).execute(status -> {
+ controller.delete("DEFAULT", "hugegraph", job.getId(), "rollback.csv", token);
+ Assert.assertTrue(this.mappings.getUploadingTokenLocks().containsKey(token));
+ status.setRollbackOnly();
+ return null;
+ });
+ Assert.assertTrue(this.mappings.getUploadingTokenLocks().containsKey(token));
+ Assert.assertEquals(job.getId(), this.mappings.get(partial.getId()).getJobId());
+ Assert.assertTrue(Files.exists(Path.of(partial.getPath())));
+ }
+
+ @Test
+ public void testExpiredUploadRevokesOnlyItsJobEvenWhenDiskCleanupFails() throws Exception {
+ JobManager job = this.seedJob();
+ JobManager other = this.seedJob();
+ FileUploadController controller = this.uploadController();
+ String token = controller.fileToken("DEFAULT", "hugegraph", job.getId(), List.of("expired.csv"))
+ .get("expired.csv");
+ String otherToken = controller.fileToken("DEFAULT", "hugegraph", other.getId(), List.of("expired.csv"))
+ .get("expired.csv");
+ controller.upload("DEFAULT", "hugegraph", job.getId(), this.file("expired.csv"),
+ "expired.csv", 20L, token, 2, 0);
+ FileMapping mapping = this.mappings.get("DEFAULT", "hugegraph", job.getId(), "expired.csv");
+ mapping.setPath(this.folder.newFolder("outside-upload-root").getAbsolutePath());
+ mapping.setUpdateTime(new Date(System.currentTimeMillis() - 10000));
+ this.mappings.update(mapping);
+ Mockito.when(this.config.get(HubbleOptions.UPLOAD_FILE_MAX_TIME_CONSUMING)).thenReturn(1L);
+ // Invoke the target synchronously: the production scheduled operation and filesystem checks are unchanged.
+ GateMappingService target = org.springframework.test.util.AopTestUtils.getUltimateTargetObject(this.mappings);
+ target.deleteUnfinishedFile();
+ Assert.assertFalse(this.mappings.getUploadingTokenLocks().containsKey(token));
+ Assert.assertTrue(this.mappings.getUploadingTokenLocks().containsKey(otherToken));
+ Assert.assertTrue(this.mappings.get(mapping.getId()) != null);
+ }
+
+ @Test(timeout = 15000)
+ public void testTokenIssuanceHoldsAdmissionThroughRegistration() throws Exception {
+ JobManager job = this.seedJob();
+ FileUploadController controller = this.uploadController();
+ CountDownLatch parentRead = new CountDownLatch(1);
+ CountDownLatch register = new CountDownLatch(1);
+ JobManagerService observed = Mockito.mock(JobManagerService.class, call -> {
+ Object result = call.getMethod().invoke(this.jobs, call.getArguments());
+ if (call.getMethod().getName().equals("get") && call.getArguments().length == 3) {
+ parentRead.countDown();
+ await(register);
+ }
+ return result;
+ });
+ ReflectionTestUtils.setField(controller, "jobService", observed);
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ Future<?> issuing = requests.submit(() -> controller.fileToken("DEFAULT", "hugegraph", job.getId(),
+ List.of("pending.csv")));
+ Assert.assertTrue(parentRead.await(5, TimeUnit.SECONDS));
+ Future<?> deleting = requests.submit(() -> this.jobs.deleteJob(job.getId()));
+ assertBlocked(deleting);
+ register.countDown();
+ issuing.get(5, TimeUnit.SECONDS);
+ deleting.get(5, TimeUnit.SECONDS);
+ Assert.assertTrue(this.mappings.getUploadingTokenLocks().isEmpty());
+ } finally {
+ register.countDown();
+ stop(requests);
+ }
+ }
+
+ @Test(timeout = 15000)
+ public void testTokenIssuanceWaitsForDeletionThenRejectsMissingParent() throws Exception {
+ JobManager job = this.seedJob();
+ CountDownLatch snapshot = new CountDownLatch(1);
+ CountDownLatch commit = new CountDownLatch(1);
+ this.mappings.afterSnapshot(() -> {
+ snapshot.countDown();
+ await(commit);
+ });
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ Future<?> deleting = requests.submit(() -> this.jobs.deleteJob(job.getId()));
+ Assert.assertTrue(snapshot.await(5, TimeUnit.SECONDS));
+ Future<?> issuing = requests.submit(() -> this.uploadController().fileToken("DEFAULT", "hugegraph",
+ job.getId(), List.of("pending.csv")));
+ assertBlocked(issuing);
+ commit.countDown();
+ deleting.get(5, TimeUnit.SECONDS);
+ Assert.assertThrows(java.util.concurrent.ExecutionException.class, () -> issuing.get(5, TimeUnit.SECONDS));
+ Assert.assertTrue(this.mappings.getUploadingTokenLocks().isEmpty());
+ } finally {
+ commit.countDown();
+ stop(requests);
+ }
+ }
+
+ @Test(timeout = 25000)
+ public void testJobDeleteRefreshesStoppedSnapshotAfterRetryReleasesTaskReservation() throws Exception {
+ // Two Spring workers let the replacement actually start while the old
+ // callback returns; the submit observer keeps the task stripe until then.
+ this.executor.setMaxPoolSize(2);
+ this.executor.setCorePoolSize(2);
+ FileMapping mapping = this.seed();
+ LoadTask oldTask = this.metadata(mapping);
+ oldTask.setStatus(LoadStatus.RUNNING);
+ CountDownLatch oldLoading = new CountDownLatch(1);
+ NativeRun oldRun = this.nativeRun(oldTask, oldLoading, null);
+ this.tasks.save(oldTask);
+ CountDownLatch oldCallbackEntered = new CountDownLatch(1);
+ CountDownLatch releaseOldCallback = new CountDownLatch(1);
+ this.taskExecutor.callbackGate(() -> {
+ oldCallbackEntered.countDown();
+ await(releaseOldCallback);
+ });
+ CountDownLatch retryReserved = new CountDownLatch(1);
+ CountDownLatch stoppedSnapshot = new CountDownLatch(1);
+ CountDownLatch newLoading = new CountDownLatch(1);
+ CountDownLatch cleanupEntered = new CountDownLatch(1);
+ CountDownLatch releaseCleanup = new CountDownLatch(1);
+ CountDownLatch newCallbackEntered = new CountDownLatch(1);
+ CountDownLatch releaseNewCallback = new CountDownLatch(1);
+ AtomicReference<NativeRun> newRun = new AtomicReference<>();
+ AtomicReference<LoadTask> newTask = new AtomicReference<>();
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ this.tasks.executeAfterCommit(oldTask);
+ Assert.assertTrue(oldLoading.await(5, TimeUnit.SECONDS));
+ this.tasks.stop(oldTask.getId());
+ Assert.assertTrue(oldCallbackEntered.await(5, TimeUnit.SECONDS));
+ oldRun.assertClosed();
+ Assert.assertNull(this.tasks.retained(oldTask.getId()));
+ this.tasks.prepareRestart(fresh -> {
+ try {
+ NativeRun owned = this.nativeRun(fresh, newLoading, null);
+ Mockito.doAnswer(call -> {
+ cleanupEntered.countDown();
+ await(releaseCleanup);
+ return call.callRealMethod();
+ }).when(owned.loader).shutdown();
+ newRun.set(owned);
+ newTask.set(fresh);
+ } catch (Exception failure) {
+ throw new AssertionError(failure);
+ }
+ });
+ LoadTaskExecutor observing = new LoadTaskExecutor() {
+ @Override
+ public void execute(LoadTask task, Runnable callback) {
+ JobAdmissionIntegrationTest.this.taskExecutor.execute(task, callback);
+ if (task != oldTask) {
+ await(newLoading);
+ }
+ }
+ };
+ ReflectionTestUtils.setField(this.tasks, "taskExecutor", observing);
+ this.tasks.afterReservation(retryReserved::countDown);
+ Future<LoadTask> retry = requests.submit(() -> this.tasks.retry(oldTask.getId(), "fixture-session"));
+ Assert.assertTrue(retryReserved.await(5, TimeUnit.SECONDS));
+ this.tasks.afterReservation(null);
+ this.tasks.afterSnapshot(() -> {
+ Assert.assertEquals(LoadStatus.STOPPED.getValue().intValue(), this.jdbc.queryForObject(
+ "SELECT load_status FROM load_task WHERE id = ?", Integer.class, oldTask.getId()).intValue());
+ stoppedSnapshot.countDown();
+ });
+ Future<?> deleting = requests.submit(() -> this.jobs.deleteJob(mapping.getJobId()));
+ Assert.assertTrue("Delete must actually snapshot STOPPED before acquiring the held stripe",
+ stoppedSnapshot.await(5, TimeUnit.SECONDS));
+ assertBlocked(deleting);
+ this.taskExecutor.callbackGate(() -> {
+ newCallbackEntered.countDown();
+ await(releaseNewCallback);
+ });
+ releaseOldCallback.countDown();
+ Assert.assertEquals(oldTask.getId(), retry.get(5, TimeUnit.SECONDS).getId());
+ Assert.assertTrue(newLoading.await(5, TimeUnit.SECONDS));
+ Assert.assertTrue("Delete must signal the replacement's actual native cleanup",
+ cleanupEntered.await(5, TimeUnit.SECONDS));
+ Assert.assertFalse(newTask.get().isFinished());
+ Assert.assertEquals(LoadStatus.RUNNING.getValue().intValue(), this.jdbc.queryForObject(
+ "SELECT load_status FROM load_task WHERE id = ?", Integer.class, oldTask.getId()).intValue());
+ assertBlocked(deleting);
+ Assert.assertEquals(1, this.count("load_task"));
+ Assert.assertEquals(1, this.count("job_manager"));
+ releaseCleanup.countDown();
+ Assert.assertTrue(newCallbackEntered.await(5, TimeUnit.SECONDS));
+ newRun.get().assertClosed();
+ Assert.assertTrue(newTask.get().isFinished());
+ assertBlocked(deleting);
+ Assert.assertEquals("Deletion must join the replacement's actual callback before SQL delete",
+ 1, this.count("load_task"));
+ releaseNewCallback.countDown();
+ deleting.get(5, TimeUnit.SECONDS);
+ Assert.assertEquals(0, this.count("load_task"));
+ Assert.assertEquals(0, this.count("job_manager"));
+ Assert.assertEquals(0, this.count("file_mapping"));
+ } finally {
+ this.taskExecutor.callbackGate(null);
+ releaseOldCallback.countDown();
+ releaseCleanup.countDown();
+ releaseNewCallback.countDown();
+ oldRun.stop.countDown();
+ NativeRun replacement = newRun.get();
+ if (replacement != null) {
+ replacement.stop.countDown();
+ }
+ requests.shutdown();
+ Assert.assertTrue(requests.awaitTermination(8, TimeUnit.SECONDS));
+ this.executor.getThreadPoolExecutor().submit(() -> { }).get(5, TimeUnit.SECONDS);
+ oldRun.close();
+ if (replacement != null) {
+ replacement.close();
+ }
+ this.tasks.prepareRestart(null);
+ this.tasks.afterReservation(null);
+ ReflectionTestUtils.setField(this.tasks, "taskExecutor", this.taskExecutor);
+ this.executor.setCorePoolSize(1);
+ this.executor.setMaxPoolSize(1);
+ }
+ }
+
+ @Test(timeout = 20000)
+ public void testRetryWaitsForStoppedRunsActualCallbackBeforeWritingRunning() throws Exception {
+ this.retryAfterCallback(false);
+ }
+
+ @Test(timeout = 20000)
+ public void testRetryRefusesFailedActualCallbackWithoutCreatingAnotherLoader() throws Exception {
+ this.retryAfterCallback(true);
+ }
+
+ private void retryAfterCallback(boolean callbackFails) throws Exception {
+ FileMapping mapping = this.seed();
+ LoadTask task = this.metadata(mapping);
+ task.setStatus(LoadStatus.RUNNING);
+ CountDownLatch oldLoading = new CountDownLatch(1);
+ NativeRun oldRun = this.nativeRun(task, oldLoading, null);
+ this.tasks.save(task);
+ CountDownLatch callbackEntered = new CountDownLatch(1);
+ CountDownLatch releaseCallback = new CountDownLatch(1);
+ CountDownLatch callbackFinished = new CountDownLatch(1);
+ this.taskExecutor.callbackGate(() -> {
+ CallbackWriteFailure.thread = Thread.currentThread();
+ callbackEntered.countDown();
+ await(releaseCallback);
+ });
+ this.taskExecutor.callbackFinished(callbackFinished);
+ ExecutorService requests = Executors.newSingleThreadExecutor();
+ AtomicInteger reconnected = new AtomicInteger();
+ AtomicReference<NativeRun> restarted = new AtomicReference<>();
+ CountDownLatch newLoading = new CountDownLatch(1);
+ CountDownLatch reserved = new CountDownLatch(1);
+ boolean trigger = false;
+ try {
+ this.tasks.executeAfterCommit(task);
+ Assert.assertTrue(oldLoading.await(5, TimeUnit.SECONDS));
+ this.tasks.stop(task.getId());
+ Assert.assertTrue(callbackEntered.await(5, TimeUnit.SECONDS));
+ Assert.assertTrue(task.isFinished());
+ oldRun.assertClosed();
+ Assert.assertNull(this.tasks.retained(task.getId()));
+ Assert.assertEquals(LoadStatus.STOPPED.getValue().intValue(), this.jdbc.queryForObject(
+ "SELECT load_status FROM load_task WHERE id = ?", Integer.class, task.getId()).intValue());
+ if (callbackFails) {
+ this.jdbc.execute("CREATE TRIGGER reject_old_callback BEFORE UPDATE ON load_task FOR EACH ROW CALL '" +
+ CallbackWriteFailure.class.getName() + "'");
+ trigger = true;
+ }
+ this.tasks.prepareRestart(fresh -> {
+ reconnected.incrementAndGet();
+ try {
+ restarted.set(this.nativeRun(fresh, newLoading, null));
+ } catch (Exception failure) {
+ throw new AssertionError(failure);
+ }
+ });
+ this.tasks.afterReservation(reserved::countDown);
+ Future<LoadTask> retry = requests.submit(() -> this.tasks.retry(task.getId(), "fixture-session"));
+ Assert.assertTrue("Retry must actually own the task reservation", reserved.await(5, TimeUnit.SECONDS));
+ assertBlocked(retry);
+ Assert.assertEquals("An unfinished old callback must precede any reconnect", 0, reconnected.get());
+ Assert.assertEquals(LoadStatus.STOPPED.getValue().intValue(), this.jdbc.queryForObject(
+ "SELECT load_status FROM load_task WHERE id = ?", Integer.class, task.getId()).intValue());
+ this.taskExecutor.callbackGate(null);
+ releaseCallback.countDown();
+ Assert.assertTrue(callbackFinished.await(5, TimeUnit.SECONDS));
+ if (callbackFails) {
+ Throwable rejected = Assert.assertThrows(java.util.concurrent.ExecutionException.class,
+ () -> retry.get(5, TimeUnit.SECONDS));
+ Throwable original = this.tasks.callbackFailure(task.getId());
+ Assert.assertNotNull(original);
+ Assert.assertSame("Restart must preserve the actual failed metadata callback cause",
+ original, rejected.getCause());
+ Assert.assertEquals(0, reconnected.get());
+ Assert.assertEquals(LoadStatus.STOPPED.getValue().intValue(), this.jdbc.queryForObject(
+ "SELECT load_status FROM load_task WHERE id = ?", Integer.class, task.getId()).intValue());
+ Assert.assertSame(original, Assert.assertThrows(RuntimeException.class,
+ () -> this.tasks.retry(task.getId(), "fixture-session")));
+ Assert.assertThrows(RuntimeException.class, () -> this.jobs.deleteJob(mapping.getJobId()));
+ } else {
+ LoadTask running = retry.get(5, TimeUnit.SECONDS);
+ Assert.assertTrue(newLoading.await(5, TimeUnit.SECONDS));
+ Assert.assertEquals(1, reconnected.get());
+ Assert.assertEquals(LoadStatus.RUNNING, running.getStatus());
+ Assert.assertEquals("Completed old callbacks cannot overwrite the restarted H2 row",
+ LoadStatus.RUNNING.getValue().intValue(), this.jdbc.queryForObject(
+ "SELECT load_status FROM load_task WHERE id = ?", Integer.class, task.getId()).intValue());
+ this.tasks.stop(task.getId());
+ this.jobs.deleteJob(mapping.getJobId());
+ Assert.assertEquals(0, this.count("load_task"));
+ restarted.get().assertClosed();
+ }
+ } finally {
+ this.taskExecutor.callbackGate(null);
+ releaseCallback.countDown();
+ oldRun.stop.countDown();
+ NativeRun newRun = restarted.get();
+ if (newRun != null) {
+ newRun.stop.countDown();
+ }
+ requests.shutdown();
+ Assert.assertTrue(requests.awaitTermination(8, TimeUnit.SECONDS));
+ this.executor.getThreadPoolExecutor().submit(() -> { }).get(5, TimeUnit.SECONDS);
+ oldRun.close();
+ if (newRun != null) {
+ newRun.close();
+ }
+ if (trigger) {
+ this.jdbc.execute("DROP TRIGGER reject_old_callback");
+ }
+ CallbackWriteFailure.thread = null;
+ this.tasks.prepareRestart(null);
+ this.tasks.afterReservation(null);
+ this.taskExecutor.callbackFinished(null);
+ }
+ }
+
+ public static class CallbackWriteFailure implements org.h2.api.Trigger {
+ private static volatile Thread thread;
+
+ @Override
+ public void fire(java.sql.Connection connection, Object[] oldRow, Object[] newRow)
+ throws java.sql.SQLException {
+ if (Thread.currentThread() == thread) {
+ throw new java.sql.SQLException("controlled actual old callback write failure", "45000");
+ }
+ }
+ }
+
+ @Test
+ public void testDirectStartDispatchErrorRetainsOneRecoveryTicket() throws Exception {
+ FileMapping mapping = this.seed();
+ Error dispatchFailure = new AssertionError("controlled direct dispatch Error");
+ RuntimeException cleanupFailure = new IllegalStateException("controlled unscheduled cleanup failure");
+ AtomicReference<LoadTask> observed = new AtomicReference<>();
+ LoadTaskExecutor rejecting = Mockito.mock(LoadTaskExecutor.class);
+ Mockito.doAnswer(call -> {
+ observed.set(call.getArgument(0));
+ throw dispatchFailure;
+ }).when(rejecting).execute(Mockito.any(LoadTask.class), Mockito.any());
+ ReflectionTestUtils.setField(this.tasks, "taskExecutor", rejecting);
+ try (MockedConstruction<HugeGraphLoader> constructed = Mockito.mockConstruction(HugeGraphLoader.class,
+ (loader, construction) -> Mockito.doThrow(cleanupFailure).when(loader).shutdown())) {
+ Assert.assertSame(dispatchFailure, Assert.assertThrows(Error.class,
+ () -> this.tasks.start(this.connection(), mapping, this.client())));
+ LoadTask task = observed.get();
+ Assert.assertNotNull(task);
+ Assert.assertTrue(task.isFinished());
+ Assert.assertTrue(task.isCleanupFailed());
+ Assert.assertSame(task, this.tasks.retained(task.getId()));
+ Assert.assertEquals(LoadStatus.FAILED, this.tasks.get(task.getId()).getStatus());
+ java.util.Set<?> runs = (java.util.Set<?>) ReflectionTestUtils.getField(this.tasks, "acceptedRuns");
+ long matching = runs.stream().filter(run -> ReflectionTestUtils.getField(run, "task") == task).count();
+ Assert.assertEquals("The dispatch and start catch paths own the same accepted run", 1L, matching);
+ Assert.assertThrows(RuntimeException.class, () -> this.jobs.deleteJob(mapping.getJobId()));
+ Assert.assertEquals(1, this.count("job_manager"));
+ Mockito.verify(constructed.constructed().get(0), Mockito.atLeastOnce()).shutdown();
+ } finally {
+ ReflectionTestUtils.setField(this.tasks, "taskExecutor", this.taskExecutor);
+ }
+ }
+
+ @Test(timeout = 20000)
+ public void testDeleteWaitsForActualCallbackAfterNativeFinished() throws Exception {
+ FileMapping mapping = this.seed();
+ LoadTask task = this.metadata(mapping);
+ task.setStatus(LoadStatus.RUNNING);
+ CountDownLatch callbackEntered = new CountDownLatch(1);
+ CountDownLatch releaseCallback = new CountDownLatch(1);
+ CountDownLatch callbackFinished = new CountDownLatch(1);
+ NativeRun nativeRun = this.nativeRun(task, null, null);
+ this.tasks.save(task);
+ this.taskExecutor.callbackGate(() -> {
+ callbackEntered.countDown();
+ await(releaseCallback);
+ });
+ this.taskExecutor.callbackFinished(callbackFinished);
+ ExecutorService requests = Executors.newSingleThreadExecutor();
+ CountDownLatch deleteSnapshot = new CountDownLatch(1);
+ this.tasks.afterSnapshot(deleteSnapshot::countDown);
+ try {
+ this.tasks.executeAfterCommit(task);
+ Assert.assertTrue(callbackEntered.await(5, TimeUnit.SECONDS));
+ Assert.assertTrue(task.isFinished());
+ nativeRun.assertClosed();
+ Future<?> deletion = requests.submit(() -> this.jobs.deleteJob(mapping.getJobId()));
+ Assert.assertTrue(deleteSnapshot.await(5, TimeUnit.SECONDS));
+ assertBlocked(deletion);
+ com.zaxxer.hikari.HikariDataSource dataSource =
+ (com.zaxxer.hikari.HikariDataSource) this.jdbc.getDataSource();
+ Assert.assertEquals("Callback waiting must leave metadata connections available", 0,
+ dataSource.getHikariPoolMXBean().getActiveConnections());
+ Assert.assertEquals("Native finished is not a callback certificate", 1, this.count("load_task"));
+ Assert.assertEquals(1, this.count("job_manager"));
+ releaseCallback.countDown();
+ deletion.get(5, TimeUnit.SECONDS);
+ Assert.assertTrue(callbackFinished.await(5, TimeUnit.SECONDS));
+ Assert.assertEquals(0, this.count("load_task"));
+ Assert.assertEquals(0, this.count("job_manager"));
+ Assert.assertEquals(0, this.count("file_mapping"));
+ } finally {
+ releaseCallback.countDown();
+ this.executor.getThreadPoolExecutor().submit(() -> { }).get(5, TimeUnit.SECONDS);
+ nativeRun.close();
+ stop(requests);
+ this.taskExecutor.callbackGate(null);
+ this.taskExecutor.callbackFinished(null);
+ }
+ }
+
+ @Test(timeout = 20000)
+ public void testReplacementUploadTokenWaitsForFailedUploadCleanup() throws Exception {
+ FileMapping parent = this.seed();
+ FileUploadController controller = this.uploadController();
+ int jobId = parent.getJobId();
+ String name = "retry.csv";
+ String token = controller.fileToken("DEFAULT", "hugegraph", jobId, List.of(name)).get(name);
+ CountDownLatch cleaning = new CountDownLatch(1);
+ CountDownLatch releaseCleanup = new CountDownLatch(1);
+ this.mappings.beforeCleanup(() -> {
+ cleaning.countDown();
+ await(releaseCleanup);
+ });
+ MockMultipartFile upload = new MockMultipartFile("file", name, "text/csv", "name\nnew\n".getBytes()) {
+ @Override
+ public void transferTo(File destination) throws IOException {
+ super.transferTo(destination);
+ // The merged result fails the final quota recheck.
+ Mockito.when(config.get(HubbleOptions.UPLOAD_TOTAL_FILE_SIZE_LIMIT)).thenReturn(1L);
+ }
+ };
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ Future<?> failing = requests.submit(() -> Assert.assertThrows(RuntimeException.class,
+ () -> controller.upload("DEFAULT", "hugegraph", jobId, upload, name, null, token, 1, 0)));
+ Assert.assertTrue(cleaning.await(5, TimeUnit.SECONDS));
+ Future<String> replacement = requests.submit(() ->
+ controller.fileToken("DEFAULT", "hugegraph", jobId, List.of(name)).get(name));
+ assertBlocked(replacement);
+ Mockito.when(this.config.get(HubbleOptions.UPLOAD_TOTAL_FILE_SIZE_LIMIT)).thenReturn(4096L);
+ releaseCleanup.countDown();
+ failing.get(5, TimeUnit.SECONDS);
+ String nextToken = replacement.get(5, TimeUnit.SECONDS);
+ MockMultipartFile next = new MockMultipartFile("file", name, "text/csv", "name\nnew\n".getBytes());
+ controller.upload("DEFAULT", "hugegraph", jobId, next, name, null, nextToken, 1, 0);
+ FileMapping completed = this.mappings.get("DEFAULT", "hugegraph", jobId, name);
+ Assert.assertNotNull(completed);
+ Assert.assertEquals(FileMappingStatus.COMPLETED, completed.getFileStatus());
+ Assert.assertTrue(Files.isRegularFile(Path.of(completed.getPath())));
+ } finally {
+ releaseCleanup.countDown();
+ this.mappings.beforeCleanup(null);
+ stop(requests);
+ }
+ }
+
+ @Test(timeout = 20000)
+ public void testDeleteRejectsPendingCallbackInsideExistingTransaction() throws Exception {
+ FileMapping mapping = this.seed();
+ LoadTask task = this.metadata(mapping);
+ task.setStatus(LoadStatus.RUNNING);
+ CountDownLatch callbackEntered = new CountDownLatch(1);
+ CountDownLatch releaseCallback = new CountDownLatch(1);
+ NativeRun nativeRun = this.nativeRun(task, null, null);
+ this.tasks.save(task);
+ this.taskExecutor.callbackGate(() -> {
+ callbackEntered.countDown();
+ await(releaseCallback);
+ });
+ try {
+ this.tasks.executeAfterCommit(task);
+ Assert.assertTrue(callbackEntered.await(5, TimeUnit.SECONDS));
+ new TransactionTemplate(this.transactionManager).executeWithoutResult(status -> {
+ Throwable failure = Assert.assertThrows(RuntimeException.class,
+ () -> this.jobs.deleteJob(mapping.getJobId()));
+ Assert.assertTrue(failure.getMessage().contains("retry after the current transaction"));
+ Assert.assertEquals(1, this.count("load_task"));
+ Assert.assertEquals(1, this.count("job_manager"));
+ Assert.assertEquals(1, this.count("file_mapping"));
+ Assert.assertEquals(LoadStatus.RUNNING, this.tasks.get(task.getId()).getStatus());
+ });
+ releaseCallback.countDown();
+ this.executor.getThreadPoolExecutor().submit(() -> { }).get(5, TimeUnit.SECONDS);
+ new TransactionTemplate(this.transactionManager).executeWithoutResult(status ->
+ this.jobs.deleteJob(mapping.getJobId()));
+ Assert.assertEquals(0, this.count("job_manager"));
+ nativeRun.assertClosed();
+ } finally {
+ releaseCallback.countDown();
+ this.executor.getThreadPoolExecutor().submit(() -> { }).get(5, TimeUnit.SECONDS);
+ nativeRun.close();
+ this.taskExecutor.callbackGate(null);
+ }
+ }
+
+ @Test(timeout = 25000)
+ public void testQueuedFirstJobDeleteSignalsEveryNativeOwnerBeforeWaiting() throws Exception {
+ FileMapping mapping = this.seed();
+ this.executor.setMaxPoolSize(8);
+ this.executor.setCorePoolSize(8);
+ CountDownLatch active = new CountDownLatch(8);
+ CountDownLatch stopRequests = new CountDownLatch(9);
+ List<NativeRun> nativeRuns = new ArrayList<>();
+ List<LoadTask> tasks = new ArrayList<>();
+ ExecutorService requests = Executors.newSingleThreadExecutor();
+ CountDownLatch deleteSnapshot = new CountDownLatch(1);
+ try {
+ for (int i = 0; i < 9; i++) {
+ LoadTask task = this.metadata(mapping);
+ task.setStatus(LoadStatus.RUNNING);
+ NativeRun nativeRun = this.nativeRun(task, active, stopRequests);
+ nativeRuns.add(nativeRun);
+ tasks.add(task);
+ this.tasks.save(task);
+ this.tasks.executeAfterCommit(task);
+ if (i == 7) {
+ Assert.assertTrue(active.await(5, TimeUnit.SECONDS));
+ }
+ }
+ LoadTask queued = tasks.get(8);
+ Assert.assertFalse(queued.isFinished());
+ Assert.assertEquals(1, this.executor.getThreadPoolExecutor().getQueue().size());
+ this.tasks.queuedFirst(queued.getId());
+ this.tasks.afterSnapshot(deleteSnapshot::countDown);
+ Future<?> deletion = requests.submit(() -> this.jobs.deleteJob(mapping.getJobId()));
+ Assert.assertTrue(deleteSnapshot.await(5, TimeUnit.SECONDS));
+ Assert.assertTrue("Every actual native stop intent must precede queued-first wait",
+ stopRequests.await(5, TimeUnit.SECONDS));
+ deletion.get(5, TimeUnit.SECONDS);
+ Assert.assertEquals(0, this.count("load_task"));
+ Assert.assertEquals(0, this.count("job_manager"));
+ for (int i = 0; i < tasks.size(); i++) {
+ Assert.assertTrue(tasks.get(i).isFinished());
+ nativeRuns.get(i).assertClosed();
+ }
+ Mockito.verify(nativeRuns.get(8).loader, Mockito.never()).load();
+ } finally {
+ for (NativeRun run : nativeRuns) {
+ run.stop.countDown();
+ }
+ requests.shutdown();
+ Assert.assertTrue(requests.awaitTermination(8, TimeUnit.SECONDS));
+ for (NativeRun run : nativeRuns) {
+ run.close();
+ }
+ this.executor.getThreadPoolExecutor().submit(() -> { }).get(5, TimeUnit.SECONDS);
+ this.executor.setCorePoolSize(1);
+ this.executor.setMaxPoolSize(1);
+ this.tasks.queuedFirst(null);
+ }
+ }
+
+ private NativeRun nativeRun(LoadTask task, CountDownLatch active, CountDownLatch stopRequests) throws Exception {
+ LoadOptions options = task.getOptions();
+ LoadContext context = Mockito.spy(LoadContext.forOffline(options, new SchemaCache(null)));
+ TaskManager manager = new TaskManager(context);
+ Class<?> unsafeType = Class.forName("sun.misc.Unsafe");
+ java.lang.reflect.Field field = unsafeType.getDeclaredField("theUnsafe");
+ field.setAccessible(true);
+ HugeGraphLoader raw = (HugeGraphLoader) unsafeType.getMethod("allocateInstance", Class.class)
+ .invoke(field.get(null), HugeGraphLoader.class);
+ ReflectionTestUtils.setField(raw, "context", context);
+ ReflectionTestUtils.setField(raw, "manager", manager);
+ ReflectionTestUtils.setField(raw, "options", options);
+ HugeGraphLoader loader = Mockito.spy(raw);
+ CountDownLatch stop = new CountDownLatch(1);
+ java.util.concurrent.atomic.AtomicBoolean requested = new java.util.concurrent.atomic.AtomicBoolean();
+ Mockito.doAnswer(call -> {
+ Object value = call.callRealMethod();
+ stop.countDown();
+ if (stopRequests != null && requested.compareAndSet(false, true)) {
+ stopRequests.countDown();
+ }
+ return value;
+ }).when(context).stopLoading();
+ Mockito.doAnswer(call -> {
+ if (active != null) {
+ active.countDown();
+ await(stop);
+ if (stopRequests != null) {
+ await(stopRequests);
+ }
+ }
+ return true;
+ }).when(loader).load();
+ task.setLoader(loader);
+ return new NativeRun(loader, context, manager, stop);
+ }
+
+ private static final class NativeRun implements AutoCloseable {
+ private final HugeGraphLoader loader;
+ private final LoadContext context;
+ private final TaskManager manager;
+ private final CountDownLatch stop;
+ private NativeRun(HugeGraphLoader loader, LoadContext context, TaskManager manager, CountDownLatch stop) {
+ this.loader = loader;
+ this.context = context;
+ this.manager = manager;
+ this.stop = stop;
+ }
+
+ private void assertClosed() {
+ Assert.assertTrue(this.context.closed());
+ Assert.assertTrue(((ExecutorService) ReflectionTestUtils.getField(this.manager, "batchService"))
+ .isTerminated());
+ Assert.assertTrue(((ExecutorService) ReflectionTestUtils.getField(this.manager, "singleService"))
+ .isTerminated());
+ Mockito.verify(this.context).close(Mockito.anyBoolean());
+ }
+
+ @Override
+ public void close() {
+ this.stop.countDown();
+ this.loader.shutdown();
+ }
+ }
+
+ private MockMultipartFile file(String name) {
+ return new MockMultipartFile("file", name, "text/csv", "name\nperson\n".getBytes());
+ }
+
+ private MockMultipartFile blockedFile(String name, CountDownLatch writing, CountDownLatch finish) {
+ return new MockMultipartFile("file", name, "text/csv", "name\nperson\n".getBytes()) {
+ @Override
+ public void transferTo(File destination) throws IOException {
+ writing.countDown();
+ await(finish);
+ super.transferTo(destination);
+ }
+ };
+ }
+
+ private JobManager seedJob() {
+ JobManager job = JobManager.builder().graphSpace("DEFAULT").graph("hugegraph")
+ .jobName(UUID.randomUUID().toString()).jobStatus(JobStatus.SETTING)
+ .createTime(new Date()).updateTime(new Date()).build();
+ this.jobs.save(job);
+ return job;
+ }
+
+ private FileMappingController mappingController() {
+ FileMappingController controller = new FileMappingController() {
+ @Override
+ protected HugeClient requireGraphSpaceWrite(String graphSpace) {
+ return client();
+ }
+ };
+ ReflectionTestUtils.setField(controller, "service", this.mappings);
+ ReflectionTestUtils.setField(controller, "jobService", this.jobs);
+ ReflectionTestUtils.setField(controller, "taskService", this.tasks);
+ return controller;
+ }
+
+ @Test(timeout = 15000)
+ public void testUploadCannotCreateMappingAfterDeleteSnapshot() throws Exception {
+ FileMapping mapping = this.seed();
+ FileUploadController controller = this.uploadController();
+ String token = controller.fileToken("DEFAULT", "hugegraph", mapping.getJobId(),
+ List.of("later.csv")).get("later.csv");
+ CountDownLatch snapshot = new CountDownLatch(1);
+ CountDownLatch complete = new CountDownLatch(1);
+ this.mappings.afterSnapshot(() -> {
+ snapshot.countDown();
+ await(complete);
+ });
+ AtomicInteger writes = new AtomicInteger();
+ MockMultipartFile file = new MockMultipartFile("file", "later.csv", "text/csv", "name\nnew\n".getBytes()) {
+ @Override
+ public void transferTo(File destination) throws IOException {
+ writes.incrementAndGet();
+ super.transferTo(destination);
+ }
+ };
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ Future<?> deleting = requests.submit(() -> this.jobs.deleteJob(mapping.getJobId()));
+ Assert.assertTrue(snapshot.await(5, TimeUnit.SECONDS));
+ CountDownLatch entered = new CountDownLatch(1);
+ Future<?> uploading = requests.submit(() -> {
+ entered.countDown();
+ return controller.upload("DEFAULT", "hugegraph", mapping.getJobId(), file,
+ "later.csv", null, token, 1, 0);
+ });
+ Assert.assertTrue(entered.await(5, TimeUnit.SECONDS));
+ assertBlocked(uploading);
+ complete.countDown();
+ deleting.get(5, TimeUnit.SECONDS);
+ Assert.assertThrows(java.util.concurrent.ExecutionException.class,
+ () -> uploading.get(5, TimeUnit.SECONDS));
+ Assert.assertEquals(0, writes.get());
+ Assert.assertEquals(0, this.count("job_manager"));
+ Assert.assertEquals(0, this.count("file_mapping"));
+ try (var files = Files.walk(this.uploads)) {
+ Assert.assertFalse(files.anyMatch(Files::isRegularFile));
+ }
+ } finally {
+ complete.countDown();
+ stop(requests);
+ }
+ }
+
+ @Test(timeout = 15000)
+ public void testDeleteWaitsForUploadingFileAndCleansCommittedMapping() throws Exception {
+ FileMapping mapping = this.seed();
+ FileUploadController controller = this.uploadController();
+ String token = controller.fileToken("DEFAULT", "hugegraph", mapping.getJobId(),
+ List.of("later.csv")).get("later.csv");
+ CountDownLatch writing = new CountDownLatch(1);
+ CountDownLatch complete = new CountDownLatch(1);
+ MockMultipartFile file = new MockMultipartFile("file", "later.csv", "text/csv", "name\nnew\n".getBytes()) {
+ @Override
+ public void transferTo(File destination) throws IOException {
+ writing.countDown();
+ await(complete);
+ super.transferTo(destination);
+ }
+ };
+ ExecutorService requests = Executors.newFixedThreadPool(2);
+ try {
+ Future<?> uploading = requests.submit(() -> controller.upload("DEFAULT", "hugegraph", mapping.getJobId(),
+ file, "later.csv", null, token, 1, 0));
+ Assert.assertTrue(writing.await(5, TimeUnit.SECONDS));
+ CountDownLatch entered = new CountDownLatch(1);
+ Future<?> deleting = requests.submit(() -> {
+ entered.countDown();
+ this.jobs.deleteJob(mapping.getJobId());
+ });
+ Assert.assertTrue(entered.await(5, TimeUnit.SECONDS));
+ assertBlocked(deleting);
+ Assert.assertEquals(1, this.count("job_manager"));
+ complete.countDown();
+ uploading.get(5, TimeUnit.SECONDS);
+ deleting.get(5, TimeUnit.SECONDS);
+ Assert.assertEquals(0, this.count("job_manager"));
+ Assert.assertEquals(0, this.count("file_mapping"));
+ try (var files = Files.walk(this.uploads)) {
+ Assert.assertFalse(files.anyMatch(Files::isRegularFile));
+ }
+ } finally {
+ complete.countDown();
+ stop(requests);
+ }
+ }
+
+ @Test
+ public void testDetachedMappingCannotStartFromStaleControllerSnapshot() throws Exception {
+ FileMapping mapping = this.seed();
+ this.jobs.deleteMappings(mapping.getJobId(), List.of(mapping));
+ Assert.assertTrue(this.mappings.get(mapping.getId()) == null ||
+ this.mappings.get(mapping.getId()).getJobId() == null);
+ try (MockedConstruction<HugeGraphLoader> constructed = Mockito.mockConstruction(HugeGraphLoader.class)) {
+ Assert.assertThrows(RuntimeException.class,
+ () -> this.tasks.start(this.connection(), mapping, this.client()));
+ Assert.assertTrue(constructed.constructed().isEmpty());
+ Assert.assertEquals(0, this.count("load_task"));
+ }
+ }
+
+ @Test
+ public void testRunningTaskPreventsMappingDetach() throws Exception {
+ FileMapping mapping = this.seed();
+ LoadTask task = this.metadata(mapping);
+ task.setStatus(LoadStatus.RUNNING);
+ this.tasks.save(task);
+ Assert.assertThrows(RuntimeException.class,
+ () -> this.jobs.deleteMappings(mapping.getJobId(), List.of(mapping)));
+ Assert.assertEquals(mapping.getJobId(), this.mappings.get(mapping.getId()).getJobId());
+ Assert.assertTrue(Files.isRegularFile(Path.of(mapping.getPath())));
+ }
+
+ @Test
+ public void testCleanupFailureBeforeGeneratedTaskIdProtectsParent() throws Exception {
+ FileMapping mapping = this.seed();
+ RuntimeException cleanup = new IllegalStateException("fixture native cleanup failed");
+ String constraint = "reject_fixture_task";
+ this.jdbc.execute("ALTER TABLE load_task ADD CONSTRAINT " + constraint +
+ " CHECK (job_id <> " + mapping.getJobId() + ")");
+ try (MockedConstruction<HugeGraphLoader> constructed = Mockito.mockConstruction(HugeGraphLoader.class,
+ (loader, context) -> Mockito.doThrow(cleanup).when(loader).shutdown())) {
+ Throwable insert = Assert.assertThrows(RuntimeException.class,
+ () -> this.tasks.start(this.connection(), mapping, this.client()));
+ Assert.assertTrue(insert instanceof org.springframework.dao.DataIntegrityViolationException);
+ Assert.assertEquals(1, constructed.constructed().size());
+ Mockito.verify(constructed.constructed().get(0)).shutdown();
+ Assert.assertEquals(0, this.count("load_task"));
+ Assert.assertThrows(RuntimeException.class, () -> this.jobs.deleteJob(mapping.getJobId()));
+ Assert.assertEquals(1, this.count("job_manager"));
+ Assert.assertEquals(1, this.count("file_mapping"));
+ Assert.assertTrue(Files.isRegularFile(Path.of(mapping.getPath())));
+ Assert.assertTrue(java.util.Arrays.asList(insert.getSuppressed()).contains(cleanup));
+ } finally {
+ this.jdbc.execute("ALTER TABLE load_task DROP CONSTRAINT " + constraint);
+ }
+ }
+
+ @Test
+ public void testRollbackCleanupFailureProtectsParentWithoutTaskRow() throws Exception {
+ FileMapping mapping = this.seed();
+ RuntimeException cleanup = new IllegalStateException("fixture rollback cleanup failed");
+ AtomicReference<LoadTask> created = new AtomicReference<>();
+ try (MockedConstruction<HugeGraphLoader> constructed = Mockito.mockConstruction(HugeGraphLoader.class,
+ (loader, context) -> Mockito.doThrow(cleanup).when(loader).shutdown())) {
+ new TransactionTemplate(this.transactionManager).execute(status -> {
+ created.set(this.tasks.start(this.connection(), mapping, this.client()));
+ status.setRollbackOnly();
+ return null;
+ });
+ Assert.assertNotNull(created.get().getId());
+ Assert.assertTrue(created.get().isCleanupFailed());
+ Assert.assertTrue(created.get().isFinished());
+ Mockito.verify(constructed.constructed().get(0)).shutdown();
+ Assert.assertEquals(0, this.count("load_task"));
+ Assert.assertThrows(RuntimeException.class, () -> this.jobs.deleteJob(mapping.getJobId()));
+ Assert.assertEquals(1, this.count("job_manager"));
+ Assert.assertEquals(1, this.count("file_mapping"));
+ }
+ }
+
+ private FileMapping seed() throws Exception {
+ JobManager job = JobManager.builder().graphSpace("DEFAULT").graph("hugegraph")
+ .jobName(UUID.randomUUID().toString()).jobStatus(JobStatus.SETTING)
+ .createTime(new Date()).updateTime(new Date()).build();
+ this.jobs.save(job);
+ Path dir = Files.createDirectories(this.uploads.resolve(UUID.randomUUID().toString()));
+ Path input = dir.resolve("people.csv");
+ Files.writeString(input, "name\nperson\n");
+ FileMapping mapping = new FileMapping();
+ mapping.setGraphSpace("DEFAULT");
+ mapping.setGraph("hugegraph");
+ mapping.setJobId(job.getId());
+ mapping.setName("people.csv");
+ mapping.setPath(input.toString());
+ mapping.setTotalLines(2L);
+ mapping.setTotalSize(12L);
+ mapping.setFileStatus(FileMappingStatus.COMPLETED);
+ mapping.setFileSetting(FileSetting.builder().columnNames(List.of("name")).format("CSV").delimiter(",")
+ .charset("UTF-8").dateFormat("yyyy-MM-dd HH:mm:ss").timeZone("GMT+8")
+ .skippedLine("(^#|^//).*").listFormat(new ListFormat()).build());
+ mapping.setVertexMappings(Collections.emptySet());
+ mapping.setEdgeMappings(Collections.emptySet());
+ mapping.setLoadParameter(new LoadParameter());
+ mapping.setCreateTime(new Date());
+ mapping.setUpdateTime(new Date());
+ this.mappings.save(mapping);
+ return mapping;
+ }
+
+ private LoadTask metadata(FileMapping mapping) {
+ return LoadTask.builder().graphSpace("DEFAULT").graph("hugegraph").jobId(mapping.getJobId())
+ .fileId(mapping.getId()).fileName(mapping.getName()).options(new LoadOptions())
+ .vertices(Collections.emptySet()).edges(Collections.emptySet()).fileTotalLines(2L)
+ .fileReadLines(0L).lastDuration(0L).currDuration(0L).createTime(new Date())
+ .status(LoadStatus.STOPPED).build();
+ }
+
+ private LoadTaskController controller() {
+ LoadTaskController controller = new LoadTaskController(this.tasks) {
+ @Override
+ protected HugeClient requireGraphSpaceWrite(String graphSpace) {
+ return client();
+ }
+
+ @Override
+ protected HugeClient authClient(String space, String graph) {
+ return client();
+ }
+
+ @Override
+ protected String getToken() {
+ return "fixture-current-token";
+ }
+
+ @Override
+ protected String getUser() {
+ return "fixture-user";
+ }
+ };
+ ReflectionTestUtils.setField(controller, "jobService", this.jobs);
+ ReflectionTestUtils.setField(controller, "fmService", this.mappings);
+ ReflectionTestUtils.setField(controller, "config", this.config);
+ return controller;
+ }
+
+ private FileUploadController uploadController() {
+ FileUploadController controller = new FileUploadController() {
+ @Override
+ protected HugeClient requireGraphSpaceWrite(String graphSpace) {
+ return client();
+ }
+ };
+ ReflectionTestUtils.setField(controller, "service", this.mappings);
+ ReflectionTestUtils.setField(controller, "jobService", this.jobs);
+ ReflectionTestUtils.setField(controller, "config", this.config);
+ return controller;
+ }
+
+ private GraphConnection connection() {
+ GraphConnection connection = new GraphConnection();
+ connection.setGraphSpace("DEFAULT");
+ connection.setGraph("hugegraph");
+ return connection;
+ }
+
+ private HugeClient client() {
+ return Mockito.mock(HugeClient.class);
+ }
+
+ private int count(String table) {
+ return this.jdbc.queryForObject("SELECT COUNT(*) FROM " + table, Integer.class);
+ }
+
+ private void clearRows() {
+ this.jdbc.update("DELETE FROM load_task");
+ this.jdbc.update("DELETE FROM file_mapping");
+ this.jdbc.update("DELETE FROM job_manager");
+ }
+
+ private static void await(CountDownLatch latch) {
+ try {
+ Assert.assertTrue(latch.await(8, TimeUnit.SECONDS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError(e);
+ }
+ }
+
+ private static void assertBlocked(Future<?> request) throws Exception {
+ try {
+ request.get(300, TimeUnit.MILLISECONDS);
+ Assert.fail("request bypassed lifecycle reservation");
+ } catch (TimeoutException expected) {
+ // The owning transaction still holds the reservation.
+ }
+ }
+
+ private static void stop(ExecutorService requests) throws Exception {
+ requests.shutdownNow();
+ Assert.assertTrue(requests.awaitTermination(5, TimeUnit.SECONDS));
+ }
+
+ public static class GateTaskService extends LoadTaskService {
+ private Runnable afterSnapshot;
+ private Integer queuedFirst;
+ private java.util.function.Consumer<LoadTask> prepareRestart;
+ private Runnable afterReservation;
+
+ public void prepareRestart(java.util.function.Consumer<LoadTask> prepare) {
+ this.prepareRestart = prepare;
+ }
+
+ public void afterReservation(Runnable action) {
+ this.afterReservation = action;
+ }
+
+ @Override
+ TaskReservation reserveTasks(List<Integer> ids) {
+ TaskReservation reservation = super.reserveTasks(ids);
+ Runnable action = this.afterReservation;
+ if (action != null) {
+ action.run();
+ }
+ return reservation;
+ }
+
+ @Override
+ public LoadTask get(int id) {
+ LoadTask task = super.get(id);
+ java.util.function.Consumer<LoadTask> prepare = this.prepareRestart;
+ if (task != null && prepare != null) {
+ LoadTask controlled = Mockito.spy(task);
+ Mockito.doAnswer(call -> {
+ prepare.accept(controlled);
+ return null;
+ }).when(controlled).reconnect("fixture-session");
+ return controlled;
+ }
+ return task;
+ }
+
+ public Throwable callbackFailure(int taskId) {
+ java.util.Set<?> accepted = (java.util.Set<?>) ReflectionTestUtils.getField(this, "acceptedRuns");
+ for (Object run : accepted) {
+ LoadTask task = (LoadTask) ReflectionTestUtils.getField(run, "task");
+ if (task.getId() != null && task.getId() == taskId) {
+ java.util.concurrent.CompletableFuture<?> completed = (java.util.concurrent.CompletableFuture<?>)
+ ReflectionTestUtils.getField(run, "completed");
+ try {
+ completed.join();
+ } catch (java.util.concurrent.CompletionException failure) {
+ return failure.getCause();
+ }
+ }
+ }
+ return null;
+ }
+
+ public void afterSnapshot(Runnable action) {
+ this.afterSnapshot = action;
+ }
+
+ public void queuedFirst(Integer id) {
+ this.queuedFirst = id;
+ }
+
+ @Override
+ public List<LoadTask> taskListByJob(int id) {
+ List<LoadTask> snapshot = super.taskListByJob(id);
+ if (this.queuedFirst != null) {
+ snapshot.sort(Comparator.comparing(task -> !task.getId().equals(this.queuedFirst)));
+ }
+ Runnable action = this.afterSnapshot;
+ this.afterSnapshot = null;
+ if (action != null) {
+ action.run();
+ }
+ return snapshot;
+ }
+
+ public LoadTask retained(int id) {
+ @SuppressWarnings("unchecked")
+ java.util.Map<Integer, LoadTask> running = (java.util.Map<Integer, LoadTask>)
+ ReflectionTestUtils.getField(this, "runningTaskContainer");
+ return running.get(id);
+ }
+ }
+
+ public static class GateMappingService extends FileMappingService {
+ private Runnable beforeCleanup;
+
+ public void beforeCleanup(Runnable action) {
+ this.beforeCleanup = action;
+ }
+
+ @Override
+ public void cleanupMappings(List<FileMapping> mappings) {
+ Runnable action = this.beforeCleanup;
+ this.beforeCleanup = null;
+ if (action != null) {
+ action.run();
+ }
+ super.cleanupMappings(mappings);
+ }
+
+ private Runnable afterSnapshot;
+ public void afterSnapshot(Runnable action) {
+ this.afterSnapshot = action;
+ }
+
+ @Override
+ public List<FileMapping> listByJob(int id) {
+ List<FileMapping> snapshot = super.listByJob(id);
+ Runnable action = this.afterSnapshot;
+ this.afterSnapshot = null;
+ if (action != null) {
+ action.run();
+ }
+ return snapshot;
+ }
+ }
+
+ public static class GateLoadTaskExecutor extends LoadTaskExecutor {
+ private volatile Runnable callbackGate;
+ private volatile CountDownLatch callbackFinished;
+
+ public void callbackGate(Runnable action) {
+ this.callbackGate = action;
+ }
+
+ public void callbackFinished(CountDownLatch signal) {
+ this.callbackFinished = signal;
+ }
+
+ @Async
+ @Override
+ public void execute(LoadTask task, Runnable callback) {
+ super.execute(task, () -> {
+ Runnable action = this.callbackGate;
+ if (action != null) {
+ action.run();
+ }
+ try {
+ callback.run();
+ } finally {
+ CountDownLatch signal = this.callbackFinished;
+ if (signal != null) {
+ signal.countDown();
+ }
+ }
+ });
+ }
+ }
+
+ @SpringBootConfiguration
+ @EnableAutoConfiguration
+ @EnableTransactionManagement
+ @EnableAsync
+ @MapperScan("org.apache.hugegraph.mapper.load")
+ public static class Configuration {
+ @Bean
+ public GateTaskService loadTaskService() {
+ return new GateTaskService();
+ }
+
+ @Bean
+ public JobManagerService jobManagerService() {
+ return new JobManagerService();
+ }
+
+ @Bean
+ public GateMappingService fileMappingService() {
+ return new GateMappingService();
+ }
+
+ @Bean
+ public GateLoadTaskExecutor loadTaskExecutor() {
+ return new GateLoadTaskExecutor();
+ }
+
+ @Bean
+ public HugeConfig hugeConfig() {
+ return Mockito.mock(HugeConfig.class);
+ }
+
+ @Bean
+ public ThreadPoolTaskExecutor taskExecutor() {
+ ThreadPoolTaskExecutor pool = new ThreadPoolTaskExecutor();
+ pool.setCorePoolSize(1);
+ pool.setMaxPoolSize(1);
+ pool.setQueueCapacity(4);
+ pool.setThreadNamePrefix("job-admission-fixture-");
+ pool.setWaitForTasksToCompleteOnShutdown(true);
+ pool.setAwaitTerminationSeconds(5);
+ return pool;
+ }
+ }
+}
diff --git a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/testutil/NativeCleanupTestSupport.java b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/testutil/NativeCleanupTestSupport.java
new file mode 100644
index 000000000..97b48efaa
--- /dev/null
+++ b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/testutil/NativeCleanupTestSupport.java
@@ -0,0 +1,171 @@
+/*
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with this
+ * work for additional information regarding copyright ownership. The ASF
+ * licenses this file to You under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance with the
+ * License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hugegraph.testutil;
+
+import java.io.OutputStream;
+import java.lang.reflect.Field;
+import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.apache.hugegraph.entity.enums.LoadStatus;
+import org.apache.hugegraph.entity.load.LoadTask;
+import org.apache.hugegraph.loader.HugeGraphLoader;
+import org.apache.hugegraph.loader.executor.LoadOptions;
+import org.apache.hugegraph.loader.mapping.LoadMapping;
+import org.apache.hugegraph.loader.task.TaskManager;
+import org.apache.hugegraph.util.JsonUtil;
+import org.mockito.Mockito;
+
+import com.sun.net.httpserver.HttpServer;
+import okhttp3.OkHttpClient;
+
+/** Real native shutdown and HTTP reset boundary; only load failure entry is controlled. */
+public final class NativeCleanupTestSupport implements AutoCloseable {
+
+ public final HugeGraphLoader loader;
+ public final RuntimeException original = new IllegalStateException("recorded load failure");
+ public final AtomicReference<Throwable> shutdownFailure = new AtomicReference<>();
+ private final HttpServer server;
+ private volatile String mode = "NONE";
+
+ public NativeCleanupTestSupport(LoadTask task, Path input, boolean rejectReset,
+ LoadStatus requestedStatus) throws Exception {
+ this.server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+ this.server.createContext("/", exchange -> {
+ Object response;
+ int status = 200;
+ if (exchange.getRequestURI().getPath().equals("/versions")) {
+ response = Map.of("versions", Map.of("core", "1.8.0", "api", "0.74"));
+ } else if (exchange.getRequestURI().getPath().endsWith("/mode")) {
+ if (exchange.getRequestMethod().equals("GET")) {
+ response = Map.of("mode", this.mode);
+ } else {
+ String value = JsonUtil.fromJson(new String(exchange.getRequestBody().readAllBytes(),
+ StandardCharsets.UTF_8), String.class);
+ if (rejectReset && value.equals("NONE")) {
+ status = 500;
+ response = Map.of("exception", "class java.lang.IllegalStateException",
+ "message", "independent mode reset failed");
+ } else {
+ this.mode = value;
+ response = Map.of();
+ }
+ }
+ } else {
+ status = 404;
+ response = Map.of("message", "unsupported fixture path");
+ }
+ byte[] body = JsonUtil.toJson(response).getBytes(StandardCharsets.UTF_8);
+ exchange.getResponseHeaders().add("Content-Type", "application/json");
+ exchange.sendResponseHeaders(status, body.length);
+ try (OutputStream output = exchange.getResponseBody()) {
+ output.write(body);
+ }
+ });
+ this.server.start();
+ LoadOptions options = new LoadOptions();
+ options.host = "127.0.0.1";
+ options.port = this.server.getAddress().getPort();
+ options.file = input.toString();
+ this.loader = Mockito.spy(new HugeGraphLoader(options, new LoadMapping(Collections.emptyList())));
+ task.setLoader(this.loader);
+ task.setOptions(options);
+ this.loader.context().setLoadingMode();
+ Mockito.doAnswer(invocation -> {
+ task.setStatus(requestedStatus);
+ this.loader.context().failLoading(this.original);
+ throw this.original;
+ }).when(this.loader).load();
+ Mockito.doAnswer(invocation -> {
+ try {
+ return invocation.callRealMethod();
+ } catch (Throwable failure) {
+ this.shutdownFailure.set(failure);
+ throw failure;
+ }
+ }).when(this.loader).shutdown();
+ }
+
+ public String mode() {
+ return this.mode;
+ }
+
+ public void assertResourcesClosed() throws Exception {
+ Assert.assertTrue(this.loader.context().closed());
+ TaskManager manager = (TaskManager) field(this.loader, HugeGraphLoader.class, "manager");
+ Assert.assertTrue(((ExecutorService) field(manager, TaskManager.class, "batchService")).isTerminated());
+ Assert.assertTrue(((ExecutorService) field(manager, TaskManager.class, "singleService")).isTerminated());
+ Assert.assertFalse(Runtime.getRuntime().removeShutdownHook(
+ (Thread) field(this.loader, HugeGraphLoader.class, "shutdownHook")));
+ OkHttpClient transport = (OkHttpClient) fieldAny(fieldAny(this.loader.context().client(), "client"), "client");
+ Assert.assertTrue(transport.dispatcher().executorService().isShutdown());
+ Assert.assertEquals(0, transport.connectionPool().connectionCount());
+ Set<Throwable> active = Collections.newSetFromMap(new IdentityHashMap<>());
+ Assert.assertTrue(containsOriginalWithoutCycle(this.shutdownFailure.get(), active));
+ }
+
+ private boolean containsOriginalWithoutCycle(Throwable failure, Set<Throwable> active) {
+ if (failure == null) {
+ return false;
+ }
+ Assert.assertTrue("cyclic throwable graph", active.add(failure));
+ boolean found = failure == this.original;
+ found |= this.containsOriginalWithoutCycle(failure.getCause(), active);
+ for (Throwable suppressed : failure.getSuppressed()) {
+ found |= this.containsOriginalWithoutCycle(suppressed, active);
+ }
+ active.remove(failure);
+ return found;
+ }
+
+ private static Object fieldAny(Object object, String name) throws Exception {
+ for (Class<?> type = object.getClass(); type != null; type = type.getSuperclass()) {
+ try {
+ return field(object, type, name);
+ } catch (NoSuchFieldException ignored) {
+ // Continue through the SDK transport hierarchy.
+ }
+ }
+ throw new NoSuchFieldException(name);
+ }
+
+ private static Object field(Object object, Class<?> type, String name) throws Exception {
+ Field field = type.getDeclaredField(name);
+ field.setAccessible(true);
+ return field.get(object);
+ }
+
+ @Override
+ public void close() throws Exception {
+ try {
+ this.loader.shutdown();
+ } finally {
+ Runtime.getRuntime().removeShutdownHook(
+ (Thread) field(this.loader, HugeGraphLoader.class, "shutdownHook"));
+ this.server.stop(0);
+ }
+ }
+}
diff --git a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/FileUploadControllerTest.java b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/FileUploadControllerTest.java
index 269d71d07..e607469d9 100644
--- a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/FileUploadControllerTest.java
+++ b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/FileUploadControllerTest.java
@@ -30,7 +30,9 @@ import java.util.Collections;
import java.util.List;
import org.apache.commons.io.FileUtils;
+import org.junit.Rule;
import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import org.springframework.mock.web.MockMultipartFile;
@@ -50,6 +52,9 @@ import org.apache.hugegraph.testutil.Assert;
public class FileUploadControllerTest {
+ @Rule
+ public TemporaryFolder folder = new TemporaryFolder();
+
@Test
public void testCheckFileValidAcceptsUppercaseTrimmedFormat()
throws Exception {
@@ -118,6 +123,11 @@ public class FileUploadControllerTest {
throws Exception {
FileUploadController controller = this.controller("csv");
FileMappingService service = Mockito.mock(FileMappingService.class);
+ Path uploadRoot = this.folder.newFolder("quota-upload").toPath();
+ Path uploadPath = uploadRoot.resolve("data.csv");
+ FileMappingService pathValidator = this.fileMappingService(uploadRoot);
+ Mockito.when(service.requirePathUnderUploadRoot(Mockito.anyString()))
+ .thenAnswer(call -> pathValidator.requirePathUnderUploadRoot(call.getArgument(0)));
JobManagerService jobService = Mockito.mock(JobManagerService.class);
JobManager job = JobManager.builder().id(1).jobSize(20L).build();
Mockito.when(jobService.get("DEFAULT", "hugegraph", 1))
@@ -131,7 +141,7 @@ public class FileUploadControllerTest {
FileMapping mapping = this.reserveUploadQuota(controller, "DEFAULT",
"hugegraph", 1,
"data.csv",
- "upload/data.csv", 10L);
+ uploadPath.toString(), 10L);
ArgumentCaptor<FileMapping> captor = ArgumentCaptor.forClass(
FileMapping.class);
@@ -140,10 +150,19 @@ public class FileUploadControllerTest {
Assert.assertEquals("DEFAULT", mapping.getGraphSpace());
Assert.assertEquals("hugegraph", mapping.getGraph());
Assert.assertEquals("data.csv", mapping.getName());
- Assert.assertEquals("upload/data.csv", mapping.getPath());
+ Assert.assertEquals(uploadPath.toString(), mapping.getPath());
+ Assert.assertTrue(Files.isDirectory(uploadPath));
Assert.assertEquals(10L, mapping.getTotalSize());
Assert.assertEquals(FileMappingStatus.UPLOADING,
mapping.getFileStatus());
+
+ Path failedPath = uploadRoot.resolve("failed.csv");
+ Mockito.doThrow(new InternalException("database failure")).when(service).save(Mockito.any());
+ Assert.assertThrows(InternalException.class, () ->
+ this.reserveUploadQuota(controller, "DEFAULT", "hugegraph", 1,
+ "failed.csv", failedPath.toString(), 10L));
+ Assert.assertFalse(Files.exists(failedPath));
+ Assert.assertTrue(Files.isDirectory(uploadPath));
}
@Test
@@ -177,6 +196,11 @@ public class FileUploadControllerTest {
throws Exception {
FileUploadController controller = this.controller("csv");
FileMappingService service = Mockito.mock(FileMappingService.class);
+ Path uploadRoot = this.folder.newFolder("quota-upload").toPath();
+ Path uploadPath = uploadRoot.resolve("data.csv");
+ FileMappingService pathValidator = this.fileMappingService(uploadRoot);
+ Mockito.when(service.requirePathUnderUploadRoot(Mockito.anyString()))
+ .thenAnswer(call -> pathValidator.requirePathUnderUploadRoot(call.getArgument(0)));
JobManagerService jobService = Mockito.mock(JobManagerService.class);
JobManager job = JobManager.builder().id(1).jobSize(20L).build();
FileMapping existing = new FileMapping("DEFAULT", "hugegraph",
@@ -199,10 +223,11 @@ public class FileUploadControllerTest {
FileMapping mapping = this.reserveUploadQuota(controller, "DEFAULT",
"hugegraph", 1,
"data.csv",
- "upload/data.csv", 10L);
+ uploadPath.toString(), 10L);
Assert.assertEquals(existing, mapping);
- Assert.assertEquals("upload/data.csv", mapping.getPath());
+ Assert.assertEquals(uploadPath.toString(), mapping.getPath());
+ Assert.assertTrue(Files.isDirectory(uploadPath));
Assert.assertEquals(10L, mapping.getTotalSize());
Assert.assertEquals(FileMappingStatus.UPLOADING,
mapping.getFileStatus());
@@ -210,6 +235,98 @@ public class FileUploadControllerTest {
Mockito.verify(service, Mockito.never()).save(Mockito.any());
}
+ @Test
+ public void testChunkReservationsReuseDirectoryAndStillEnforceQuota()
+ throws Exception {
+ FileUploadController controller = this.controller("csv");
+ Path uploadRoot = this.folder.newFolder("chunk-upload").toPath();
+ Path uploadPath = uploadRoot.resolve("data.csv");
+ FileMappingService service = Mockito.spy(this.fileMappingService(uploadRoot));
+ JobManagerService jobService = Mockito.mock(JobManagerService.class);
+ JobManager job = JobManager.builder().id(1).jobSize(20L).build();
+ Mockito.when(jobService.get("DEFAULT", "hugegraph", 1)).thenReturn(job);
+ List<FileMapping> mappings = new ArrayList<>();
+ Path otherPath = Files.createDirectory(uploadRoot.resolve("other.csv"));
+ FileMapping other = new FileMapping("DEFAULT", "hugegraph", "other.csv", otherPath.toString());
+ other.setId(3);
+ other.setFileStatus(FileMappingStatus.UPLOADING);
+ other.setTotalSize(5L);
+ mappings.add(other);
+ Mockito.doReturn(mappings).when(service).listByJob(1);
+ Mockito.doAnswer(call -> mappings.size() == 1 ? null : mappings.get(1))
+ .when(service).get("DEFAULT", "hugegraph", 1, "data.csv");
+ Mockito.doAnswer(call -> {
+ FileMapping mapping = call.getArgument(0);
+ mapping.setId(2);
+ mappings.add(mapping);
+ return null;
+ }).when(service).save(Mockito.any());
+ Mockito.doNothing().when(service).update(Mockito.any());
+ this.setField(controller, "service", service);
+ this.setField(controller, "jobService", jobService);
+
+ for (int i = 0; i < 3; i++) {
+ this.reserveUploadQuota(controller, "DEFAULT", "hugegraph", 1,
+ "data.csv", uploadPath.toString(), 10L);
+ MockMultipartFile chunk = new MockMultipartFile(
+ "file", "data.csv", "text/csv", ("part" + i).getBytes());
+ service.uploadFile(chunk, "data.csv", i, uploadPath.toString());
+ }
+ // One quota scan per chunk and one physical alias scan for the first chunk.
+ Mockito.verify(service, Mockito.times(4)).listByJob(1);
+ Mockito.verify(service, Mockito.times(6))
+ .requirePathUnderUploadRoot(Mockito.anyString());
+ Mockito.verify(service).requirePathUnderUploadRoot(otherPath.toString());
+ Path outsidePath = uploadRoot.getParent().resolve("outside.csv");
+ Assert.assertThrows(InternalException.class, () ->
+ this.reserveUploadQuota(controller, "DEFAULT", "hugegraph", 1,
+ "data.csv", outsidePath.toString(), 10L));
+ Assert.assertTrue(service.tryMergePartFiles(uploadPath.toString(), 3));
+ Assert.assertEquals("part0part1part2", Files.readString(uploadPath));
+
+ // Increasing the declared size and a concurrently increased job size are still checked.
+ Assert.assertThrows(ExternalException.class, () ->
+ this.reserveUploadQuota(controller, "DEFAULT", "hugegraph", 1,
+ "data.csv", uploadPath.toString(), 1025L));
+ job.setJobSize(2040L);
+ Assert.assertThrows(ExternalException.class, () ->
+ this.reserveUploadQuota(controller, "DEFAULT", "hugegraph", 1,
+ "data.csv", uploadPath.toString(), 10L));
+ Assert.assertEquals(10L, mappings.get(1).getTotalSize());
+ }
+
+ @Test
+ public void testFirstReservationRejectsNativeAliases() throws Exception {
+ Path uploadRoot = this.folder.newFolder("alias-upload").toPath();
+ for (String[] names : new String[][]{{"data.csv", "DATA.csv"},
+ {"caf\u00e9.csv", "cafe\u0301.csv"}}) {
+ Path original = Files.createDirectory(uploadRoot.resolve(names[0]));
+ Path alias = uploadRoot.resolve(names[1]);
+ // Exercise case/Unicode aliases where the native filesystem actually aliases them.
+ if (!Files.exists(alias)) {
+ Files.createSymbolicLink(alias, original);
+ }
+ Assert.assertTrue(Files.isSameFile(original, alias));
+ FileUploadController controller = this.controller("csv");
+ FileMappingService service = Mockito.spy(this.fileMappingService(uploadRoot));
+ FileMapping existing = new FileMapping("DEFAULT", "hugegraph", names[0], original.toString());
+ existing.setFileStatus(FileMappingStatus.UPLOADING);
+ Mockito.doReturn(null).when(service).get("DEFAULT", "hugegraph", 1, names[1]);
+ Mockito.doReturn(Collections.singletonList(existing)).when(service).listByJob(1);
+ JobManagerService jobService = Mockito.mock(JobManagerService.class);
+ Mockito.when(jobService.get("DEFAULT", "hugegraph", 1))
+ .thenReturn(JobManager.builder().id(1).jobSize(0L).build());
+ this.setField(controller, "service", service);
+ this.setField(controller, "jobService", jobService);
+
+ Assert.assertThrows(ExternalException.class, () ->
+ this.reserveUploadQuota(controller, "DEFAULT", "hugegraph", 1,
+ names[1], alias.toString(), 10L));
+ Mockito.verify(service, Mockito.never()).save(Mockito.any());
+ Assert.assertTrue(Files.isDirectory(original));
+ }
+ }
+
@Test
public void testCheckTotalAndIndexValidRejectsIndexAtTotal()
throws Exception {
@@ -414,7 +531,7 @@ public class FileUploadControllerTest {
String.class,
int.class,
JobManager.class,
- org.springframework.web.multipart.MultipartFile.class,
+ org.springframework.web.multipart.MultipartFile.class,
String.class);
method.setAccessible(true);
try {
diff --git a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/JobManagerServiceTest.java b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/JobManagerServiceTest.java
index c2bb8d8a3..95db52f72 100644
--- a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/JobManagerServiceTest.java
+++ b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/JobManagerServiceTest.java
@@ -21,6 +21,16 @@ package org.apache.hugegraph.unit;
import java.lang.reflect.Field;
import java.util.Collections;
import java.util.Arrays;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.hugegraph.exception.ExternalException;
+
+import org.apache.hugegraph.loader.HugeGraphLoader;
+import org.apache.hugegraph.loader.executor.LoadContext;
+import org.apache.hugegraph.loader.executor.LoadOptions;
+import org.apache.hugegraph.mapper.load.LoadTaskMapper;
import org.junit.Test;
import org.mockito.Mockito;
@@ -176,8 +186,199 @@ public class JobManagerServiceTest {
Mockito.verify(taskService).start(connection, mapping, client);
}
- private JobManagerService service() {
- return new JobManagerService();
+ @Test
+ public void testDeleteJobAllowsSucceededTaskAfterRunningSnapshot() throws Exception {
+ this.deleteCompletedTask(LoadStatus.SUCCEED);
+ }
+
+ @Test
+ public void testDeleteJobAllowsFailedLoadAfterRunningSnapshot() throws Exception {
+ this.deleteCompletedTask(LoadStatus.FAILED);
+ }
+
+ @Test
+ public void testDeleteJobAllowsStoppedTaskAfterRunningSnapshot() throws Exception {
+ this.deleteCompletedTask(LoadStatus.STOPPED);
+ }
+
+ private void deleteCompletedTask(LoadStatus terminal) throws Exception {
+ DeleteFixture fixture = this.deleteFixture(terminal, null);
+
+ fixture.jobs.deleteJob(7);
+
+ Assert.assertEquals(terminal, fixture.task.getStatus());
+ Assert.assertTrue(fixture.task.isFinished());
+ Mockito.verify(fixture.loader).shutdown();
+ Mockito.verify(fixture.tasks).deleteById(9);
+ Mockito.verify(fixture.mapper).deleteById(7);
+ Assert.assertTrue(fixture.running.isEmpty());
+ }
+
+ @Test(timeout = 10000)
+ public void testDeleteJobPreservesTaskWhenCleanupFailsDuringStopWait() throws Exception {
+ DeleteFixture fixture = this.deleteFixture(LoadStatus.FAILED,
+ new IllegalStateException("release failed"));
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch stopped = new CountDownLatch(1);
+ Mockito.when(fixture.loader.load()).thenAnswer(invocation -> {
+ entered.countDown();
+ Assert.assertTrue(stopped.await(5, TimeUnit.SECONDS));
+ return false;
+ });
+ Mockito.doAnswer(invocation -> {
+ stopped.countDown();
+ return null;
+ }).when(fixture.context).stopLoading();
+ LoadTask snapshot = LoadTask.builder().id(9).status(LoadStatus.RUNNING).build();
+ Mockito.doReturn(Collections.singletonList(snapshot)).when(fixture.service).taskListByJob(7);
+ Mockito.when(fixture.tasks.updateById(Mockito.any(LoadTask.class))).thenReturn(1);
+ Thread runner = new Thread(fixture.task, "delete-job-cleanup-race-test");
+ runner.start();
+ try {
+ Assert.assertTrue(entered.await(5, TimeUnit.SECONDS));
+
+ Assert.assertThrows(RuntimeException.class, () -> fixture.jobs.deleteJob(7));
+
+ Mockito.verify(fixture.tasks, Mockito.never()).deleteById(Mockito.anyInt());
+ Mockito.verify(fixture.mapper, Mockito.never()).deleteById(Mockito.anyInt());
+ Assert.assertSame(fixture.task, fixture.running.get(9));
+ Assert.assertEquals(LoadStatus.FAILED, fixture.task.getStatus());
+ } finally {
+ stopped.countDown();
+ runner.join(5000);
+ Assert.assertFalse(runner.isAlive());
+ }
+ }
+
+ @Test
+ public void testDeleteJobPreservesTaskWhenCleanupFailedAfterSnapshot() throws Exception {
+ DeleteFixture fixture = this.deleteFixture(LoadStatus.FAILED,
+ new IllegalStateException("release failed"));
+
+ Assert.assertThrows(RuntimeException.class, () -> fixture.jobs.deleteJob(7));
+
+ Mockito.verify(fixture.tasks, Mockito.never()).deleteById(Mockito.anyInt());
+ Mockito.verify(fixture.mapper, Mockito.never()).deleteById(Mockito.anyInt());
+ Mockito.verify(fixture.mappings, Mockito.never()).listByJob(Mockito.anyInt());
+ Assert.assertEquals(LoadStatus.FAILED, fixture.task.getStatus());
+ Assert.assertSame(fixture.task, fixture.running.get(9));
+ }
+
+ @Test
+ public void testDeleteJobDoesNotAssumeDatabaseTerminalCleanupSucceeded() throws Exception {
+ DeleteFixture fixture = this.deleteFixture(LoadStatus.FAILED, null);
+ fixture.running.clear();
+ LoadTask snapshot = LoadTask.builder().id(9).status(LoadStatus.RUNNING).build();
+ Mockito.doReturn(Collections.singletonList(snapshot)).when(fixture.service).taskListByJob(7);
+ Mockito.when(fixture.tasks.selectById(9))
+ .thenReturn(LoadTask.builder().id(9).status(LoadStatus.FAILED).build());
+
+ Assert.assertThrows(ExternalException.class, () -> fixture.jobs.deleteJob(7));
+
+ Mockito.verify(fixture.tasks, Mockito.never()).deleteById(Mockito.anyInt());
+ Mockito.verify(fixture.mapper, Mockito.never()).deleteById(Mockito.anyInt());
+ }
+
+ @Test
+ public void testOrdinaryStopStillRejectsCompletedTasks() throws Exception {
+ for (LoadStatus terminal : new LoadStatus[]{LoadStatus.SUCCEED, LoadStatus.FAILED, LoadStatus.STOPPED}) {
+ DeleteFixture fixture = this.deleteFixture(terminal, null);
+ fixture.task.run();
+
+ Assert.assertThrows(ExternalException.class, () -> fixture.service.stop(9));
+
+ Assert.assertEquals(terminal, fixture.task.getStatus());
+ Mockito.verify(fixture.tasks, Mockito.never()).updateById(Mockito.any(LoadTask.class));
+ }
+ }
+
+ @Test
+ public void testDeleteJobDoesNotRemoveMetadataWhenStopPersistenceFails() throws Exception {
+ DeleteFixture fixture = this.deleteFixture(LoadStatus.PAUSED, null);
+ fixture.running.clear();
+ LoadTask paused = LoadTask.builder().id(9).status(LoadStatus.PAUSED).build();
+ Mockito.doReturn(Collections.singletonList(paused)).when(fixture.service).taskListByJob(7);
+ Mockito.when(fixture.tasks.selectById(9)).thenReturn(paused);
+ IllegalStateException failure = new IllegalStateException("persist stop failed");
+ Mockito.when(fixture.tasks.updateById(Mockito.any(LoadTask.class))).thenThrow(failure);
+
+ Assert.assertSame(failure, Assert.assertThrows(IllegalStateException.class,
+ () -> fixture.jobs.deleteJob(7)));
+
+ Mockito.verify(fixture.tasks, Mockito.never()).deleteById(Mockito.anyInt());
+ Mockito.verify(fixture.mapper, Mockito.never()).deleteById(Mockito.anyInt());
+ Mockito.verify(fixture.mappings, Mockito.never()).listByJob(Mockito.anyInt());
+ }
+
+ private DeleteFixture deleteFixture(LoadStatus terminal, RuntimeException cleanup) throws Exception {
+ DeleteFixture fixture = new DeleteFixture();
+ fixture.jobs = this.service();
+ fixture.service = Mockito.spy(new LoadTaskService());
+ org.springframework.transaction.PlatformTransactionManager transactions =
+ Mockito.mock(org.springframework.transaction.PlatformTransactionManager.class);
+ Mockito.when(transactions.getTransaction(Mockito.any()))
+ .thenReturn(new org.springframework.transaction.support.SimpleTransactionStatus());
+ this.setField(fixture.service, "transactionManager", transactions);
+ fixture.mapper = Mockito.mock(JobManagerMapper.class);
+ fixture.tasks = Mockito.mock(LoadTaskMapper.class);
+ fixture.mappings = Mockito.mock(FileMappingService.class);
+ fixture.loader = Mockito.mock(HugeGraphLoader.class);
+ fixture.context = Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS);
+ Mockito.when(fixture.loader.context()).thenReturn(fixture.context);
+ fixture.task = LoadTask.builder().id(9).jobId(7).loader(fixture.loader)
+ .options(new LoadOptions()).status(LoadStatus.RUNNING).lastDuration(0L).build();
+ Mockito.when(fixture.loader.load()).thenAnswer(invocation -> {
+ if (terminal == LoadStatus.STOPPED) {
+ fixture.task.setStatus(LoadStatus.STOPPED);
+ }
+ return terminal == LoadStatus.SUCCEED;
+ });
+ if (cleanup != null) {
+ Mockito.doThrow(cleanup).when(fixture.loader).shutdown();
+ }
+ fixture.running = new ConcurrentHashMap<>();
+ fixture.running.put(9, fixture.task);
+ this.setField(fixture.service, "runningTaskContainer", fixture.running);
+ this.setField(fixture.service, "mapper", fixture.tasks);
+ this.setField(fixture.jobs, "mapper", fixture.mapper);
+ this.setField(fixture.jobs, "taskService", fixture.service);
+ this.setField(fixture.jobs, "fileMappingService", fixture.mappings);
+ Mockito.when(fixture.mapper.selectById(7)).thenReturn(JobManager.builder().id(7).build());
+ Mockito.when(fixture.mapper.deleteById(7)).thenReturn(1);
+ Mockito.when(fixture.tasks.deleteById(9)).thenReturn(1);
+ Mockito.when(fixture.mappings.listByJob(7)).thenReturn(Collections.emptyList());
+ LoadTask snapshot = LoadTask.builder().id(9).status(LoadStatus.RUNNING).build();
+ Mockito.when(fixture.tasks.selectById(9)).thenAnswer(invocation ->
+ LoadTask.builder().id(9).status(fixture.task.getStatus()).build());
+ Mockito.doAnswer(invocation -> {
+ // The real runner finishes cleanup after the database snapshot,
+ // before deleteJob reaches the real stop implementation.
+ fixture.task.run();
+ return Collections.singletonList(snapshot);
+ }).when(fixture.service).taskListByJob(7);
+ return fixture;
+ }
+
+ private static class DeleteFixture {
+ private JobManagerService jobs;
+ private LoadTaskService service;
+ private JobManagerMapper mapper;
+ private LoadTaskMapper tasks;
+ private FileMappingService mappings;
+ private HugeGraphLoader loader;
+ private LoadContext context;
+ private LoadTask task;
+ private ConcurrentHashMap<Integer, LoadTask> running;
+ }
+
+ private JobManagerService service() throws Exception {
+ JobManagerService service = new JobManagerService();
+ org.springframework.transaction.PlatformTransactionManager transactions =
+ Mockito.mock(org.springframework.transaction.PlatformTransactionManager.class);
+ Mockito.when(transactions.getTransaction(Mockito.any()))
+ .thenReturn(new org.springframework.transaction.support.SimpleTransactionStatus());
+ this.setField(service, "transactionManager", transactions);
+ return service;
}
private void setField(Object object, String name, Object value)
diff --git a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/LoadTaskCheckpointTest.java b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/LoadTaskCheckpointTest.java
new file mode 100644
index 000000000..9285789ce
--- /dev/null
+++ b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/LoadTaskCheckpointTest.java
@@ -0,0 +1,259 @@
+/*
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with this
+ * work for additional information regarding copyright ownership. The ASF
+ * licenses this file to You under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance with the
+ * License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hugegraph.unit;
+
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.lang.reflect.Field;
+import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.zip.GZIPInputStream;
+
+import org.apache.hugegraph.entity.GraphConnection;
+import org.apache.hugegraph.entity.enums.LoadStatus;
+import org.apache.hugegraph.entity.load.FileMapping;
+import org.apache.hugegraph.entity.load.LoadTask;
+import org.apache.hugegraph.loader.HugeGraphLoader;
+import org.apache.hugegraph.loader.executor.LoadOptions;
+import org.apache.hugegraph.loader.progress.LoadProgress;
+import org.apache.hugegraph.util.JsonUtil;
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import com.sun.net.httpserver.HttpServer;
+
+/** Real CSV/parser/semaphore/insert/cleanup pipeline against a controlled HTTP endpoint. */
+public class LoadTaskCheckpointTest {
+
+ @Rule
+ public TemporaryFolder folder = new TemporaryFolder();
+
+ @Test(timeout = 20000)
+ public void testPauseAtBatchSemaphoreThenResumeDoesNotSkipUnsubmittedRows() throws Exception {
+ try (InsertServer server = new InsertServer()) {
+ server.holdFirstBatch = true;
+ LoadTask task = this.task(server);
+ HugeGraphLoader first = task.getLoader();
+ Object manager = field(first, "manager");
+ Semaphore semaphore = (Semaphore) field(manager, "batchSemaphore");
+ ExecutorService runner = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> running = runner.submit(task);
+ Assert.assertTrue(server.firstBatch.await(5, TimeUnit.SECONDS));
+ long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+ while (!semaphore.hasQueuedThreads() && System.nanoTime() < deadline) {
+ Thread.sleep(10);
+ }
+ Assert.assertTrue("The actual parser must wait on the native batch semaphore",
+ semaphore.hasQueuedThreads());
+ task.setStatus(LoadStatus.PAUSED);
+ task.context().stopLoading();
+ server.releaseBatch.countDown();
+ running.get(5, TimeUnit.SECONDS);
+ Assert.assertEquals(LoadStatus.PAUSED, task.getStatus());
+ Assert.assertTrue(first.context().closed());
+ Assert.assertTrue("The parsed waiting batch was never inserted", server.rows.size() < 10);
+ this.resume(task);
+ Assert.assertEquals(LoadStatus.SUCCEED, task.getStatus());
+ Assert.assertEquals("Resume must include every unsubmitted row", 10, server.rows.size());
+ this.assertNoCheckpoint(task, first);
+ Assert.assertTrue(task.context().closed());
+ } finally {
+ server.releaseBatch.countDown();
+ runner.shutdownNow();
+ Assert.assertTrue(runner.awaitTermination(5, TimeUnit.SECONDS));
+ first.shutdown();
+ task.getLoader().shutdown();
+ }
+ }
+ }
+
+ @Test(timeout = 20000)
+ public void testFailedSingleInsertThenResumeDoesNotSkipFailedRow() throws Exception {
+ try (InsertServer server = new InsertServer()) {
+ server.rejectName = "person2";
+ LoadTask task = this.task(server);
+ HugeGraphLoader first = task.getLoader();
+ task.run();
+ Assert.assertEquals(LoadStatus.FAILED, task.getStatus());
+ Assert.assertEquals(9, server.rows.size());
+ Assert.assertFalse(server.rows.contains("person2"));
+ Assert.assertTrue(first.context().closed());
+ server.rejectName = null;
+ this.resume(task);
+ Assert.assertEquals(LoadStatus.SUCCEED, task.getStatus());
+ Assert.assertEquals("The failed row must be read and inserted on resume", 10, server.rows.size());
+ Assert.assertTrue(server.rows.contains("person2"));
+ this.assertNoCheckpoint(task, first);
+ Assert.assertTrue(task.context().closed());
+ }
+ }
+
+ private LoadTask task(InsertServer server) throws Exception {
+ Path input = this.folder.newFile("people.csv").toPath();
+ StringBuilder csv = new StringBuilder();
+ for (int i = 0; i < 10; i++) {
+ csv.append("person").append(i).append('\n');
+ }
+ Files.writeString(input, csv);
+ Path mapping = this.folder.newFile("mapping.json").toPath();
+ Files.writeString(mapping, JsonUtil.toJson(Map.of("vertices", List.of(Map.of(
+ "label", "person", "input", Map.of("type", "FILE", "path", input.toString(),
+ "format", "CSV", "header", List.of("name"), "charset", "UTF-8"))))));
+ LoadOptions options = new LoadOptions();
+ options.host = "127.0.0.1";
+ options.port = server.server.getAddress().getPort();
+ options.file = mapping.toString();
+ options.batchSize = server.holdFirstBatch ? 4 : 1;
+ options.batchInsertThreads = 1;
+ options.singleInsertThreads = 1;
+ options.retryTimes = 0;
+ GraphConnection connection = new GraphConnection();
+ connection.setId(1);
+ FileMapping file = new FileMapping().setId(1).setJobId(1).setName("people.csv").setTotalLines(10);
+ return new LoadTask(options, connection, file);
+ }
+
+ private void resume(LoadTask task) {
+ // The same native reconnect path used after LoadTaskService joins the old callback.
+ task.getOptions().incrementalMode = true;
+ task.reconnect("current-token");
+ task.setStatus(LoadStatus.RUNNING);
+ task.run();
+ }
+
+ private void assertNoCheckpoint(LoadTask task, HugeGraphLoader loader) {
+ Assert.assertFalse("Hubble must not persist parsed offsets as inserted-row checkpoints",
+ Files.exists(Path.of(LoadProgress.format(task.getOptions(), loader.context().timestamp()))));
+ }
+
+ private static Object field(Object target, String name) throws Exception {
+ Field field = target.getClass().getDeclaredField(name);
+ field.setAccessible(true);
+ return field.get(target);
+ }
+
+ private static class InsertServer implements AutoCloseable {
+
+ private final HttpServer server;
+ private final ExecutorService handlers = Executors.newCachedThreadPool();
+ private final Set<String> rows = ConcurrentHashMap.newKeySet();
+ private final CountDownLatch firstBatch = new CountDownLatch(1);
+ private final CountDownLatch releaseBatch = new CountDownLatch(1);
+ private volatile boolean holdFirstBatch;
+ private volatile String rejectName;
+ private String mode = "NONE";
+
+ private InsertServer() throws Exception {
+ this.server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+ this.server.setExecutor(this.handlers);
+ this.server.createContext("/", exchange -> {
+ Object response = Map.of();
+ int status = 200;
+ String path = exchange.getRequestURI().getPath();
+ if (path.equals("/versions")) {
+ response = Map.of("versions", Map.of("core", "1.8.0", "api", "0.74"));
+ } else if (path.endsWith("/mode")) {
+ if (exchange.getRequestMethod().equals("PUT")) {
+ this.mode = JsonUtil.fromJson(new String(exchange.getRequestBody().readAllBytes(),
+ StandardCharsets.UTF_8), String.class);
+ } else {
+ response = Map.of("mode", this.mode);
+ }
+ } else if (path.endsWith("/schema/propertykeys")) {
+ response = Map.of("propertykeys", List.of(Map.of("id", 1, "name", "name", "data_type", "TEXT",
+ "cardinality", "SINGLE", "aggregate_type", "NONE", "properties", List.of())));
+ } else if (path.endsWith("/schema/vertexlabels")) {
+ response = Map.of("vertexlabels", List.of(Map.of("id", 1, "name", "person",
+ "id_strategy", "PRIMARY_KEY", "primary_keys", List.of("name"),
+ "properties", List.of("name"), "nullable_keys", List.of(), "enable_label_index", true)));
+ } else if (path.endsWith("/schema/edgelabels")) {
+ response = Map.of("edgelabels", List.of());
+ } else if (path.endsWith("/vertices/batch") || path.endsWith("/vertices")) {
+ if (path.endsWith("/batch") && this.holdFirstBatch) {
+ this.firstBatch.countDown();
+ try {
+ if (!this.releaseBatch.await(10, TimeUnit.SECONDS)) {
+ throw new IllegalStateException("Batch release deadline expired");
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException(e);
+ }
+ }
+ InputStream input = exchange.getRequestBody();
+ if ("gzip".equals(exchange.getRequestHeaders().getFirst("Content-Encoding"))) {
+ input = new GZIPInputStream(input);
+ }
+ String body = new String(input.readAllBytes(), StandardCharsets.UTF_8);
+ List<?> vertices = path.endsWith("/batch") ? JsonUtil.fromJson(body, List.class) :
+ Collections.singletonList(JsonUtil.fromJson(body, Map.class));
+ List<String> names = new ArrayList<>();
+ for (Object value : vertices) {
+ Map<?, ?> vertex = (Map<?, ?>) value;
+ names.add((String) ((Map<?, ?>) vertex.get("properties")).get("name"));
+ }
+ if (this.rejectName != null && names.contains(this.rejectName)) {
+ status = 500;
+ response = Map.of("exception", "class java.lang.IllegalStateException",
+ "message", "controlled insert failure");
+ } else {
+ this.rows.addAll(names);
+ response = path.endsWith("/batch") ? names : Map.of("id", names.get(0), "label", "person",
+ "properties", Map.of("name", names.get(0)));
+ }
+ } else {
+ status = 404;
+ response = Map.of("message", "Unsupported test endpoint " + path);
+ }
+ byte[] body = JsonUtil.toJson(response).getBytes(StandardCharsets.UTF_8);
+ exchange.getResponseHeaders().add("Content-Type", "application/json");
+ exchange.sendResponseHeaders(status, body.length);
+ try (OutputStream output = exchange.getResponseBody()) {
+ output.write(body);
+ }
+ });
+ this.server.start();
+ }
+
+ @Override
+ public void close() throws Exception {
+ this.releaseBatch.countDown();
+ this.server.stop(0);
+ this.handlers.shutdownNow();
+ Assert.assertTrue(this.handlers.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+}
diff --git a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/LoadTaskServiceTest.java b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/LoadTaskServiceTest.java
index 68c0ab58e..54130b8db 100644
--- a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/LoadTaskServiceTest.java
+++ b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/LoadTaskServiceTest.java
@@ -18,12 +18,20 @@
package org.apache.hugegraph.unit;
+import java.io.OutputStream;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
+import java.lang.reflect.UndeclaredThrowableException;
+import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
import java.util.Arrays;
+import java.util.Collections;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.RejectedExecutionException;
+import org.junit.Rule;
import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
import org.apache.hugegraph.entity.GraphConnection;
import org.apache.hugegraph.entity.enums.LoadStatus;
@@ -33,20 +41,349 @@ import org.apache.hugegraph.entity.load.ListFormat;
import org.apache.hugegraph.entity.load.LoadParameter;
import org.apache.hugegraph.entity.load.LoadTask;
import org.apache.hugegraph.handler.LoadTaskExecutor;
+import org.apache.hugegraph.loader.HugeGraphLoader;
+import org.apache.hugegraph.loader.executor.LoadContext;
import org.apache.hugegraph.loader.executor.LoadOptions;
+import org.apache.hugegraph.loader.mapping.LoadMapping;
+import org.apache.hugegraph.loader.util.LoadUtil;
+import org.apache.hugegraph.loader.exception.LoadException;
import org.apache.hugegraph.loader.source.file.FileSource;
import org.apache.hugegraph.service.load.LoadTaskService;
import org.apache.hugegraph.mapper.load.LoadTaskMapper;
import org.apache.hugegraph.testutil.Assert;
+import org.apache.hugegraph.testutil.NativeCleanupTestSupport;
import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
+import org.mockito.MockedConstruction;
+import org.springframework.transaction.PlatformTransactionManager;
+import org.springframework.transaction.support.SimpleTransactionStatus;
import org.springframework.transaction.support.TransactionSynchronization;
import org.springframework.transaction.support.TransactionSynchronizationManager;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
+import com.sun.net.httpserver.HttpServer;
public class LoadTaskServiceTest {
+ @Rule
+ public TemporaryFolder folder = new TemporaryFolder();
+
+ @Test
+ public void testRunnerClosesBeforePublishingCompletion() {
+ for (boolean success : new boolean[]{true, false}) {
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ LoadContext context = Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS);
+ Mockito.when(loader.context()).thenReturn(context);
+ Mockito.when(loader.load()).thenReturn(success);
+ LoadTask task = LoadTask.builder().id(9).loader(loader).options(new LoadOptions())
+ .status(LoadStatus.RUNNING).lastDuration(0L).build();
+ Mockito.doAnswer(invocation -> {
+ Assert.assertFalse(task.isFinished());
+ Assert.assertEquals(LoadStatus.RUNNING, task.getStatus());
+ return null;
+ }).when(loader).shutdown();
+
+ task.run();
+
+ Mockito.verify(loader).shutdown();
+ Assert.assertTrue(task.isFinished());
+ Assert.assertEquals(success ? LoadStatus.SUCCEED : LoadStatus.FAILED, task.getStatus());
+ }
+ }
+
+ @Test
+ public void testCancelledLoadFailureDoesNotOverwriteRequestedStatus() throws Exception {
+ HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+ server.createContext("/versions", exchange -> {
+ byte[] body = "{\"versions\":{\"core\":\"1.8.0\",\"api\":\"0.74\"}}"
+ .getBytes(StandardCharsets.UTF_8);
+ exchange.getResponseHeaders().add("Content-Type", "application/json");
+ exchange.sendResponseHeaders(200, body.length);
+ try (OutputStream output = exchange.getResponseBody()) {
+ output.write(body);
+ }
+ });
+ server.start();
+ try {
+ for (LoadStatus status : new LoadStatus[]{LoadStatus.PAUSED, LoadStatus.STOPPED}) {
+ LoadOptions options = new LoadOptions();
+ options.host = "127.0.0.1";
+ options.port = server.getAddress().getPort();
+ options.file = this.folder.newFile("cancel-" + status + ".json").getAbsolutePath();
+ HugeGraphLoader loader = Mockito.spy(new HugeGraphLoader(
+ options, new LoadMapping(Collections.emptyList())));
+ IllegalStateException failure = new IllegalStateException("load failed");
+ LoadTask task = LoadTask.builder().id(9).loader(loader).options(options)
+ .status(LoadStatus.RUNNING).lastDuration(0L).build();
+ Mockito.doAnswer(invocation -> {
+ task.setStatus(status);
+ loader.context().failLoading(failure);
+ throw failure;
+ }).when(loader).load();
+ Field field = HugeGraphLoader.class.getDeclaredField("shutdownHook");
+ field.setAccessible(true);
+ Thread hook = (Thread) field.get(loader);
+ try {
+ task.run();
+
+ Assert.assertTrue(loader.context().closed());
+ Assert.assertFalse(Runtime.getRuntime().removeShutdownHook(hook));
+ Assert.assertTrue(task.isFinished());
+ Assert.assertEquals(status, task.getStatus());
+ Assert.assertEquals(0, failure.getSuppressed().length);
+ } finally {
+ loader.shutdown();
+ Runtime.getRuntime().removeShutdownHook(hook);
+ }
+ }
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ @Test
+ public void testCancelledLoadWithIndependentCleanupFailureIsFailed() {
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ LoadContext context = Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS);
+ Mockito.when(loader.context()).thenReturn(context);
+ IllegalStateException failure = new IllegalStateException("load failed");
+ IllegalStateException cleanup = new IllegalStateException("cleanup failed");
+ Mockito.when(loader.load()).thenThrow(failure);
+ Mockito.doThrow(cleanup).when(loader).shutdown();
+ LoadTask task = LoadTask.builder().id(9).loader(loader).options(new LoadOptions())
+ .status(LoadStatus.STOPPED).lastDuration(0L).build();
+
+ task.run();
+
+ Assert.assertTrue(task.isFinished());
+ Assert.assertEquals(LoadStatus.FAILED, task.getStatus());
+ Assert.assertEquals(1, failure.getSuppressed().length);
+ Assert.assertSame(cleanup, failure.getSuppressed()[0]);
+ }
+
+ @Test
+ public void testCancelledRecordedLoadFailureRethrowKeepsRequestedStatus() {
+ for (Throwable original : new Throwable[]{new IllegalStateException("load failed"),
+ new Exception("load failed"),
+ new AssertionError("load failed"),
+ new UndeclaredThrowableException(
+ new RuntimeException(new IllegalStateException("wrapped")))}) {
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ LoadContext context = Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS);
+ Mockito.when(loader.context()).thenReturn(context);
+ RuntimeException loadFailure = LoadUtil.targetRuntimeException(original);
+ RuntimeException recorded = original instanceof RuntimeException ? (RuntimeException) original :
+ new LoadException("recorded failure", original);
+ RuntimeException rethrown = original instanceof RuntimeException ? (RuntimeException) original :
+ new LoadException("recorded failure", original);
+ Mockito.when(loader.load()).thenThrow(loadFailure);
+ Mockito.doThrow(recorded).when(context).throwIfFailed();
+ Mockito.doThrow(rethrown).when(loader).shutdown();
+ LoadTask task = LoadTask.builder().id(9).loader(loader).options(new LoadOptions())
+ .status(LoadStatus.PAUSED).lastDuration(0L).build();
+
+ task.run();
+
+ Assert.assertEquals(LoadStatus.PAUSED, task.getStatus());
+ Assert.assertTrue(task.isFinished());
+ Assert.assertEquals(0, loadFailure.getSuppressed().length);
+ }
+ }
+
+ @Test
+ public void testCancelledRethrowWithNewSuppressedCleanupFailureIsFailed() {
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ LoadContext context = Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS);
+ Mockito.when(loader.context()).thenReturn(context);
+ IllegalStateException failure = new IllegalStateException("load failed");
+ IllegalStateException cleanup = new IllegalStateException("close failed");
+ Mockito.when(loader.load()).thenThrow(failure);
+ Mockito.doThrow(failure).when(context).throwIfFailed();
+ Mockito.doAnswer(invocation -> {
+ failure.addSuppressed(cleanup);
+ throw failure;
+ }).when(loader).shutdown();
+ LoadTask task = LoadTask.builder().id(9).loader(loader).options(new LoadOptions())
+ .status(LoadStatus.STOPPED).lastDuration(0L).build();
+
+ task.run();
+
+ Assert.assertEquals(LoadStatus.FAILED, task.getStatus());
+ Assert.assertTrue(task.isFinished());
+ Assert.assertSame(cleanup, failure.getSuppressed()[0]);
+ }
+
+ @Test
+ public void testAfterCommitDispatchErrorClosesLoaderAndPreservesError() throws Exception {
+ LoadTaskService service = new LoadTaskService();
+ LoadTaskMapper mapper = Mockito.mock(LoadTaskMapper.class);
+ LoadTaskExecutor executor = Mockito.mock(LoadTaskExecutor.class);
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ LoadTask task = LoadTask.builder().id(9).loader(loader).options(new LoadOptions())
+ .status(LoadStatus.RUNNING).build();
+ Error failure = new AssertionError("dispatch failed");
+ IllegalStateException cleanup = new IllegalStateException("close failed");
+ Mockito.doThrow(failure).when(executor).execute(Mockito.eq(task), Mockito.any());
+ Mockito.doThrow(cleanup).when(loader).shutdown();
+ Mockito.when(mapper.update(Mockito.isNull(), Mockito.any())).thenReturn(1);
+ this.setField(service, "mapper", mapper);
+ this.setField(service, "taskExecutor", executor);
+ ConcurrentHashMap<Integer, LoadTask> running = new ConcurrentHashMap<>();
+ this.setField(service, "runningTaskContainer", running);
+ PlatformTransactionManager transactions = Mockito.mock(PlatformTransactionManager.class);
+ Mockito.when(transactions.getTransaction(Mockito.any())).thenReturn(new SimpleTransactionStatus());
+ this.setField(service, "transactionManager", transactions);
+ Method dispatch = LoadTaskService.class.getDeclaredMethod("executeAfterCommit", LoadTask.class);
+ dispatch.setAccessible(true);
+ TransactionSynchronizationManager.initSynchronization();
+ try {
+ dispatch.invoke(service, task);
+ TransactionSynchronization synchronization =
+ TransactionSynchronizationManager.getSynchronizations().get(0);
+ try {
+ synchronization.afterCommit();
+ Assert.fail("Expected dispatch Error");
+ } catch (Error actual) {
+ Assert.assertSame(failure, actual);
+ }
+ synchronization.afterCompletion(TransactionSynchronization.STATUS_COMMITTED);
+ Mockito.verify(loader).shutdown();
+ Mockito.verify(mapper).update(Mockito.isNull(), Mockito.any());
+ Assert.assertTrue(task.isFinished());
+ Assert.assertEquals(LoadStatus.FAILED, task.getStatus());
+ Assert.assertSame(cleanup, failure.getSuppressed()[0]);
+ Assert.assertTrue(task.isCleanupFailed());
+ Assert.assertSame(task, running.get(9));
+ Assert.assertThrows(RuntimeException.class, () -> service.retry(9, "current-token"));
+ Assert.assertThrows(RuntimeException.class, () -> service.resume(9, "current-token"));
+ Assert.assertSame(failure, Assert.assertThrows(Error.class, () -> service.remove(9)));
+ Mockito.verify(mapper, Mockito.never()).deleteById(Mockito.anyInt());
+ } finally {
+ TransactionSynchronizationManager.clearSynchronization();
+ }
+ }
+
+ @Test
+ public void testCleanupFailureWithoutLoadFailureStillMarksTaskFailed() {
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ LoadContext context = Mockito.mock(LoadContext.class, Mockito.RETURNS_DEEP_STUBS);
+ Mockito.when(loader.context()).thenReturn(context);
+ Mockito.when(loader.load()).thenReturn(true);
+ Mockito.doThrow(new IllegalStateException("cleanup failed")).when(loader).shutdown();
+ LoadTask task = LoadTask.builder().id(9).loader(loader).options(new LoadOptions())
+ .status(LoadStatus.PAUSED).lastDuration(0L).build();
+
+ task.run();
+
+ Assert.assertTrue(task.isFinished());
+ Assert.assertEquals(LoadStatus.FAILED, task.getStatus());
+ }
+
+ @Test
+ public void testResumeAndRetryCloseLoaderOnUpdateFailure() throws Exception {
+ for (boolean resume : new boolean[]{true, false}) {
+ LoadTaskService service = new LoadTaskService();
+ LoadTaskMapper mapper = Mockito.mock(LoadTaskMapper.class);
+ LoadTaskExecutor executor = Mockito.mock(LoadTaskExecutor.class);
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ LoadTask task = Mockito.spy(LoadTask.builder().id(9).loader(loader)
+ .options(new LoadOptions())
+ .status(resume ? LoadStatus.PAUSED : LoadStatus.STOPPED)
+ .build());
+ Mockito.doNothing().when(task).reconnect("current-token");
+ Mockito.when(mapper.selectById(9)).thenReturn(task);
+ IllegalStateException failure = new IllegalStateException("update failed");
+ IllegalStateException cleanup = new IllegalStateException("cleanup failed");
+ Mockito.when(mapper.updateById(task)).thenThrow(failure);
+ Mockito.doThrow(cleanup).when(loader).shutdown();
+ ConcurrentHashMap<Integer, LoadTask> running = new ConcurrentHashMap<>();
+ this.setField(service, "mapper", mapper);
+ this.setField(service, "taskExecutor", executor);
+ this.setField(service, "runningTaskContainer", running);
+
+ try {
+ if (resume) {
+ service.resume(9, "current-token");
+ } else {
+ service.retry(9, "current-token");
+ }
+ Assert.fail("Expected update failure");
+ } catch (IllegalStateException actual) {
+ Assert.assertSame(failure, actual);
+ Assert.assertEquals(1, actual.getSuppressed().length);
+ Assert.assertSame(cleanup, actual.getSuppressed()[0]);
+ }
+ Mockito.verify(loader).shutdown();
+ Assert.assertTrue(task.isFinished());
+ Assert.assertTrue(task.isCleanupFailed());
+ Assert.assertSame(task, running.get(9));
+ Assert.assertEquals(LoadStatus.FAILED, task.getStatus());
+ Assert.assertThrows(RuntimeException.class, () -> service.retry(9, "current-token"));
+ Assert.assertThrows(RuntimeException.class, () -> service.resume(9, "current-token"));
+ Mockito.verify(task).reconnect("current-token");
+ Mockito.verify(executor, Mockito.never()).execute(Mockito.any(), Mockito.any());
+ }
+ }
+
+ @Test
+ public void testResumeAndRetryCloseLoaderOnDispatchRejection() throws Exception {
+ for (boolean resume : new boolean[]{true, false}) {
+ LoadTaskService service = new LoadTaskService();
+ LoadTaskMapper mapper = Mockito.mock(LoadTaskMapper.class);
+ LoadTaskExecutor executor = Mockito.mock(LoadTaskExecutor.class);
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ LoadTask task = Mockito.spy(LoadTask.builder().id(9).loader(loader)
+ .options(new LoadOptions())
+ .status(resume ? LoadStatus.PAUSED : LoadStatus.STOPPED)
+ .build());
+ Mockito.doNothing().when(task).reconnect("current-token");
+ Mockito.when(mapper.selectById(9)).thenReturn(task);
+ Mockito.when(mapper.updateById(task)).thenReturn(1);
+ RejectedExecutionException rejected = new RejectedExecutionException("queue full");
+ Mockito.doThrow(rejected).when(executor).execute(Mockito.eq(task), Mockito.any());
+ ConcurrentHashMap<Integer, LoadTask> running = new ConcurrentHashMap<>();
+ this.setField(service, "mapper", mapper);
+ this.setField(service, "taskExecutor", executor);
+ this.setField(service, "runningTaskContainer", running);
+ PlatformTransactionManager transactions = Mockito.mock(PlatformTransactionManager.class);
+ Mockito.when(transactions.getTransaction(Mockito.any()))
+ .thenReturn(new SimpleTransactionStatus());
+ Mockito.when(mapper.update(Mockito.isNull(), Mockito.any())).thenReturn(1);
+ this.setField(service, "transactionManager", transactions);
+
+ Throwable actual = Assert.assertThrows(RejectedExecutionException.class, () -> {
+ if (resume) {
+ service.resume(9, "current-token");
+ } else {
+ service.retry(9, "current-token");
+ }
+ });
+ Assert.assertSame(rejected, actual);
+ Mockito.verify(loader).shutdown();
+ Assert.assertTrue(task.isFinished());
+ Assert.assertEquals(LoadStatus.FAILED, task.getStatus());
+ Assert.assertTrue(running.isEmpty());
+ Mockito.verify(mapper).update(Mockito.isNull(), Mockito.any());
+ }
+ }
+
+ @Test
+ public void testPreviewAlwaysClosesLoader() {
+ for (boolean success : new boolean[]{true, false}) {
+ try (MockedConstruction<HugeGraphLoader> construction = Mockito.mockConstruction(
+ HugeGraphLoader.class, (loader, ignored) -> {
+ if (success) {
+ Mockito.when(loader.load()).thenReturn(true);
+ } else {
+ Mockito.when(loader.load()).thenThrow(new IllegalStateException("load failed"));
+ }
+ })) {
+ new LoadTaskService().loader(new LoadOptions());
+ Mockito.verify(construction.constructed().get(0)).shutdown();
+ }
+ }
+ }
+
@Test
public void testListByIdsAppliesGraphScope() throws Exception {
LoadTaskService service = new LoadTaskService();
@@ -179,6 +516,100 @@ public class LoadTaskServiceTest {
Mockito.verify(executor).execute(Mockito.eq(task), Mockito.any());
}
+ @Test
+ public void testRetryAndResumeRejectUnfinishedOrCleanupFailedRuntime() throws Exception {
+ for (boolean resume : new boolean[]{false, true}) {
+ for (boolean cleanupFailed : new boolean[]{false, true}) {
+ LoadTaskService service = new LoadTaskService();
+ LoadTaskMapper mapper = Mockito.mock(LoadTaskMapper.class);
+ LoadTaskExecutor executor = Mockito.mock(LoadTaskExecutor.class);
+ LoadTask persisted = Mockito.mock(LoadTask.class);
+ Mockito.when(persisted.getStatus()).thenReturn(LoadStatus.FAILED);
+ Mockito.when(persisted.getOptions()).thenReturn(new LoadOptions());
+ Mockito.when(mapper.selectById(9)).thenReturn(persisted);
+ Mockito.when(mapper.updateById(persisted)).thenReturn(1);
+ LoadTask retained = LoadTask.builder().id(9).status(LoadStatus.FAILED)
+ .finished(cleanupFailed).cleanupFailed(cleanupFailed).build();
+ ConcurrentHashMap<Integer, LoadTask> running = new ConcurrentHashMap<>();
+ running.put(9, retained);
+ this.setField(service, "mapper", mapper);
+ this.setField(service, "taskExecutor", executor);
+ this.setField(service, "runningTaskContainer", running);
+
+ Assert.assertThrows(RuntimeException.class, () -> {
+ if (resume) {
+ service.resume(9, "current-session-token");
+ } else {
+ service.retry(9, "current-session-token");
+ }
+ });
+
+ Assert.assertSame(retained, running.get(9));
+ Mockito.verify(persisted, Mockito.never()).reconnect(Mockito.anyString());
+ Mockito.verify(mapper, Mockito.never()).updateById(Mockito.any(LoadTask.class));
+ Mockito.verify(executor, Mockito.never()).execute(Mockito.any(), Mockito.any());
+ }
+ }
+ }
+
+ @Test
+ public void testNativeResetFailureAndPurePauseStopCloseResources() throws Exception {
+ for (LoadStatus status : new LoadStatus[]{LoadStatus.PAUSED, LoadStatus.STOPPED}) {
+ for (boolean rejectReset : new boolean[]{false, true}) {
+ LoadTask task = LoadTask.builder().id(9).status(LoadStatus.RUNNING).lastDuration(0L).build();
+ try (NativeCleanupTestSupport nativeCleanup = new NativeCleanupTestSupport(
+ task, this.folder.newFile("native-" + status + rejectReset + ".json").toPath(),
+ rejectReset, status)) {
+ task.run();
+
+ Assert.assertEquals(rejectReset ? LoadStatus.FAILED : status, task.getStatus());
+ Assert.assertEquals(rejectReset, task.isCleanupFailed());
+ if (rejectReset) {
+ Assert.assertNotSame(nativeCleanup.original, nativeCleanup.shutdownFailure.get());
+ } else {
+ Assert.assertSame(nativeCleanup.original, nativeCleanup.shutdownFailure.get());
+ }
+ Assert.assertEquals(rejectReset ? "LOADING" : "NONE", nativeCleanup.mode());
+ nativeCleanup.assertResourcesClosed();
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testPauseAndStopKeepCleanupFailedRuntimeAndLoader() throws Exception {
+ for (boolean pause : new boolean[]{false, true}) {
+ LoadTaskService service = new LoadTaskService();
+ LoadTaskMapper mapper = Mockito.mock(LoadTaskMapper.class);
+ HugeGraphLoader loader = Mockito.mock(HugeGraphLoader.class);
+ LoadContext context = Mockito.mock(LoadContext.class);
+ Mockito.when(loader.context()).thenReturn(context);
+ Mockito.when(context.closed()).thenReturn(true);
+ LoadTask task = LoadTask.builder().id(9).status(LoadStatus.RUNNING).loader(loader).build();
+ Mockito.doAnswer(invocation -> {
+ task.setStatus(LoadStatus.FAILED);
+ task.setCleanupFailed(true);
+ task.setFinished(true);
+ return null;
+ }).when(context).stopLoading();
+ Mockito.when(mapper.updateById(task)).thenReturn(1);
+ ConcurrentHashMap<Integer, LoadTask> running = new ConcurrentHashMap<>();
+ running.put(9, task);
+ this.setField(service, "mapper", mapper);
+ this.setField(service, "runningTaskContainer", running);
+
+ if (pause) {
+ service.pause(9);
+ } else {
+ service.stop(9);
+ }
+
+ Assert.assertSame(task, running.get(9));
+ Assert.assertSame(loader, task.getLoader());
+ Assert.assertTrue(task.isCleanupFailed());
+ }
+ }
+
@Test
public void testLoadOptionsUseTokenWithoutCredentialPassword()
throws Exception {
diff --git a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/LoaderScopeControllerTest.java b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/LoaderScopeControllerTest.java
index 9c3901d08..a719e10a9 100644
--- a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/LoaderScopeControllerTest.java
+++ b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/LoaderScopeControllerTest.java
@@ -15,6 +15,7 @@
* License for the specific language governing permissions and limitations
* under the License.
*/
+
package org.apache.hugegraph.unit;
import java.util.Collections;
@@ -82,6 +83,8 @@ public class LoaderScopeControllerTest {
JobManagerService jobService = Mockito.mock(JobManagerService.class);
LoadTaskService taskService = Mockito.mock(LoadTaskService.class);
FileMapping mapping = FileMapping.builder().id(11).build();
+ Mockito.when(jobService.get("space-a", "graph-a", 7))
+ .thenReturn(JobManager.builder().id(7).graphSpace("space-a").graph("graph-a").build());
Mockito.when(service.listByJob("space-a", "graph-a", 7))
.thenReturn(Collections.singletonList(mapping));
Mockito.when(taskService.taskListByJob(7))
@@ -93,6 +96,7 @@ public class LoaderScopeControllerTest {
controller.clear("space-a", "graph-a", 7);
+ Mockito.verify(jobService).get("space-a", "graph-a", 7);
Mockito.verify(jobService).deleteMappings(
7, Collections.singletonList(mapping));
Mockito.verify(service, Mockito.never()).deleteDiskFile(mapping);
diff --git a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/UnitTestSuite.java b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/UnitTestSuite.java
index 07800995c..b745912fd 100644
--- a/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/UnitTestSuite.java
+++ b/hugegraph-hubble/hubble-be/src/test/java/org/apache/hugegraph/unit/UnitTestSuite.java
@@ -77,6 +77,7 @@ import org.junit.runners.Suite;
HubbleOptionsTest.class,
IngestControllerTest.class,
IngestTransactionIntegrationTest.class,
+ org.apache.hugegraph.service.load.JobAdmissionIntegrationTest.class,
LangChainControllerSecurityTest.class,
LegacyFacadeRemovalTest.class,
MessageSourceHandlerTest.class,
@@ -87,6 +88,7 @@ import org.junit.runners.Suite;
JobManagerServiceTest.class,
K8sTokenEndpointSecurityTest.class,
LoadTaskServiceTest.class,
+ LoadTaskCheckpointTest.class,
LoaderScopeControllerTest.class,
LoginAttemptGuardTest.class,
OltpAlgoControllerTest.class,
diff --git a/hugegraph-loader/README.md b/hugegraph-loader/README.md
index c3e69388d..bcd19f1cd 100644
--- a/hugegraph-loader/README.md
+++ b/hugegraph-loader/README.md
@@ -7,6 +7,10 @@
hugegraph-loader is a customizable command line utility for loading small to medium size graph datasets into the HugeGraph database from multiple data sources with various input formats.
+Embedded callers must call `shutdown()` after `load()`, including failed imports. Successful shutdown waits for this Loader's
+workers, closes its clients and unregisters its JVM shutdown hook. Closing a Loader whose `load()` was never called
+leaves the existing import checkpoint unchanged.
+
## 1. Features
- Multiple data sources, such as local file(path), HDFS file(path), MySQL
diff --git a/hugegraph-loader/src/main/java/org/apache/hugegraph/loader/HugeGraphLoader.java b/hugegraph-loader/src/main/java/org/apache/hugegraph/loader/HugeGraphLoader.java
index 42cccff7e..7325e50f6 100644
--- a/hugegraph-loader/src/main/java/org/apache/hugegraph/loader/HugeGraphLoader.java
+++ b/hugegraph-loader/src/main/java/org/apache/hugegraph/loader/HugeGraphLoader.java
@@ -20,9 +20,11 @@ package org.apache.hugegraph.loader;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
+import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -90,6 +92,9 @@ public final class HugeGraphLoader {
private final LoadMapping mapping;
private final TaskManager manager;
private final LoadOptions options;
+ private final Thread shutdownHook;
+ private volatile boolean loadStarted;
+ private final boolean saveProgress;
public static class InputTaskItem {
@@ -137,25 +142,32 @@ public final class HugeGraphLoader {
}
public HugeGraphLoader(LoadOptions options) {
- this(options, LoadMapping.of(options.file));
+ this(options, true);
+ }
+
+ /** Embedded callers may close resources without persisting parsed offsets. */
+ public HugeGraphLoader(LoadOptions options, boolean saveProgress) {
+ this(options, LoadMapping.of(options.file), saveProgress);
// Set concurrency
GlobalExecutorManager.setBatchThreadCount(options.batchInsertThreads);
GlobalExecutorManager.setSingleThreadCount(options.singleInsertThreads);
}
public HugeGraphLoader(LoadOptions options, LoadMapping mapping) {
+ this(options, mapping, true);
+ }
+
+ public HugeGraphLoader(LoadOptions options, LoadMapping mapping, boolean saveProgress) {
+ this.saveProgress = saveProgress;
this.context = new LoadContext(options);
this.options = options;
this.mapping = mapping;
this.manager = new TaskManager(this.context);
- this.addShutdownHook();
- }
-
- private void addShutdownHook() {
- Runtime.getRuntime().addShutdownHook(new Thread(() -> {
+ this.shutdownHook = new Thread(() -> {
LOG.info("Shutdown hook was triggered");
this.stopThenShutdown();
- }));
+ });
+ Runtime.getRuntime().addShutdownHook(this.shutdownHook);
}
public LoadContext context() {
@@ -203,6 +215,7 @@ public final class HugeGraphLoader {
}
public boolean load() {
+ this.loadStarted = true;
this.options.dumpParams();
try {
@@ -903,6 +916,20 @@ public final class HugeGraphLoader {
* TODO: How to distinguish load task finished normally or abnormally
*/
private synchronized void stopThenShutdown() {
+ try {
+ this.drainThenClose();
+ } finally {
+ if (this.context.closed() && this.shutdownHook != null) {
+ try {
+ Runtime.getRuntime().removeShutdownHook(this.shutdownHook);
+ } catch (IllegalStateException ignored) {
+ // JVM shutdown already started; this hook can finish closing.
+ }
+ }
+ }
+ }
+
+ private void drainThenClose() {
if (this.context.closed()) {
return;
}
@@ -927,9 +954,19 @@ public final class HugeGraphLoader {
if (tasksStopped) {
try {
this.context.unsetLoadingMode();
- } finally {
- this.context.close();
+ } catch (Throwable resetFailure) {
+ try {
+ this.context.close(this.loadStarted && this.saveProgress);
+ } catch (Throwable closeFailure) {
+ if (closeFailure instanceof Error && !(resetFailure instanceof Error)) {
+ suppressCleanupFailure(closeFailure, resetFailure);
+ throw (Error) closeFailure;
+ }
+ suppressCleanupFailure(resetFailure, closeFailure);
+ }
+ throw resetFailure;
}
+ this.context.close(this.loadStarted && this.saveProgress);
}
} finally {
if (interrupted) {
@@ -939,6 +976,28 @@ public final class HugeGraphLoader {
}
}
+ private static void suppressCleanupFailure(Throwable primary, Throwable secondary) {
+ Set<Throwable> seen = Collections.newSetFromMap(new IdentityHashMap<>());
+ if (!containsFailure(secondary, primary, seen)) {
+ primary.addSuppressed(secondary);
+ }
+ }
+
+ private static boolean containsFailure(Throwable failure, Throwable target, Set<Throwable> seen) {
+ if (failure == null || !seen.add(failure)) {
+ return false;
+ }
+ if (failure == target || containsFailure(failure.getCause(), target, seen)) {
+ return true;
+ }
+ for (Throwable suppressed : failure.getSuppressed()) {
+ if (containsFailure(suppressed, target, seen)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
private void cleanupEmptyProgress() {
Map<String, InputProgress> inputProgressMap = this.context.newProgress().inputProgress();
inputProgressMap.entrySet().removeIf(entry -> entry.getValue().loadedItems().isEmpty());
diff --git a/hugegraph-loader/src/main/java/org/apache/hugegraph/loader/executor/LoadContext.java b/hugegraph-loader/src/main/java/org/apache/hugegraph/loader/executor/LoadContext.java
index bbd935143..1023b967e 100644
--- a/hugegraph-loader/src/main/java/org/apache/hugegraph/loader/executor/LoadContext.java
+++ b/hugegraph-loader/src/main/java/org/apache/hugegraph/loader/executor/LoadContext.java
@@ -259,6 +259,10 @@ public final class LoadContext implements Cloneable {
}
public void close() {
+ this.close(true);
+ }
+
+ public void close(boolean saveProgress) {
if (this.closed) {
return;
}
@@ -284,7 +288,7 @@ public final class LoadContext implements Cloneable {
this.throwIfFailed();
this.newProgress.plusVertexLoaded(this.summary.vertexLoaded());
this.newProgress.plusEdgeLoaded(this.summary.edgeLoaded());
- if (this.client != null) {
+ if (this.client != null && saveProgress) {
try {
this.newProgress.write(this);
} catch (IOException e) {
diff --git a/hugegraph-loader/src/test/java/org/apache/hugegraph/loader/test/unit/LoadContextTest.java b/hugegraph-loader/src/test/java/org/apache/hugegraph/loader/test/unit/LoadContextTest.java
index 1684d411a..6841945c2 100644
--- a/hugegraph-loader/src/test/java/org/apache/hugegraph/loader/test/unit/LoadContextTest.java
+++ b/hugegraph-loader/src/test/java/org/apache/hugegraph/loader/test/unit/LoadContextTest.java
@@ -31,8 +31,12 @@ import java.nio.file.Files;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
+import java.util.IdentityHashMap;
import java.util.LinkedHashMap;
+import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import org.apache.hugegraph.driver.GraphsManager;
import org.apache.hugegraph.driver.HugeClient;
import org.apache.hugegraph.driver.HugeClientBuilder;
import org.apache.hugegraph.exception.ServerException;
@@ -65,6 +69,8 @@ import org.junit.rules.TemporaryFolder;
import com.sun.net.httpserver.HttpServer;
+import okhttp3.OkHttpClient;
+
public class LoadContextTest {
@Rule
@@ -677,6 +683,360 @@ public class LoadContextTest {
return loader;
}
+ @Test
+ public void testUnstartedShutdownPreservesCheckpoint() throws Exception {
+ try (ModeServer server = new ModeServer()) {
+ LoadOptions options = new LoadOptions();
+ options.host = "127.0.0.1";
+ options.port = server.server.getAddress().getPort();
+ options.file = this.folder.newFile("unstarted.json").getAbsolutePath();
+ options.incrementalMode = true;
+ File checkpoint = new File(LoadProgress.format(options, "2000-01-01"));
+ Files.createDirectories(checkpoint.toPath().getParent());
+ String saved = "{\"vertex_progress\":7,\"edge_progress\":0,\"input_progress\":{}}";
+ Files.writeString(checkpoint.toPath(), saved);
+ HugeGraphLoader loader = new HugeGraphLoader(options, new LoadMapping(Collections.emptyList()));
+ Assert.assertEquals(7L, loader.context().oldProgress().vertexLoaded());
+
+ loader.shutdown();
+
+ Assert.assertTrue(loader.context().closed());
+ Assert.assertEquals(saved, Files.readString(checkpoint.toPath()));
+ Assert.assertFalse(new File(LoadProgress.format(options, loader.context().timestamp())).exists());
+ }
+ }
+
+ @Test
+ public void testEmbeddedShutdownHookDoesNotReplaceExistingCheckpoint() throws Exception {
+ try (ModeServer server = new ModeServer()) {
+ LoadOptions options = new LoadOptions();
+ options.host = "127.0.0.1";
+ options.port = server.server.getAddress().getPort();
+ options.file = this.folder.newFile("embedded.json").getAbsolutePath();
+ options.incrementalMode = true;
+ File checkpoint = new File(LoadProgress.format(options, "2000-01-01"));
+ Files.createDirectories(checkpoint.toPath().getParent());
+ String saved = "{\"vertex_progress\":7,\"edge_progress\":0,\"input_progress\":{}}";
+ Files.writeString(checkpoint.toPath(), saved);
+ HugeGraphLoader loader = new HugeGraphLoader(options, new LoadMapping(Collections.emptyList()), false);
+ Assert.assertEquals(7L, loader.context().oldProgress().vertexLoaded());
+ set(loader, "loadStarted", true);
+ loader.context().newProgress().plusVertexLoaded(99);
+
+ ((Thread) field(loader, "shutdownHook")).run();
+
+ Assert.assertTrue(loader.context().closed());
+ Assert.assertEquals(saved, Files.readString(checkpoint.toPath()));
+ Assert.assertFalse(new File(LoadProgress.format(options, loader.context().timestamp())).exists());
+ }
+ }
+
+ @Test
+ public void testDefaultLoaderShutdownStillWritesProgress() throws Exception {
+ try (ModeServer server = new ModeServer()) {
+ LoadOptions options = new LoadOptions();
+ options.host = "127.0.0.1";
+ options.port = server.server.getAddress().getPort();
+ options.file = this.folder.newFile("standalone.json").getAbsolutePath();
+ HugeGraphLoader loader = new HugeGraphLoader(options, new LoadMapping(Collections.emptyList()));
+ set(loader, "loadStarted", true);
+ loader.context().newProgress().plusVertexLoaded(7);
+
+ loader.shutdown();
+
+ File checkpoint = new File(LoadProgress.format(options, loader.context().timestamp()));
+ Assert.assertTrue(loader.context().closed());
+ Assert.assertEquals(7L, LoadProgress.read(checkpoint).vertexLoaded());
+ }
+ }
+
+ @Test
+ public void testClosedLoaderUnregistersOwnHook() throws Exception {
+ for (boolean failed : new boolean[]{false, true}) {
+ try (ModeServer server = new ModeServer()) {
+ LoadOptions options = new LoadOptions();
+ options.host = "127.0.0.1";
+ options.port = server.server.getAddress().getPort();
+ options.file = this.folder.newFile("hook-" + failed + ".json").getAbsolutePath();
+ HugeGraphLoader loader = new HugeGraphLoader(options, new LoadMapping(Collections.emptyList()));
+ Field field = HugeGraphLoader.class.getDeclaredField("shutdownHook");
+ field.setAccessible(true);
+ Thread hook = (Thread) field.get(loader);
+ RuntimeException failure = new IllegalStateException("loading failed");
+ try {
+ if (failed) {
+ loader.context().failLoading(failure);
+ try {
+ loader.shutdown();
+ Assert.fail("shutdown must retain the loading failure");
+ } catch (RuntimeException actual) {
+ Assert.assertSame(failure, actual);
+ }
+ } else {
+ loader.shutdown();
+ }
+ Assert.assertTrue(loader.context().closed());
+ Assert.assertFalse(Runtime.getRuntime().removeShutdownHook(hook));
+ loader.shutdown();
+ } finally {
+ Runtime.getRuntime().removeShutdownHook(hook);
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testLoaderResetFailurePreservesIndependentAndLoadingFailures() throws Exception {
+ this.assertLoaderShutdownFailure(true);
+ }
+
+ @Test
+ public void testLoaderSuccessfulResetPreservesLoadingFailureIdentity() throws Exception {
+ this.assertLoaderShutdownFailure(false);
+ }
+
+ private void assertLoaderShutdownFailure(boolean rejectReset) throws Exception {
+ try (ModeServer server = new ModeServer()) {
+ server.mode = GraphMode.NONE;
+ LoadOptions options = new LoadOptions();
+ options.host = "127.0.0.1";
+ options.port = server.server.getAddress().getPort();
+ options.file = this.folder.newFile("reset-" + rejectReset + ".json").getAbsolutePath();
+ HugeGraphLoader loader = new HugeGraphLoader(options, new LoadMapping(Collections.emptyList()));
+ Thread hook = (Thread) field(loader, "shutdownHook");
+ LoadContext context = loader.context();
+ Object manager = field(loader, "manager");
+ ExecutorService batch = (ExecutorService) field(manager, "batchService");
+ ExecutorService single = (ExecutorService) field(manager, "singleService");
+ try {
+ // Start both actual native pools; shutdown must drain and terminate them.
+ batch.submit(() -> { }).get();
+ single.submit(() -> { }).get();
+ context.setLoadingMode();
+ OkHttpClient transport = (OkHttpClient) field(field(context.client(), "client"), "client");
+ RuntimeException original = new IllegalStateException("recorded loading failure");
+ context.failLoading(original);
+ if (rejectReset) {
+ server.failure = "Independent reset denied";
+ }
+ RuntimeException actual = null;
+ try {
+ loader.shutdown();
+ Assert.fail("Shutdown must report the recorded failure");
+ } catch (RuntimeException failure) {
+ actual = failure;
+ }
+ if (rejectReset) {
+ Assert.assertTrue("Reset failure must remain primary", actual instanceof LoadException);
+ Assert.assertNotSame(original, actual);
+ Assert.assertTrue(actual.getCause() instanceof ServerException);
+ Assert.assertEquals("Independent reset denied", actual.getCause().getMessage());
+ Assert.assertEquals(1, actual.getSuppressed().length);
+ Assert.assertSame(original, actual.getSuppressed()[0]);
+ Assert.assertEquals(GraphMode.LOADING, server.mode);
+ Assert.assertEquals(Boolean.TRUE, field(context, "loadingModeSet"));
+ } else {
+ Assert.assertSame(original, actual);
+ Assert.assertEquals(0, actual.getSuppressed().length);
+ Assert.assertEquals(GraphMode.NONE, server.mode);
+ Assert.assertEquals(Boolean.FALSE, field(context, "loadingModeSet"));
+ }
+ assertAcyclic(actual, Collections.newSetFromMap(new IdentityHashMap<>()));
+ Assert.assertTrue(context.closed());
+ Assert.assertTrue(batch.isTerminated());
+ Assert.assertTrue(single.isTerminated());
+ Assert.assertTrue(transport.dispatcher().executorService().isShutdown());
+ Assert.assertEquals(0, transport.connectionPool().connectionCount());
+ Assert.assertFalse(Runtime.getRuntime().removeShutdownHook(hook));
+ Assert.assertEquals(1, server.modeReads);
+ Assert.assertEquals(2, server.modeWrites);
+ loader.shutdown();
+ Assert.assertEquals(1, server.modeReads);
+ Assert.assertEquals(2, server.modeWrites);
+ } finally {
+ server.failure = null;
+ try {
+ loader.shutdown();
+ } catch (RuntimeException ignored) {
+ // Retain the original assertion while closing the actual native resources.
+ }
+ Runtime.getRuntime().removeShutdownHook(hook);
+ }
+ }
+ }
+
+ @Test
+ public void testLoaderPromotesFatalCloseOverOrdinaryResetFailure() throws Exception {
+ this.assertFatalLoaderShutdown(false, false, false);
+ }
+
+ @Test
+ public void testLoaderPreservesFatalResetAndRecordedLoadingFailure() throws Exception {
+ this.assertFatalLoaderShutdown(true, false, false);
+ }
+
+ @Test
+ public void testLoaderDoesNotSuppressFatalFailureOnItself() throws Exception {
+ this.assertFatalLoaderShutdown(true, true, false);
+ }
+
+ @Test
+ public void testLoaderDoesNotCreateFatalResetCauseCycle() throws Exception {
+ this.assertFatalLoaderShutdown(true, false, true);
+ }
+
+ private void assertFatalLoaderShutdown(boolean fatalReset, boolean sameError,
+ boolean causeCycle) throws Exception {
+ try (ModeServer server = new ModeServer()) {
+ server.mode = GraphMode.NONE;
+ LoadOptions options = new LoadOptions();
+ options.host = "127.0.0.1";
+ options.port = server.server.getAddress().getPort();
+ options.file = this.folder.newFile().getAbsolutePath();
+ HugeGraphLoader loader = new HugeGraphLoader(options, new LoadMapping(Collections.emptyList()));
+ LoadContext context = loader.context();
+ Thread hook = (Thread) field(loader, "shutdownHook");
+ // Use real SDK clients; only the fatal accessor / post-close boundary is controlled.
+ BoundaryClient primary = new BoundaryClient(server);
+ HugeClient secondary = server.client();
+ context.client().close();
+ set(context, "client", primary);
+ set(context, "indirectClient", secondary);
+ OkHttpClient transport = (OkHttpClient) field(field(primary, "client"), "client");
+ OkHttpClient secondaryTransport = (OkHttpClient) field(field(secondary, "client"), "client");
+ Object manager = field(loader, "manager");
+ ExecutorService batch = (ExecutorService) field(manager, "batchService");
+ ExecutorService single = (ExecutorService) field(manager, "singleService");
+ Error fatal = new AssertionError("fatal SDK cleanup boundary");
+ RuntimeException original = causeCycle ? new IllegalStateException("loading failed", fatal) :
+ new IllegalStateException("recorded loading failure");
+ try {
+ batch.submit(() -> { }).get();
+ single.submit(() -> { }).get();
+ context.setLoadingMode();
+ if (fatalReset) {
+ primary.resetFailure = fatal;
+ if (!sameError) {
+ context.failLoading(original);
+ }
+ } else {
+ server.failure = "Independent reset denied";
+ }
+ primary.closeFailure = !fatalReset || sameError ? fatal : null;
+ Throwable actual = null;
+ try {
+ loader.shutdown();
+ } catch (RuntimeException | Error failure) {
+ actual = failure;
+ }
+ Assert.assertSame("First fatal Error must escape with its identity", fatal, actual);
+ if (causeCycle || sameError) {
+ Assert.assertEquals(0, fatal.getSuppressed().length);
+ } else {
+ Assert.assertEquals(1, fatal.getSuppressed().length);
+ if (fatalReset) {
+ Assert.assertSame(original, fatal.getSuppressed()[0]);
+ } else {
+ Assert.assertTrue(fatal.getSuppressed()[0] instanceof LoadException);
+ Assert.assertEquals("Independent reset denied",
+ fatal.getSuppressed()[0].getCause().getMessage());
+ }
+ }
+ if (fatalReset && !sameError) {
+ try {
+ context.throwIfFailed();
+ Assert.fail("Context must retain the original recorded failure");
+ } catch (RuntimeException failure) {
+ Assert.assertSame(original, failure);
+ }
+ }
+ assertAcyclic(fatal, Collections.newSetFromMap(new IdentityHashMap<>()));
+ Assert.assertTrue(context.closed());
+ Assert.assertTrue(batch.isTerminated());
+ Assert.assertTrue(single.isTerminated());
+ Assert.assertTrue(transport.dispatcher().executorService().isShutdown());
+ Assert.assertEquals(0, transport.connectionPool().connectionCount());
+ Assert.assertTrue(secondaryTransport.dispatcher().executorService().isShutdown());
+ Assert.assertEquals(0, secondaryTransport.connectionPool().connectionCount());
+ Assert.assertEquals(1, primary.closeCalls);
+ Assert.assertEquals(GraphMode.LOADING, server.mode);
+ Assert.assertEquals(Boolean.TRUE, field(context, "loadingModeSet"));
+ Assert.assertFalse(Runtime.getRuntime().removeShutdownHook(hook));
+ int reads = server.modeReads;
+ int writes = server.modeWrites;
+ loader.shutdown();
+ Assert.assertEquals(reads, server.modeReads);
+ Assert.assertEquals(writes, server.modeWrites);
+ Assert.assertEquals(1, primary.closeCalls);
+ } finally {
+ server.failure = null;
+ primary.resetFailure = null;
+ primary.closeFailure = null;
+ try {
+ loader.shutdown();
+ } catch (RuntimeException | Error ignored) {
+ // Preserve assertions while tearing down the real native resources.
+ }
+ primary.close();
+ secondary.close();
+ Runtime.getRuntime().removeShutdownHook(hook);
+ }
+ }
+ }
+
+ private static class BoundaryClient extends HugeClient {
+
+ private Error resetFailure;
+ private Error closeFailure;
+ private int closeCalls;
+
+ private BoundaryClient(ModeServer server) {
+ super(HugeClient.builder("http://127.0.0.1:" + server.server.getAddress().getPort(), "hugegraph"));
+ }
+
+ @Override
+ public GraphsManager graphs() {
+ if (this.resetFailure != null) {
+ throw this.resetFailure;
+ }
+ return super.graphs();
+ }
+
+ @Override
+ public void close() {
+ this.closeCalls++;
+ super.close();
+ if (this.closeFailure != null) {
+ throw this.closeFailure;
+ }
+ }
+ }
+
+ private static void assertAcyclic(Throwable failure, Set<Throwable> path) {
+ if (failure == null) {
+ return;
+ }
+ Assert.assertTrue("Failure graph must not contain cycles", path.add(failure));
+ assertAcyclic(failure.getCause(), path);
+ for (Throwable suppressed : failure.getSuppressed()) {
+ assertAcyclic(suppressed, path);
+ }
+ path.remove(failure);
+ }
+
+ private static Object field(Object target, String name) throws Exception {
+ for (Class<?> type = target.getClass(); type != null; type = type.getSuperclass()) {
+ try {
+ Field field = type.getDeclaredField(name);
+ field.setAccessible(true);
+ return field.get(target);
+ } catch (NoSuchFieldException ignored) {
+ // SDK transport fields are declared in its existing superclass.
+ }
+ }
+ throw new NoSuchFieldException(name);
+ }
+
private static class ModeServer implements AutoCloseable {
private final HttpServer server;