diff --git a/.github/workflows/pd-store-ci.yml b/.github/workflows/pd-store-ci.yml index 231751e303..47a8abaefc 100644 --- a/.github/workflows/pd-store-ci.yml +++ b/.github/workflows/pd-store-ci.yml @@ -431,6 +431,11 @@ jobs: mvn test -pl hugegraph-store/hg-store-test -am \ -P store-core-test -Djacoco.sessionId=store-core-test + - name: Run server lifecycle test + run: | + mvn test -pl hugegraph-store/hg-store-test -am \ + -P store-server-test -Djacoco.sessionId=store-server-test + - name: Generate aggregate coverage report run: | mvn verify -pl hugegraph-store/hg-store-test -am -P jacoco \ @@ -450,6 +455,8 @@ jobs: --require-test-report \ "$TEST_REPORT_DIR/TEST-org.apache.hugegraph.store.core.CoreSuiteTest.xml" \ --require-test-report \ + "$TEST_REPORT_DIR/TEST-org.apache.hugegraph.store.service.ServerSuiteTest.xml" \ + --require-test-report \ "hugegraph-store/hg-store-node/target/surefire-reports/TEST-org.apache.hugegraph.store.business.StoredRowIngressTest.xml" \ --require-covered-group hg-store-common \ --require-covered-group hg-store-client \ @@ -460,6 +467,7 @@ jobs: --require-session store-rocksdb-test \ --require-session store-raftcore-test \ --require-session store-core-test \ + --require-session store-server-test \ --require-session store-node-test \ "$REPORT_FILE" \ hg-store-grpc hg-store-common hg-store-client hg-store-rocksdb hg-store-core diff --git a/README.md b/README.md index 00cacdb561..de996d16bf 100644 --- a/README.md +++ b/README.md @@ -100,7 +100,7 @@ HugeGraph supports both **standalone** and **distributed** deployments: See the [backend evolution guide](hugegraph-server/README.md#backend-evolution-and-compatibility) for lifecycle and historical compatibility guidance. -To use ToplingDB with the RocksDB backend, follow the [three-step switching guide](docs/toplingdb.md). +To use ToplingDB with the RocksDB backend, follow the [three-step switching guide](docs/storage-lifecycle.md#storage-provider-selection-and-toplingdb). ### Deployment Mode Comparison @@ -317,9 +317,7 @@ gremlin> :> g.V().limit(5) For comprehensive documentation, visit the [HugeGraph Documentation](https://hugegraph.apache.org/docs/). -See [standalone RocksDB snapshot recovery](docs/rocksdb-recovery.md) before restoring data or mounting store directories. - -For an existing deployment, read the [RocksDB upgrade guidance](docs/rocksdb-upgrade.md) before upgrading the storage runtime. +Read the [storage lifecycle and recovery guide](docs/storage-lifecycle.md) before restoring data, mounting store directories or upgrading the storage runtime. diff --git a/docs/rocksdb-recovery.md b/docs/rocksdb-recovery.md deleted file mode 100644 index c8ade0a86a..0000000000 --- a/docs/rocksdb-recovery.md +++ /dev/null @@ -1,42 +0,0 @@ -# Standalone RocksDB snapshot recovery - -Standalone RocksDB snapshot restore retains its source checkpoint until the replacement is installed and reopened. Before changing data, it records the source checkpoint and WAL location in a sibling `.resume-pending` file, together with the restore operation and checkpoint generation. After an I/O failure or process interruption, a subsequent HugeGraph open retries the interrupted installation before native recovery. Missing sources, incomplete metadata or a changed WAL configuration stop opening instead of replaying uncertain data. - -On Linux, recovery synchronizes the marker parent directory before replacing live data and after removing the marker, before consuming the checkpoint. A directory-sync failure stops that step and preserves the checkpoint; after installation has succeeded, an unlink-sync failure leaves the restored live data intact. Other platforms retain their existing filesystem durability behavior. This marker ordering does not guarantee that every copied data or WAL file survives a power loss. - -Recovery metadata publication requires same-directory hard-link support in the data directory's parent filesystem. The complete temporary record is written and forced before its final name is published without overwriting an existing marker. A pre-publication I/O or unsupported-link failure leaves live data and the checkpoint untouched; resolve the filesystem or permission fault, reopen normally and reissue restore. Existing damaged or old-format markers are still preserved and rejected. - -Quiesce graph requests, background tasks and open query iterators before starting a restore. The recovery lease coordinates database opens and file installation; it does not drain active application work. - -Preserve the checkpoint, pending marker and configured paths after a failure. Restore access or free space, then retry normal startup with the same runtime and configuration. Do not delete the pending marker to bypass the guard. Successful native reopening clears it and then attempts checkpoint cleanup, preserving the existing consume-on-success behavior. - -A sibling `.resume-lock` serializes cooperating opens and restores. The OS lock remains held through native close and reopen, and is released when the database finally closes. The lock file remains on disk; its presence alone does not prove an active owner. Do not delete it to force another opener through. Within one JVM, physical-file ownership also covers path aliases. An unrelated database can still open while another database is closing. A failed native or lock-descriptor close retains recovery ownership and reports the failure; retrying close cannot bypass it. Preserve the recovery files and stop the process before retrying startup after such a failure. Older binaries and unrelated writers do not honor this protocol and must not access these directories concurrently. - -Mount the parent data root, such as `rocksdb-data`, rather than an individual store directory such as `data/g`. Java opening rejects a store that is itself a volume mount, including same-filesystem bind mounts on Linux. Such aliases can hide the sibling guards, and directory replacement cannot replace a mount point. This Java check is per database: it does not promise a pre-JVM scan of every graph or that no sibling database has initialized before another open fails. Keep mount layout unchanged throughout startup and recovery. Checkpoint/live-data overlap and a checkpoint containing the WAL directory are rejected by physical directory identity, including parent bind aliases; separate path strings do not make overlapping trees safe. Recovery refuses the operation if a required physical identity cannot be determined. - -Independent WAL, WAL inside data, and data inside a WAL root use in-place log replacement. Preserve WAL symlink configuration across retries. Local fault tests cover interrupted operations and reopening; they do not establish power-cut durability. This mechanism does not provide an atomic whole-graph restore or an HStore multi-partition snapshot protocol. - -Each non-consuming restore creates a unique checkpoint directory. Pending recovery binds the original directory and immutable SST filesystem identities, plus checkpoint metadata and WAL contents, so replacing a checkpoint at the same path is rejected before data replacement. Keep the original checkpoint in place: copying it elsewhere and back may change its identity. Filesystems without file identities use content hashes instead. Unrelated writers must not modify checkpoint files; this protocol does not protect against malicious in-place changes to immutable SST files. - -WAL staging uses a sibling `.resume-staging-` directory recorded by the pending operation. Retired WAL remains there until successful native reopening; repeated attempts reuse that same directory. Cleanup removes only the current operation's directory, never all matching prefixes. Unknown staging directories and markers from older formats remain preserved for diagnosis rather than being guessed or silently upgraded. - -After a Store-level restore failure closes its native owner, reopen the database to retry the recorded pending operation. Calling restore again on the already closed Store does not reopen it. Recovery retries the recorded unique checkpoint; it does not create a new non-consuming copy. - -Non-consuming Store restore first creates all temporary checkpoints, then restores each database. If either phase fails, it attempts to remove only the exact temporary copies created by that invocation, including an unfinished checkpoint created before its path could be returned. Cleanup requires the same database recovery ownership and a confirmed absent pending marker. Any existing or unreadable marker preserves the copy and emits a diagnostic; lock contention also preserves it. Original user checkpoints and unrelated temporary directories are never part of this cleanup. This does not collect historical orphan directories or copies left by a process killed before writing its pending marker. Restore remains non-atomic across databases. - -Pending records include a SHA-256 checksum over all metadata property names and values, excluding the checksum itself. Recovery verifies it before reconstructing any WAL alias, then uses the filesystem's actual canonical WAL path to validate the destination. This detects incomplete or changed metadata; it is not authentication against someone who can edit the record and recompute its checksum, and it does not freeze the surrounding filesystem. Records from older versions without this checksum are rejected and preserved, not silently upgraded. Finish an outstanding restore with its producing version before upgrading; otherwise preserve the checkpoint, marker and paths for diagnosis. Do not add a checksum manually to bypass this check. - -Before creating a non-consuming copy, restore holds the database recovery ownership and requires a confirmed absent pending marker. An existing or unreadable marker rejects the new copy before opening its checkpoint source; reopen the database to retry the recorded recovery instead. Repeated requests therefore do not allocate more UUID copies for an already pending database. If a later database is pending, copies already created for earlier databases are still cleaned by the same invocation's failure handling. - - -The database directory itself must not be a symbolic link, including a dangling link; put a stable alias on its parent directory instead. This keeps the -pending marker and recovery lock at the same physical database identity while its contents are replaced. Parent-directory symbolic links and bind mounts -remain supported. Data and WAL aliases of the same physical directory are treated as one directory, so checkpoint WAL records are not deleted as a separate source. - -Simple WAL symbolic links remain supported, including an external stable link followed by `..`. Recovery rejects a WAL path combining `..` with a link -inside the data directory, or a symbolic-link target chain containing another link inside data: data replacement would delete an unrecorded part of that -path. Reconfigure such layouts to use a stable external WAL directory before restoring. Rejection happens before writing new pending metadata or deleting data. -If marker existence cannot be determined, opening fails with the recovery material preserved rather than bypassing a possible pending restore. - -After a failed restore has closed the native handle, release the remaining Store sessions with normal close calls before reopening the same Store to retry -its recorded pending checkpoint. A native close failure still retains its recovery ownership and blocks unsafe competing opens. diff --git a/docs/rocksdb-upgrade.md b/docs/rocksdb-upgrade.md deleted file mode 100644 index 5101b1b24c..0000000000 --- a/docs/rocksdb-upgrade.md +++ /dev/null @@ -1,37 +0,0 @@ -# Standard RocksDB runtime upgrade - -The standard runtime aligns RocksDB JNI to 8.10.2. Validate existing data before rolling out this dependency change. Keep a restorable backup made by the old runtime; do not assume that data opened or modified by 8.10.2 can be reopened by an older runtime. Rollback requires restoring the backup with the matching old application and dependencies, rather than replacing the JAR in an upgraded data directory. - -All components use JRaft 1.3.14: Server upgrades from 1.3.11, and PD/Store from 1.3.13. Older JRaft versions call compressed-cache methods removed from RocksDB 8.x, preventing Raft storage initialization, including Server when `raft.mode=true`. Treat the JNI and JRaft changes as one service upgrade and validate existing Raft logs with the complete new package. - -## RocksDB compatibility test - -The permanent [compatibility test](../hugegraph-server/hugegraph-test/src/test/rocksdb-compatibility/run.sh) uses Maven effective POMs to compare the actual Server, PD and Store RocksDB versions. CI compares PR base/head or push before/after, fails if a version cannot be resolved, and runs only changed version pairs, deduplicated across components. POM or JRaft changes alone do not run the JNI test. No historical version matrix is maintained. - -Run from two complete source checkouts with Java 11+, Maven and Python 3: - -```bash -bash hugegraph-server/hugegraph-test/src/test/rocksdb-compatibility/run.sh \ - /path/to/before /path/to/after /tmp/new-rocksdb-compatibility-run -``` - -Each pair creates synthetic SST and synchronous WAL data with the old JNI, reads/modifies/deletes/adds data with the new JNI, then verifies it in a separate JVM. Original data, commands, JAR hashes, native versions and failure evidence are retained in the new output directory. This test never opens business data, runs at application startup, or enters a production binary distribution. When changing this test or its CI, verify both changed/unchanged version selection and run the changed version pairs. Service-created data, Raft logs, cluster upgrades and Topling acceptance remain separate checks. - -## Store monitoring changes - -The JNI upgrade removes registration of these unavailable ticker constants: - -- `BLOCK_CACHE_INDEX_BYTES_EVICT`, `BLOCK_CACHE_FILTER_BYTES_EVICT` -- `NO_FILE_CLOSES`, `RATE_LIMIT_DELAY_MILLIS`, `NO_ITERATORS` -- `NUMBER_FILTERED_DELETES` -- `BLOCK_CACHE_COMPRESSED_MISS`, `BLOCK_CACHE_COMPRESSED_HIT`, - `BLOCK_CACHE_COMPRESSED_ADD`, `BLOCK_CACHE_COMPRESSED_ADD_FAILURES` -- `WRITE_TIMEDOUT` - -It also removes these histogram constants: - -- `STALL_L0_SLOWDOWN_COUNT`, `STALL_MEMTABLE_COMPACTION_COUNT`, - `STALL_L0_NUM_FILES_COUNT` -- `HARD_RATE_LIMIT_DELAY_COUNT`, `SOFT_RATE_LIMIT_DELAY_COUNT` - -Store meter names use `rocks.stats.` plus the lowercase constant name. For histograms, the removed meters include `.max`, `.mean`, `.min`, `.summary` (with its quantile tags), `.summary.sum` and `.summary.count`. Exporters may normalize names further. Update dashboards and alerts that reference these series; their disappearance is not a zero reading. No replacement series or semantic equivalence is asserted here. See [RocksDBMetricsConst](../hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/metrics/RocksDBMetricsConst.java) and [meter registration](../hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/metrics/RocksDBMetrics.java). diff --git a/docs/storage-lifecycle.md b/docs/storage-lifecycle.md new file mode 100644 index 0000000000..374a4572b3 --- /dev/null +++ b/docs/storage-lifecycle.md @@ -0,0 +1,580 @@ +# Storage lifecycle, recovery and runtime compatibility + +This guide covers Store RPC ownership and shutdown coordination, standalone RocksDB +snapshot recovery, and storage runtime upgrades. Use it when changing request cleanup, +deploying storage components, recovering an interrupted restore, or upgrading JNI. + +Store RPC lifecycle and standalone RocksDB recovery are separate contracts. The +standalone recovery protocol does not provide HStore multi-partition recovery or an +atomic restore across all databases in a graph. + +## Guide map + +- [RPC and Scan ownership](#rpc-and-scan-ownership): final batches, half-close, + cancellation, iterator cleanup and callback dispatch. +- [Shutdown coordination](#shutdown-coordination): admission, worker drain and + application callbacks before database teardown. +- [Standalone snapshot recovery](#standalone-snapshot-recovery): pending records, + checkpoint identity, WAL paths and safe retry. +- [Storage provider selection and ToplingDB](#storage-provider-selection-and-toplingdb): + JNI preparation, provider configuration and persistent paths. +- [Runtime compatibility and upgrades](#runtime-compatibility-and-upgrades): backup, + JNI/JRaft compatibility, the compatibility runner and monitoring changes. + +## RPC and Scan ownership + +### Half-close, feedback and cancellation + +Request half-close is separate from transport cancellation. +An aggregate query finishes work already supported by feedback credit after half-close; +a later transport cancellation or deadline still cancels its workers. Graph partition +scans deliver their permitted pages and final batch, or return an explicit error when +half-close leaves insufficient feedback credit. Batch-scan RPCs accept one initial +query. + +### Client response and iterator ownership + +Response parsing runs outside the client state lock. Accepted final batches remain +visible across response completion; errors and iterator close discard unpublished data. +Early query close cancels the transport without waiting for a blocked request send. +Normal half-close stays serialized with sends. Closing a composite iterator closes all +started children once. A closed multi-partition iterator is terminal: later probes do +not reopen another partition, including after a child cleanup failure. + +Accepting a response refreshes the watchdog before parsing starts. Completion preserves +an accepted batch while it is being parsed; cancellation and errors can discard it. +The state lock protects these transitions rather than response extraction. + +### Server worker and cleanup ownership + +Each scan/query owns its workers and iterators until cleanup finishes. Iterator cleanup +failures remain visible, including automatic native iterator close and empty-partition +selection. Final aggregate success is sent only after plan and iterator cleanup +succeeds. + +Cancellation signals the currently owned reader before waiting for its iterator lock. +Publishing, interrupting and clearing that thread use the same ownership lock so a +reused executor thread is not mistaken for the cancelled reader. Half-close does not +interrupt the reader. Queued cancelled tasks still execute their cleanup/accounting +paths, and iterator close waits for an active read to leave the iterator. + +Unary query and count requests share the query admission, worker and cleanup owner +registry. Count keeps partition work parallel on the service executor. Iterator cleanup +failure prevents successful completion and remains visible to request drain. + +One-shot scans run on the gRPC caller thread. Transport cancellation and Store shutdown +signal that caller directly, so interruptible waits can exit without waiting for a +queued call-executor callback. Interruption is cooperative: a RocksDB native operation +does not guarantee that it will return on Java thread interruption. The call remains +registered until its read and iterator cleanup finish; cancellation does not permit +another thread to close a native iterator that is still in use. + +### Callback executor configuration + +The gRPC application callback queue is always unbounded (`Integer.MAX_VALUE`), matching +its existing default. `thread.pool.grpc.queue` remains accepted for configuration +compatibility but no longer limits this callback queue. Bounded callback dispatch can +reject cancellation or completion callbacks in gRPC 1.55.3, leaving request cleanup +registered after the transport has closed. `thread.pool.grpc.core` and `.max` still +configure the same executor; an unbounded queue normally keeps dispatch at the core +thread count. Bound expensive scan/query work and admission at their worker queues; +sustained overload can otherwise increase callback backlog and memory use. + +## Shutdown coordination + +RPC and Scan drain establishes completion of their request owners. It does not by itself +establish completion of unrelated background maintenance jobs. + +The RPC boundary stops admission to every Store RPC before cancelling active scans and +queries. It runs queued request cleanup and waits for terminal application callbacks +before Spring destroys their databases. This reuses the same per-call owners; it does +not create a second shutdown registry. Interrupted cleanup waits and native teardown +preserve interruption without abandoning these owners. + +The Store context-close listener stops RPC, scan and query admission, requests transport +shutdown, and drains scan/query workers before waiting for terminal gRPC callbacks and +scan cleanup. Queued cancelled scans still run their cleanup paths. + +Shutdown terminates unfinished aggregate responses with `UNAVAILABLE`, rather than +reporting normal completion for a truncated result. Terminal delivery is serialized +with batch delivery and does not release the request's cleanup ownership: shutdown +still waits for the workers and iterators even if the response callback fails. + +Transport termination alone is insufficient evidence that an application callback or +native iterator owner has finished. Preserve the application cleanup barrier when +changing the server framework, executor configuration or shutdown ordering. + +Drain waits retain ownership until cleanup finishes and restore the shutdown thread's +interrupted status. They must not close databases under a worker solely because a wait +was interrupted. A blocked owner can keep shutdown waiting; investigate that owner +rather than bypassing the database lifetime guard. + +The engine closes partition databases after their partition engines shut down, then +releases graph database and schema resources. Keep this ordering when modifying Raft or +database teardown. Request drain and database release have different responsibilities; +reuse their existing owners instead of adding a parallel registry. + +## Standalone snapshot recovery + +### Installation and pending metadata + +Standalone RocksDB snapshot restore retains its source checkpoint until the replacement +is installed and reopened. Before changing data, it records the source checkpoint and +WAL location in a sibling `.resume-pending` file, together with the restore +operation and checkpoint generation. After an I/O failure or process interruption, a +subsequent HugeGraph open retries the interrupted installation before native recovery. +Missing sources, incomplete metadata or a changed WAL configuration stop opening instead +of replaying uncertain data. + +### Filesystem durability boundaries + +On Linux, recovery synchronizes the marker parent directory before replacing live data +and after removing the marker, before consuming the checkpoint. A directory-sync failure +stops that step and preserves the checkpoint; after installation has succeeded, an +unlink-sync failure leaves the restored live data intact. Other platforms retain their +existing filesystem durability behavior. This marker ordering does not guarantee that +every copied data or WAL file survives a power loss. + +### Publishing the pending record + +Recovery metadata publication requires same-directory hard-link support in the data +directory's parent filesystem. The complete temporary record is written and forced +before its final name is published without overwriting an existing marker. A +pre-publication I/O or unsupported-link failure leaves live data and the checkpoint +untouched; resolve the filesystem or permission fault, reopen normally and reissue +restore. Existing damaged or old-format markers are still preserved and rejected. + +### Preparing application work + +Quiesce graph requests, background tasks and open query iterators before starting a +restore. The recovery lease coordinates database opens and file installation; it does +not drain active application work. + +### Retrying an interrupted restore + +Preserve the checkpoint, pending marker and configured paths after a failure. Restore +access or free space, then retry normal startup with the same runtime and configuration. +Do not delete the pending marker to bypass the guard. Successful native reopening clears +it and then attempts checkpoint cleanup, preserving the existing consume-on-success +behavior. + +The normal retry sequence is: + +1. Preserve the original checkpoint, pending record, WAL location and mount layout. +2. Correct the access, space or filesystem fault without changing their identities. +3. Release the remaining Store sessions with normal close calls when the failed restore + has already closed the native handle. +4. Reopen with the same runtime and configuration so startup retries the recorded operation. +5. Allow successful native reopening to clear the record and perform checkpoint cleanup. + +If native close itself failed, use the failure handling below instead of attempting a +competing open. + +### Recovery ownership and native close failures + +A sibling `.resume-lock` serializes cooperating opens and restores. The OS +lock remains held through native close and reopen, and is released when the database +finally closes. The lock file remains on disk; its presence alone does not prove an +active owner. Do not delete it to force another opener through. Within one JVM, +physical-file ownership also covers path aliases. An unrelated database can still open +while another database is closing. A failed native or lock-descriptor close retains +recovery ownership and reports the failure; retrying close cannot bypass it. Preserve +the recovery files and stop the process before retrying startup after such a failure. +Older binaries and unrelated writers do not honor this protocol and must not access +these directories concurrently. + +### Mount layout and physical path identity + +Mount the parent data root, such as `rocksdb-data`, rather than an individual store +directory such as `data/g`. Java opening rejects a store that is itself a volume mount, +including same-filesystem bind mounts on Linux. Such aliases can hide the sibling +guards, and directory replacement cannot replace a mount point. This Java check is per +database: it does not promise a pre-JVM scan of every graph or that no sibling database +has initialized before another open fails. Keep mount layout unchanged throughout +startup and recovery. Checkpoint/live-data overlap and a checkpoint containing the WAL +directory are rejected by physical directory identity, including parent bind aliases; +separate path strings do not make overlapping trees safe. Recovery refuses the operation +if a required physical identity cannot be determined. + +### WAL replacement and recovery scope + +Independent WAL, WAL inside data, and data inside a WAL root use in-place log +replacement. Preserve WAL symlink configuration across retries. Local fault tests cover +interrupted operations and reopening; they do not establish power-cut durability. This +mechanism does not provide an atomic whole-graph restore or an HStore multi-partition +snapshot protocol. + +### Checkpoint generation and identity + +Each non-consuming restore creates a unique checkpoint directory. Pending recovery binds +the original directory and immutable SST filesystem identities, plus checkpoint metadata +and WAL contents, so replacing a checkpoint at the same path is rejected before data +replacement. Keep the original checkpoint in place: copying it elsewhere and back may +change its identity. Filesystems without file identities use content hashes instead. +Unrelated writers must not modify checkpoint files; this protocol does not protect +against malicious in-place changes to immutable SST files. + +### WAL staging and retained material + +WAL staging uses a sibling `.resume-staging-` directory recorded +by the pending operation. Retired WAL remains there until successful native reopening; +repeated attempts reuse that same directory. Cleanup removes only the current +operation's directory, never all matching prefixes. Unknown staging directories and +markers from older formats remain preserved for diagnosis rather than being guessed or +silently upgraded. + +### Store-level retry + +After a Store-level restore failure closes its native owner, reopen the database to +retry the recorded pending operation. Calling restore again on the already closed Store +does not reopen it. Recovery retries the recorded unique checkpoint; it does not create +a new non-consuming copy. + +### Non-consuming restore and temporary-copy cleanup + +Non-consuming Store restore first creates all temporary checkpoints, then restores each +database. If either phase fails, it attempts to remove only the exact temporary copies +created by that invocation, including an unfinished checkpoint created before its path +could be returned. Cleanup requires the same database recovery ownership and a confirmed +absent pending marker. Any existing or unreadable marker preserves the copy and emits a +diagnostic; lock contention also preserves it. Original user checkpoints and unrelated +temporary directories are never part of this cleanup. This does not collect historical +orphan directories or copies left by a process killed before writing its pending marker. +Restore remains non-atomic across databases. + +### Metadata integrity and runtime changes + +Pending records include a SHA-256 checksum over all metadata property names and values, +excluding the checksum itself. Recovery verifies it before reconstructing any WAL alias, +then uses the filesystem's actual canonical WAL path to validate the destination. This +detects incomplete or changed metadata; it is not authentication against someone who can +edit the record and recompute its checksum, and it does not freeze the surrounding +filesystem. Records from older versions without this checksum are rejected and +preserved, not silently upgraded. Finish an outstanding restore with its producing +version before upgrading; otherwise preserve the checkpoint, marker and paths for +diagnosis. Do not add a checksum manually to bypass this check. + +### Pending recovery blocks new checkpoint copies + +Before creating a non-consuming copy, restore holds the database recovery ownership and +requires a confirmed absent pending marker. An existing or unreadable marker rejects the +new copy before opening its checkpoint source; reopen the database to retry the recorded +recovery instead. Repeated requests therefore do not allocate more UUID copies for an +already pending database. If a later database is pending, copies already created for +earlier databases are still cleaned by the same invocation's failure handling. + +### Database directory aliases + +The database directory itself must not be a symbolic link, including a dangling link; +put a stable alias on its parent directory instead. This keeps the pending marker and +recovery lock at the same physical database identity while its contents are replaced. +Parent-directory symbolic links and bind mounts remain supported. Data and WAL aliases +of the same physical directory are treated as one directory, so checkpoint WAL records +are not deleted as a separate source. + +### WAL symbolic links + +Simple WAL symbolic links remain supported, including an external stable link followed +by `..`. Recovery rejects a WAL path combining `..` with a link inside the data +directory, or a symbolic-link target chain containing another link inside data: data +replacement would delete an unrecorded part of that path. Reconfigure such layouts to +use a stable external WAL directory before restoring. Rejection happens before writing +new pending metadata or deleting data. If marker existence cannot be determined, opening +fails with the recovery material preserved rather than bypassing a possible pending +restore. + +## Runtime compatibility and upgrades + +The standard runtime uses RocksDB JNI 8.10.2. Validate existing data before upgrading +the storage runtime. Keep a restorable backup made by the old runtime; do not +assume that data opened or modified by 8.10.2 can be reopened by an older runtime. +Rollback requires restoring the backup with the matching old application and +dependencies, rather than replacing the JAR in an upgraded data directory. + +All components use JRaft 1.3.14. Older JRaft versions call compressed-cache methods removed from RocksDB 8.x, preventing +Raft storage initialization, including Server when `raft.mode=true`. Treat the JNI and +JRaft changes as one service upgrade and validate existing Raft logs with the complete +new package. + +Use the complete matching application package for an upgrade or backup-based rollback. +Validate service-created data and Raft logs separately from the synthetic JNI check. A +successful JNI compatibility run does not establish that a cluster rolling upgrade or +downgrade is safe. + +### RocksDB compatibility runner + +The permanent [compatibility +test](../hugegraph-server/hugegraph-test/src/test/rocksdb-compatibility/run.sh) uses +Maven effective POMs to compare the actual Server, PD and Store RocksDB versions. CI +compares PR base/head or push before/after, fails if a version cannot be resolved, and +runs only changed version pairs, deduplicated across components. POM or JRaft changes +alone do not run the JNI test. No historical version matrix is maintained. + +Run from two complete source checkouts with Java 11+, Maven and Python 3: + +```bash +bash hugegraph-server/hugegraph-test/src/test/rocksdb-compatibility/run.sh \ + /path/to/before /path/to/after /tmp/new-rocksdb-compatibility-run +``` + +Each pair creates synthetic SST and synchronous WAL data with the old JNI, +reads/modifies/deletes/adds data with the new JNI, then verifies it in a separate JVM. +Original data, commands, JAR hashes, native versions and failure evidence are retained +in the new output directory. This test never opens business data. It does not run at application startup or enter a +production binary distribution. When changing this test or its CI, +verify both changed/unchanged version selection and run the changed version pairs. +Service-created data, Raft logs, cluster upgrades and Topling acceptance remain separate +checks. + +The runner requires a new output directory and refuses environment overrides that inject +Java options or native libraries. Before running, unset `JAVA_TOOL_OPTIONS`, +`JDK_JAVA_OPTIONS`, `_JAVA_OPTIONS`, `LD_PRELOAD` and `DYLD_INSERT_LIBRARIES`. Keep each +run's synthetic data and evidence separate; the script does not reuse a previous output +directory. + +### Monitoring compatibility + +The JNI upgrade removes registration of these unavailable ticker constants: + +- `BLOCK_CACHE_INDEX_BYTES_EVICT`, `BLOCK_CACHE_FILTER_BYTES_EVICT` +- `NO_FILE_CLOSES`, `RATE_LIMIT_DELAY_MILLIS`, `NO_ITERATORS` +- `NUMBER_FILTERED_DELETES` +- `BLOCK_CACHE_COMPRESSED_MISS`, `BLOCK_CACHE_COMPRESSED_HIT`, + `BLOCK_CACHE_COMPRESSED_ADD`, `BLOCK_CACHE_COMPRESSED_ADD_FAILURES` +- `WRITE_TIMEDOUT` + +It also removes these histogram constants: + +- `STALL_L0_SLOWDOWN_COUNT`, `STALL_MEMTABLE_COMPACTION_COUNT`, + `STALL_L0_NUM_FILES_COUNT` +- `HARD_RATE_LIMIT_DELAY_COUNT`, `SOFT_RATE_LIMIT_DELAY_COUNT` + +Store meter names use `rocks.stats.` plus the lowercase constant name. For histograms, +the removed meters include `.max`, `.mean`, `.min`, `.summary` (with its quantile tags), +`.summary.sum` and `.summary.count`. Exporters may normalize names further. Update +dashboards and alerts that reference these series; their disappearance is not a zero +reading. No replacement series or semantic equivalence is asserted here. See +[RocksDBMetricsConst](../hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/metrics/RocksDBMetricsConst.java) +and [meter +registration](../hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/metrics/RocksDBMetrics.java). + +## Storage provider selection and ToplingDB + +HugeGraph defaults to standard RocksDB. Server, PD and Store can explicitly select a +prepared ToplingDB runtime using their normal distributions. Prepare and configure each +component separately; this guide covers the shipped launch scripts, not custom IDE or +embedded launchers. + +| Component | Business provider setting | Start command | +| --- | --- | --- | +| Server with the RocksDB backend | `rocksdb.provider=topling` in each graph properties file | `bin/start-hugegraph.sh` | +| PD | `rocksdb.provider: topling` in `conf/application.yml` | `bin/start-hugegraph-pd.sh` | +| Store | `rocksdb.provider: topling` in `conf/application-pd.yml` | `bin/start-hugegraph-store.sh` | + +For a distributed graph, Server uses `backend=hstore`; PD and Store own its persistent +storage. A standalone Server keeps `backend=rocksdb`. The launcher selection and the +business provider must agree. A mismatch fails before opening that component's database. + +### Standalone Server setup + +![Switch RocksDB to ToplingDB](images/switch-rocksdb-to-toplingdb.png) + +Use a stopped Server distribution on Linux x86_64 with a Java 17 JDK. If your +distribution lacks these scripts, [build it from this source](#build-and-prepare) first. +Prepare the runtime once, select it, then initialize and start the Server. The graph +APIs and normal start/stop commands stay the same; Topling's native options are +configured in `conf/toplingdb.yaml`. Use new, empty, service-writable data and WAL +directories outside the distribution for this example. + +**1. Prepare the runtime.** Supply your ToplingDB JNI JAR and its trusted SHA-256: + +```bash +export TOPLING_JNI_JAR=/absolute/path/to/rocksdbjni-topling.jar +export TOPLING_JNI_SHA256='' +bash bin/prepare-topling.sh +``` + +**2. Select Topling.** Edit `conf/graphs/hugegraph.properties` and every other graph loaded from the configured graphs directory: + +```properties +backend=rocksdb +rocksdb.provider=topling +rocksdb.data_path=/srv/hugegraph/topling/server/data +rocksdb.wal_path=/srv/hugegraph/topling/server/wal +``` + +Create those external directories with permissions for the service account. The paths +above are examples; choose persistent locations for your deployment. + +**3. Initialize and start.** From the same shell, select the matching JNI runtime and initialize the new store once: + +```bash +export TOPLINGDB_ROCKSDB_PROVIDER=topling +bash bin/init-store.sh +bash bin/start-hugegraph.sh +``` + +Keep `TOPLINGDB_ROCKSDB_PROVIDER=topling` in the service environment for subsequent +starts; initialization is only needed for a new store. For PD and Store, prepare each +component separately, then use its YAML provider setting and start command from the +table above. A Server connected to distributed storage keeps `backend=hstore`; PD and +Store select their own storage runtime. + +### Build and prepare + +Build the normal distributions from the repository root using Java 17 and Maven 3.6.3 or +later: + +```bash +mvn clean package -Dmaven.test.skip=true -Dmaven.javadoc.skip=true +``` + +Obtain a trusted ToplingDB JNI JAR and its SHA-256 through your approved artifact +channel. The runtime is not downloaded by the launcher. It must implement the +`TOPLINGDB_EASY_MIGRATE_CONF` configuration entry point. Preparation requires a Java 17 +JDK, Linux x86_64, `sha256sum`, `unzip`, `od`, `ldd` and GNU `mv`. Native dependencies depend on the selected JNI JAR; builds requiring glibc 2.38 or +newer and libaio need those libraries on the deployment host. Preparation rejects +unresolved dependencies. + +In each stopped, unpacked component directory: + +```bash +export TOPLING_JNI_JAR=/absolute/path/to/rocksdbjni-topling.jar +export TOPLING_JNI_SHA256='' +bash bin/prepare-topling.sh +``` + +The script verifies the copied JAR's hash, the Topling Java marker, one Linux x86_64 JNI +entry, the ELF architecture and native dependencies. Before installation, it opens and +closes a disposable database with the selected JNI and checks that a probe configuration +changes the persisted write-buffer size. A marker or matching hash alone is +insufficient: builds without the EasyMigrate configuration entry point are rejected. The +probe validates the configuration entry point, not existing-data compatibility or every +setting in your production profile. Its checksum receipt binds the probe to the +installed JAR and native library. Launchers reject missing or stale receipts; prepare a +fresh distribution when upgrading from an older preparation script. It installs +`topling/rocksdbjni.jar`, its matching native library and optional web resources outside +the standard `lib` directory. On Ubuntu with libaio's t64 name, the compatibility +symlink stays inside that component's runtime directory. The published runtime +directories use 0755 and regular files use 0644, so a different service account can read +the prepared JNI. That account also needs traversal permission on the component's parent +directories; preparation does not change those parents. + +Preparation refuses an existing `topling` directory. To replace a runtime, stop the +service and prepare a fresh distribution; it never changes database files. + +### Select and start + +Set the provider in the component's business configuration. For PD and Store, edit the +existing `rocksdb` mapping, retaining its other options: + +```yaml +rocksdb: + provider: topling +``` + +For a standalone Server, edit every selected graph configuration: + +```properties +backend=rocksdb +rocksdb.provider=topling +``` + +Use separate, initially empty data directories for validation. This integration does not +convert or certify an existing database for a different provider. + +For persistent use, create service-writable directories outside source checkouts, build +outputs and unpacked distributions, and configure all applicable paths: + +| Component | Configuration file | Setting | Example external path | +| --- | --- | --- | --- | +| Server data | Graph properties | `rocksdb.data_path` | `/srv/hugegraph/topling/server/data` | +| Server WAL | Graph properties | `rocksdb.wal_path` | `/srv/hugegraph/topling/server/wal` | +| PD metadata | `conf/application.yml` | `pd.data-path` | `/srv/hugegraph/topling/pd` | +| Store data | `conf/application.yml` | `app.data-path` | `/srv/hugegraph/topling/store/data` | +| Store Raft | `conf/application.yml` | `app.raft-path` | `/srv/hugegraph/topling/store/raft` | + +Set both Server data and WAL paths. Retaining the default distribution-relative WAL path +can lose that directory when rebuilding or replacing the distribution. Preserve the +external data, WAL and Raft roots across runtime updates. Start each component with +explicit runtime selection: + +```bash +export TOPLINGDB_ROCKSDB_PROVIDER=topling +# Run the start command for this component from the table above. +``` + +The component uses its own `conf/toplingdb.yaml`. An explicit +`TOPLINGDB_EASY_MIGRATE_CONF=/absolute/path/config.yaml` overrides that native +configuration. The shipped profile disables the native HTTP server and keeps +`memtable_as_log_index=false`, which is required by the Java write path. Tune the native +profile for your machine before production use; it does not replace the component's +business configuration. + +Server's `init-store.sh` and `dump-store.sh` use the same explicit runtime selection. PD +and Store put the prepared JNI JAR before their Spring Boot dependencies. Standard +launch remains the default and does not require Topling files or preparation tools. +Launch rejects inherited recognizable RocksDB JNI preloads in default/standard mode and +competing JNI in Topling mode. Unrelated allocator or tracing preloads remain intact; +switching from a previously selected Topling runtime removes only that runtime. Store +adds verified jemalloc without replacing or duplicating caller preloads; the selected +JNI stays first. Do not preload a second RocksDB JNI alongside the selected runtime. + +### Validate and stop + +Before adopting a JAR, validate each component's real writes, reads, normal stop and +restart using a dedicated database directory. Confirm both the Java classes and the +loaded JNI library come from the prepared runtime; a standard-JNI unit test or a +successful startup alone is insufficient. For distributed storage, validate PD and Store +together, then perform a Server `hstore` graph operation. + +At request/task completion, explicitly committed writes are preserved; unfinished writes +are rolled back and cached backend leases are released. Shared schema and element caches +retain their invalidation listeners until the graph closes. + +Stop incoming work and use the component's normal stop script. PD drains its scheduled +metadata work, joins Raft and closes the metadata database and native options. Do not +use a forced kill as evidence of normal resource cleanup. + +To return to standard RocksDB, stop the component, unset `TOPLINGDB_ROCKSDB_PROVIDER` +and set its business provider back to `rocksdb`. Use the corresponding standard-provider +data directory; do not test a provider switch by reusing the other provider's database. + +## Implementation navigation + +Use the current implementation and its tests when reviewing a lifetime change. These +entry points separate client response ownership, server worker ownership, callback +admission and standalone recovery installation. + +### RPC and shutdown entry points + +- [CommonKvStreamObserver](../hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserver.java) + owns client stream state and response delivery. +- [KvBatchScanner](../hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvBatchScanner.java) + sends batch scan feedback and distinguishes early cancellation from half-close. +- [ScanResponseObserver](../hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/ScanResponseObserver.java) + coordinates scan readers, cancellation and iterator cleanup. +- [GrpcShutdownBarrier](../hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GrpcShutdownBarrier.java) + guards admission and waits for terminal application callbacks. +- [ContextClosedListener](../hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/listener/ContextClosedListener.java) + orders RPC shutdown, worker drain and callback completion. +- [GRpcServerConfig](../hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GRpcServerConfig.java) + configures application callback dispatch. +- [HgStoreEngine](../hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java) + orders partition shutdown and database release. + +### Recovery and compatibility entry points + +- [RocksDBSnapshotRestore](../hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBSnapshotRestore.java) + implements pending recovery and filesystem installation. +- [RocksDBSessions](../hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBSessions.java) + manages database sessions and recovery ownership. +- [Snapshot restore tests](../hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBSnapshotRestoreTest.java) + exercise restore failures, path identities and retry behavior. +- [Compatibility runner](../hugegraph-server/hugegraph-test/src/test/rocksdb-compatibility/run.sh) + selects changed JNI version pairs and runs isolated synthetic data checks. +- [CI guidance](ci.md) describes how to select relevant checks for a change. + +Keep operational contracts here and deployment topology in the [Store operations +guide](../hugegraph-store/docs/operations-guide.md) and [deployment +guide](../hugegraph-store/docs/deployment-guide.md). diff --git a/docs/toplingdb.md b/docs/toplingdb.md deleted file mode 100644 index 2e12590c97..0000000000 --- a/docs/toplingdb.md +++ /dev/null @@ -1,149 +0,0 @@ -# Switch RocksDB to ToplingDB - -HugeGraph defaults to standard RocksDB. Server, PD and Store can explicitly select a prepared ToplingDB runtime using their normal distributions. -Prepare and configure each component separately; this guide covers the shipped launch scripts, not custom IDE or embedded launchers. - -| Component | Business provider setting | Start command | -| --- | --- | --- | -| Server with the RocksDB backend | `rocksdb.provider=topling` in each graph properties file | `bin/start-hugegraph.sh` | -| PD | `rocksdb.provider: topling` in `conf/application.yml` | `bin/start-hugegraph-pd.sh` | -| Store | `rocksdb.provider: topling` in `conf/application-pd.yml` | `bin/start-hugegraph-store.sh` | - -For a distributed graph, Server uses `backend=hstore`; PD and Store own its persistent storage. A standalone Server keeps `backend=rocksdb`. -The launcher selection and the business provider must agree. A mismatch fails before opening that component's database. - -## Quick start: standalone Server - -![Switch RocksDB to ToplingDB](images/switch-rocksdb-to-toplingdb.png) - -Use a stopped Server distribution on Linux x86_64 with a Java 17 JDK. If your distribution lacks these scripts, [build it from this source](#build-and-prepare) first. -Prepare the runtime once, select it, then initialize and start the Server. -The graph APIs and normal start/stop commands stay the same; Topling's native options are configured in `conf/toplingdb.yaml`. -Use new, empty, service-writable data and WAL directories outside the distribution for this example. - -**1. Prepare the runtime.** Supply your ToplingDB JNI JAR and its trusted SHA-256: - -```bash -export TOPLING_JNI_JAR=/absolute/path/to/rocksdbjni-topling.jar -export TOPLING_JNI_SHA256='' -bash bin/prepare-topling.sh -``` - -**2. Select Topling.** Edit `conf/graphs/hugegraph.properties` and every other graph loaded from the configured graphs directory: - -```properties -backend=rocksdb -rocksdb.provider=topling -rocksdb.data_path=/srv/hugegraph/topling/server/data -rocksdb.wal_path=/srv/hugegraph/topling/server/wal -``` - -Create those external directories with permissions for the service account. The paths above are examples; choose persistent locations for your deployment. - -**3. Initialize and start.** From the same shell, select the matching JNI runtime and initialize the new store once: - -```bash -export TOPLINGDB_ROCKSDB_PROVIDER=topling -bash bin/init-store.sh -bash bin/start-hugegraph.sh -``` - -Keep `TOPLINGDB_ROCKSDB_PROVIDER=topling` in the service environment for subsequent starts; initialization is only needed for a new store. -For PD and Store, prepare each component separately, then use its YAML provider setting and start command from the table above. -A Server connected to distributed storage keeps `backend=hstore`; PD and Store select their own storage runtime. - -## Build and prepare - -Build the normal distributions from the repository root using Java 17 and Maven 3.6.3 or later: - -```bash -mvn clean package -Dmaven.test.skip=true -Dmaven.javadoc.skip=true -``` - -Obtain a trusted ToplingDB JNI JAR and its SHA-256 through your approved artifact channel. The runtime is not downloaded by the launcher. It must implement the `TOPLINGDB_EASY_MIGRATE_CONF` configuration entry point. -Preparation requires a Java 17 JDK, Linux x86_64, `sha256sum`, `unzip`, `od`, `ldd` and GNU `mv`. The currently exercised JNI requires glibc 2.38 or newer and libaio. -Native dependency requirements depend on the selected JAR; preparation rejects unresolved dependencies. - -In each stopped, unpacked component directory: - -```bash -export TOPLING_JNI_JAR=/absolute/path/to/rocksdbjni-topling.jar -export TOPLING_JNI_SHA256='' -bash bin/prepare-topling.sh -``` - -The script verifies the copied JAR's hash, the Topling Java marker, one Linux x86_64 JNI entry, the ELF architecture and native dependencies. -Before installation, it opens and closes a disposable database with the selected JNI and checks that a probe configuration changes the persisted write-buffer size. -A marker or matching hash alone is insufficient: incompatible builds, including the older `frocksdbjni-8.10.2-topling-1.0.jar` without EasyMigrate support, are rejected. -The probe validates the configuration entry point, not existing-data compatibility or every setting in your production profile. -Its checksum receipt binds the probe to the installed JAR and native library. Launchers reject missing or stale receipts; prepare a fresh distribution when upgrading from an older preparation script. -It installs `topling/rocksdbjni.jar`, its matching native library and optional web resources outside the standard `lib` directory. -On Ubuntu with libaio's t64 name, the compatibility symlink stays inside that component's runtime directory. -The published runtime directories use 0755 and regular files use 0644, so a different service account can read the prepared JNI. -That account also needs traversal permission on the component's parent directories; preparation does not change those parents. - -Preparation refuses an existing `topling` directory. To replace a runtime, stop the service and prepare a fresh distribution; it never changes database files. - -## Select and start - -Set the provider in the component's business configuration. For PD and Store, edit the existing `rocksdb` mapping, retaining its other options: - -```yaml -rocksdb: - provider: topling -``` - -For a standalone Server, edit every selected graph configuration: - -```properties -backend=rocksdb -rocksdb.provider=topling -``` - -Use separate, initially empty data directories for validation. This integration does not convert or certify an existing database for a different provider. - -For persistent use, create service-writable directories outside source checkouts, build outputs and unpacked distributions, and configure all applicable paths: - -| Component | Configuration file | Setting | Example external path | -| --- | --- | --- | --- | -| Server data | Graph properties | `rocksdb.data_path` | `/srv/hugegraph/topling/server/data` | -| Server WAL | Graph properties | `rocksdb.wal_path` | `/srv/hugegraph/topling/server/wal` | -| PD metadata | `conf/application.yml` | `pd.data-path` | `/srv/hugegraph/topling/pd` | -| Store data | `conf/application.yml` | `app.data-path` | `/srv/hugegraph/topling/store/data` | -| Store Raft | `conf/application.yml` | `app.raft-path` | `/srv/hugegraph/topling/store/raft` | - -Set both Server data and WAL paths. Retaining the default distribution-relative WAL path can lose that directory when rebuilding or replacing the distribution. -Preserve the external data, WAL and Raft roots across runtime updates. -Start each component with explicit runtime selection: - -```bash -export TOPLINGDB_ROCKSDB_PROVIDER=topling -# Run the start command for this component from the table above. -``` - -The component uses its own `conf/toplingdb.yaml`. An explicit `TOPLINGDB_EASY_MIGRATE_CONF=/absolute/path/config.yaml` overrides that native configuration. -The shipped profile disables the native HTTP server and keeps `memtable_as_log_index=false`, which is required by the Java write path. -Tune the native profile for your machine before production use; it does not replace the component's business configuration. - -Server's `init-store.sh` and `dump-store.sh` use the same explicit runtime selection. PD and Store put the prepared JNI JAR before their Spring Boot dependencies. -Standard launch remains the default and does not require Topling files or preparation tools. -Launch rejects inherited recognizable RocksDB JNI preloads in default/standard mode and competing JNI in Topling mode. -Unrelated allocator or tracing preloads remain intact; switching from a previously selected Topling runtime removes only that runtime. Store adds verified jemalloc without replacing or duplicating caller preloads; the selected JNI stays first. -Do not preload a second RocksDB JNI alongside the selected runtime. - -## Validate and stop - -Before adopting a JAR, validate each component's real writes, reads, normal stop and restart using a dedicated database directory. -Confirm both the Java classes and the loaded JNI library come from the prepared runtime; a standard-JNI unit test or a successful startup alone is insufficient. -For distributed storage, validate PD and Store together, then perform a Server `hstore` graph operation. - -At request/task completion, explicitly committed writes are preserved; unfinished writes are rolled back and cached backend leases are released. -Shared schema and element caches retain their invalidation listeners until the graph closes. - -Stop incoming work and use the component's normal stop script. PD drains its scheduled metadata work, joins Raft and closes the metadata database and native options. -Do not use a forced kill as evidence of normal resource cleanup. - -To return to standard RocksDB, stop the component, unset `TOPLINGDB_ROCKSDB_PROVIDER` and set its business provider back to `rocksdb`. -Use the corresponding standard-provider data directory; do not test a provider switch by reusing the other provider's database. - -Dedicated Topling distributions, Docker/Compose packaging and broader lifecycle/recovery changes are separate follow-ups. diff --git a/hugegraph-server/hugegraph-dist/src/assembly/travis/test-check-jacoco-report.sh b/hugegraph-server/hugegraph-dist/src/assembly/travis/test-check-jacoco-report.sh index 59b4caf32c..bf7fd64530 100755 --- a/hugegraph-server/hugegraph-dist/src/assembly/travis/test-check-jacoco-report.sh +++ b/hugegraph-server/hugegraph-dist/src/assembly/travis/test-check-jacoco-report.sh @@ -430,8 +430,8 @@ def validation_command(job): def reports_for_option(job, option): pattern = re.escape(option) + ( - r'\s+\\?\s*"\$TEST_REPORT_DIR/' - r'(TEST-[A-Za-z0-9_.]+SuiteTest[.]xml)"' + r'\s+\\?\s*"[^"]*/' + r'(TEST-[A-Za-z0-9_.]+Test[.]xml)"' ) return set(re.findall(pattern, validation_command(job))) @@ -492,9 +492,12 @@ assert_order(store_job, [ "-P store-client-test -Djacoco.sessionId=store-client-test", "-P store-rocksdb-test -Djacoco.sessionId=store-rocksdb-test", "-P store-raftcore-test -Djacoco.sessionId=store-raftcore-test", + "-P store-core-test -Djacoco.sessionId=store-core-test", + "-P store-server-test -Djacoco.sessionId=store-server-test", "mvn verify", "--require-session store-common-test", "--require-session store-client-test", "--require-session store-rocksdb-test", - "--require-session store-raftcore-test", "codecov/codecov-action", + "--require-session store-raftcore-test", "--require-session store-core-test", + "--require-session store-server-test", "codecov/codecov-action", ]) assert store_job.count("mvn clean") == 1 assert "hugegraph-store/hg-store-test/target/site/jacoco/jacoco.xml" in store_job @@ -504,7 +507,7 @@ assert "mvn verify -pl hugegraph-store/hg-store-test -am -P jacoco \\ " \ "-DskipTests -Deditorconfig.skip=true -ntp" in " ".join(store_job.split()) assert selected_profiles(store_job, "store") == { "store-common-test", "store-client-test", "store-rocksdb-test", - "store-raftcore-test", "store-core-test", + "store-raftcore-test", "store-core-test", "store-server-test", } assert reports_for_option(store_job, "--require-test-report") == { "TEST-org.apache.hugegraph.store.common.CommonSuiteTest.xml", @@ -512,8 +515,11 @@ assert reports_for_option(store_job, "--require-test-report") == { "TEST-org.apache.hugegraph.store.rocksdb.RocksDbSuiteTest.xml", "TEST-org.apache.hugegraph.store.raftcore.RaftSuiteTest.xml", "TEST-org.apache.hugegraph.store.core.CoreSuiteTest.xml", + "TEST-org.apache.hugegraph.store.service.ServerSuiteTest.xml", + "TEST-org.apache.hugegraph.store.business.StoredRowIngressTest.xml", } assert not reports_for_option(store_job, "--require-suite-report") +assert values_for_option(store_job, "--require-session") == (selected_profiles(store_job, "store") | {"store-node-test"}) assert values_for_option(store_job, "--require-covered-group") == { "hg-store-common", "hg-store-client", "hg-store-rocksdb", "hg-store-core", } diff --git a/hugegraph-store/README.md b/hugegraph-store/README.md index 58943bb94b..d038c9d52d 100644 --- a/hugegraph-store/README.md +++ b/hugegraph-store/README.md @@ -22,7 +22,7 @@ HugeGraph Store is a distributed storage backend for HugeGraph that provides hig ### Technology Stack -- **Storage Engine**: RocksDB 8.10.2 (see the [upgrade guidance](../docs/rocksdb-upgrade.md)) +- **Storage Engine**: RocksDB 8.10.2 (see the [upgrade guidance](../docs/storage-lifecycle.md#runtime-compatibility-and-upgrades)) - **Consensus Protocol**: Apache JRaft (Ant Financial's Raft implementation) - **RPC Framework**: gRPC + Protocol Buffers - **Deployment**: Java 17 (currently supported release), Docker/Kubernetes support @@ -114,6 +114,15 @@ The assembled distribution will be available at: hugegraph-store/apache-hugegraph-store-/lib/hg-store-node-.jar ``` +The Node module keeps a regular JAR for development dependencies and attaches its +Spring Boot executable as `hg-store-node--exec.jar`. The distribution +uses the executable JAR with the filename shown above, so the start script is +unchanged. To run the Store server tests through the package phase: + +```bash +mvn clean package -pl hugegraph-store/hg-store-test -am -Pstore-server-test +``` + ### Configuration Extract the distribution package and edit `conf/application.yml`: @@ -390,6 +399,9 @@ For Docker and Kubernetes deployment details, see [Deployment Guide](docs/deploy ## Documentation +See [RPC and Scan lifecycle](../docs/storage-lifecycle.md#rpc-and-scan-ownership) for request ownership, +cancellation and feedback semantics. + Comprehensive documentation for HugeGraph Store: | Documentation | Description | diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvBatchScanner.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvBatchScanner.java index 5b711844ca..ab1a021ff3 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvBatchScanner.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvBatchScanner.java @@ -40,6 +40,7 @@ import org.apache.hugegraph.store.grpc.common.ScanOrderType; import org.apache.hugegraph.store.grpc.stream.HgStoreStreamGrpc; import org.apache.hugegraph.store.grpc.stream.KvStream; +import org.apache.hugegraph.store.grpc.stream.ScanCancelRequest; import org.apache.hugegraph.store.grpc.stream.ScanReceiptRequest; import org.apache.hugegraph.store.grpc.stream.ScanStreamBatchReq; @@ -167,12 +168,16 @@ public void dataArrived(Supplier> supplier) throws * Data reception ended */ public void dataComplete() { - close(); + close(false); } // Flow is closed @Override public void close() { + close(true); + } + + private void close(boolean cancel) { try { if (notifier.unregisterScanner(this) < 0) { notifier.dataArrived(this, NO_DATA); // Task finished, wake up the queue @@ -183,6 +188,12 @@ public void close() { synchronized (this.sender) { try { if (running) { + running = false; + if (cancel) { + sender.onNext(ScanStreamBatchReq.newBuilder() + .setHeader(Header.newBuilder().setGraph(graphName)) + .setCancelRequest(ScanCancelRequest.getDefaultInstance()).build()); + } sender.onCompleted(); } } catch (Exception e) { @@ -344,7 +355,7 @@ public void onCompleted() { private void close() { if (scanner != null) { - scanner.close(); + scanner.dataComplete(); } } } diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserver.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserver.java index 7abcbf88d0..21da9a39fe 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserver.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserver.java @@ -18,17 +18,16 @@ package org.apache.hugegraph.store.client.query; import java.util.Iterator; +import java.util.Objects; import java.util.concurrent.BlockingQueue; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; import java.util.function.Function; import org.apache.hugegraph.store.client.type.HgStoreClientException; import io.grpc.stub.StreamObserver; -import lombok.Data; import lombok.Setter; import lombok.extern.slf4j.Slf4j; @@ -47,7 +46,7 @@ public class CommonKvStreamObserver implements StreamObserver { private Consumer requestSender; /** - * Handling the case that server has no results, close channel + * Close the request stream: true means serialized half-close; false means transport cancellation. */ @Setter private Consumer transferComplete; @@ -61,19 +60,18 @@ public class CommonKvStreamObserver implements StreamObserver { */ private final Function stateWatcher; - /** - * It can be ended by the client to stop receiving redundant data. - */ - private final AtomicBoolean closed = new AtomicBoolean(false); + private final Object stateLock = new Object(); + // Only outgoing sends and normal half-close share this lock. Cancellation bypasses it. + private final Object requestLock = new Object(); + private volatile ResultState terminal; + private volatile int parsing; + private volatile boolean closed; + private volatile long current = System.nanoTime(); + private String queryId; @Setter private long timeout = 1800 * 1000; - /** - * Monitor internal state - */ - private final ResultStateWatcher watcher = new ResultStateWatcher(); - public CommonKvStreamObserver(Function> valueExtractor, Function stateWatcher) { this.queue = new LinkedBlockingQueue<>(); @@ -85,55 +83,89 @@ public CommonKvStreamObserver(Function> valueExtractor, * Send requests */ public void sendRequest() { - if (!isServerFinished() && !closed.get()) { + synchronized (this.requestLock) { + synchronized (this.stateLock) { + if (this.closed || this.terminal != null) { + return; + } + this.current = System.nanoTime(); + } this.requestSender.accept(true); - this.watcher.setState(ResultState.WAITING); } } public boolean isServerFinished() { - return this.watcher.getState() == ResultState.FINISHED - || this.watcher.getState() == ResultState.ERROR; + ResultState state = this.terminal; + return this.closed || state == ResultState.ERROR || + (state == ResultState.FINISHED && this.parsing == 0); } @Override public void onNext(R value) { - watcher.setState(ResultState.INNER_BUSY); + synchronized (this.stateLock) { + if (this.closed || this.terminal != null) { + return; + } + this.parsing++; + this.current = System.nanoTime(); + } + + Iterator iterator = null; + ResultState result = ResultState.ERROR; try { - var state = stateWatcher.apply(value); - log.debug("observer state: {}", state); - - switch (state) { - case IDLE: - case FINISHED: - if (!this.closed.get()) { - queue.offer(this.valueExtractor.apply(value)); + ResultState responseState = this.stateWatcher.apply(value); + if (responseState == ResultState.IDLE || responseState == ResultState.FINISHED) { + iterator = Objects.requireNonNull(this.valueExtractor.apply(value), "response iterator"); + result = responseState; + } else { + iterator = new ErrorMessageIterator<>(responseState.getMessage()); + result = ResultState.ERROR; + } + } catch (Throwable e) { + log.error("handling server data for query {}, got error: ", this.queryId, e); + result = ResultState.ERROR; + iterator = new ErrorMessageIterator<>(e.getMessage()); + if (e instanceof Error) { + throw (Error) e; + } + } finally { + synchronized (this.stateLock) { + // Completion can arrive while parsing. Publish before releasing the parsing count. + // An error or client close instead discards responses still being parsed. + if (!this.closed && this.terminal != ResultState.ERROR) { + if (iterator != null) { + this.queue.offer(iterator); + } + // A failed accepted response overrides completion received during parsing. + if (result == ResultState.ERROR || + (this.terminal == null && result != ResultState.IDLE)) { + this.terminal = result; } - // this.stop(); - break; - default: - queue.offer(new ErrorMessageIterator<>(state.getMessage())); - break; + } + this.parsing--; + this.current = System.nanoTime(); } - watcher.setState(state); - // sendRequest(); - } catch (Exception e) { - log.error("handling server data, got error: ", e); - queue.offer(new ErrorMessageIterator<>(e.getMessage())); } } public Iterator consume() { try { - while (!Thread.currentThread().isInterrupted() && (!this.queue.isEmpty() || - !isServerFinished())) { + while (!Thread.currentThread().isInterrupted()) { + // Read terminal state before the queue, including a concurrently published final batch. + if (isServerFinished() && this.queue.isEmpty()) { + return null; + } var iterator = this.queue.poll(200, TimeUnit.MILLISECONDS); if (iterator != null) { sendRequest(); return iterator; } - if ((System.nanoTime() - watcher.current) / 1000_000 > this.timeout) { + // Read terminal state before the queue: a final batch is published with that state. + if (isServerFinished() && this.queue.isEmpty()) { + return null; + } + if ((System.nanoTime() - this.current) / 1000_000 > this.timeout) { throw new HgStoreClientException("iterator timeout"); } @@ -146,48 +178,52 @@ public Iterator consume() { } /** - * Send onComplete, stop receiving data + * Stop feedback and invoke the request stream's completion or cancellation callback. */ public void clear() { - if (!this.closed.get()) { - this.closed.set(true); - this.transferComplete.accept(true); + boolean finished; + synchronized (this.stateLock) { + if (this.closed) { + return; + } + finished = this.terminal == ResultState.FINISHED && this.parsing == 0; + this.closed = true; + this.queue.clear(); + } + if (finished) { + synchronized (this.requestLock) { + this.transferComplete.accept(true); + } + } else { + // The transport cancellation hook is thread-safe even during a blocked send. + this.transferComplete.accept(false); } - this.queue.clear(); } @Override public void onError(Throwable t) { + synchronized (this.stateLock) { + if (this.closed || this.terminal != null) { + return; + } + this.queue.offer(new ErrorMessageIterator<>(t.getMessage())); + this.terminal = ResultState.ERROR; + this.current = System.nanoTime(); + } log.error("StreamObserver got error:", t); - this.queue.offer(new ErrorMessageIterator<>(t.getMessage())); - this.watcher.setState(ResultState.ERROR); } @Override public void onCompleted() { - if (watcher.getState() != ResultState.ERROR) { - watcher.setState(ResultState.FINISHED); + synchronized (this.stateLock) { + if (!this.closed && this.terminal == null) { + this.terminal = ResultState.FINISHED; + this.current = System.nanoTime(); + } } } public void setWatcherQueryId(String queryId) { - this.watcher.setQueryId(queryId); - } - - @Data - private static class ResultStateWatcher { - - private long current = System.nanoTime(); - private volatile ResultState state = ResultState.IDLE; - - private String queryId; - - public void setState(ResultState state) { - log.debug("query Id: {}, COST_STAT: {} -> {}, cost {} ms", this.queryId, this.state, - state, - +(System.nanoTime() - current) * 1.0 / 1000000); - this.state = state; - this.current = System.nanoTime(); - } + this.queryId = queryId; } } diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIterator.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIterator.java index 0a4373c7f3..22be95b1ce 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIterator.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIterator.java @@ -17,7 +17,12 @@ package org.apache.hugegraph.store.client.query; +import java.util.ArrayList; +import java.util.Collections; +import java.util.IdentityHashMap; import java.util.Iterator; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.List; import java.util.NoSuchElementException; @@ -27,40 +32,74 @@ public class MultiStreamIterator implements HgKvIterator { private HgKvIterator currentIterator = null; + private final List> iterators; private final Iterator> listIterator; + private final AtomicBoolean closed = new AtomicBoolean(); public MultiStreamIterator(List> iterators) { - this.listIterator = iterators.iterator(); + this.iterators = new ArrayList<>(iterators); + this.listIterator = this.iterators.iterator(); } @Override public byte[] key() { + checkOpen(); return currentIterator.key(); } @Override public byte[] value() { + checkOpen(); return currentIterator.value(); } @Override public void close() { - //Todo is syntax correct? - if (currentIterator != null && currentIterator.hasNext()) { - currentIterator.close(); + if (!this.closed.compareAndSet(false, true)) { + return; + } + Throwable failure = null; + Set> visited = Collections.newSetFromMap(new IdentityHashMap<>()); + for (HgKvIterator iterator : this.iterators) { + if (!visited.add(iterator)) { + continue; + } + try { + iterator.close(); + } catch (RuntimeException | Error error) { + if (failure == null) { + failure = error; + } else if (failure != error) { + failure.addSuppressed(error); + } + } + } + if (failure instanceof Error) { + throw (Error) failure; + } + if (failure != null) { + throw (RuntimeException) failure; } } @Override public byte[] position() { + checkOpen(); return currentIterator.position(); } @Override public void seek(byte[] position) { + checkOpen(); this.currentIterator.seek(position); } + private void checkOpen() { + if (this.closed.get()) { + throw new IllegalStateException("Iterator is closed"); + } + } + private void getNextIterator() { if (currentIterator != null && currentIterator.hasNext()) { return; @@ -76,13 +115,16 @@ private void getNextIterator() { @Override public boolean hasNext() { + if (this.closed.get()) { + return false; + } getNextIterator(); return currentIterator != null && currentIterator.hasNext(); } @Override public E next() { - if (currentIterator == null || !currentIterator.hasNext()) { + if (this.closed.get() || currentIterator == null || !currentIterator.hasNext()) { throw new NoSuchElementException(); } return currentIterator.next(); diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryExecutor.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryExecutor.java index 9f6fd0dc8b..bacc0aa383 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryExecutor.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryExecutor.java @@ -61,6 +61,8 @@ import com.google.protobuf.ByteString; +import io.grpc.stub.ClientCallStreamObserver; +import io.grpc.stub.ClientResponseObserver; import lombok.extern.slf4j.Slf4j; @Slf4j @@ -237,10 +239,35 @@ public BaseElement next() { ); - var reqStream = stub.query(observer); observer.setWatcherQueryId(request.getQueryId() + '-' + address); - observer.setRequestSender(r -> reqStream.onNext(request)); - observer.setTransferComplete(r -> reqStream.onCompleted()); + stub.query(new ClientResponseObserver() { + @Override + public void beforeStart(ClientCallStreamObserver requestStream) { + observer.setRequestSender(ignored -> requestStream.onNext(request)); + observer.setTransferComplete(finished -> { + if (finished) { + requestStream.onCompleted(); + } else { + requestStream.cancel("Query iterator closed", null); + } + }); + } + + @Override + public void onNext(QueryResponse response) { + observer.onNext(response); + } + + @Override + public void onError(Throwable error) { + observer.onError(error); + } + + @Override + public void onCompleted() { + observer.onCompleted(); + } + }); observer.setTimeout(this.timeout); var itr = new StreamKvIterator<>(b -> observer.clear(), observer::consume); diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java index 73dbfceb7a..ea6f1b5870 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java @@ -398,14 +398,19 @@ public void monitorMemory() { @Override public void shutdown() { log.info("HeartbeatService shutdown"); - stopProducers(); - joinHeartbeat(storeHeartbeatThread); - joinHeartbeat(partitionHeartbeatThread); - // Producers may have submitted callbacks that create partition owners. - // Never hold this monitor while joining the producer threads. - synchronized (stateChangeLock) { - if (Thread.currentThread().isInterrupted()) { - throw new IllegalStateException("Interrupted while stopping store state callbacks"); + boolean interrupted = Thread.interrupted(); + try { + stopProducers(); + interrupted |= joinHeartbeat(storeHeartbeatThread); + interrupted |= joinHeartbeat(partitionHeartbeatThread); + // Producers may have submitted callbacks that create partition owners. + // Never hold this monitor while joining the producer threads. + synchronized (stateChangeLock) { + // Acquiring the monitor waits for any in-flight state callback to finish. + } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); } } } @@ -434,19 +439,24 @@ private void stopProducers() { } } - private static void joinHeartbeat(Thread thread) { + private static boolean joinHeartbeat(Thread thread) { if (thread == null) { - return; + return false; } if (thread == Thread.currentThread()) { throw new IllegalStateException("Cannot close the store from its heartbeat thread"); } - try { - thread.join(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IllegalStateException("Interrupted while stopping store heartbeat", e); + boolean interrupted = false; + boolean joined = false; + while (!joined) { + try { + thread.join(); + joined = true; + } catch (InterruptedException e) { + interrupted = true; + } } + return interrupted; } @Override diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java index 4753c85264..4299e79cee 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java @@ -207,35 +207,42 @@ public void shutdown() { if (rpcServer == null) { return; } - closing.set(true); - heartbeatService.shutdown(); - metricService.shutdown(); -// Use sequential processing for safer shutdown - partitionEngines.values().forEach(pe -> { - try { - Node raftNode = pe.getRaftNode(); - if (raftNode.isLeader(false)) { - Status status = raftNode.transferLeadershipTo(PeerId.ANY_PEER); - if (!status.isOk()) { - log.warn("transfer leader error: {}", status); + boolean interrupted = Thread.interrupted(); + try { + closing.set(true); + heartbeatService.shutdown(); + interrupted |= Thread.interrupted(); + metricService.shutdown(); + // Use sequential processing so each partition drains before its database closes. + for (PartitionEngine pe : partitionEngines.values()) { + try { + Node raftNode = pe.getRaftNode(); + if (raftNode.isLeader(false)) { + Status status = raftNode.transferLeadershipTo(PeerId.ANY_PEER); + if (!status.isOk()) { + log.warn("transfer leader error: {}", status); + } } + } catch (Exception e) { + log.error("transfer leader error: ", e); } - } catch (Exception e) { - log.error("transfer leader error: ", e); + pe.shutdown(); + interrupted |= Thread.interrupted(); + businessHandler.closeDB(pe.getGroupId()); + } + partitionEngines.clear(); + rpcServer.shutdown(); + // HgStoreEngine.init skips RPC initialization while this field is non-null. + rpcServer = null; + try { + RocksDBFactory.getInstance().releaseAllGraphDB(); + } finally { + BusinessHandlerImpl.closeSchemaResources(); } - pe.shutdown(); - businessHandler.closeDB(pe.getGroupId()); - }); - partitionEngines.clear(); - rpcServer.shutdown(); - // HgStoreEngine.init function check rpcServer whether is null, skipped if the instance - // exists even shut down. - rpcServer = null; - // close all db session - try { - RocksDBFactory.getInstance().releaseAllGraphDB(); } finally { - BusinessHandlerImpl.closeSchemaResources(); + if (interrupted) { + Thread.currentThread().interrupt(); + } } } diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/PartitionEngine.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/PartitionEngine.java index 78ac781375..85e43fea02 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/PartitionEngine.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/PartitionEngine.java @@ -95,7 +95,6 @@ import com.alipay.sofa.jraft.util.Endpoint; import com.alipay.sofa.jraft.util.ThreadId; import com.alipay.sofa.jraft.util.Utils; -import com.alipay.sofa.jraft.util.internal.ThrowUtil; import com.google.protobuf.CodedInputStream; import lombok.Getter; @@ -565,16 +564,28 @@ public void shutdown() { if (!this.started) { return; } - if (this.raftGroupService != null) { - this.raftGroupService.shutdown(); - try { - this.raftGroupService.join(); - } catch (final InterruptedException e) { - ThrowUtil.throwException(e); + boolean interrupted = false; + try { + if (this.raftGroupService != null) { + this.raftGroupService.shutdown(); + boolean terminated = false; + while (!terminated) { + try { + this.raftGroupService.join(); + terminated = true; + } catch (InterruptedException e) { + // Storage remains in use until Raft and its callbacks have stopped. + interrupted = true; + } + } + } + this.started = false; + log.info("PartitionEngine shutdown successfully: {}.", this); + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); } } - this.started = false; - log.info("PartitionEngine shutdown successfully: {}.", this); } /** diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/InnerKeyFilter.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/InnerKeyFilter.java index 368032f2ce..75572ec609 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/InnerKeyFilter.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/InnerKeyFilter.java @@ -31,13 +31,13 @@ public class InnerKeyFilter implements ScanIterator { final boolean codeFilter; ScanIterator iterator; T current = null; + private boolean initialized; public InnerKeyFilter(ScanIterator iterator) { this.iterator = iterator; this.codeFrom = Integer.MIN_VALUE; this.codeTo = Integer.MAX_VALUE; this.codeFilter = false; - moveNext(); } public InnerKeyFilter(ScanIterator iterator, boolean codeFilter) { @@ -45,7 +45,6 @@ public InnerKeyFilter(ScanIterator iterator, boolean codeFilter) { this.codeFrom = Integer.MIN_VALUE; this.codeTo = Integer.MAX_VALUE; this.codeFilter = codeFilter; - moveNext(); } public InnerKeyFilter(ScanIterator iterator, int codeFrom, int codeTo) { @@ -53,7 +52,14 @@ public InnerKeyFilter(ScanIterator iterator, int codeFrom, int codeTo) { this.codeFrom = codeFrom; this.codeTo = codeTo; this.codeFilter = true; - moveNext(); + } + + private void init() { + if (!this.initialized) { + // Publish the wrapper to its owner before prefetch can close or fail on the raw iterator. + this.initialized = true; + moveNext(); + } } private void moveNext() { @@ -76,16 +82,19 @@ private void moveNext() { @Override public boolean hasNext() { + init(); return current != null; } @Override public boolean isValid() { + init(); return iterator.isValid(); } @Override public T next() { + init(); T column = current; if (!codeFilter) // Remove the image ID and hash suffix @@ -103,11 +112,13 @@ public T next() { @Override public void close() { + this.initialized = true; iterator.close(); } @Override public long count() { + init(); return iterator.count(); } } diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/MultiPartitionIterator.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/MultiPartitionIterator.java index 72cc472b21..54fd80be72 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/MultiPartitionIterator.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/MultiPartitionIterator.java @@ -18,13 +18,13 @@ package org.apache.hugegraph.store.business; import java.nio.ByteBuffer; +import java.util.ArrayList; import java.util.Collections; import java.util.LinkedList; import java.util.List; import java.util.NoSuchElementException; import java.util.Queue; import java.util.function.BiFunction; -import java.util.stream.Collectors; import org.apache.hugegraph.rocksdb.access.ScanIterator; @@ -45,6 +45,8 @@ public class MultiPartitionIterator implements ScanIterator { private Integer curPartitionId; private Integer positionPartitionId; private byte[] positionKey; + private RuntimeException cleanupFailure; + private boolean closed; private MultiPartitionIterator(List partitionIds, BiFunction supplier) { @@ -75,35 +77,31 @@ public static int toInt(byte[] bytes) { } private ScanIterator getIterator() { - if (this.partitions.isEmpty()) { - return null; - } - ScanIterator buf = null; - while (!partitions.isEmpty()) { - this.curPartitionId = partitions.poll(); + while (!this.partitions.isEmpty()) { + this.curPartitionId = this.partitions.poll(); if (!this.inPosition(this.curPartitionId)) { continue; } - buf = supplier.apply(this.curPartitionId, getPositionKey(this.curPartitionId)); - if (buf == null) { + ScanIterator child = this.supplier.apply(this.curPartitionId, + getPositionKey(this.curPartitionId)); + if (child == null) { continue; } - if (buf.hasNext()) { - break; + try { + if (child.hasNext()) { + return child; + } + } catch (RuntimeException | Error failure) { + closeAfterFailure(child, failure); + throw failure; } + closeCreatedIterator(child); } - if (buf == null) { - return null; - } - if (!buf.hasNext()) { - buf.close(); - buf = null; - } - return buf; + return null; } private void init() { - if (this.iterator == null) { + if (!this.closed && this.iterator == null) { this.iterator = this.getIterator(); } } @@ -128,8 +126,7 @@ public T next() { } T t = this.iterator.next(); if (!this.iterator.hasNext()) { - this.iterator.close(); - this.iterator = null; + closeCurrentIterator(); } return t; } @@ -137,11 +134,16 @@ public T next() { @Override public long count() { long count = 0; - this.iterator = this.getIterator(); - while (this.iterator != null) { - count += this.iterator.count(); - // this.iterator.close(); - this.iterator = this.getIterator(); + while (this.hasNext()) { + try { + count += this.iterator.count(); + } catch (RuntimeException | Error failure) { + ScanIterator child = this.iterator; + this.iterator = null; + closeAfterFailure(child, failure); + throw failure; + } + closeCurrentIterator(); } return count; } @@ -172,8 +174,33 @@ public void seek(byte[] position) { @Override public void close() { - if (this.iterator != null) { - this.iterator.close(); + this.closed = true; + try { + closeCurrentIterator(); + } catch (RuntimeException | Error failure) { + // closeCreatedIterator retains the failure even after ownership is released here. + throw this.cleanupFailure; + } + if (this.cleanupFailure != null) { + throw this.cleanupFailure; + } + } + + private void closeCurrentIterator() { + ScanIterator child = this.iterator; + this.iterator = null; + if (child != null) { + closeCreatedIterator(child); + } + } + + private void closeAfterFailure(ScanIterator child, Throwable failure) { + try { + closeCreatedIterator(child); + } catch (RuntimeException | Error closeFailure) { + if (closeFailure != failure) { + failure.addSuppressed(closeFailure); + } } } @@ -205,10 +232,50 @@ private byte[] getPositionKey(int partitionId) { * @return iteration list */ public List getIterators() { - return this.partitions.stream() - .map(id -> supplier.apply(id, getPositionKey(id))) - .filter(ScanIterator::hasNext) - .collect(Collectors.toList()); + List opened = new ArrayList<>(); + if (this.closed) { + return opened; + } + try { + for (int id : this.partitions) { + ScanIterator child = this.supplier.apply(id, getPositionKey(id)); + if (child == null) { + continue; + } + opened.add(child); + if (!child.hasNext()) { + opened.remove(opened.size() - 1); + closeCreatedIterator(child); + } + } + // Ownership transfers to the caller only when all partitions were opened. + return opened; + } catch (RuntimeException | Error failure) { + for (ScanIterator child : opened) { + try { + closeCreatedIterator(child); + } catch (RuntimeException | Error closeFailure) { + if (closeFailure != failure) { + failure.addSuppressed(closeFailure); + } + } + } + throw failure; + } + } + + private void closeCreatedIterator(ScanIterator child) { + try { + child.close(); + } catch (RuntimeException | Error failure) { + if (this.cleanupFailure == null) { + this.cleanupFailure = new IllegalStateException( + "partition iterator cleanup failed", failure); + } else { + this.cleanupFailure.addSuppressed(failure); + } + throw failure; + } } } diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java index f7f013b89f..cba162a58b 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java @@ -41,6 +41,7 @@ public final class OrderedMultiPartitionIterator implements ScanIterator { private boolean initialized; private boolean closed; + private Throwable cleanupFailure; private Integer currentPartitionId; private OrderedMultiPartitionIterator(List partitionIds, @@ -142,22 +143,29 @@ public void seek(byte[] position) { @Override public void close() { if (this.closed) { + rethrowCleanupFailure(); return; } this.closed = true; - Throwable failure = null; + Throwable failure = this.cleanupFailure; for (SourceEntry source : this.sources) { try { this.closeSource(source); } catch (RuntimeException | Error e) { if (failure == null) { failure = e; - } else { + } else if (failure != e) { failure.addSuppressed(e); } } } this.queue.clear(); + this.cleanupFailure = failure; + rethrowCleanupFailure(); + } + + private void rethrowCleanupFailure() { + Throwable failure = this.cleanupFailure; if (failure instanceof RuntimeException) { throw (RuntimeException) failure; } @@ -197,14 +205,23 @@ private void closeSource(SourceEntry source) { return; } source.closed = true; - source.iterator.close(); + try { + source.iterator.close(); + } catch (RuntimeException | Error failure) { + if (this.cleanupFailure == null) { + this.cleanupFailure = failure; + } + throw failure; + } } private void closeAfterFailure(Throwable failure) { try { this.close(); } catch (RuntimeException | Error closeFailure) { - failure.addSuppressed(closeFailure); + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } } } diff --git a/hugegraph-store/hg-store-dist/pom.xml b/hugegraph-store/hg-store-dist/pom.xml index 12acb0ba9e..c75fc2ad6d 100644 --- a/hugegraph-store/hg-store-dist/pom.xml +++ b/hugegraph-store/hg-store-dist/pom.xml @@ -42,6 +42,8 @@ org.apache.hugegraph hg-store-node ${revision} + exec + runtime diff --git a/hugegraph-store/hg-store-dist/src/assembly/descriptor/server-assembly.xml b/hugegraph-store/hg-store-dist/src/assembly/descriptor/server-assembly.xml index 2e761200ef..47170a6c22 100644 --- a/hugegraph-store/hg-store-dist/src/assembly/descriptor/server-assembly.xml +++ b/hugegraph-store/hg-store-dist/src/assembly/descriptor/server-assembly.xml @@ -68,11 +68,12 @@ /lib + ${artifact.artifactId}-${artifact.version}.jar false runtime false - org.apache.hugegraph:${executable.jar.name}:jar:* + org.apache.hugegraph:${executable.jar.name}:jar:exec:* diff --git a/hugegraph-store/hg-store-node/pom.xml b/hugegraph-store/hg-store-node/pom.xml index 56cff46dba..9c952507af 100644 --- a/hugegraph-store/hg-store-node/pom.xml +++ b/hugegraph-store/hg-store-node/pom.xml @@ -169,6 +169,7 @@ repackage + exec org.apache.hugegraph.store.node.StoreNodeApplication diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/BatchScanIterator.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/BatchScanIterator.java index f7e02481fe..4d26c7d839 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/BatchScanIterator.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/BatchScanIterator.java @@ -41,6 +41,7 @@ public final class BatchScanIterator implements ScanIterator { private final AtomicBoolean closed = new AtomicBoolean(); private ScanIterator iterator; private boolean hasNext = false; + private Throwable cleanupFailure; private long curCount; private long curLimit; @@ -65,13 +66,14 @@ private ScanIterator getIterator() { do { buf = this.batchSupplier.get().getValue(); + this.iterator = buf; if (buf == null) { break; } if (!buf.hasNext()) { - buf.close(); + closeIterator(); buf = null; } @@ -99,13 +101,17 @@ private ScanIterator getIterator() { @Override public boolean hasNext() { + rethrowCleanupFailure(); + if (this.closed.get()) { + return false; + } if (this.iterator == null) { this.iterator = this.getIterator(); } else if (!this.iterator.hasNext()) { - this.iterator.close(); + closeIterator(); this.iterator = this.getIterator(); } else if (this.curCount == this.curLimit) { - this.iterator.close(); + closeIterator(); this.iterator = this.getIterator(); } @@ -135,10 +141,31 @@ public T next() { @Override public void close() { if (!this.closed.getAndSet(true)) { - if (this.iterator != null) { - this.iterator.close(); + closeIterator(); + } + rethrowCleanupFailure(); + } + + private void closeIterator() { + ScanIterator current = this.iterator; + this.iterator = null; + if (current != null) { + try { + current.close(); + } catch (RuntimeException | Error failure) { + this.cleanupFailure = failure; } } + rethrowCleanupFailure(); + } + + private void rethrowCleanupFailure() { + if (this.cleanupFailure instanceof RuntimeException) { + throw (RuntimeException) this.cleanupFailure; + } + if (this.cleanupFailure != null) { + throw (Error) this.cleanupFailure; + } } @Override diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/FusingScanIterator.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/FusingScanIterator.java index 56a1cc8c58..afdcb2e59c 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/FusingScanIterator.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/FusingScanIterator.java @@ -33,6 +33,8 @@ final class FusingScanIterator implements ScanIterator { private long accumulator; private Supplier supplier; private ScanIterator iterator; + private Throwable cleanupFailure; + private boolean closed; private byte[] position = EMPTY_BYTES; private FusingScanIterator() { @@ -47,14 +49,13 @@ public static FusingScanIterator maxOf(long maxThreshold, } private ScanIterator getIterator() { - ScanIterator buf = this.supplier.get(); - if (buf == null) { - return null; + while ((this.iterator = this.supplier.get()) != null) { + if (this.iterator.hasNext()) { + return this.iterator; + } + closeCurrent(); } - if (!buf.hasNext()) { - buf = null; - } - return buf; + return null; } private void init() { @@ -65,7 +66,8 @@ private void init() { @Override public boolean hasNext() { - if (this.isThresholdExceeded()) { + rethrowCleanupFailure(); + if (this.closed || this.isThresholdExceeded()) { return false; } if (this.iterator == null) { @@ -104,16 +106,36 @@ public T next() { position = this.iterator.position(); this.accumulator++; if (!this.iterator.hasNext() || this.isThresholdExceeded()) { - this.iterator.close(); - this.iterator = null; + closeCurrent(); } return t; } @Override public void close() { - if (this.iterator != null) { - this.iterator.close(); + this.closed = true; + closeCurrent(); + } + + private void closeCurrent() { + ScanIterator current = this.iterator; + this.iterator = null; + if (current != null) { + try { + current.close(); + } catch (RuntimeException | Error failure) { + this.cleanupFailure = failure; + } + } + rethrowCleanupFailure(); + } + + private void rethrowCleanupFailure() { + if (this.cleanupFailure instanceof RuntimeException) { + throw (RuntimeException) this.cleanupFailure; + } + if (this.cleanupFailure != null) { + throw (Error) this.cleanupFailure; } } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GRpcServerConfig.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GRpcServerConfig.java index 27cb69a1de..1d22764aba 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GRpcServerConfig.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GRpcServerConfig.java @@ -38,9 +38,10 @@ public class GRpcServerConfig extends GRpcServerBuilderConfigurer { @Override public void configure(ServerBuilder serverBuilder) { AppConfig.ThreadPoolGrpc grpc = appConfig.getThreadPoolGrpc(); + // Terminal callbacks must still be queued when all dispatch threads are busy. serverBuilder.executor( HgExecutorUtil.createExecutor(EXECUTOR_NAME, grpc.getCore(), grpc.getMax(), - grpc.getQueue()) + Integer.MAX_VALUE) ); } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GrpcShutdownBarrier.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GrpcShutdownBarrier.java new file mode 100644 index 0000000000..a9a6a9f95a --- /dev/null +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GrpcShutdownBarrier.java @@ -0,0 +1,110 @@ +/* + * 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.store.node.grpc; + +import java.util.HashSet; +import java.util.Set; + +import org.lognet.springboot.grpc.GRpcGlobalInterceptor; +import org.springframework.stereotype.Component; + +import io.grpc.ForwardingServerCallListener; +import io.grpc.Metadata; +import io.grpc.ServerCall; +import io.grpc.ServerCallHandler; +import io.grpc.ServerInterceptor; +import io.grpc.Status; +import lombok.extern.slf4j.Slf4j; + +/** Waits for terminal application callbacks, not just transport termination. */ +@Slf4j +@Component +@GRpcGlobalInterceptor +public class GrpcShutdownBarrier implements ServerInterceptor { + + private final Set calls = new HashSet<>(); + private boolean closing; + + public synchronized void stopAcceptingCalls() { + this.closing = true; + } + + @Override + public ServerCall.Listener interceptCall( + ServerCall call, Metadata headers, + ServerCallHandler next) { + Object token = new Object(); + synchronized (this) { + if (this.closing) { + call.close(Status.UNAVAILABLE.withDescription("Store is stopping"), new Metadata()); + return new ServerCall.Listener() { }; + } + this.calls.add(token); + } + try { + return new ForwardingServerCallListener.SimpleForwardingServerCallListener( + next.startCall(call, headers)) { + @Override + public void onCancel() { + try { + super.onCancel(); + } finally { + finished(token); + } + } + + @Override + public void onComplete() { + try { + super.onComplete(); + } finally { + finished(token); + } + } + }; + } catch (RuntimeException | Error e) { + finished(token); + throw e; + } + } + + private synchronized void finished(Object token) { + this.calls.remove(token); + this.notifyAll(); + } + + public synchronized void awaitCallbacks() { + boolean interrupted = false; + try { + while (!this.calls.isEmpty()) { + try { + this.wait(5000); + if (!this.calls.isEmpty()) { + log.warn("Waiting for {} RPC callbacks before closing databases", this.calls.size()); + } + } catch (InterruptedException e) { + interrupted = true; + } + } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + } +} diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreStreamImpl.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreStreamImpl.java index 7d01fa3db4..24a7d3c4d2 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreStreamImpl.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreStreamImpl.java @@ -17,9 +17,19 @@ package org.apache.hugegraph.store.node.grpc; +import java.util.Map; +import java.util.function.Function; +import java.util.function.Consumer; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.hugegraph.store.business.BusinessHandler; +import org.apache.hugegraph.store.grpc.Graphpb.ScanPartitionRequest; +import org.apache.hugegraph.store.grpc.Graphpb.ScanResponse; +import org.apache.hugegraph.store.node.grpc.scan.ScanResponseObserver; import org.apache.hugegraph.store.grpc.state.ScanState; import org.apache.hugegraph.store.grpc.stream.HgStoreStreamGrpc; import org.apache.hugegraph.store.grpc.stream.KvPageRes; @@ -31,6 +41,8 @@ import org.lognet.springboot.grpc.GRpcService; import org.springframework.beans.factory.annotation.Autowired; +import io.grpc.Context; +import io.grpc.Status; import io.grpc.stub.StreamObserver; import lombok.extern.slf4j.Slf4j; @@ -47,6 +59,191 @@ public class HgStoreStreamImpl extends HgStoreStreamGrpc.HgStoreStreamImplBase { private AppConfig appConfig; private HgStoreWrapperEx wrapper; private ThreadPoolExecutor executor; + private boolean closing; + private final Map scans = new ConcurrentHashMap<>(); + + /** Close admission before cancelling the gRPC server. */ + public synchronized void stopAcceptingScans() { + this.closing = true; + } + + /** Cancel streams, then drain queued cleanup tasks. Never discard the queue. */ + public void shutdownScans() { + stopAcceptingScans(); + for (Runnable cancel : this.scans.values().toArray(new Runnable[0])) { + try { + cancel.run(); + } catch (RuntimeException | Error failure) { + log.warn("Failed to cancel scan; continuing other cancellations", failure); + } + } + ThreadPoolExecutor current = getRealExecutor(); + if (current != null) { + current.shutdown(); + } + } + + private synchronized void checkAcceptingScans() { + if (this.closing) { + throw Status.UNAVAILABLE.withDescription("Store scans are stopping") + .asRuntimeException(); + } + } + + /** Wait independently of executor termination: failed native release is sticky. */ + public void awaitScanCleanup() { + boolean interrupted = false; + long nextLog = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + synchronized (this.scans) { + while (!this.scans.isEmpty()) { + try { + this.scans.wait(5000L); + } catch (InterruptedException e) { + interrupted = true; + } + if (!this.scans.isEmpty() && System.nanoTime() - nextLog >= 0) { + log.warn("Still waiting for {} scans to clean up before closing databases", + this.scans.size()); + for (ScanLifecycle scan : this.scans.keySet()) { + Throwable failure = scan.cleanupFailure(); + if (failure != null) { + log.warn("Scan cleanup failed; database close stays blocked", failure); + } + } + nextLog = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + } + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + + private StreamObserver register( + ScanLifecycle lifecycle, Function> factory) { + Context context = Context.current(); + StreamObserver observer = factory.apply(lifecycle); + Runnable cancel = () -> { + if (lifecycle.tryEnterCancellation()) { + try { + observer.onError(Status.CANCELLED.asRuntimeException()); + } finally { + lifecycle.leave(); + } + } + }; + this.scans.put(lifecycle, cancel); + Context.CancellationListener listener = ignored -> cancel.run(); + lifecycle.onFinished(() -> { + this.scans.remove(lifecycle); + context.removeListener(listener); + synchronized (this.scans) { + this.scans.notifyAll(); + } + }); + context.addListener(listener, Runnable::run); + return new StreamObserver() { + @Override + public void onNext(T value) { + if (lifecycle.tryEnter()) { + try { + observer.onNext(value); + } finally { + lifecycle.leave(); + } + } + } + + @Override + public void onError(Throwable failure) { + if (lifecycle.tryEnterCancellation()) { + try { + observer.onError(failure); + } finally { + lifecycle.leave(); + } + } + } + + @Override + public void onCompleted() { + if (lifecycle.tryEnter()) { + try { + observer.onCompleted(); + } finally { + lifecycle.leave(); + } + } + } + }; + } + + public StreamObserver scanGraphPartition( + StreamObserver response, + BusinessHandler handler) { + synchronized (this) { + checkAcceptingScans(); + return register(new ScanLifecycle(), lifecycle -> + new ScanResponseObserver<>( + lifecycle.response(response), handler, + task -> lifecycle.execute(getExecutor(), task), lifecycle::failedCleanup)); + } + } + + private void oneShot(StreamObserver response, Consumer action) { + ScanLifecycle lifecycle = new ScanLifecycle(); + Thread caller = Thread.currentThread(); + boolean alreadyInterrupted = caller.isInterrupted(); + AtomicBoolean interruptedByCancellation = new AtomicBoolean(); + Context context = Context.current(); + Runnable cancel = () -> { + synchronized (lifecycle) { + if (lifecycle.tryCancel()) { + // Wake interruptible reads; native reads retain ownership until they return. + if (!caller.isInterrupted()) { + interruptedByCancellation.set(true); + caller.interrupt(); + } + } + } + }; + Context.CancellationListener listener = ignored -> cancel.run(); + synchronized (this) { + checkAcceptingScans(); + lifecycle.enter(); + this.scans.put(lifecycle, cancel); + lifecycle.onFinished(() -> { + this.scans.remove(lifecycle); + synchronized (this.scans) { + this.scans.notifyAll(); + } + }); + } + try { + context.addListener(listener, Runnable::run); + action.accept(lifecycle); + } finally { + try { + // An iterator may close automatically during hasNext(), after which the + // scan exits without entering its ordinary response-completion path. + if (lifecycle.isCancelled()) { + lifecycle.response(response).onError(Status.CANCELLED + .withDescription("Store scan cancelled before completion") + .asRuntimeException()); + } + } finally { + context.removeListener(listener); + // A cancellation snapshot must not interrupt this executor thread after reuse. + synchronized (lifecycle) { + lifecycle.finishWithoutResponse(); + lifecycle.leave(); + if (!alreadyInterrupted && interruptedByCancellation.get()) { + Thread.interrupted(); + } + } + } + } + } private HgStoreWrapperEx getWrapper() { if (this.wrapper == null) { @@ -60,20 +257,16 @@ private HgStoreWrapperEx getWrapper() { return this.wrapper; } - public ThreadPoolExecutor getRealExecutor() { - return executor; + public synchronized ThreadPoolExecutor getRealExecutor() { + return this.executor; } - public ThreadPoolExecutor getExecutor() { + public synchronized ThreadPoolExecutor getExecutor() { + checkAcceptingScans(); if (this.executor == null) { - synchronized (this) { - if (this.executor == null) { - AppConfig.ThreadPoolScan scan = this.appConfig.getThreadPoolScan(); - this.executor = - HgExecutorUtil.createExecutor("hg-scan", scan.getCore(), scan.getMax(), - scan.getQueue()); - } - } + AppConfig.ThreadPoolScan scan = this.appConfig.getThreadPoolScan(); + this.executor = HgExecutorUtil.createExecutor("hg-scan", scan.getCore(), + scan.getMax(), scan.getQueue()); } return this.executor; } @@ -95,27 +288,39 @@ public ScanState getState() { } @Override - public StreamObserver scan(StreamObserver response) { - return ScanStreamResponse.of(response, getWrapper(), getExecutor(), appConfig); + public synchronized StreamObserver scan(StreamObserver response) { + checkAcceptingScans(); + return register(new ScanLifecycle(), lifecycle -> + new ScanStreamResponse(response, getWrapper(), getExecutor(), appConfig, lifecycle)); } @Override public void scanOneShot(ScanStreamReq request, StreamObserver response) { - ScanOneShotResponse.scanOneShot(request, response, getWrapper()); + oneShot(response, lifecycle -> ScanOneShotResponse.scanOneShot(request, response, getWrapper(), lifecycle)); } @Override - public StreamObserver scanBatch(StreamObserver response) { - return ScanBatchResponse3.of(response, getWrapper(), getExecutor()); + public synchronized StreamObserver scanBatch(StreamObserver response) { + checkAcceptingScans(); + return register(new ScanLifecycle(), lifecycle -> + ScanBatchResponse3.of(response, getWrapper(), getExecutor(), lifecycle)); } @Override public StreamObserver scanBatch2(StreamObserver response) { - return ScanBatchResponseFactory.of(response, getWrapper(), getExecutor()); + StreamObserver observer; + synchronized (this) { + checkAcceptingScans(); + observer = register(new ScanLifecycle(), lifecycle -> + ScanBatchResponseFactory.of(response, getWrapper(), getExecutor(), lifecycle)); + } + ScanBatchResponseFactory.getInstance().checkStreamActive(); + return observer; } @Override public void scanBatchOneShot(ScanStreamBatchReq request, StreamObserver response) { - ScanBatchOneShotResponse.scanOneShot(request, response, getWrapper()); + oneShot(response, lifecycle -> + ScanBatchOneShotResponse.scanOneShot(request, response, getWrapper(), lifecycle)); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ParallelScanIterator.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ParallelScanIterator.java index 56ce1f45f3..65823da8ab 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ParallelScanIterator.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ParallelScanIterator.java @@ -27,6 +27,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.ReentrantLock; import java.util.function.Supplier; +import java.util.function.Consumer; import org.apache.hugegraph.pd.common.KVPair; import org.apache.hugegraph.rocksdb.access.RocksDBSession; @@ -67,11 +68,33 @@ public class ParallelScanIterator implements ScanIterator { private int maxInQueue = maxWorkThreads * 2; private volatile boolean finished; private List current = null; + private final Consumer cleanupReporter; + private volatile Throwable scanFailure; + private volatile Throwable cleanupFailure; + + private static void rethrow(Throwable failure) { + if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } + if (failure != null) { + throw (Error) failure; + } + } + + private synchronized void failedCleanup(Throwable failure) { + if (this.cleanupFailure == null) { + this.cleanupFailure = failure; + } + this.cleanupReporter.accept(failure); + this.scanFailure = this.cleanupFailure; + requestStop(); + } private ParallelScanIterator(Supplier> iteratorSupplier, Supplier limitSupplier, ScanQueryRequest query, - ThreadPoolExecutor executor) { + ThreadPoolExecutor executor, Consumer cleanupReporter) { + this.cleanupReporter = cleanupReporter; this.executor = executor; this.batchSupplier = iteratorSupplier; this.limitSupplier = limitSupplier; @@ -88,7 +111,12 @@ private ParallelScanIterator(Supplier> iter this.maxInQueue = maxWorkThreads * 2; // Edge sorted requires a larger queue queue = new LinkedBlockingQueue<>(maxInQueue * 2); - createScanner(); + try { + createScanner(); + } catch (RuntimeException e) { + close(); + throw e; + } } public static ParallelScanIterator of( @@ -98,7 +126,13 @@ public static ParallelScanIterator of( ThreadPoolExecutor executor) { HgAssert.isArgumentNotNull(iteratorSupplier, "iteratorSupplier"); HgAssert.isArgumentNotNull(limitSupplier, "limitSupplier"); - return new ParallelScanIterator(iteratorSupplier, limitSupplier, query, executor); + return of(iteratorSupplier, limitSupplier, query, executor, ignored -> { }); + } + + static ParallelScanIterator of(Supplier> iteratorSupplier, + Supplier limitSupplier, ScanQueryRequest query, + ThreadPoolExecutor executor, Consumer cleanupReporter) { + return new ParallelScanIterator(iteratorSupplier, limitSupplier, query, executor, cleanupReporter); } @Override @@ -126,6 +160,7 @@ public boolean hasNext() { log.error("Wait data timeout!!!, scanner is {}/{}", scanners.size(), pauseScanners.size()); } + rethrow(this.scanFailure); return current != null && current != NO_DATA; } @@ -144,20 +179,30 @@ public List next() { return t; } + void requestStop() { + this.finished = true; + } + @Override public void close() { - finished = true; + requestStop(); + List pending; synchronized (scanners) { - scanners.forEach(scanner -> { - scanner.close(); - }); + pending = new ArrayList<>(scanners); + } + // Stop active scanner reads before close() waits for their iterator locks. + for (KVScanner scanner : pending) { + scanner.interruptReader(); + } + // Never hold a registry lock while waiting for a scanner's iterator lock. + for (KVScanner scanner : pending) { + scanner.close(); } synchronized (pauseScanners) { - pauseScanners.forEach(s -> { - s.close(); - }); + pauseScanners.clear(); } queue.clear(); + rethrow(this.cleanupFailure); } /** @@ -179,10 +224,15 @@ private void createScanner() { */ private void wakeUpScanner() { synchronized (pauseScanners) { - if (!pauseScanners.isEmpty()) { + if (!finished && !pauseScanners.isEmpty()) { KVScanner scanner = pauseScanners.poll(); if (scanner != null) { - executor.execute(() -> scanner.scanKV()); + try { + executor.execute(() -> scanner.scanKV()); + } catch (java.util.concurrent.RejectedExecutionException e) { + scanner.close(); + throw e; + } } } } @@ -195,7 +245,11 @@ private void wakeUpScanner() { */ private void suspendScanner(KVScanner scanner) { synchronized (pauseScanners) { - pauseScanners.add(scanner); + if (!finished) { + pauseScanners.add(scanner); + } else { + scanner.close(); + } } } @@ -218,30 +272,43 @@ private void quitScanner(KVScanner scanner) { */ private boolean putData(List data) { try { - this.queue.put(data); + while (!finished && !this.queue.offer(data, 100, TimeUnit.MILLISECONDS)) { + // Cancellation must release a producer whose client stopped consuming. + } } catch (InterruptedException e) { log.error("exception ", e); this.finished = true; return false; } - return this.queue.size() < maxInQueue; + return !finished && this.queue.size() < maxInQueue; } private boolean putData(List data, boolean hasNext) { + boolean locked = false; try { - queueLock.lock(); - this.queue.put(data); + while (!finished && !(locked = queueLock.tryLock(100, TimeUnit.MILLISECONDS))) { + // An ordered producer may be waiting behind a cancelled scanner. + } + if (!locked) { + return false; + } + while (!finished && !this.queue.offer(data, 100, TimeUnit.MILLISECONDS)) { + // Recheck cancellation while the output queue is full. + } } catch (InterruptedException e) { - log.error("exception ", e); + Thread.currentThread().interrupt(); this.finished = true; return false; } finally { - if (!hasNext) { + if (locked && finished) { + while (queueLock.isHeldByCurrentThread()) { + queueLock.unlock(); + } + } else if (locked && !hasNext) { queueLock.unlock(); } } - // Data not ended, thread continues to execute - return hasNext || this.queue.size() < maxInQueue; + return !finished && (hasNext || this.queue.size() < maxInQueue); } private synchronized KVPair getIterator() { @@ -300,18 +367,45 @@ public int size() { class KVScanner { private final ReentrantLock iteratorLock = new ReentrantLock(); + private final Object readerLock = new Object(); + private Thread reader; private ScanIterator iterator = null; private QueryCondition query = null; private long limit; private long counter; private volatile boolean closed = false; + private void interruptReader() { + synchronized (this.readerLock) { + if (this.reader != null && this.reader != Thread.currentThread()) { + this.reader.interrupt(); + } + } + } + + private boolean startReader() { + synchronized (this.readerLock) { + if (this.reader != null) { + return false; + } + this.reader = Thread.currentThread(); + return true; + } + } + + private void finishReader() { + synchronized (this.readerLock) { + if (this.reader == Thread.currentThread()) { + this.reader = null; + } + } + } + private ScanIterator getIterator() { // Iterator has no data, or the point has reached the limit, switch to a new iterator. if (iterator == null || !iterator.hasNext() || counter >= limit) { - if (iterator != null) { - iterator.close(); - } + closeIterator(); + rethrow(cleanupFailure); KVPair pair = ParallelScanIterator.this.getIterator(); query = pair.getKey(); iterator = pair.getValue(); @@ -325,15 +419,18 @@ public void scanKV() { boolean canNext = true; ArrayList dataList = new ArrayList<>(batchSize); dataList.ensureCapacity(batchSize); + if (!startReader()) { + return; + } iteratorLock.lock(); try { long entriesSize = 0, bodySize = 0; - while (canNext && !closed) { + while (canNext && !closed && !finished) { iterator = this.getIterator(); if (iterator == null) { break; } - while (iterator.hasNext() && entriesSize < batchSize && + while (!closed && !finished && iterator.hasNext() && entriesSize < batchSize && bodySize < maxBodySize && counter < limit && !closed) { KV kv = KV.of(iterator.next()); @@ -363,11 +460,20 @@ public void scanKV() { putData(dataList); } } - } catch (Exception e) { - log.error("exception {}", e); + } catch (RuntimeException | Error e) { + scanFailure = e; + requestStop(); } finally { + // putData(..., true) intentionally keeps this lock across batches to + // serialize ordered output. The scanner must always release it before + // suspending or quitting, including when the next iterator is empty. + while (queueLock.isHeldByCurrentThread()) { + queueLock.unlock(); + } iteratorLock.unlock(); - if (iterator != null && counter < limit && !closed) { + // Clear ownership before a paused scanner can be scheduled again. + finishReader(); + if (iterator != null && counter < limit && !closed && !finished) { suspendScanner(this); } else { quitScanner(this); @@ -379,12 +485,22 @@ public void close() { closed = true; iteratorLock.lock(); try { - if (iterator != null) { - iterator.close(); - } + closeIterator(); } finally { iteratorLock.unlock(); } } + + private void closeIterator() { + ScanIterator current = this.iterator; + this.iterator = null; + if (current != null) { + try { + current.close(); + } catch (RuntimeException | Error failure) { + failedCleanup(failure); + } + } + } } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchOneShotResponse.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchOneShotResponse.java index 1234a5b3de..bedb7eaf86 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchOneShotResponse.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchOneShotResponse.java @@ -30,6 +30,7 @@ import com.google.protobuf.ByteString; +import io.grpc.Context; import io.grpc.Status; import io.grpc.stub.StreamObserver; import lombok.extern.slf4j.Slf4j; @@ -51,10 +52,19 @@ public class ScanBatchOneShotResponse { public static void scanOneShot(ScanStreamBatchReq request, StreamObserver responseObserver, HgStoreWrapperEx wrapper) { + scanOneShot(request, responseObserver, wrapper, new ScanLifecycle()); + } + + static void scanOneShot(ScanStreamBatchReq request, StreamObserver output, + HgStoreWrapperEx wrapper, ScanLifecycle lifecycle) { + StreamObserver responseObserver = lifecycle.response(output); String graph = request.getHeader().getGraph(); ScanQueryRequest queryRequest = request.getQueryRequest(); - ScanIterator iterator = getIterator(graph, queryRequest, wrapper); + if ((Context.current().isCancelled() || lifecycle.isCancelled())) { + return; + } + ScanIterator iterator = getIterator(graph, queryRequest, wrapper, lifecycle::failedCleanup); KvPageRes.Builder resBuilder = KvPageRes.newBuilder(); Kv.Builder kvBuilder = Kv.newBuilder(); @@ -68,9 +78,11 @@ public static void scanOneShot(ScanStreamBatchReq request, } int count = 0; + boolean cleanupAttempted = false; try { - while (iterator.hasNext()) { + while (!(Context.current().isCancelled() || lifecycle.isCancelled()) && + !Thread.currentThread().isInterrupted() && iterator.hasNext()) { if (++count > limit) { break; @@ -86,6 +98,19 @@ public static void scanOneShot(ScanStreamBatchReq request, } + if ((Context.current().isCancelled() || lifecycle.isCancelled())) { + return; + } + if (Thread.currentThread().isInterrupted()) { + responseObserver.onError(HgGrpc.toErr(Status.Code.CANCELLED, "Scanning interrupted")); + return; + } + cleanupAttempted = true; + if (!lifecycle.close(iterator)) { + responseObserver.onError(Status.INTERNAL.withDescription("Failed to close scan iterator") + .withCause(lifecycle.cleanupFailure()).asRuntimeException()); + return; + } responseObserver.onNext(resBuilder.build()); responseObserver.onCompleted(); @@ -94,7 +119,10 @@ public static void scanOneShot(ScanStreamBatchReq request, responseObserver.onError( HgGrpc.toErr(Status.Code.INTERNAL, msg, t)); } finally { - iterator.close(); + if (!cleanupAttempted && !lifecycle.close(iterator)) { + responseObserver.onError(Status.INTERNAL.withDescription("Failed to close scan iterator") + .withCause(lifecycle.cleanupFailure()).asRuntimeException()); + } } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse.java index f4485e1527..25606b5b02 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse.java @@ -21,6 +21,7 @@ import java.util.List; import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; @@ -33,6 +34,7 @@ import org.apache.hugegraph.store.node.util.HgGrpc; import org.apache.hugegraph.store.node.util.PropertyUtil; +import io.grpc.Status; import io.grpc.stub.StreamObserver; import lombok.extern.slf4j.Slf4j; @@ -59,7 +61,7 @@ public class ScanBatchResponse implements StreamObserver { private final Object stateLock = new Object(); private final Lock iteratorLock = new ReentrantLock(); // Currently traversing iterator - private ScanIterator iterator; + private volatile ScanIterator iterator; // Next send sequence number private volatile int seqNo; // Client consumed sequence number @@ -72,10 +74,21 @@ public class ScanBatchResponse implements StreamObserver { // Last read data time private long activeTime; private volatile State state; + private final AtomicBoolean cancelled = new AtomicBoolean(); + private final ScanLifecycle lifecycle; + private boolean iteratorClosed; + private volatile boolean inputCompleted; + private boolean queryReceived; public ScanBatchResponse(StreamObserver response, HgStoreWrapperEx wrapper, ThreadPoolExecutor executor) { - this.sender = response; + this(response, wrapper, executor, new ScanLifecycle()); + } + + ScanBatchResponse(StreamObserver response, HgStoreWrapperEx wrapper, + ThreadPoolExecutor executor, ScanLifecycle lifecycle) { + this.lifecycle = lifecycle; + this.sender = lifecycle.response(response); this.wrapper = wrapper; this.executor = executor; this.iterator = null; @@ -92,45 +105,70 @@ public ScanBatchResponse(StreamObserver response, HgStoreWrapperEx wra */ @Override public void onNext(ScanStreamBatchReq request) { + if (this.cancelled.get()) { + return; + } switch (request.getQueryCase()) { case QUERY_REQUEST: // query conditions - executor.execute(() -> { + synchronized (this.stateLock) { + this.queryReceived = true; + } + submit(() -> { startQuery(request.getHeader().getGraph(), request.getQueryRequest()); }); break; case RECEIPT_REQUEST: // Message asynchronous response this.clientSeqNo = request.getReceiptRequest().getTimes(); if (seqNo - clientSeqNo < maxInFlightCount) { + boolean send = false; + boolean done; synchronized (stateLock) { + done = state == State.DONE; if (state == State.IDLE) { state = State.DOING; - executor.execute(() -> { - sendEntries(); - }); - } else if (state == State.DONE) { - sendNoDataEntries(); + send = true; } } + if (send) { + submit(this::sendEntries); + } else if (done) { + sendNoDataEntries(); + } } break; case CANCEL_REQUEST: // close stream - closeQuery(); + this.lifecycle.tryCancel(); + closeQuery(Status.CANCELLED.withDescription("Batch scan cancelled").asRuntimeException()); break; default: - sender.onError( - HgGrpc.toErr("Unsupported sub-request: [ " + request + " ]")); + closeQuery(HgGrpc.toErr("Unsupported sub-request: [ " + request + " ]")); } } @Override public void onError(Throwable t) { log.error("onError ", t); - closeQuery(); + this.lifecycle.tryCancel(); + closeQuery(t); } @Override public void onCompleted() { - closeQuery(); + boolean finish; + boolean resume = false; + synchronized (this.stateLock) { + this.inputCompleted = true; + finish = !this.queryReceived || this.state == State.DONE; + if (!finish && this.state == State.IDLE && this.iterator != null) { + this.state = State.DOING; + resume = true; + } + } + if (finish) { + closeQuery(); + } else if (resume) { + submit(this::sendEntries); + } } /** @@ -138,18 +176,43 @@ public void onCompleted() { * * @param request */ + private void submit(Runnable task) { + if (this.cancelled.get()) { + return; + } + try { + this.lifecycle.execute(this.executor, task); + } catch (io.grpc.StatusRuntimeException e) { + closeQuery(e); + } + } + private void startQuery(String graphName, ScanQueryRequest request) { - this.query = request; - this.limit = request.getLimit(); - this.count = 0; - this.iterator = getParallelIterator(graphName, request, this.wrapper, executor); - synchronized (stateLock) { - if (state == State.IDLE) { + this.iteratorLock.lock(); + try { + if (this.cancelled.get() || this.query != null) { + return; + } + this.query = request; + this.limit = request.getLimit(); + this.count = 0; + this.iterator = getParallelIterator(graphName, request, this.wrapper, executor, this.lifecycle); + if (this.cancelled.get()) { + closeIter(); + return; + } + synchronized (stateLock) { + if (state != State.IDLE) { + return; + } state = State.DOING; - executor.execute(() -> { - sendEntries(); - }); } + submit(this::sendEntries); + } catch (RuntimeException e) { + closeQuery(e); + log.warn("Failed to start scan", e); + } finally { + this.iteratorLock.unlock(); } } @@ -157,25 +220,51 @@ private void startQuery(String graphName, ScanQueryRequest request) { * Generate iterator */ private void closeQuery() { + closeQuery(null); + } + + private void closeQuery(Throwable error) { + if (error instanceof java.util.concurrent.RejectedExecutionException) { + error = ScanLifecycle.rejected(this.executor, error); + } + if (!this.cancelled.compareAndSet(false, true)) { + return; + } setStateDone(); + ScanIterator current = this.iterator; + if (current instanceof ParallelScanIterator) { + ((ParallelScanIterator) current).requestStop(); + } try { closeIter(); - this.sender.onCompleted(); - } catch (Exception e) { - log.error("exception ", e); + if (this.lifecycle.cleanupFailure() != null) { + error = Status.INTERNAL.withDescription("Failed to close batch scan iterator") + .withCause(this.lifecycle.cleanupFailure()).asRuntimeException(); + } + if (error == null) { + this.sender.onCompleted(); + } else { + this.sender.onError(error); + } + } catch (RuntimeException | Error e) { + log.error("Failed to terminate batch scan response", e); + } finally { + int active = ScanBatchResponseFactory.getInstance().removeStreamObserver(this); + log.info("ScanBatchResponse closeQuery, active count is {}", active); } - int active = ScanBatchResponseFactory.getInstance().removeStreamObserver(this); - log.info("ScanBatchResponse closeQuery, active count is {}", active); } private void closeIter() { + this.iteratorLock.lock(); try { - if (this.iterator != null) { - this.iterator.close(); - this.iterator = null; + if (this.iterator != null && !this.iteratorClosed) { + this.iteratorClosed = true; + if (this.lifecycle.close(this.iterator)) { + this.iterator = null; + } } - } catch (Exception e) { - + } finally { + this.iteratorLock.unlock(); } } @@ -209,24 +298,38 @@ private void sendEntries() { this.sender.onNext(dataBuilder.build()); this.activeTime = System.currentTimeMillis(); } - if (!iterator.hasNext() || this.count >= limit || state == State.DONE) { + if (this.cancelled.get()) { + return; + } + if (this.count >= limit || !iterator.hasNext()) { closeIter(); + if (this.lifecycle.cleanupFailure() != null) { + closeQuery(this.lifecycle.cleanupFailure()); + return; + } this.sender.onNext(KvStream.newBuilder().setOver(true).build()); setStateDone(); + if (this.inputCompleted) { + closeQuery(); + } } else { - setStateIdle(); + boolean halfClosed; + synchronized (this.stateLock) { + halfClosed = this.inputCompleted; + if (!halfClosed && this.state != State.DONE) { + this.state = State.IDLE; + } + } + if (halfClosed) { + closeQuery(Status.FAILED_PRECONDITION + .withDescription("Batch scan half-closed before sufficient receipts") + .asRuntimeException()); + } } } catch (Throwable e) { if (this.state != State.DONE) { log.error(" send data exception: ", e); - setStateIdle(); - if (this.sender != null) { - try { - this.sender.onError(e); - } catch (Exception ex) { - log.warn("Error when call sender.onError {}", e.getMessage()); - } - } + closeQuery(e); } } finally { iteratorLock.unlock(); @@ -234,9 +337,13 @@ private void sendEntries() { } private void sendNoDataEntries() { + this.iteratorLock.lock(); try { - this.sender.onNext(KvStream.newBuilder().setOver(true).build()); - } catch (Exception e) { + if (!this.cancelled.get()) { + this.sender.onNext(KvStream.newBuilder().setOver(true).build()); + } + } finally { + this.iteratorLock.unlock(); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse3.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse3.java index d0c52b372e..3c39ac40f5 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse3.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse3.java @@ -61,7 +61,13 @@ public static StreamObserver of(StreamObserver responseObserver, HgStoreWrapperEx wrapper, ThreadPoolExecutor executor) { HgAssert.isArgumentNotNull(responseObserver, "responseObserver"); HgAssert.isArgumentNotNull(wrapper, "wrapper"); - return new Broker(responseObserver, wrapper, executor); + return of(responseObserver, wrapper, executor, new ScanLifecycle()); + } + + static StreamObserver of(StreamObserver responseObserver, + HgStoreWrapperEx wrapper, ThreadPoolExecutor executor, + ScanLifecycle lifecycle) { + return new Broker(lifecycle.response(responseObserver), wrapper, executor, lifecycle); } private enum OrderState { @@ -81,10 +87,12 @@ private static class Broker implements StreamObserver { private final HgStoreWrapperEx wrapper; private final ThreadPoolExecutor executor; private final OrderManager manager = new OrderManager(); + private final ScanLifecycle lifecycle; private String graph; Broker(StreamObserver responseObserver, HgStoreWrapperEx wrapper, - ThreadPoolExecutor executor) { + ThreadPoolExecutor executor, ScanLifecycle lifecycle) { + this.lifecycle = lifecycle; this.responseObserver = responseObserver; this.wrapper = wrapper; this.executor = executor; @@ -101,23 +109,27 @@ public void onNext(ScanStreamBatchReq request) { this.manager.receipt(request.getReceiptRequest().getTimes()); break; case CANCEL_REQUEST: - this.manager.finished(); + this.lifecycle.tryCancel(); + this.onError(Status.CANCELLED.withDescription("Batch scan cancelled").asRuntimeException()); break; default: - responseObserver.onError( - HgGrpc.toErr("Unsupported sub-request: [ " + request + " ]")); + this.onError(HgGrpc.toErr("Unsupported sub-request: [ " + request + " ]")); } } @Override public void onError(Throwable t) { log.warn(t.getMessage()); + this.lifecycle.tryCancel(); this.manager.breakdown(); + this.responseObserver.onError(t); } @Override public void onCompleted() { - this.manager.finished(); + if (this.manager.inputCompleted()) { + this.responseObserver.onCompleted(); + } } private void handleHeader(ScanStreamBatchReq request) { @@ -145,9 +157,9 @@ private void makeADeal(ScanQueryRequest request) { OrderWorker worker = new OrderWorker( request.getLimit(), request.getPageSize(), - ScanUtil.getIterator(this.graph, request, this.wrapper), + ScanUtil.getIterator(this.graph, request, this.wrapper, this.lifecycle::failedCleanup), deliverer, - this.executor); + this.executor, this.lifecycle); this.manager.deal(worker, deliverer); } @@ -160,16 +172,19 @@ private static class OrderManager { OrderState state = OrderState.NEW; OrderWorker worker; OrderDeliverer deliverer; + boolean cancelled; synchronized void deal(OrderWorker worker, OrderDeliverer deliverer) { if (log.isDebugEnabled()) { log.debug("Receiving query request."); } - if (this.state == OrderState.NEW) { + if (!this.cancelled && this.state == OrderState.NEW) { this.worker = worker; this.deliverer = deliverer; this.worker.hereWeGo(); this.state = OrderState.WORKING; + } else { + worker.closeIterator(); } } @@ -177,17 +192,24 @@ synchronized void receipt(int receiptTimes) { if (log.isDebugEnabled()) { log.debug("Receiving receipt request."); } - this.worker.setReceipt(receiptTimes); + if (this.worker != null) { + this.worker.setReceipt(receiptTimes); + } } - synchronized void finished() { - if (log.isDebugEnabled()) { - log.debug("Receiving finished request."); + synchronized boolean inputCompleted() { + if (this.cancelled) { + return false; } - this.breakdown(); + if (this.worker == null) { + return true; + } + this.worker.inputCompleted(); + return false; } synchronized void breakdown() { + this.cancelled = true; if (this.worker != null) { this.worker.breakdown(); } @@ -240,9 +262,12 @@ void error(String msg) { } void error(String msg, Throwable t) { + this.error(Status.INTERNAL, msg, t); + } + + void error(Status status, String msg, Throwable t) { if (!finishFlag.getAndSet(true)) { - this.responseObserver.onError(HgGrpc.toErr(Status.INTERNAL, - msg, t)); + this.responseObserver.onError(HgGrpc.toErr(status, msg, t)); } } } @@ -261,11 +286,17 @@ private static class OrderWorker { private final AtomicInteger curTimes = new AtomicInteger(); private final ThreadPoolExecutor executor; private final long limit; + private final ScanLifecycle lifecycle; + private final AtomicBoolean iteratorClosed = new AtomicBoolean(); + private final Object readerLock = new Object(); + private Thread reader; private long packageSize; private long counter; + private volatile boolean inputCompleted; OrderWorker(long limit, long packageSize, ScanIterator iterator, OrderDeliverer deliverer, - ThreadPoolExecutor executor) { + ThreadPoolExecutor executor, ScanLifecycle lifecycle) { + this.lifecycle = lifecycle; this.limit = limit; this.packageSize = packageSize; this.iterator = iterator; @@ -297,7 +328,13 @@ void hereWeGo() { return; } - executor.execute(() -> working()); + try { + this.lifecycle.execute(executor, this::working); + } catch (io.grpc.StatusRuntimeException e) { + this.completeFlag.set(true); + closeIterator(); + this.deliverer.error(e.getStatus(), "Store scan task rejected", e); + } Thread.yield(); } @@ -318,18 +355,51 @@ void continueWorking() { } } + void inputCompleted() { + this.inputCompleted = true; + synchronized (this.iterator) { + this.iterator.notify(); + } + } + void breakdown() { this.breakdown.set(true); + // Interrupt before waiting for the iterator monitor held during hasNext(). + synchronized (this.readerLock) { + if (this.reader != null && this.reader != Thread.currentThread()) { + this.reader.interrupt(); + } + } synchronized (this.iterator) { this.iterator.notify(); } } + private void closeIterator() { + if (this.iteratorClosed.compareAndSet(false, true) && + !this.lifecycle.close(this.iterator)) { + this.deliverer.error("Failed to close scan iterator", this.lifecycle.cleanupFailure()); + } + } + private void working() { if (this.isWorking.getAndSet(true)) { return; } + boolean startReading; + synchronized (this.readerLock) { + startReading = !this.breakdown.get() && !this.completeFlag.get(); + if (startReading) { + this.reader = Thread.currentThread(); + } + } + if (!startReading) { + this.completeFlag.set(true); + closeIterator(); + return; + } + this.workingLock.lock(); try { @@ -338,7 +408,8 @@ private void working() { Kv.Builder kvBuilder = Kv.newBuilder(); long packageCount = 0; - while (iterator.hasNext()) { + while (!this.breakdown.get() && !Thread.currentThread().isInterrupted() && + iterator.hasNext()) { if (++this.counter > limit) { this.completeFlag.set(true); break; @@ -353,25 +424,26 @@ private void working() { deliverer.deliver(dataBuilder, curTimes.incrementAndGet(), false); Thread.yield(); - if (!this.checkContinue()) { + if (!this.breakdown.get() && !this.checkContinue()) { long start = System.currentTimeMillis(); - iterator.wait( - HgStoreConst.SCAN_WAIT_CLIENT_TAKING_TIME_OUT_SECONDS * - 1000); - - if (System.currentTimeMillis() - start - >= - HgStoreConst.SCAN_WAIT_CLIENT_TAKING_TIME_OUT_SECONDS * 1000) { - throw new TimeoutException("Waiting continue more than " - + - HgStoreConst.SCAN_WAIT_CLIENT_TAKING_TIME_OUT_SECONDS + - " seconds."); + while (!this.breakdown.get() && !this.checkContinue()) { + if (this.inputCompleted) { + throw Status.FAILED_PRECONDITION.withDescription( + "Batch scan half-closed before sufficient receipts") + .asRuntimeException(); + } + long remaining = HgStoreConst.SCAN_WAIT_CLIENT_TAKING_TIME_OUT_SECONDS * + 1000L - (System.currentTimeMillis() - start); + if (remaining <= 0) { + throw new TimeoutException("Waiting continue more than " + + HgStoreConst.SCAN_WAIT_CLIENT_TAKING_TIME_OUT_SECONDS + + " seconds."); + } + iterator.wait(remaining); } - if (this.breakdown.get()) { break; } - } packageCount = 1; @@ -390,15 +462,24 @@ private void working() { } + if (!this.breakdown.get() && Thread.currentThread().isInterrupted()) { + throw new InterruptedException("Scanning interrupted"); + } this.completeFlag.set(true); - deliverer.deliver(dataBuilder, curTimes.incrementAndGet(), true); + closeIterator(); + if (!this.breakdown.get() && this.lifecycle.cleanupFailure() == null) { + deliverer.deliver(dataBuilder, curTimes.incrementAndGet(), true); + } } + } catch (io.grpc.StatusRuntimeException e) { + this.deliverer.error(e.getStatus(), e.getMessage(), e); } catch (InterruptedException e) { log.error("Interrupted waiting of iterator, canceled while.", e); - this.deliverer.error("Failed to finish scanning, cause by InterruptedException."); + this.deliverer.error(Status.CANCELLED, + "Failed to finish scanning, cause by InterruptedException.", e); } catch (TimeoutException t) { log.info(t.getMessage()); this.deliverer.error("Sever waiting exceeded [" @@ -408,8 +489,16 @@ private void working() { log.error("Failed to do while for scanning, cause by:", t); this.deliverer.error("Failed to finish scanning ", t); } finally { - this.workingLock.unlock(); - this.iterator.close(); + try { + this.workingLock.unlock(); + closeIterator(); + } finally { + synchronized (this.readerLock) { + if (this.reader == Thread.currentThread()) { + this.reader = null; + } + } + } } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponseFactory.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponseFactory.java index 43abfd97ce..dbb5406387 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponseFactory.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponseFactory.java @@ -37,12 +37,19 @@ public static ScanBatchResponseFactory getInstance() { public static StreamObserver of(StreamObserver responseObserver, HgStoreWrapperEx wrapper, ThreadPoolExecutor executor) { - StreamObserver observer = new ScanBatchResponse(responseObserver, wrapper, executor); - getInstance().addStreamObserver(observer); + StreamObserver observer = of(responseObserver, wrapper, executor, new ScanLifecycle()); getInstance().checkStreamActive(); return observer; } + static StreamObserver of(StreamObserver responseObserver, + HgStoreWrapperEx wrapper, ThreadPoolExecutor executor, + ScanLifecycle lifecycle) { + StreamObserver observer = new ScanBatchResponse(responseObserver, wrapper, executor, lifecycle); + getInstance().addStreamObserver(observer); + return observer; + } + public int addStreamObserver(StreamObserver observer) { streamObservers.add(observer); return streamObservers.size(); diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanLifecycle.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanLifecycle.java new file mode 100644 index 0000000000..6cce65ad02 --- /dev/null +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanLifecycle.java @@ -0,0 +1,192 @@ +/* + * 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.store.node.grpc; + +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +import org.apache.hugegraph.rocksdb.access.ScanIterator; + +import io.grpc.Status; +import io.grpc.stub.StreamObserver; + +/** Per-call response and cleanup accounting; a failed release remains registered. */ +final class ScanLifecycle { + + // 0 = open, 1 = cancellation won, 2 = response terminated. + private final AtomicInteger terminal = new AtomicInteger(); + private final Object responseLock = new Object(); + private int activities; + private volatile Throwable cleanupFailure; + private Runnable finished = () -> { }; + + synchronized void onFinished(Runnable action) { + this.finished = action; + } + + synchronized void enter() { + this.activities++; + } + + synchronized boolean tryEnter() { + if (this.terminal.get() != 0) { + return false; + } + this.activities++; + return true; + } + + synchronized boolean tryEnterCancellation() { + if (!tryCancel()) { + return false; + } + this.activities++; + return true; + } + + synchronized void leave() { + this.activities--; + checkFinished(); + } + + synchronized void failedCleanup(Throwable failure) { + if (this.cleanupFailure == null) { + this.cleanupFailure = failure; + } + } + + Throwable cleanupFailure() { + return this.cleanupFailure; + } + + private synchronized void checkFinished() { + if (this.terminal.get() == 2 && this.activities == 0 && this.cleanupFailure == null) { + this.finished.run(); + } + } + + boolean tryCancel() { + return this.terminal.compareAndSet(0, 1); + } + + boolean isCancelled() { + return this.terminal.get() == 1; + } + + void finishWithoutResponse() { + this.terminal.set(2); + checkFinished(); + } + + boolean close(ScanIterator iterator) { + if (iterator == null) { + return true; + } + try { + iterator.close(); + return true; + } catch (RuntimeException | Error failure) { + failedCleanup(failure); + return false; + } + } + + void execute(ThreadPoolExecutor executor, Runnable task) { + enter(); + AtomicBoolean released = new AtomicBoolean(); + Runnable release = () -> { + if (released.compareAndSet(false, true)) { + leave(); + } + }; + try { + executor.execute(() -> { + try { + task.run(); + } finally { + release.run(); + } + }); + } catch (RuntimeException | Error failure) { + release.run(); + if (failure instanceof RejectedExecutionException) { + throw rejected(executor, failure); + } + throw failure; + } + } + + static RuntimeException rejected(ThreadPoolExecutor executor, Throwable failure) { + Status status = executor.isShutdown() ? Status.UNAVAILABLE : Status.RESOURCE_EXHAUSTED; + return status.withDescription("Store scan task rejected").withCause(failure) + .asRuntimeException(); + } + + StreamObserver response(StreamObserver delegate) { + return new StreamObserver() { + @Override + public void onNext(T value) { + enter(); + try { + synchronized (responseLock) { + if (terminal.get() == 0) { + delegate.onNext(value); + } + } + } finally { + leave(); + } + } + + @Override + public void onError(Throwable failure) { + enter(); + try { + synchronized (responseLock) { + if (terminal.getAndSet(2) != 2) { + delegate.onError(failure); + } + } + } finally { + leave(); + } + } + + @Override + public void onCompleted() { + enter(); + try { + synchronized (responseLock) { + int previous = terminal.getAndSet(2); + if (previous == 0) { + delegate.onCompleted(); + } else if (previous == 1) { + delegate.onError(Status.CANCELLED + .withDescription("Store scan cancelled before completion") + .asRuntimeException()); + } + } + } finally { + leave(); + } + } + }; + } +} diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java index 5b77a4449a..e800cc1166 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java @@ -27,6 +27,7 @@ import com.google.protobuf.ByteString; +import io.grpc.Context; import io.grpc.Status; import io.grpc.stub.StreamObserver; import lombok.extern.slf4j.Slf4j; @@ -48,6 +49,12 @@ public class ScanOneShotResponse { public static void scanOneShot(ScanStreamReq request, StreamObserver responseObserver, HgStoreWrapperEx wrapper) { + scanOneShot(request, responseObserver, wrapper, new ScanLifecycle()); + } + + static void scanOneShot(ScanStreamReq request, StreamObserver output, + HgStoreWrapperEx wrapper, ScanLifecycle lifecycle) { + StreamObserver responseObserver = lifecycle.response(output); KvPageRes.Builder resBuilder = KvPageRes.newBuilder(); Kv.Builder kvBuilder = Kv.newBuilder(); @@ -57,12 +64,17 @@ public static void scanOneShot(ScanStreamReq request, responseObserver.onError(HgGrpc.toErr("limit<=0, please to invoke stream scan.")); return; } - ScanIterator iterator = ScanUtil.getIterator(request, wrapper); + if ((Context.current().isCancelled() || lifecycle.isCancelled())) { + return; + } + ScanIterator iterator = ScanUtil.getIterator(request, wrapper, lifecycle::failedCleanup); int count = 0; + boolean cleanupAttempted = false; try { - while (iterator.hasNext()) { + while (!(Context.current().isCancelled() || lifecycle.isCancelled()) && + !Thread.currentThread().isInterrupted() && iterator.hasNext()) { if (++count > limit) { break; @@ -78,6 +90,19 @@ public static void scanOneShot(ScanStreamReq request, } + if ((Context.current().isCancelled() || lifecycle.isCancelled())) { + return; + } + if (Thread.currentThread().isInterrupted()) { + responseObserver.onError(HgGrpc.toErr(Status.Code.CANCELLED, "Scanning interrupted")); + return; + } + cleanupAttempted = true; + if (!lifecycle.close(iterator)) { + responseObserver.onError(Status.INTERNAL.withDescription("Failed to close scan iterator") + .withCause(lifecycle.cleanupFailure()).asRuntimeException()); + return; + } responseObserver.onNext( resBuilder.setVersion(ScanUtil.responseVersion(request)) .build()); @@ -87,7 +112,10 @@ public static void scanOneShot(ScanStreamReq request, String msg = "an exception occurred during data scanning"; responseObserver.onError(HgGrpc.toErr(Status.INTERNAL, msg, t)); } finally { - iterator.close(); + if (!cleanupAttempted && !lifecycle.close(iterator)) { + responseObserver.onError(Status.INTERNAL.withDescription("Failed to close scan iterator") + .withCause(lifecycle.cleanupFailure()).asRuntimeException()); + } } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java index e572862823..8628e0c92b 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java @@ -61,6 +61,21 @@ public class ScanStreamResponse implements StreamObserver { private final int waitTime; private final HgChannel channel; private ScanIterator iterator; + private final Object cancellationLock = new Object(); + private final Object responseLock = new Object(); + private Thread worker; + private final ScanLifecycle lifecycle; + + private void cancel() { + this.isStop.set(true); + this.channel.close(); + Thread current = Thread.currentThread(); + synchronized (this.cancellationLock) { + if (this.worker != null && this.worker != current) { + this.worker.interrupt(); + } + } + } private long limit = 0; private int times = 0; private long pageSize = 0; @@ -72,7 +87,14 @@ public class ScanStreamResponse implements StreamObserver { ScanStreamResponse(StreamObserver responseObserver, HgStoreWrapperEx wrapper, ThreadPoolExecutor executor, AppConfig appConfig) { - this.responseObserver = responseObserver; + this(responseObserver, wrapper, executor, appConfig, new ScanLifecycle()); + } + + ScanStreamResponse(StreamObserver responseObserver, + HgStoreWrapperEx wrapper, ThreadPoolExecutor executor, + AppConfig appConfig, ScanLifecycle lifecycle) { + this.lifecycle = lifecycle; + this.responseObserver = lifecycle.response(responseObserver); this.wrapper = wrapper; this.executor = executor; this.config = appConfig; @@ -91,6 +113,9 @@ public static ScanStreamResponse of(StreamObserver responseObserver, @Override public void onNext(ScanStreamReq request) { + if (this.isStop.get()) { + return; + } try { this.responseVersion = Math.max( this.responseVersion, ScanUtil.responseVersion(request)); @@ -100,22 +125,20 @@ public void onNext(ScanStreamReq request) { next(request); } } catch (Exception e) { - responseObserver.onError(e); + this.failServer(e); } } @Override public void onError(Throwable t) { - this.isStop.set(true); - this.finishServer(); + this.cancelServer(false); log.warn("onError from client [ graph: {} , table: {}]; Reason: {}]", graph, table, t.getMessage()); } @Override public void onCompleted() { - this.isStop.set(true); - this.finishServer(); + this.cancelServer(false); } private void initIterator(ScanStreamReq request) { @@ -123,7 +146,6 @@ private void initIterator(ScanStreamReq request) { if (this.isStarted.getAndSet(true)) { return; } - this.iterator = getIterator(request, this.wrapper); this.graph = request.getHeader().getGraph(); this.table = request.getTable(); this.limit = request.getLimit(); @@ -143,7 +165,15 @@ private void initIterator(ScanStreamReq request) { Kv.Builder kvBuilder = Kv.newBuilder(); int pageCount = 0; try { - while (iterator.hasNext()) { + synchronized (this.cancellationLock) { + if (this.isStop.get()) { + return; + } + this.worker = Thread.currentThread(); + } + this.iterator = getIterator(request, this.wrapper, this.lifecycle::failedCleanup); + while (!this.isStop.get() && !Thread.currentThread().isInterrupted() && + iterator.hasNext()) { if (limit > 0 && ++this.total > limit) { break; } @@ -164,32 +194,45 @@ private void initIterator(ScanStreamReq request) { } dataBuilder.addData(toKv(kvBuilder, iterator.next(), iterator.position())); } + if (this.isStop.get()) { + return; + } + if (Thread.currentThread().isInterrupted()) { + this.failServer(HgGrpc.toErr(Status.Code.CANCELLED, "Scanning interrupted")); + return; + } this.channel.send(dataBuilder); } catch (Throwable t) { + if (this.isStop.get()) { + return; + } String msg = "an exception occurred while scanning data:"; - StatusRuntimeException ex = - HgGrpc.toErr(Status.INTERNAL, msg + t.getMessage(), t); - responseObserver.onError(ex); + Status status = t instanceof InterruptedException || + Thread.currentThread().isInterrupted() ? + Status.CANCELLED : Status.INTERNAL; + StatusRuntimeException ex = HgGrpc.toErr(status, msg + t.getMessage(), t); + this.failServer(ex); } finally { try { - this.iterator.close(); + if (!this.lifecycle.close(this.iterator)) { + this.failServer(Status.INTERNAL.withDescription("Failed to close scan iterator") + .withCause(this.lifecycle.cleanupFailure()) + .asRuntimeException()); + } + } finally { this.channel.close(); - } catch (Exception e) { - + synchronized (this.cancellationLock) { + this.worker = null; + } } } }; - this.executor.execute(scanning); + this.lifecycle.execute(this.executor, scanning); } catch (Exception e) { - StatusRuntimeException ex = HgGrpc.toErr(Status.INTERNAL, null, e); - responseObserver.onError(ex); - try { - this.iterator.close(); - this.channel.close(); - } catch (Exception exception) { - - } + StatusRuntimeException ex = e instanceof StatusRuntimeException ? + (StatusRuntimeException) e : HgGrpc.toErr(Status.INTERNAL, null, e); + this.failServer(ex); } /*** Scanning loop end ***/ @@ -205,19 +248,7 @@ private Kv toKv(Kv.Builder kvBuilder, RocksDBSession.BackendColumn col, } private void close() { - this.isStop.set(true); - this.channel.close(); - if (!this.finishFlag.get()) { - responseObserver.onNext(KvPageRes.newBuilder() - .addAllData(Collections.EMPTY_LIST) - .setOver(true) - .setTimes(++times) - .setVersion(this.responseVersion) - .build() - ); - } - - this.finishServer(); + this.cancelServer(true); } private void next(ScanStreamReq request) { @@ -225,12 +256,19 @@ private void next(ScanStreamReq request) { KvPageRes.Builder resBuilder; try { + if (this.isStop.get()) { + return; + } resBuilder = this.channel.receive(); times++; } catch (Exception e) { + if (this.isStop.get()) { + return; + } String msg = "failed to poll a page of data, cause by:"; + Status status = Thread.currentThread().isInterrupted() ? Status.CANCELLED : Status.INTERNAL; log.error(msg, e); - responseObserver.onError(HgGrpc.toErr(msg + e.getMessage())); + this.failServer(HgGrpc.toErr(status, msg + e.getMessage(), e)); return; } boolean isOver = false; @@ -239,31 +277,65 @@ private void next(ScanStreamReq request) { isOver = true; resBuilder = KvPageRes.newBuilder().addAllData(Collections.EMPTY_LIST); } - if (!this.finishFlag.get()) { - responseObserver.onNext(resBuilder.setOver(isOver) - .setTimes(times) - .setVersion(this.responseVersion) - .build()); - } + this.sendPage(resBuilder.setOver(isOver) + .setTimes(times) + .setVersion(this.responseVersion) + .build()); if (isOver) { this.finishServer(); } } + private void cancelServer(boolean sendFinalPage) { + // Cancellation can make the worker throw; establish the normal terminal first. + boolean complete = this.finishFlag.compareAndSet(false, true); + this.cancel(); + if (complete) { + synchronized (this.responseLock) { + if (sendFinalPage) { + this.responseObserver.onNext(KvPageRes.newBuilder() + .setOver(true) + .setTimes(++times) + .setVersion(this.responseVersion) + .build()); + } + this.responseObserver.onCompleted(); + } + } + } + + private void sendPage(KvPageRes page) { + synchronized (this.responseLock) { + if (!this.finishFlag.get()) { + this.responseObserver.onNext(page); + } + } + } + + private void failServer(Throwable failure) { + // Claim the terminal signal before waking a receiver on the closed channel. + boolean reportFailure = this.finishFlag.compareAndSet(false, true); + this.cancel(); + if (reportFailure) { + synchronized (this.responseLock) { + this.responseObserver.onError(failure); + } + } + } + private void finishServer() { if (!this.finishFlag.getAndSet(true)) { - responseObserver.onCompleted(); + synchronized (this.responseLock) { + this.responseObserver.onCompleted(); + } } } private void timeoutSever() { - if (!this.finishFlag.getAndSet(true)) { - String msg = "server wait time exceeds the threshold[" + waitTime + - "] seconds."; - responseObserver.onError( - HgGrpc.toErr(Status.Code.DEADLINE_EXCEEDED, msg)); - } + String msg = "server wait time exceeds the threshold[" + waitTime + + "] seconds."; + this.failServer(HgGrpc.toErr(Status.Code.DEADLINE_EXCEEDED, msg)); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java index b7678f91d9..8f873cbaeb 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java @@ -29,6 +29,7 @@ import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Supplier; +import java.util.function.Consumer; import java.util.stream.Collectors; import javax.annotation.concurrent.NotThreadSafe; @@ -58,6 +59,11 @@ class ScanUtil { private final static Map tableKeyMap = new HashMap<>(); static ScanIterator getIterator(ScanStreamReq request, HgStoreWrapperEx wrapper) { + return getIterator(request, wrapper, ignored -> { }); + } + + static ScanIterator getIterator(ScanStreamReq request, HgStoreWrapperEx wrapper, + Consumer cleanupFailure) { boolean ordered = isOrdered(request); if (ordered && !request.getPosition().isEmpty()) { throw new IllegalArgumentException( @@ -109,16 +115,19 @@ static ScanIterator getIterator(ScanStreamReq request, HgStoreWrapperEx wrapper) throw new IllegalArgumentException("Property selection requires a vertex or edge table: " + table); } iter = new SelectIterator(iter, properties, BusinessHandlerImpl.getGraphSupplier(graph), isVertex); - } catch (RuntimeException failure) { + } catch (RuntimeException | Error failure) { try { iter.close(); - } catch (RuntimeException closeFailure) { - failure.addSuppressed(closeFailure); + } catch (RuntimeException | Error closeFailure) { + cleanupFailure.accept(closeFailure); + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } } throw failure; } } - iter.seek(request.getPosition().toByteArray()); + seek(iter, request.getPosition().toByteArray(), cleanupFailure); return iter; } @@ -137,6 +146,11 @@ static int responseVersion(ScanStreamReq request) { } static ScanIterator getIterator(ScanQuery sq, HgStoreWrapperEx wrapper) { + return getIterator(sq, wrapper, ignored -> { }); + } + + private static ScanIterator getIterator(ScanQuery sq, HgStoreWrapperEx wrapper, + Consumer cleanupFailure) { if (log.isDebugEnabled()) { log.debug("{}", sq); } @@ -161,12 +175,29 @@ static ScanIterator getIterator(ScanQuery sq, HgStoreWrapperEx wrapper) { iter = new EmptyIterator(); } - iter.seek(sq.position); + seek(iter, sq.position, cleanupFailure); return iter; } + private static void seek(ScanIterator iterator, byte[] position, + Consumer cleanupFailure) { + try { + iterator.seek(position); + } catch (RuntimeException | Error failure) { + try { + iterator.close(); + } catch (RuntimeException | Error closing) { + cleanupFailure.accept(closing); + if (failure != closing) { + failure.addSuppressed(closing); + } + } + throw failure; + } + } + static ScanQuery toSq(ScanStreamReq request) { ScanQuery res = ScanQuery.of(); @@ -192,7 +223,12 @@ static ScanQuery toSq(ScanStreamReq request) { static ScanIterator getIterator(String graph, ScanQueryRequest request, HgStoreWrapperEx wrapper) { - ScanIteratorSupplier supplier = new ScanIteratorSupplier(graph, request, wrapper); + return getIterator(graph, request, wrapper, ignored -> { }); + } + + static ScanIterator getIterator(String graph, ScanQueryRequest request, + HgStoreWrapperEx wrapper, Consumer cleanupFailure) { + ScanIteratorSupplier supplier = new ScanIteratorSupplier(graph, request, wrapper, cleanupFailure); return BatchScanIterator.of(supplier, supplier.getLimitSupplier()); } @@ -200,10 +236,11 @@ static ScanIterator getIterator(String graph, ScanQueryRequest request, * Support for multi-iterators with parallel reading */ static ScanIterator getParallelIterator(String graph, ScanQueryRequest request, - HgStoreWrapperEx wrapper, ThreadPoolExecutor executor) { - ScanIteratorSupplier supplier = new ScanIteratorSupplier(graph, request, wrapper); + HgStoreWrapperEx wrapper, ThreadPoolExecutor executor, + ScanLifecycle lifecycle) { + ScanIteratorSupplier supplier = new ScanIteratorSupplier(graph, request, wrapper, lifecycle::failedCleanup); return ParallelScanIterator.of(supplier, supplier.getLimitSupplier(), - request, executor); + request, executor, lifecycle::failedCleanup); } @NotThreadSafe @@ -216,6 +253,7 @@ private static class ScanIteratorSupplier implements private final long perKeyMax; private final long skipDegree; private final HgStoreWrapperEx wrapper; + private final Consumer cleanupFailure; private long perKeyLimit; private List sqs = new LinkedList<>(); private Iterator sqIterator; @@ -223,7 +261,9 @@ private static class ScanIteratorSupplier implements private ScanQueryProducer scanQueryProducer; private Iterator scanQueryIterator; - ScanIteratorSupplier(String graph, ScanQueryRequest request, HgStoreWrapperEx wrapper) { + ScanIteratorSupplier(String graph, ScanQueryRequest request, HgStoreWrapperEx wrapper, + Consumer cleanupFailure) { + this.cleanupFailure = cleanupFailure; this.graph = graph; this.perKeyLimit = request.getPerKeyLimit(); this.perKeyMax = request.getPerKeyMax(); @@ -311,7 +351,7 @@ public KVPair get1() { ScanQuery query = null; if (this.sqIterator != null && this.sqIterator.hasNext()) { query = this.sqIterator.next(); - iterator = getIterator(query, this.wrapper); + iterator = getIterator(query, this.wrapper, this.cleanupFailure); } else { this.sqs.clear(); this.sqIterator = null; @@ -369,7 +409,7 @@ public ScanIterator get() { if (index + 1 > queries.length) { return null; } - return getIterator(queries[index++], wrapper); + return getIterator(queries[index++], wrapper, cleanupFailure); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryObserver.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryObserver.java index 0ba569cb92..973ab459f3 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryObserver.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryObserver.java @@ -21,13 +21,18 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.HashSet; import java.util.List; +import java.util.Set; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import org.apache.hugegraph.backend.BackendColumn; import org.apache.hugegraph.rocksdb.access.RocksDBSession; @@ -47,6 +52,8 @@ import com.google.protobuf.ByteString; +import io.grpc.Context; +import io.grpc.Status; import io.grpc.stub.StreamObserver; import lombok.extern.slf4j.Slf4j; @@ -65,18 +72,193 @@ public class AggregativeQueryObserver implements StreamObserver { private volatile ScanIterator iterator = null; private QueryPlan plan = null; private String queryId; + private final Consumer completion; + private final Set workers = new HashSet<>(); + private int pendingTasks; + private boolean finished; + private volatile boolean requestCompleted; + private final Object responseLock = new Object(); + private boolean responseFinished; + private boolean completeResponse = true; + private QueryResponse finalResponse; + private boolean errorReported; + private final AtomicReference cleanupFailure = new AtomicReference<>(); public AggregativeQueryObserver(StreamObserver sender, ExecutorService threadPool, long timeout, int batchSize) { + this(sender, threadPool, timeout, batchSize, ignored -> { }); + } + + AggregativeQueryObserver(StreamObserver sender, + ExecutorService threadPool, long timeout, int batchSize, + Consumer completion) { this.sender = sender; this.threadPool = threadPool; this.batchSize = batchSize; this.timeout = timeout; + this.completion = completion; + } + + ScanIterator getIterator(QueryRequest request) { + return QueryUtil.getIterator(request); + } + + ScanIterator getCountIterator(QueryRequest request) { + return new QueryUtil().getHandler().scanAll(request.getGraph(), request.getTable()); + } + + void runUnary(QueryRequest request, boolean countOnly) { + synchronized (this) { + if (this.finished || this.clientCanceled.get()) { + return; + } + this.queryId = request.getQueryId(); + this.pendingTasks = 1; + this.workers.add(Thread.currentThread()); + } + Context context = Context.current(); + Context.CancellationListener cancellation = ignored -> onTransportCancel(); + Throwable failure = null; + try { + // Unary callbacks occupy the call executor while reading. Context cancellation + // must interrupt that read without waiting for the queued onCancel callback. + context.addListener(cancellation, Runnable::run); + if (this.clientCanceled.get()) { + return; + } + this.iterator = countOnly ? getCountIterator(request) : getIterator(request); + QueryResponse.Builder response = getBuilder(); + Kv.Builder kv = getKvBuilder(); + long count = countOnly ? countPartitions() : 0; + while (!countOnly && !this.clientCanceled.get() && this.iterator.hasNext()) { + Object next = this.iterator.next(); + if (next != null) { + RocksDBSession.BackendColumn column = (RocksDBSession.BackendColumn) next; + response.addData(kv.setKey(ByteString.copyFrom(column.name)) + .setValue(column.value == null ? ByteString.EMPTY : + ByteString.copyFrom(column.value)).build()); + } + } + if (countOnly) { + List counts = new ArrayList<>(); + for (int i = 0; i < request.getFunctionsCount(); i++) { + counts.add(new AtomicLong(count)); + } + response.addData(kv.setKey(ByteString.copyFrom(KvSerializer.toBytes(List.of()))) + .setValue(ByteString.copyFrom(KvSerializer.toBytes(counts))).build()); + } + synchronized (this.responseLock) { + this.finalResponse = response.setQueryId(this.queryId).setIsOk(true) + .setIsFinished(true).build(); + } + } catch (RuntimeException | Error e) { + failure = e; + synchronized (this.responseLock) { + this.finalResponse = errorResponse(getBuilder(), this.queryId, e); + } + } finally { + try { + context.removeListener(cancellation); + } finally { + workerFinished(failure); + } + } + } + + private long countIterator(ScanIterator source) { + long count = 0; + while (!this.clientCanceled.get() && source.hasNext()) { + source.next(); + count++; + } + return count; + } + + private long countPartitions() { + if (!(this.iterator instanceof MultiPartitionIterator)) { + return countIterator(this.iterator); + } + List children = ((MultiPartitionIterator) this.iterator).getIterators(); + CountDownLatch completed = new CountDownLatch(children.size()); + AtomicLong total = new AtomicLong(); + AtomicReference failure = new AtomicReference<>(); + synchronized (this) { + this.pendingTasks += children.size(); + } + int submitted = 0; + try { + for (ScanIterator child : children) { + this.threadPool.execute(() -> { + taskStarted(); + try { + total.addAndGet(countIterator(child)); + } catch (RuntimeException | Error e) { + recordCountFailure(failure, e); + } finally { + try { + cleanup(child::close); + } finally { + workerFinished(); + completed.countDown(); + } + } + }); + submitted++; + } + } catch (RuntimeException | Error e) { + recordCountFailure(failure, e); + for (int i = submitted; i < children.size(); i++) { + ScanIterator child = children.get(i); + try { + cleanup(child::close); + } finally { + taskFinished(); + completed.countDown(); + } + } + } + boolean interrupted = false; + for (;;) { + try { + completed.await(); + break; + } catch (InterruptedException e) { + interrupted = true; + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + Throwable error = failure.get(); + if (error instanceof RuntimeException) { + throw (RuntimeException) error; + } + if (error instanceof Error) { + throw (Error) error; + } + return total.get(); + } + + private static void recordCountFailure(AtomicReference failure, Throwable error) { + synchronized (failure) { + if (failure.get() == null) { + failure.set(error); + } else if (failure.get() != error) { + failure.get().addSuppressed(error); + } + } + } + + QueryPlan buildPlan(QueryRequest request) { + return QueryUtil.buildPlan(request); } @Override - public void onNext(QueryRequest request) { + public synchronized void onNext(QueryRequest request) { + if (this.clientCanceled.get() || this.finished || this.requestCompleted) { + return; + } if (this.queryId == null) { log.debug("got request: {}", request); this.queryId = request.getQueryId(); @@ -85,9 +267,23 @@ public void onNext(QueryRequest request) { // the first request, start the sending thread if (iterator == null) { long current = System.nanoTime(); - iterator = QueryUtil.getIterator(request); - plan = QueryUtil.buildPlan(request); - threadPool.submit(this::sendData); + this.pendingTasks = 1; + try { + iterator = getIterator(request); + plan = buildPlan(request); + threadPool.execute(this::sendData); + } catch (RuntimeException | Error e) { + // The framework must report the original synchronous failure, not normal completion. + synchronized (this.responseLock) { + this.completeResponse = false; + } + taskFinished(); + Throwable failure = this.cleanupFailure.get(); + if (failure != null && failure != e) { + e.addSuppressed(failure); + } + throw e; + } log.debug("query id: {}, init data cost: {} ms", queryId, (System.nanoTime() - current) * 1.0 / 1000000); } else { @@ -98,18 +294,162 @@ public void onNext(QueryRequest request) { @Override public void onError(Throwable t) { - // Stop calculating when channel got error - this.clientCanceled.set(true); + onTransportCancel(); log.error("AggregativeQueryService, query id: {}, got error", this.queryId, t); } + void onTransportCancel() { + // Both transport cancellation and request onError can report the same event. + synchronized (this.responseLock) { + if (this.responseFinished) { + return; + } + this.completeResponse = false; + this.responseFinished = true; + } + cancel(); + } + @Override - public void onCompleted() { - // client my be cancelled earlier + public synchronized void onCompleted() { + // Half-close ends feedback, not work already supported by its existing credit. + this.requestCompleted = true; + if (this.pendingTasks == 0 && !this.finished) { + this.finished = true; + finishResponse(); + } + } + + void cancelForShutdown() { + try { + synchronized (this.responseLock) { + if (!this.responseFinished) { + this.responseFinished = true; + this.completeResponse = false; + this.sender.onError(Status.UNAVAILABLE.withDescription("Store queries are stopping") + .asRuntimeException()); + } + } + } finally { + cancel(); + } + } + + public void cancel() { this.clientCanceled.set(true); + synchronized (this) { + for (Thread worker : this.workers) { + worker.interrupt(); + } + // An RPC with no first request has no iterator or cleanup task. + if (this.pendingTasks == 0 && !this.finished) { + this.finished = true; + finishResponse(); + } + } + } + + private synchronized void taskStarted() { + this.workers.add(Thread.currentThread()); + } + + private void workerFinished() { + workerFinished(null); + } + + private void workerFinished(Throwable failure) { + synchronized (this) { + this.workers.remove(Thread.currentThread()); + } + taskFinished(failure); + } + + private void taskFinished() { + taskFinished(null); + } + + private void taskFinished(Throwable primaryFailure) { + boolean cleanup; + synchronized (this) { + cleanup = --this.pendingTasks == 0; + if (cleanup) { + this.finished = true; + } + } + if (cleanup) { + if (this.plan != null) { + cleanup(this.plan::clear); + } + if (this.iterator != null) { + cleanup(this.iterator::close); + } + Throwable releaseFailure = this.cleanupFailure.get(); + if (primaryFailure != null && releaseFailure != null && primaryFailure != releaseFailure) { + primaryFailure.addSuppressed(releaseFailure); + } + finishResponse(); + } + } + + private void cleanup(Runnable release) { + try { + release.run(); + } catch (RuntimeException | Error failure) { + synchronized (this.cleanupFailure) { + Throwable first = this.cleanupFailure.get(); + if (first == null) { + this.cleanupFailure.set(failure); + } else if (first != failure) { + first.addSuppressed(failure); + } + } + log.error("Aggregate query {} failed to release resources; Store shutdown remains blocked", + this.queryId, failure); + } + } + + private void finishResponse() { + try { + synchronized (this.responseLock) { + if (this.completeResponse && !this.responseFinished) { + this.responseFinished = true; + Throwable failure = this.cleanupFailure.get(); + if (this.finalResponse != null && !this.finalResponse.getIsOk() && + !this.clientCanceled.get()) { + this.sender.onNext(this.finalResponse); + } else if (failure != null && !this.errorReported) { + this.sender.onNext(errorResponse(getBuilder(), this.queryId, failure)); + } else if (failure == null && this.finalResponse != null && + !this.clientCanceled.get()) { + this.sender.onNext(this.finalResponse); + } + this.sender.onCompleted(); + } + } + } finally { + // A terminal RPC does not prove that its native resources were released. + if (this.cleanupFailure.get() == null) { + this.completion.accept(this); + } + } + } + + private void sendResponse(QueryResponse response) { + synchronized (this.responseLock) { + if (!this.responseFinished && !this.clientCanceled.get()) { + if (response.getIsOk() && response.getIsFinished()) { + // The client treats this batch as success; publish it only after cleanup. + this.finalResponse = response; + } else { + this.errorReported |= !response.getIsOk(); + this.sender.onNext(response); + } + } + } } public void sendData() { + taskStarted(); try { long lastSend = System.currentTimeMillis(); var responseBuilder = getBuilder(); @@ -118,12 +458,18 @@ public void sendData() { while (!this.clientCanceled.get()) { // produces more result than consumer, just waiting if (sendCount.get() - consumeCount.get() >= RESULT_COUNT) { + if (this.requestCompleted) { + sendResponse(errorResponse(getBuilder(), queryId, + new IllegalStateException( + "Request completed without enough feedback"))); + cancel(); + return; + } // read timeout, takes long time not to read data if (System.currentTimeMillis() - lastSend > timeout) { - this.sender.onNext(errorResponse(getBuilder(), queryId, - new RuntimeException( - "sending-timeout, server closed"))); - this.sender.onCompleted(); + sendResponse(errorResponse(getBuilder(), queryId, + new RuntimeException("sending-timeout, server closed"))); + cancel(); return; } @@ -141,23 +487,33 @@ public void sendData() { } else { try { builder.setQueryId(queryId); - sender.onNext(builder.build()); + sendResponse(builder.build()); this.sendCount.incrementAndGet(); lastSend = System.currentTimeMillis(); } catch (Exception e) { log.error("send data got error: ", e); + cancel(); break; } } - if (builder.getIsFinished() || !builder.getIsOk()) { + if (!builder.getIsOk()) { + // Report the internal error before cancelling remaining partition work. + cancel(); + break; + } + if (builder.getIsFinished()) { break; } } + } catch (Exception e) { + try { + sendResponse(errorResponse(getBuilder(), queryId, e)); + } finally { + cancel(); + } } finally { - this.plan.clear(); - this.iterator.close(); - this.sender.onCompleted(); + workerFinished(); } } @@ -188,32 +544,64 @@ private QueryResponse.Builder readBatchData(QueryResponse.Builder builder, return builder; } else if (this.plan.hasIteratorResult()) { checkIterator = false; - AtomicReference exception = new AtomicReference<>(); + AtomicReference exception = new AtomicReference<>(); if (this.iterator instanceof MultiPartitionIterator) { var iterators = ((MultiPartitionIterator) this.iterator).getIterators(); CountDownLatch latch = new CountDownLatch(iterators.size()); - for (var itr2 : iterators) { - threadPool.execute(() -> { + synchronized (this) { + this.pendingTasks += iterators.size(); + } + int submitted = 0; + try { + for (var itr2 : iterators) { + threadPool.execute(() -> { + taskStarted(); + try { + execute(itr2); + } catch (RuntimeException | Error e) { + exception.compareAndSet(null, e); + } finally { + try { + cleanup(itr2::close); + Throwable failure = this.cleanupFailure.get(); + if (failure != null) { + exception.compareAndSet(null, new RuntimeException( + "partition iterator cleanup failed", failure)); + } + } finally { + latch.countDown(); + workerFinished(); + } + } + }); + submitted++; + } + } catch (RuntimeException | Error failure) { + // Rejection must also release iterators whose tasks were not accepted. + for (int i = submitted; i < iterators.size(); i++) { try { - execute(itr2); - } catch (Exception e) { - exception.set(e); + cleanup(iterators.get(i)::close); } finally { - // MultiPartitionIterator close() not working - itr2.close(); latch.countDown(); + taskFinished(); } - }); + } + throw failure; + } + if (!latch.await(timeout, TimeUnit.MILLISECONDS)) { + throw new TimeoutException("partition query timed out"); } - latch.await(timeout, TimeUnit.MILLISECONDS); if (exception.get() != null) { - throw exception.get(); + throw new RuntimeException("partition query failed", exception.get()); } } else { // can't be parallel, but has agg like stage execute(this.iterator); } + if (this.clientCanceled.get()) { + return builder.setIsOk(false).setIsFinished(false); + } try { // last empty element itr = (ScanIterator) plan.execute(PipelineResult.EMPTY); @@ -227,7 +615,7 @@ private QueryResponse.Builder readBatchData(QueryResponse.Builder builder, builder.clear(); List batchResult = new ArrayList<>(); - while (itr.hasNext() && !this.clientCanceled.get()) { + while (!this.clientCanceled.get() && itr.hasNext()) { if (count >= batchSize) { break; } @@ -265,6 +653,9 @@ private QueryResponse.Builder readBatchData(QueryResponse.Builder builder, builder.addAllData(batchResult); } catch (Exception e) { + if (e instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } log.error("readBatchData got error: ", e); return builder.setIsOk(false).setIsFinished(false).setMessage("Store Server Error: " + Arrays.toString( @@ -273,7 +664,7 @@ private QueryResponse.Builder readBatchData(QueryResponse.Builder builder, if (checkIterator) { // check the iterator - finish = !itr.hasNext(); + finish = this.clientCanceled.get() || !itr.hasNext(); } log.debug("query id: {}, finished batch, with size :{}, finish:{}, cost: {} ms", queryId, count, @@ -320,7 +711,7 @@ public void close() { private void execute(ScanIterator itr) { long recordCount = 0; long current = System.nanoTime(); - while (itr.hasNext() && !this.clientCanceled.get()) { + while (!this.clientCanceled.get() && itr.hasNext()) { try { recordCount++; executePipeline(itr.next()); diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryService.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryService.java index 3eb81be851..7b024e69cc 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryService.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryService.java @@ -17,24 +17,20 @@ package org.apache.hugegraph.store.node.grpc.query; -import java.util.ArrayList; -import java.util.List; +import java.util.HashSet; +import java.util.Set; import java.util.concurrent.ThreadPoolExecutor; -import java.util.concurrent.atomic.AtomicLong; -import org.apache.hugegraph.rocksdb.access.RocksDBSession; import org.apache.hugegraph.store.HgStoreEngine; import org.apache.hugegraph.store.consts.PoolNames; -import org.apache.hugegraph.store.grpc.common.Kv; import org.apache.hugegraph.store.grpc.query.QueryRequest; import org.apache.hugegraph.store.grpc.query.QueryResponse; import org.apache.hugegraph.store.grpc.query.QueryServiceGrpc; -import org.apache.hugegraph.store.query.KvSerializer; import org.apache.hugegraph.store.util.ExecutorUtil; import org.lognet.springboot.grpc.GRpcService; -import com.google.protobuf.ByteString; - +import io.grpc.Status; +import io.grpc.stub.ServerCallStreamObserver; import io.grpc.stub.StreamObserver; import lombok.Getter; import lombok.extern.slf4j.Slf4j; @@ -43,6 +39,9 @@ @GRpcService public class AggregativeQueryService extends QueryServiceGrpc.QueryServiceImplBase { + private final Set queries = new HashSet<>(); + private boolean closing; + private final int batchSize; private final Long timeout; @@ -62,6 +61,65 @@ public AggregativeQueryService() { 10000, true); } + AggregativeQueryService(ThreadPoolExecutor threadPool, long timeout, int batchSize) { + this.threadPool = threadPool; + this.timeout = timeout; + this.batchSize = batchSize; + } + + public synchronized void stopAcceptingQueries() { + this.closing = true; + } + + /** Cancel queries, but let every queued parent and partition release its iterator. */ + public void shutdownQueries() { + AggregativeQueryObserver[] active; + synchronized (this) { + this.closing = true; + active = this.queries.toArray(new AggregativeQueryObserver[0]); + } + for (AggregativeQueryObserver query : active) { + try { + query.cancelForShutdown(); + } catch (RuntimeException | Error failure) { + // Response callbacks cannot skip cancellation or cleanup waits for other queries. + log.error("Failed to cancel aggregate query response; continuing shutdown", failure); + } + } + boolean interrupted = false; + synchronized (this) { + try { + while (!this.queries.isEmpty()) { + try { + this.wait(5000); + if (!this.queries.isEmpty()) { + log.warn("Waiting for {} aggregate queries to release iterators", + this.queries.size()); + } + } catch (InterruptedException e) { + interrupted = true; + } + } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + } + // Parents can still enqueue partition cleanup until the last query is finished. + this.threadPool.shutdown(); + } + + private synchronized void finished(AggregativeQueryObserver observer) { + this.queries.remove(observer); + this.notifyAll(); + } + + AggregativeQueryObserver newObserver(StreamObserver sender) { + return new AggregativeQueryObserver(sender, this.threadPool, this.timeout, + this.batchSize, this::finished); + } + /** * Generate error response. * @@ -79,73 +137,26 @@ public static QueryResponse errorResponse(QueryResponse.Builder builder, String } @Override - public StreamObserver query(StreamObserver observer) { - return new AggregativeQueryObserver(observer, threadPool, timeout, batchSize); + public synchronized StreamObserver query(StreamObserver observer) { + if (this.closing) { + throw Status.UNAVAILABLE.withDescription("Store queries are stopping").asRuntimeException(); + } + AggregativeQueryObserver query = newObserver(observer); + this.queries.add(query); + if (observer instanceof ServerCallStreamObserver) { + ((ServerCallStreamObserver) observer).setOnCancelHandler( + query::onTransportCancel); + } + return query; } @Override public void query0(QueryRequest request, StreamObserver observer) { - - var itr = QueryUtil.getIterator(request); - var builder = QueryResponse.newBuilder(); - var kvBuilder = Kv.newBuilder(); - - try { - while (itr.hasNext()) { - var column = (RocksDBSession.BackendColumn) itr.next(); - if (column != null) { - builder.addData(kvBuilder.setKey(ByteString.copyFrom(column.name)) - .setValue(column.value == null ? ByteString.EMPTY : - ByteString.copyFrom(column.value)) - .build()); - } - } - builder.setQueryId(request.getQueryId()); - builder.setIsOk(true); - builder.setIsFinished(true); - observer.onNext(builder.build()); - } catch (Exception e) { - observer.onNext(errorResponse(builder, request.getQueryId(), e)); - } - observer.onCompleted(); + ((AggregativeQueryObserver) query(observer)).runUnary(request, false); } - /** - * Query data count - * - * @param request query request object - * @param observer Observer object for receiving query response results - */ @Override public void count(QueryRequest request, StreamObserver observer) { - - log.debug("query id : {}, simple count of table: {}", request.getQueryId(), - request.getTable()); - var builder = QueryResponse.newBuilder(); - var kvBuilder = Kv.newBuilder(); - - try { - - var handler = new QueryUtil().getHandler(); - long start = System.currentTimeMillis(); - long count = handler.count(request.getGraph(), request.getTable()); - log.debug("query id: {}, count of cost: {} ms", request.getQueryId(), - System.currentTimeMillis() - start); - List array = new ArrayList<>(); - for (int i = 0; i < request.getFunctionsList().size(); i++) { - array.add(new AtomicLong(count)); - } - - kvBuilder.setKey(ByteString.copyFrom(KvSerializer.toBytes(List.of()))); - kvBuilder.setValue(ByteString.copyFrom(KvSerializer.toBytes(array))); - builder.addData(kvBuilder.build()); - builder.setQueryId(request.getQueryId()); - builder.setIsOk(true); - builder.setIsFinished(true); - observer.onNext(builder.build()); - } catch (Exception e) { - observer.onNext(errorResponse(builder, request.getQueryId(), e)); - } - observer.onCompleted(); + ((AggregativeQueryObserver) query(observer)).runUnary(request, true); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/model/QueryPlan.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/model/QueryPlan.java index 228850d65c..d3289c2492 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/model/QueryPlan.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/model/QueryPlan.java @@ -118,10 +118,25 @@ public String toString() { } public void clear() { + Throwable failure = null; for (var stage : stages) { - stage.close(); + try { + stage.close(); + } catch (RuntimeException | Error e) { + if (failure == null) { + failure = e; + } else if (failure != e) { + failure.addSuppressed(e); + } + } } this.stages.clear(); + if (failure instanceof Error) { + throw (Error) failure; + } + if (failure != null) { + throw (RuntimeException) failure; + } } public boolean isEmpty() { diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/GraphStoreImpl.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/GraphStoreImpl.java index 6583e59745..31bd69c40f 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/GraphStoreImpl.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/GraphStoreImpl.java @@ -77,7 +77,7 @@ public ThreadPoolExecutor getExecutor() { @Override public StreamObserver scanPartition( StreamObserver ro) { - return new ScanResponseObserver(ro, getHandler(), getExecutor()); + return this.storeStream.scanGraphPartition(ro, getHandler()); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/ScanResponseObserver.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/ScanResponseObserver.java index b4d10f7c43..2bb7ee44f0 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/ScanResponseObserver.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/ScanResponseObserver.java @@ -18,249 +18,285 @@ package org.apache.hugegraph.store.node.grpc.scan; import java.util.ArrayList; -import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Consumer; import org.apache.hugegraph.store.business.BusinessHandler; import org.apache.hugegraph.store.business.GraphStoreIterator; -import org.apache.hugegraph.store.grpc.Graphpb.Error; +import org.apache.hugegraph.store.grpc.Graphpb; import org.apache.hugegraph.store.grpc.Graphpb.ErrorType; import org.apache.hugegraph.store.grpc.Graphpb.ResponseHeader; import org.apache.hugegraph.store.grpc.Graphpb.ScanPartitionRequest; -import org.apache.hugegraph.store.grpc.Graphpb.ScanPartitionRequest.Request; import org.apache.hugegraph.store.grpc.Graphpb.ScanPartitionRequest.ScanType; import org.apache.hugegraph.store.grpc.Graphpb.ScanResponse; import com.google.protobuf.Descriptors; +import io.grpc.Status; import io.grpc.stub.StreamObserver; -import lombok.extern.slf4j.Slf4j; -@Slf4j -public class ScanResponseObserver implements - StreamObserver { +public class ScanResponseObserver implements StreamObserver { private static final int BATCH_SIZE = 100000; - private static final int MAX_PAGE = 8; // - private static final Error ok = Error.newBuilder().setType(ErrorType.OK).build(); - private static final ResponseHeader okHeader = - ResponseHeader.newBuilder().setError(ok).build(); + private static final int MAX_PAGE = 8; + private static final ResponseHeader OK = ResponseHeader.newBuilder() + .setError(Graphpb.Error.newBuilder().setType(ErrorType.OK)).build(); private final BusinessHandler handler; - private final AtomicInteger nextSeqNo = new AtomicInteger(0); - private final AtomicInteger cltSeqNo = new AtomicInteger(0); - private final ThreadPoolExecutor executor; - private final AtomicBoolean readOver = new AtomicBoolean(false); - private final LinkedBlockingQueue packages = - new LinkedBlockingQueue(MAX_PAGE * 2); - private final Descriptors.FieldDescriptor vertexField = - ScanResponse.getDescriptor().findFieldByNumber(3); - private final Descriptors.FieldDescriptor edgeField = - ScanResponse.getDescriptor().findFieldByNumber(4); - private final ReentrantLock readLock = new ReentrantLock(); - private final ReentrantLock sendLock = new ReentrantLock(); - private StreamObserver sender; + private final StreamObserver sender; + private final Consumer execute; + private final Consumer failedCleanup; + private final AtomicInteger nextSeqNo = new AtomicInteger(); + private final AtomicInteger cltSeqNo = new AtomicInteger(); + private final AtomicBoolean readOver = new AtomicBoolean(); + private final AtomicBoolean closed = new AtomicBoolean(); + private final AtomicBoolean halfClosed = new AtomicBoolean(); + private final AtomicBoolean reading = new AtomicBoolean(); + private final AtomicBoolean sending = new AtomicBoolean(); + private final LinkedBlockingQueue packages = new LinkedBlockingQueue<>(MAX_PAGE * 2); + private final Object iteratorLock = new Object(); + private final Object readerLock = new Object(); + private Thread reader; + private GraphStoreIterator iter; private ScanPartitionRequest scanReq; - private GraphStoreIterator iter; - private volatile long leftCount; - private volatile Future sendTask; - private volatile Future readTask; + private boolean iteratorClosed; + private long leftCount; - /* - * November 1, 2022 - * 1. onNext needs to be processed asynchronously to prevent the grpc call from being blocked. - * 2. Do not read iterators or send data do not produce thread waiting. - * 3. Before sending, try to prepare the data to be sent as much as possible. - * */ - - /* - * November 2, 2022 - * 1. Read the thread of rocksdb iterator read - * 2. Perform data conversion and send to the blocking queue thread offer - * 3. Thread for reading data from the blocking queue and sending, including waking up the - * reading and sending threads when no data is read - * */ + public ScanResponseObserver(StreamObserver sender, + BusinessHandler handler, ThreadPoolExecutor executor) { + this(sender, handler, executor::execute, failure -> { }); + } public ScanResponseObserver(StreamObserver sender, BusinessHandler handler, - ThreadPoolExecutor executor) { + Consumer execute, Consumer failedCleanup) { this.sender = sender; this.handler = handler; - this.executor = executor; - } - - private boolean readCondition() { - return packages.remainingCapacity() != 0 && !readOver.get(); - } - - private boolean readTaskCondition() { - return readCondition() && (readTask == null || readTask.isDone()); - } - - private boolean sendCondition() { - return nextSeqNo.get() - cltSeqNo.get() < MAX_PAGE; + this.execute = execute; + this.failedCleanup = failedCleanup; } - private boolean sendTaskCondition() { - return sendCondition() && (sendTask == null || sendTask.isDone()); - } - - private void offer(Iterable data, boolean isVertex) { - ScanResponse.Builder builder = ScanResponse.newBuilder(); - builder.setHeader(okHeader).setSeqNo(nextSeqNo.get()); - if (isVertex) { - builder = builder.setField(vertexField, data); - } else { - builder = builder.setField(edgeField, data); + @Override + public void onNext(ScanPartitionRequest request) { + if (this.closed.get() || this.halfClosed.get()) { + return; } - ScanResponse response = builder.build(); - packages.offer(response); - startSend(); - } - - private void startRead() { - if (readTaskCondition()) { - if (readLock.tryLock()) { - if (readTaskCondition()) { - readTask = executor.submit(rr); + try { + if (request.hasScanRequest() && !request.hasReplyRequest()) { + synchronized (this.iteratorLock) { + if (this.closed.get() || this.halfClosed.get()) { + return; + } + if (this.iter != null) { + throw Status.INVALID_ARGUMENT.withDescription("Scan request was already initialized") + .asRuntimeException(); + } + this.scanReq = request; + long limit = request.getScanRequest().getLimit(); + this.leftCount = limit > 0 ? limit : Long.MAX_VALUE; + this.iter = this.handler.scan(request); } - readLock.unlock(); + startRead(); + } else { + this.cltSeqNo.incrementAndGet(); + startSend(); } + } catch (RuntimeException | Error failure) { + terminate(failure); } } - private void startSend() { - if (sendTaskCondition()) { - if (sendLock.tryLock()) { - if (sendTaskCondition()) { - sendTask = executor.submit(sr); - } - sendLock.unlock(); - } - } + @Override + public void onError(Throwable failure) { + terminate(failure); } @Override - public void onNext(ScanPartitionRequest scanReq) { - if (scanReq.hasScanRequest() && !scanReq.hasReplyRequest()) { - this.scanReq = scanReq; - Request request = scanReq.getScanRequest(); - long rl = request.getLimit(); - leftCount = rl > 0 ? rl : Long.MAX_VALUE; - iter = handler.scan(scanReq); - if (!iter.hasNext()) { - close(); - sender.onCompleted(); - } else { - readTask = executor.submit(rr); - } + public void onCompleted() { + boolean initialized; + synchronized (this.iteratorLock) { + this.halfClosed.set(true); + initialized = this.iter != null; + } + if (!initialized) { + terminate(null); } else { - cltSeqNo.getAndIncrement(); startSend(); } } - @Override - public void onError(Throwable t) { - close(); - log.warn("receive client error:", t); + private boolean readCondition() { + return !this.closed.get() && !this.readOver.get() && this.packages.remainingCapacity() > 0; } - @Override - public void onCompleted() { - close(); + private boolean sendCondition() { + return !this.closed.get() && this.nextSeqNo.get() - this.cltSeqNo.get() < MAX_PAGE; } - private void close() { + private boolean sendWorkAvailable() { + return !this.closed.get() && + (sendCondition() || this.halfClosed.get() || + (this.readOver.get() && this.packages.isEmpty())); + } + + private void startRead() { + if (!readCondition() || !this.reading.compareAndSet(false, true)) { + return; + } try { - nextSeqNo.set(0); - if (sendTask != null) { - sendTask.cancel(true); - } - if (readTask != null) { - readTask.cancel(true); + this.execute.accept(() -> { + try { + synchronized (this.readerLock) { + if (this.closed.get()) { + return; + } + this.reader = Thread.currentThread(); + } + read(); + } catch (RuntimeException | Error failure) { + terminate(failure); + } finally { + synchronized (this.readerLock) { + if (this.reader == Thread.currentThread()) { + this.reader = null; + } + } + this.reading.set(false); + if (readCondition()) { + startRead(); + } + } + }); + } catch (RuntimeException | Error failure) { + this.reading.set(false); + terminate(failure); + } + } + + private void read() { + synchronized (this.iteratorLock) { + while (readCondition()) { + ArrayList data = new ArrayList<>(BATCH_SIZE); + while (!this.closed.get() && this.leftCount > 0 && + data.size() < BATCH_SIZE && this.iter.hasNext()) { + data.add(this.iter.next()); + this.leftCount--; + } + if (this.closed.get()) { + return; + } + boolean ended = this.leftCount == 0 || !this.iter.hasNext(); + if (!data.isEmpty()) { + Descriptors.FieldDescriptor field = ScanResponse.getDescriptor().findFieldByNumber( + this.scanReq.getScanRequest().getScanType() == ScanType.SCAN_VERTEX ? 3 : 4); + this.packages.add(ScanResponse.newBuilder().setHeader(OK) + .setSeqNo(this.nextSeqNo.get()).setField(field, data).build()); + } + if (ended) { + closeIterator(); + this.readOver.set(true); + } + startSend(); } - readOver.set(true); - iter.close(); - } catch (Exception e) { - log.warn("on Complete with error:", e); } } - Runnable rr = new Runnable() { - @Override - public void run() { - try { - if (readCondition()) { - synchronized (iter) { - while (readCondition()) { - Request r = scanReq.getScanRequest(); - ScanType t = r.getScanType(); - boolean isVertex = t.equals(ScanType.SCAN_VERTEX); - ArrayList data = new ArrayList<>(BATCH_SIZE); - int count = 0; - while (iter.hasNext() && leftCount > -1) { - count++; - leftCount--; - T next = (T) iter.next(); - data.add(next); - if (count >= BATCH_SIZE) { - offer(data, isVertex); - // data.clear(); - break; - } - } - if (!(iter.hasNext() && leftCount > -1)) { - if (data.size() > 0 && - data.size() < BATCH_SIZE) { - offer(data, isVertex); - } - readOver.set(true); - data = null; - //log.warn("scan complete , count: {},time: {}", - // sum, System.currentTimeMillis()); - return; + private void startSend() { + if (!sendWorkAvailable() || !this.sending.compareAndSet(false, true)) { + return; + } + try { + this.execute.accept(() -> { + try { + while (!this.closed.get()) { + if (this.readOver.get() && this.packages.isEmpty()) { + terminate(null); + return; + } + if (!sendCondition()) { + // A half-closed request stream cannot grant further page credit. + if (this.halfClosed.get() && !this.packages.isEmpty()) { + terminate(Status.FAILED_PRECONDITION.withDescription( + "Partition scan needs more response credit after client half-close") + .asRuntimeException()); } + return; } + ScanResponse response = this.packages.poll(10, TimeUnit.MILLISECONDS); + if (response != null) { + this.sender.onNext(response.toBuilder().setSeqNo(this.nextSeqNo.get()).build()); + this.nextSeqNo.incrementAndGet(); + startRead(); + } else if (this.readOver.get() && this.packages.isEmpty()) { + terminate(null); + return; + } else { + return; + } + } + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + terminate(failure); + } catch (RuntimeException | Error failure) { + terminate(failure); + } finally { + this.sending.set(false); + if (sendWorkAvailable() && (!this.packages.isEmpty() || this.readOver.get())) { + startSend(); } } - } catch (Exception e) { - log.warn("read data with error: ", e); - sender.onError(e); - } + }); + } catch (RuntimeException | Error failure) { + this.sending.set(false); + terminate(failure); } - }; + } - Runnable sr = () -> { - while (sendCondition()) { - ScanResponse response; + private void closeIterator() { + synchronized (this.iteratorLock) { + if (this.iter == null || this.iteratorClosed) { + return; + } + this.iteratorClosed = true; try { - if (readOver.get()) { - if ((response = packages.poll()) == null) { - sender.onCompleted(); - } else { - sender.onNext(response); - nextSeqNo.incrementAndGet(); - } - } else { - response = packages.poll(10, - TimeUnit.MILLISECONDS); - if (response != null) { - sender.onNext(response); - nextSeqNo.incrementAndGet(); - startRead(); - } else { - break; - } - } + this.iter.close(); + } catch (RuntimeException | Error failure) { + this.failedCleanup.accept(failure); + throw failure; + } + } + } - } catch (InterruptedException e) { - break; + private void terminate(Throwable failure) { + if (!this.closed.compareAndSet(false, true)) { + return; + } + this.readOver.set(true); + this.packages.clear(); + if (failure != null) { + synchronized (this.readerLock) { + // Keep ownership protected until interrupt is issued; the pool may reuse this thread. + if (this.reader != null && this.reader != Thread.currentThread()) { + this.reader.interrupt(); + } + } + } + try { + closeIterator(); + } catch (RuntimeException | Error cleanup) { + if (failure == null) { + failure = cleanup; + } else if (failure != cleanup) { + failure.addSuppressed(cleanup); } } - }; + if (failure == null) { + this.sender.onCompleted(); + } else { + this.sender.onError(failure); + } + } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/listener/ContextClosedListener.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/listener/ContextClosedListener.java index 285ec3b968..a46a4a9b2c 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/listener/ContextClosedListener.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/listener/ContextClosedListener.java @@ -17,64 +17,121 @@ package org.apache.hugegraph.store.node.listener; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ThreadPoolExecutor; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; +import org.apache.hugegraph.store.node.grpc.GrpcShutdownBarrier; import org.apache.hugegraph.store.node.grpc.HgStoreStreamImpl; +import org.apache.hugegraph.store.node.grpc.query.AggregativeQueryService; import org.apache.hugegraph.store.node.task.TTLCleaner; +import org.lognet.springboot.grpc.context.GRpcServerInitializedEvent; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.ApplicationListener; import org.springframework.context.event.ContextClosedEvent; +import org.springframework.context.event.EventListener; import org.springframework.stereotype.Service; +import io.grpc.Server; import lombok.extern.slf4j.Slf4j; @Slf4j @Service public class ContextClosedListener implements ApplicationListener { + private final List grpcServers = new CopyOnWriteArrayList<>(); + @Autowired HgStoreStreamImpl storeStream; @Autowired + AggregativeQueryService queryService; + @Autowired TTLCleaner cleaner; + @Autowired + GrpcShutdownBarrier grpcBarrier; + + @EventListener + public void onServerInitialized(GRpcServerInitializedEvent event) { + this.grpcServers.add(event.getServer()); + } @Override public void onApplicationEvent(ContextClosedEvent event) { + // Spring invokes HgStoreNodeService.destroy() after this event. Raft and + // its databases must remain available until local workers have stopped. + this.grpcBarrier.stopAcceptingCalls(); + if (storeStream != null) { + storeStream.stopAcceptingScans(); + } + if (queryService != null) { + queryService.stopAcceptingQueries(); + } + this.grpcServers.forEach(Server::shutdownNow); + if (cleaner != null) { + // TTL failure ownership and worker drain are added by the native shutdown PR. + if (cleaner.getExecutor() != null) { + cleaner.getExecutor().shutdownNow(); + } + if (cleaner.getScheduler() != null) { + cleaner.getScheduler().shutdownNow(); + } + } + if (storeStream != null) { + // Cancelled queued scans must run their finally blocks to release iterators. + storeStream.shutdownScans(); + awaitWorkers(storeStream.getRealExecutor(), "scan workers"); + } + if (queryService != null) { + queryService.shutdownQueries(); + awaitWorkers(queryService.getThreadPool(), "aggregate query workers"); + } + boolean interrupted = false; try { - log.info("closing scan threads...."); - if (storeStream != null) { - ThreadPoolExecutor executor = storeStream.getRealExecutor(); - if (executor != null) { + for (Server server : this.grpcServers) { + while (!server.isTerminated()) { try { - executor.shutdownNow(); - } catch (Exception e) { + if (!server.awaitTermination(5, TimeUnit.SECONDS)) { + log.warn("Still waiting for gRPC callbacks before closing databases"); + } + } catch (InterruptedException e) { + interrupted = true; } } } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + this.grpcBarrier.awaitCallbacks(); + if (storeStream != null) { + storeStream.awaitScanCleanup(); + } + log.info("closed gRPC callbacks, scan and aggregate query workers"); + } - if (cleaner != null) { - ThreadPoolExecutor cleanerExecutor = cleaner.getExecutor(); - if (cleanerExecutor != null) { - try { - cleanerExecutor.shutdownNow(); - } catch (Exception e) { - - } - } - ScheduledExecutorService scheduler = cleaner.getScheduler(); - if (scheduler != null) { - try { - scheduler.shutdownNow(); - } catch (Exception e) { - + private static void awaitWorkers(ExecutorService executor, String name) { + if (executor == null) { + return; + } + boolean interrupted = false; + try { + while (!executor.isTerminated()) { + try { + if (!executor.awaitTermination(5, TimeUnit.SECONDS)) { + log.warn("Still waiting for {} to stop before closing databases", name); } + } catch (InterruptedException e) { + // An interrupted shutdown thread must not close databases underneath + // a worker that still owns a native iterator. + interrupted = true; } } - } catch (Exception e) { - log.error("ContextClosedListener: ", e); } finally { - log.info("closed scan threads"); + if (interrupted) { + Thread.currentThread().interrupt(); + } } } - } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/util/HgChannel.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/util/HgChannel.java index c5b81fe5ca..debca4bea3 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/util/HgChannel.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/util/HgChannel.java @@ -37,6 +37,8 @@ @Slf4j public final class HgChannel { + // Recheck close without interrupting a borrowed caller thread. + private static final long CLOSE_POLL_NANOS = TimeUnit.MILLISECONDS.toNanos(100); private final BlockingQueue> queue; private final AtomicBoolean closed = new AtomicBoolean(false); private final long timeoutSeconds; @@ -69,14 +71,27 @@ public boolean send(T t) { if (this.closed.get()) { return false; } - boolean flag; + long timeoutNanos = Math.max(0L, TimeUnit.SECONDS.toNanos(this.timeoutSeconds)); + long started = System.nanoTime(); + long remaining = timeoutNanos; + Supplier supplier = () -> t; try { - flag = this.queue.offer(() -> t, timeoutSeconds, TimeUnit.SECONDS); + do { + if (this.closed.get()) { + return false; + } + if (this.queue.offer(supplier, Math.max(0L, Math.min(remaining, CLOSE_POLL_NANOS)), + TimeUnit.NANOSECONDS)) { + return true; + } + remaining = timeoutNanos - (System.nanoTime() - started); + } while (remaining > 0L); + return false; } catch (InterruptedException e) { + Thread.currentThread().interrupt(); log.error("failed to send a item to chan. cause by: ", t); throw new RuntimeException(e); } - return flag; } } @@ -105,11 +120,29 @@ public T receive(Consumer timeoutCallBack) { if (this.closed.get()) { s = this.queue.poll(); } else { + long timeoutNanos = Math.max(0L, TimeUnit.SECONDS.toNanos(this.timeoutSeconds)); + long started = System.nanoTime(); + long remaining = timeoutNanos; + s = null; try { - s = this.queue.poll(timeoutSeconds, TimeUnit.SECONDS); - } catch (Throwable t) { - log.error("Failed to receive a item from chan. cause by: ", t); - throw new RuntimeException(t); + do { + if (this.closed.get()) { + return null; + } + s = this.queue.poll(Math.max(0L, Math.min(remaining, CLOSE_POLL_NANOS)), + TimeUnit.NANOSECONDS); + if (s != null) { + break; + } + remaining = timeoutNanos - (System.nanoTime() - started); + } while (remaining > 0L); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + log.error("Failed to receive a item from chan. cause by: ", e); + throw new RuntimeException(e); + } + if (s == null && this.closed.get()) { + return null; } if (s == null) { if (timeoutCallBack == null) { @@ -136,14 +169,7 @@ public boolean isClosed() { * @throws RuntimeException when fail to close the chan */ public void close() { - if (this.closed.get()) { - return; - } this.closed.set(true); - this.queue.offer(() -> null); - Thread.yield(); - this.queue.poll(); - } } diff --git a/hugegraph-store/hg-store-rocksdb/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIterator.java b/hugegraph-store/hg-store-rocksdb/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIterator.java index dca8179308..7e23370cf6 100644 --- a/hugegraph-store/hg-store-rocksdb/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIterator.java +++ b/hugegraph-store/hg-store-rocksdb/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIterator.java @@ -41,6 +41,7 @@ public class RocksDBScanIterator implements ScanIterator { private final AtomicBoolean closed = new AtomicBoolean(false); private final RocksDBSession.RefCounter iterReference; private final Consumer closeOp; + private Throwable closeFailure; private byte[] key; private boolean matched; @@ -225,13 +226,26 @@ private boolean filter(byte[] v) { } @Override - public void close() { + public synchronized void close() { if (!this.closed.getAndSet(true)) { - if (this.rawIt.isOwningHandle()) { - this.rawIt.close(); + try { + if (this.rawIt.isOwningHandle()) { + this.rawIt.close(); + } + this.closeOp.accept(true); + this.iterReference.release(); + } catch (RuntimeException | Error failure) { + this.closeFailure = failure; + throw failure; } - this.closeOp.accept(true); - this.iterReference.release(); + } + // Automatic exhaustion can initiate close before its owner enters finally. + // A repeated close must wait for that release and retain any failure, without retrying JNI. + if (this.closeFailure instanceof RuntimeException) { + throw (RuntimeException) this.closeFailure; + } + if (this.closeFailure instanceof Error) { + throw (Error) this.closeFailure; } } diff --git a/hugegraph-store/hg-store-test/pom.xml b/hugegraph-store/hg-store-test/pom.xml index 9804435aba..bf6ee0bb70 100644 --- a/hugegraph-store/hg-store-test/pom.xml +++ b/hugegraph-store/hg-store-test/pom.xml @@ -178,6 +178,22 @@ + + org.apache.hugegraph + hg-store-node + ${revision} + compile + + + org.springframework + spring-web + + + org.springframework + spring-webmvc + + + diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIteratorCloseTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIteratorCloseTest.java new file mode 100644 index 0000000000..c66c8532b4 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIteratorCloseTest.java @@ -0,0 +1,444 @@ +/* + * 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.rocksdb.access; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.FutureTask; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.LockSupport; +import java.util.function.Consumer; + +import org.apache.hugegraph.rocksdb.access.RocksDBSession.BackendColumn; +import org.apache.hugegraph.store.business.InnerKeyFilter; +import org.apache.hugegraph.store.business.MultiPartitionIterator; +import org.apache.hugegraph.store.term.Bits; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.mockito.InOrder; +import org.rocksdb.Options; +import org.rocksdb.RocksDB; +import org.rocksdb.RocksIterator; + +public class RocksDBScanIteratorCloseTest { + + @Rule + public TemporaryFolder directory = new TemporaryFolder(); + + @Test + @SuppressWarnings("unchecked") + public void testRealNativeEmptyIteratorRetainsOneTimeCallbackFailure() throws Exception { + RocksDB.loadLibrary(); + try (Options options = new Options().setCreateIfMissing(true); + RocksDB database = RocksDB.open(options, this.directory.newFolder("empty").getAbsolutePath())) { + RocksIterator raw = database.newIterator(); + RocksDBSession.RefCounter reference = mock(RocksDBSession.RefCounter.class); + Consumer closeOp = mock(Consumer.class); + IllegalStateException failure = new IllegalStateException("close callback failed once"); + doThrow(failure).doNothing().when(closeOp).accept(true); + RocksDBScanIterator iterator = new RocksDBScanIterator<>(raw, null, null, + ScanIterator.Trait.SCAN_ANY, reference, closeOp); + + assertSame(failure, assertThrows(IllegalStateException.class, iterator::hasNext)); + assertFalse("JNI must release the native iterator before callback failure", raw.isOwningHandle()); + assertSame(failure, assertThrows(IllegalStateException.class, iterator::close)); + assertSame(failure, assertThrows(IllegalStateException.class, iterator::close)); + verify(closeOp).accept(true); + verify(reference, never()).release(); + } + } + + @Test + @SuppressWarnings("unchecked") + public void testRealNativeExhaustedIteratorRetainsOneTimeReferenceReleaseError() throws Exception { + RocksDB.loadLibrary(); + try (Options options = new Options().setCreateIfMissing(true); + RocksDB database = RocksDB.open(options, this.directory.newFolder("exhausted").getAbsolutePath())) { + database.put(new byte[]{1}, new byte[]{2}); + RocksIterator raw = database.newIterator(); + RocksDBSession.RefCounter reference = mock(RocksDBSession.RefCounter.class); + Consumer closeOp = mock(Consumer.class); + AssertionError failure = new AssertionError("reference release failed once"); + doThrow(failure).doNothing().when(reference).release(); + RocksDBScanIterator iterator = new RocksDBScanIterator<>(raw, null, null, + ScanIterator.Trait.SCAN_ANY, reference, closeOp); + + assertTrue(iterator.hasNext()); + iterator.next(); + assertSame(failure, assertThrows(AssertionError.class, iterator::hasNext)); + assertFalse("JNI must release the native iterator before reference failure", raw.isOwningHandle()); + assertSame(failure, assertThrows(AssertionError.class, iterator::close)); + assertSame(failure, assertThrows(AssertionError.class, iterator::close)); + InOrder order = inOrder(closeOp, reference); + order.verify(closeOp).accept(true); + order.verify(reference).release(); + } + } + + @Test + public void testEmptyIteratorRetainsOneTimeNativeCloseFailure() { + Fixture fixture = new Fixture(); + IllegalStateException failure = new IllegalStateException("native close failed once"); + doThrow(failure).doNothing().when(fixture.raw).close(); + + assertAutomaticFailureIsRetained(fixture, failure); + verify(fixture.closeOp, never()).accept(true); + verify(fixture.reference, never()).release(); + } + + @Test + public void testExhaustedIteratorRetainsOneTimeNativeCloseError() { + Fixture fixture = new Fixture(); + when(fixture.raw.isValid()).thenReturn(true, false); + when(fixture.raw.key()).thenReturn(new byte[]{1}); + when(fixture.raw.value()).thenReturn(new byte[]{2}); + AssertionError failure = new AssertionError("native close failed once"); + doThrow(failure).doNothing().when(fixture.raw).close(); + + assertTrue(fixture.iterator.hasNext()); + fixture.iterator.next(); + assertAutomaticFailureIsRetained(fixture, failure); + verify(fixture.closeOp, never()).accept(true); + verify(fixture.reference, never()).release(); + } + + @Test + public void testAutomaticCloseRetainsOneTimeCallbackFailure() { + Fixture fixture = new Fixture(); + IllegalStateException failure = new IllegalStateException("close callback failed once"); + doThrow(failure).doNothing().when(fixture.closeOp).accept(true); + + assertAutomaticFailureIsRetained(fixture, failure); + InOrder order = inOrder(fixture.raw, fixture.closeOp); + order.verify(fixture.raw).close(); + order.verify(fixture.closeOp).accept(true); + verify(fixture.reference, never()).release(); + } + + @Test + public void testAutomaticCloseRetainsOneTimeReferenceReleaseError() { + Fixture fixture = new Fixture(); + AssertionError failure = new AssertionError("reference release failed once"); + doThrow(failure).doNothing().when(fixture.reference).release(); + + assertAutomaticFailureIsRetained(fixture, failure); + InOrder order = inOrder(fixture.raw, fixture.closeOp, fixture.reference); + order.verify(fixture.raw).close(); + order.verify(fixture.closeOp).accept(true); + order.verify(fixture.reference).release(); + } + + @Test + public void testSuccessfulCloseIsIdempotentAndPreservesReleaseOrder() { + Fixture fixture = new Fixture(); + + fixture.iterator.close(); + fixture.iterator.close(); + assertFalse(fixture.iterator.hasNext()); + InOrder order = inOrder(fixture.raw, fixture.closeOp, fixture.reference); + order.verify(fixture.raw).close(); + order.verify(fixture.closeOp).accept(true); + order.verify(fixture.reference).release(); + verify(fixture.raw).close(); + verify(fixture.closeOp).accept(true); + verify(fixture.reference).release(); + } + + @Test(timeout = 5000) + public void testConcurrentCloseWaitsForSuccessfulRelease() throws Exception { + assertConcurrentCloseWaits(null); + } + + @Test(timeout = 5000) + public void testConcurrentCloseWaitsAndReportsTheSameFailure() throws Exception { + assertConcurrentCloseWaits(new IllegalStateException("close callback failed once")); + } + + @Test + public void testAllInnerKeyConstructorsDeferReadingUntilOwnershipTransfer() { + for (int constructor = 0; constructor < 3; constructor++) { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true); + when(raw.next()).thenReturn(column(1, 15)); + InnerKeyFilter filter = filter(raw, constructor); + + verify(raw, never()).hasNext(); + verify(raw, never()).next(); + assertTrue(filter.hasNext()); + verify(raw).hasNext(); + verify(raw).next(); + filter.close(); + verify(raw).close(); + } + } + + @Test + public void testPlainInnerKeyFilterStripsGraphAndCodeAndPreservesPrefetch() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true, true, false); + when(raw.next()).thenReturn(column(1, 5), column(2, 19)); + InnerKeyFilter filter = new InnerKeyFilter<>(raw); + + assertTrue(filter.hasNext()); + assertTrue(filter.hasNext()); + assertArrayEquals(new byte[]{1}, filter.next().name); + assertTrue(filter.hasNext()); + assertArrayEquals(new byte[]{2}, filter.next().name); + assertFalse(filter.hasNext()); + verify(raw, times(2)).next(); + verify(raw, times(3)).hasNext(); + filter.close(); + } + + @Test + public void testCodeFilterKeepsHashSuffix() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true, false); + when(raw.next()).thenReturn(column(3, 15)); + InnerKeyFilter filter = new InnerKeyFilter<>(raw, true); + + assertArrayEquals(new byte[]{3, 0, 15}, filter.next().name); + assertFalse(filter.hasNext()); + filter.close(); + } + + @Test + public void testCodeRangeFilterPreservesInclusiveLowerAndExclusiveUpperBounds() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true, true, true, true, false); + when(raw.next()).thenReturn(column(1, 5), column(2, 10), column(3, 20), column(4, 19)); + InnerKeyFilter filter = new InnerKeyFilter<>(raw, 10, 20); + + assertTrue(filter.hasNext()); + assertArrayEquals(new byte[]{2, 0, 10}, filter.next().name); + assertTrue(filter.hasNext()); + assertArrayEquals(new byte[]{4, 0, 19}, filter.next().name); + assertFalse(filter.hasNext()); + verify(raw, times(4)).next(); + filter.close(); + } + + @Test + public void testInnerKeyFilterNextBeforeHasNextPreservesTheFirstColumn() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true, false); + when(raw.next()).thenReturn(column(7, 15)); + InnerKeyFilter filter = new InnerKeyFilter<>(raw); + + assertArrayEquals(new byte[]{7}, filter.next().name); + assertFalse(filter.hasNext()); + verify(raw).next(); + filter.close(); + } + + @Test + public void testInnerKeyFilterIsValidUsesTheExistingPrefetchedDelegatePosition() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true); + when(raw.next()).thenReturn(column(1, 15)); + when(raw.isValid()).thenAnswer(invocation -> { + verify(raw).next(); + return false; + }); + InnerKeyFilter filter = new InnerKeyFilter<>(raw); + + assertFalse(filter.isValid()); + assertTrue(filter.hasNext()); + verify(raw).next(); + filter.close(); + } + + @Test + public void testInnerKeyFilterCountKeepsTheExistingPrefetchedDelegatePosition() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true); + when(raw.next()).thenReturn(column(1, 15)); + when(raw.count()).thenAnswer(invocation -> { + verify(raw).next(); + return 7L; + }); + InnerKeyFilter filter = new InnerKeyFilter<>(raw); + + assertEquals(7L, filter.count()); + verify(raw).next(); + filter.close(); + } + + @Test + public void testInnerKeyFilterCloseBeforeUseDoesNotPrefetch() { + for (int constructor = 0; constructor < 3; constructor++) { + ScanIterator raw = mock(ScanIterator.class); + InnerKeyFilter filter = filter(raw, constructor); + + filter.close(); + assertFalse(filter.hasNext()); + verify(raw, never()).hasNext(); + verify(raw, never()).next(); + verify(raw).close(); + } + } + + @Test + public void testInnerKeyFilterPreservesPartitionResumeKey() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true, false); + when(raw.next()).thenReturn(column(9, 15)); + AtomicInteger supplied = new AtomicInteger(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), (id, position) -> { + assertEquals(Integer.valueOf(2), id); + assertArrayEquals(new byte[]{9}, position); + supplied.incrementAndGet(); + return new InnerKeyFilter<>(raw); + }); + iterator.seek(ByteBuffer.allocate(Integer.BYTES + 1).putInt(2).put((byte) 9).array()); + + assertTrue(iterator.hasNext()); + BackendColumn result = iterator.next(); + assertArrayEquals(new byte[]{9}, result.name); + assertArrayEquals(ByteBuffer.allocate(Integer.BYTES).putInt(2).array(), iterator.position()); + assertFalse(iterator.hasNext()); + assertEquals(1, supplied.get()); + verify(raw).close(); + } + + private static InnerKeyFilter filter(ScanIterator raw, int constructor) { + if (constructor == 0) { + return new InnerKeyFilter<>(raw); + } + if (constructor == 1) { + return new InnerKeyFilter<>(raw, true); + } + return new InnerKeyFilter<>(raw, 10, 20); + } + + private static BackendColumn column(int key, int code) { + byte[] name = new byte[Short.BYTES + 1 + Short.BYTES]; + Bits.putShort(name, 0, 1); + name[Short.BYTES] = (byte) key; + Bits.putShort(name, name.length - Short.BYTES, code); + return BackendColumn.of(name, new byte[]{(byte) key}); + } + + private static void assertAutomaticFailureIsRetained(Fixture fixture, Throwable failure) { + assertSame(failure, assertThrows(Throwable.class, fixture.iterator::hasNext)); + assertSame(failure, assertThrows(Throwable.class, fixture.iterator::close)); + assertSame(failure, assertThrows(Throwable.class, fixture.iterator::close)); + verify(fixture.raw).close(); + } + + private static void assertConcurrentCloseWaits(Throwable failure) throws Exception { + Fixture fixture = new Fixture(); + CountDownLatch nativeCloseStarted = new CountDownLatch(1); + CountDownLatch allowNativeClose = new CountDownLatch(1); + doAnswer(invocation -> { + nativeCloseStarted.countDown(); + allowNativeClose.await(); + return null; + }).when(fixture.raw).close(); + if (failure != null) { + doThrow(failure).doNothing().when(fixture.closeOp).accept(true); + } + FutureTask first = new FutureTask<>(() -> { + fixture.iterator.close(); + return null; + }); + FutureTask second = new FutureTask<>(() -> { + fixture.iterator.close(); + return null; + }); + Thread firstThread = new Thread(first, "test-first-native-close"); + Thread secondThread = new Thread(second, "test-second-native-close"); + firstThread.setDaemon(true); + secondThread.setDaemon(true); + try { + firstThread.start(); + assertTrue(nativeCloseStarted.await(1, TimeUnit.SECONDS)); + secondThread.start(); + awaitBlockedClose(secondThread); + assertFalse("second close cannot claim release before native close finishes", second.isDone()); + verify(fixture.reference, never()).release(); + allowNativeClose.countDown(); + if (failure == null) { + first.get(2, TimeUnit.SECONDS); + second.get(2, TimeUnit.SECONDS); + verify(fixture.reference).release(); + } else { + assertSame(failure, assertThrows(ExecutionException.class, + () -> first.get(2, TimeUnit.SECONDS)).getCause()); + assertSame(failure, assertThrows(ExecutionException.class, + () -> second.get(2, TimeUnit.SECONDS)).getCause()); + verify(fixture.reference, never()).release(); + } + verify(fixture.raw).close(); + verify(fixture.closeOp).accept(true); + } finally { + allowNativeClose.countDown(); + firstThread.join(1000); + secondThread.join(1000); + } + } + + private static void awaitBlockedClose(Thread thread) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (System.nanoTime() < deadline) { + if (thread.getState() == Thread.State.BLOCKED) { + return; + } + if (thread.getState() == Thread.State.TERMINATED) { + break; + } + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); + } + fail("concurrent close must wait for the first resource release"); + } + + private static final class Fixture { + + private final RocksIterator raw = mock(RocksIterator.class); + private final RocksDBSession.RefCounter reference = mock(RocksDBSession.RefCounter.class); + @SuppressWarnings("unchecked") + private final Consumer closeOp = mock(Consumer.class); + private final RocksDBScanIterator iterator; + + private Fixture() { + when(this.raw.isOwningHandle()).thenReturn(true); + this.iterator = new RocksDBScanIterator<>(this.raw, null, null, + ScanIterator.Trait.SCAN_ANY, this.reference, this.closeOp); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java index d746299873..7a89ddf241 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java @@ -23,6 +23,7 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.List; +import java.util.Map; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -35,6 +36,7 @@ import org.apache.hugegraph.pd.common.PDException; import org.apache.hugegraph.pd.grpc.Metapb; import org.apache.hugegraph.pd.grpc.Pdpb.ErrorType; +import org.apache.hugegraph.store.business.BusinessHandler; import org.apache.hugegraph.store.meta.Partition; import org.apache.hugegraph.store.meta.PartitionManager; import org.apache.hugegraph.store.meta.Store; @@ -46,6 +48,8 @@ import org.junit.Assert; import org.junit.Test; import org.mockito.Mockito; +import org.rocksdb.Options; +import org.rocksdb.RocksDB; import com.alipay.sofa.jraft.rpc.RpcServer; @@ -72,6 +76,11 @@ public void testFatalHeartbeatExitsAfterShutdownHooks() throws Exception { } } + @Test + public void testInterruptedEngineShutdownFinishesNativeTeardown() throws Exception { + assertFatalExit("interrupted-shutdown", 0); + } + private static void assertFatalExit(String phase, int code) throws Exception { Path directory = Files.createTempDirectory(Path.of("target"), "heartbeat-fatal-"); Path marker = directory.resolve("hook-complete"); @@ -99,7 +108,9 @@ private static void assertFatalExit(String phase, int code) throws Exception { } Assert.assertTrue("Fatal heartbeat did not exit; child log: " + output + "\n" + Files.readString(output), exited); - Assert.assertEquals(Files.readString(output), "recovery".equals(phase) ? 0 : 255, process.exitValue()); + Assert.assertEquals(Files.readString(output), + "recovery".equals(phase) || "interrupted-shutdown".equals(phase) ? 0 : 255, + process.exitValue()); Assert.assertEquals("joined", Files.readString(marker)); if ("recovery".equals(phase)) { String log = Files.readString(output); @@ -122,6 +133,17 @@ public static void main(String[] args) throws Exception { runInvalidShardRecovery(marker); return; } + if ("interrupted-shutdown".equals(phase)) { + try { + runInterruptedEngineShutdown(marker); + } catch (Exception | AssertionError failure) { + failure.printStackTrace(); + System.exit(1); + } + // RocksDBFactory owns a process-wide scheduler; this isolated probe has finished. + System.exit(0); + return; + } CountDownLatch producersStarted = new CountDownLatch(1); PdProvider provider = Mockito.mock(PdProvider.class); Mockito.when(provider.getClusterStats()).thenReturn(Metapb.ClusterStats.newBuilder() @@ -223,6 +245,62 @@ private static void runInvalidShardRecovery(Path marker) throws Exception { service.onStateChanged(Metapb.StoreState.Up); } + private static void runInterruptedEngineShutdown(Path marker) throws Exception { + RocksDB.loadLibrary(); + Options nativeOptions = new Options().setCreateIfMissing(true); + RocksDB nativeDatabase = RocksDB.open(nativeOptions, marker.resolveSibling("native-store").toString()); + Constructor constructor = HgStoreEngine.class.getDeclaredConstructor(); + constructor.setAccessible(true); + HgStoreEngine engine = constructor.newInstance(); + Fixture heartbeat = new Fixture(false); + heartbeat.start(); + Assert.assertTrue(heartbeat.entered.await(5, TimeUnit.SECONDS)); + RpcServer rpc = Mockito.mock(RpcServer.class); + BusinessHandler handler = Mockito.mock(BusinessHandler.class); + PartitionEngine partition = Mockito.mock(PartitionEngine.class); + Mockito.when(partition.getGroupId()).thenReturn(1); + Mockito.when(partition.getRaftNode()).thenReturn(Mockito.mock(com.alipay.sofa.jraft.Node.class)); + Mockito.doAnswer(invocation -> { + Assert.assertFalse(Thread.currentThread().isInterrupted()); + Thread.currentThread().interrupt(); + return null; + }).when(partition).shutdown(); + Mockito.doAnswer(invocation -> { + Assert.assertFalse("Raft's restored interrupt must not escape into native cleanup", + Thread.currentThread().isInterrupted()); + Assert.assertFalse(heartbeat.storeThread.isAlive()); + Assert.assertFalse(heartbeat.partitionThread.isAlive()); + nativeDatabase.close(); + return null; + }).when(handler).closeDB(1); + setField(HgStoreEngine.class, engine, "heartbeatService", heartbeat); + setField(HgStoreEngine.class, engine, "metricService", Mockito.mock(HgMetricService.class)); + setField(HgStoreEngine.class, engine, "rpcServer", rpc); + setField(HgStoreEngine.class, engine, "businessHandler", handler); + Field partitionsField = HgStoreEngine.class.getDeclaredField("partitionEngines"); + partitionsField.setAccessible(true); + @SuppressWarnings("unchecked") + Map partitions = (Map) partitionsField.get(engine); + partitions.put(1, partition); + try { + Thread.currentThread().interrupt(); + engine.shutdown(); + Assert.assertTrue("restore the listener's interrupt after all teardown", + Thread.currentThread().isInterrupted()); + Assert.assertTrue(partitions.isEmpty()); + Mockito.verify(handler).closeDB(1); + Assert.assertFalse("engine teardown must reach actual native database close", + nativeDatabase.isOwningHandle()); + Mockito.verify(rpc).shutdown(); + Files.writeString(marker, "joined"); + } finally { + Thread.interrupted(); + heartbeat.shutdown(); + nativeDatabase.close(); + nativeOptions.close(); + } + } + private static void setField(Class owner, Object object, String name, Object value) throws Exception { Field field = owner.getDeclaredField(name); field.setAccessible(true); @@ -307,20 +385,16 @@ public void testShutdownWaitsForDerivedStateCallbackAndRejectsLateUpdates() thro } @Test - public void testInterruptedShutdownDoesNotAllowDatabaseClose() throws Exception { + public void testInterruptedShutdownKeepsWaitingForBothProducers() throws Exception { Fixture service = new Fixture(true); CountDownLatch closing = new CountDownLatch(1); AtomicBoolean databaseClosed = new AtomicBoolean(); - AtomicBoolean interruptedFailure = new AtomicBoolean(); + AtomicBoolean interruptPreserved = new AtomicBoolean(); Thread closer = new Thread(() -> { closing.countDown(); - try { - service.shutdown(); - // HgStoreEngine only closes native owners after this call succeeds. - databaseClosed.set(true); - } catch (IllegalStateException expected) { - interruptedFailure.set(Thread.currentThread().isInterrupted()); - } + service.shutdown(); + databaseClosed.set(true); + interruptPreserved.set(Thread.currentThread().isInterrupted()); }); try { service.start(); @@ -328,10 +402,18 @@ public void testInterruptedShutdownDoesNotAllowDatabaseClose() throws Exception closer.start(); Assert.assertTrue(closing.await(5, TimeUnit.SECONDS)); closer.interrupt(); + closer.join(100); + Assert.assertTrue("interruption must not abandon producer drain", closer.isAlive()); + Assert.assertFalse(databaseClosed.get()); + service.storeRelease.countDown(); + closer.join(100); + Assert.assertTrue("both producers must drain", closer.isAlive()); + Assert.assertFalse(databaseClosed.get()); + service.partitionRelease.countDown(); closer.join(5000); Assert.assertFalse(closer.isAlive()); - Assert.assertTrue(interruptedFailure.get()); - Assert.assertFalse(databaseClosed.get()); + Assert.assertTrue(databaseClosed.get()); + Assert.assertTrue(interruptPreserved.get()); } finally { service.storeRelease.countDown(); service.partitionRelease.countDown(); diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java index 78445e5934..f47b571f2d 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java @@ -19,6 +19,8 @@ import org.apache.hugegraph.store.client.grpc.AbstractGrpcClientTest; import org.apache.hugegraph.store.client.grpc.KvPageScannerTest; +import org.apache.hugegraph.store.client.query.CommonKvStreamObserverTest; +import org.apache.hugegraph.store.client.query.MultiStreamIteratorTest; import org.junit.runner.RunWith; import org.junit.runners.Suite; @@ -32,7 +34,9 @@ NodeTxExecutorTest.class, KvPageScannerTest.class, NodeTxSessionProxyTest.class, - OrderedKvIteratorTest.class + OrderedKvIteratorTest.class, + CommonKvStreamObserverTest.class, + MultiStreamIteratorTest.class }) public class ClientSuiteTest { } diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserverTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserverTest.java new file mode 100644 index 0000000000..71e88f7fa7 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserverTest.java @@ -0,0 +1,485 @@ +/* + * 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.store.client.query; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.hugegraph.store.HgKvIterator; +import org.junit.Test; + +public class CommonKvStreamObserverTest { + + @Test + public void testCompletedEmptyStreamSkipsTimedPoll() throws Exception { + RecordingQueue queue = new RecordingQueue(); + CommonKvStreamObserver, Integer> observer = observer(queue); + observer.onCompleted(); + StreamKvIterator stream = stream(observer); + assertFalse(stream.hasNext()); + stream.close(); + assertEquals("completed empty stream must not enter timed poll", 0, queue.emptyTimedPolls.get()); + } + + @Test + public void testCompletedSingleStreamAndCloseSkipEmptyTimedPoll() throws Exception { + RecordingQueue queue = new RecordingQueue(); + CommonKvStreamObserver, Integer> observer = observer(queue); + observer.onNext(Collections.singletonList(1)); + observer.onCompleted(); + StreamKvIterator stream = stream(observer); + assertTrue(stream.hasNext()); + assertEquals(Integer.valueOf(1), stream.next()); + assertFalse(stream.hasNext()); + assertFalse(stream.hasNext()); + stream.close(); + assertEquals("exhausted terminal stream must not enter empty timed poll", 0, queue.emptyTimedPolls.get()); + } + + @Test + public void testCompletedThreeStreamsMergeAndCloseSkipEmptyTimedPoll() throws Exception { + List queues = new ArrayList<>(); + MultiStreamIterator merged = merged(queues); + List actual = drain(merged); + merged.close(); + assertEquals(Arrays.asList(1, 2, 3), actual); + assertEquals("merged terminal streams and close must not enter empty timed poll", 0, + queues.stream().mapToInt(queue -> queue.emptyTimedPolls.get()).sum()); + } + + @Test(timeout = 5000) + public void testFinalBatchPublishedAtTerminalBoundaryIsNotLost() throws Exception { + CountDownLatch boundary = new CountDownLatch(1); + CountDownLatch published = new CountDownLatch(1); + GatedQueue queue = new GatedQueue(boundary, published); + CommonKvStreamObserver, Integer> observer = observer(queue); + StreamKvIterator stream = stream(observer); + AtomicReference failure = new AtomicReference<>(); + Thread producer = new Thread(() -> { + try { + assertTrue("consumer must reach terminal/queue boundary", boundary.await(2, TimeUnit.SECONDS)); + observer.onNext(Collections.singletonList(7)); + observer.onCompleted(); + } catch (Throwable error) { + failure.set(error); + } finally { + published.countDown(); + } + }, "publish-final-client-batch"); + producer.start(); + try { + assertTrue("final batch must remain visible when terminal state changes", stream.hasNext()); + assertEquals(Integer.valueOf(7), stream.next()); + assertFalse(stream.hasNext()); + assertEquals(null, failure.get()); + } finally { + boundary.countDown(); + published.countDown(); + producer.join(2000); + assertFalse("producer must finish", producer.isAlive()); + stream.close(); + } + } + + @Test(timeout = 5000) + public void testAcceptedResponseRefreshesTimeoutBeforeParsing() throws Exception { + CountDownLatch parsing = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + AtomicInteger polls = new AtomicInteger(); + RecordingQueue queue = new RecordingQueue() { + private static final long serialVersionUID = 1L; + + @Override + public Iterator poll(long timeout, TimeUnit unit) throws InterruptedException { + if (polls.getAndIncrement() == 0) { + // Exercise the timeout check while the accepted response is still being parsed. + return null; + } + release.countDown(); + return super.poll(timeout, unit); + } + }; + CommonKvStreamObserver, Integer> observer = parsingObserver(parsing, release); + observer.setTimeout(1000); + Field queueField = CommonKvStreamObserver.class.getDeclaredField("queue"); + queueField.setAccessible(true); + queueField.set(observer, queue); + Field current = CommonKvStreamObserver.class.getDeclaredField("current"); + current.setAccessible(true); + current.setLong(observer, System.nanoTime() - TimeUnit.SECONDS.toNanos(2)); + AtomicReference failure = new AtomicReference<>(); + Thread producer = new Thread(() -> { + try { + observer.onNext(Collections.singletonList(7)); + } catch (Throwable error) { + failure.set(error); + } + }); + producer.start(); + try { + assertTrue(parsing.await(2, TimeUnit.SECONDS)); + assertEquals(Integer.valueOf(7), observer.consume().next()); + producer.join(2000); + assertFalse(producer.isAlive()); + assertEquals(null, failure.get()); + assertEquals(null, observer.consume()); + } finally { + release.countDown(); + producer.join(2000); + observer.clear(); + } + } + + @Test(timeout = 5000) + public void testCompletionWaitsForResponseBeingParsed() throws Exception { + CountDownLatch parsing = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CommonKvStreamObserver, Integer> observer = parsingObserver(parsing, release); + Thread producer = new Thread(() -> observer.onNext(Collections.singletonList(7))); + producer.start(); + try { + assertTrue(parsing.await(2, TimeUnit.SECONDS)); + observer.onCompleted(); + assertFalse("completion must wait for the accepted batch", observer.isServerFinished()); + release.countDown(); + producer.join(2000); + assertFalse(producer.isAlive()); + assertTrue(observer.isServerFinished()); + assertEquals(Integer.valueOf(7), observer.consume().next()); + assertEquals(null, observer.consume()); + } finally { + release.countDown(); + producer.join(2000); + observer.clear(); + } + } + + @Test(timeout = 5000) + public void testCloseDiscardsResponseBeingParsed() throws Exception { + CountDownLatch parsing = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CommonKvStreamObserver, Integer> observer = parsingObserver(parsing, release); + AtomicInteger cancellations = new AtomicInteger(); + observer.setTransferComplete(finished -> { + assertFalse(finished); + cancellations.incrementAndGet(); + }); + Thread producer = new Thread(() -> observer.onNext(Collections.singletonList(7))); + producer.start(); + try { + assertTrue(parsing.await(2, TimeUnit.SECONDS)); + observer.clear(); + observer.clear(); + assertEquals(null, observer.consume()); + release.countDown(); + producer.join(2000); + assertFalse(producer.isAlive()); + observer.onCompleted(); + observer.onError(new IllegalStateException("late error")); + assertEquals(null, observer.consume()); + assertEquals(1, cancellations.get()); + } finally { + release.countDown(); + producer.join(2000); + } + } + + @Test(timeout = 5000) + public void testCancellationBypassesBlockedRequestSender() throws Exception { + CountDownLatch sending = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch cancelled = new CountDownLatch(1); + CommonKvStreamObserver, Integer> observer = observer(new RecordingQueue()); + observer.setRequestSender(ignored -> { + sending.countDown(); + try { + release.await(); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new AssertionError(error); + } + }); + AtomicInteger cancellations = new AtomicInteger(); + observer.setTransferComplete(finished -> { + assertFalse(finished); + cancellations.incrementAndGet(); + cancelled.countDown(); + release.countDown(); + }); + Thread sender = new Thread(observer::sendRequest); + Thread closer = new Thread(observer::clear); + sender.start(); + try { + assertTrue(sending.await(2, TimeUnit.SECONDS)); + closer.start(); + assertTrue("cancel must not wait for the blocked sender", cancelled.await(2, TimeUnit.SECONDS)); + sender.join(1000); + closer.join(1000); + assertFalse(sender.isAlive()); + assertFalse(closer.isAlive()); + assertEquals(1, cancellations.get()); + } finally { + release.countDown(); + sender.join(1000); + closer.join(1000); + } + } + + @Test + public void testTerminalCallbacksAreIdempotentAndDoNotResurrectStream() throws Exception { + CommonKvStreamObserver, Integer> observer = observer(new RecordingQueue()); + observer.onError(new IllegalStateException("first error")); + observer.onCompleted(); + observer.onError(new IllegalStateException("duplicate error")); + observer.onNext(Collections.singletonList(7)); + assertTrue(observer.isServerFinished()); + Iterator error = observer.consume(); + assertTrue(error instanceof ErrorMessageIterator); + assertEquals(null, observer.consume()); + + CommonKvStreamObserver, Integer> completed = observer(new RecordingQueue()); + AtomicInteger halfCloses = new AtomicInteger(); + completed.setTransferComplete(finished -> { + assertTrue(finished); + halfCloses.incrementAndGet(); + }); + completed.onCompleted(); + completed.onCompleted(); + completed.onError(new IllegalStateException("late error")); + completed.onNext(Collections.singletonList(7)); + completed.clear(); + completed.clear(); + assertEquals(1, halfCloses.get()); + assertEquals(null, completed.consume()); + } + + @Test(timeout = 5000) + public void testErrorDiscardsResponseBeingParsed() throws Exception { + CountDownLatch parsing = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CommonKvStreamObserver, Integer> observer = parsingObserver(parsing, release); + Thread producer = new Thread(() -> observer.onNext(Collections.singletonList(7))); + producer.start(); + try { + assertTrue(parsing.await(2, TimeUnit.SECONDS)); + observer.onError(new IllegalStateException("transport error")); + observer.onCompleted(); + assertTrue("error does not wait for a discarded batch", observer.isServerFinished()); + assertTrue(observer.consume() instanceof ErrorMessageIterator); + assertEquals(null, observer.consume()); + release.countDown(); + producer.join(2000); + assertFalse(producer.isAlive()); + assertEquals(null, observer.consume()); + } finally { + release.countDown(); + producer.join(2000); + observer.clear(); + } + } + + @Test + public void testParserErrorReleasesAcceptedResponse() { + CommonKvStreamObserver, Integer> observer = new CommonKvStreamObserver<>(values -> { + throw new AssertionError("parser error"); + }, ignored -> ResultState.FINISHED); + try { + observer.onNext(Collections.singletonList(7)); + throw new AssertionError("parser Error must propagate"); + } catch (AssertionError error) { + assertEquals("parser error", error.getMessage()); + } + assertTrue(observer.isServerFinished()); + assertTrue(observer.consume() instanceof ErrorMessageIterator); + assertEquals(null, observer.consume()); + } + + @Test(timeout = 5000) + public void testParserErrorOverridesConcurrentCompletion() throws Exception { + assertParserFailureOverridesCompletion(new AssertionError("parser error")); + } + + @Test(timeout = 5000) + public void testParserExceptionOverridesConcurrentCompletion() throws Exception { + assertParserFailureOverridesCompletion(new IllegalStateException("parser exception")); + } + + private static void assertParserFailureOverridesCompletion(Throwable failure) throws Exception { + CountDownLatch parsing = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CommonKvStreamObserver, Integer> observer = new CommonKvStreamObserver<>(values -> { + parsing.countDown(); + await(release); + if (failure instanceof Error) { + throw (Error) failure; + } + throw (RuntimeException) failure; + }, ignored -> ResultState.FINISHED); + observer.setRequestSender(ignored -> { }); + observer.setTransferComplete(ignored -> { }); + AtomicReference propagated = new AtomicReference<>(); + Thread producer = new Thread(() -> { + try { + observer.onNext(Collections.singletonList(7)); + } catch (Throwable error) { + propagated.set(error); + } + }); + producer.start(); + try { + assertTrue(parsing.await(2, TimeUnit.SECONDS)); + observer.onCompleted(); + assertFalse(observer.isServerFinished()); + release.countDown(); + producer.join(2000); + assertFalse(producer.isAlive()); + assertEquals(failure instanceof Error ? failure : null, propagated.get()); + assertTrue(observer.isServerFinished()); + assertTrue("parser failure must remain visible to the consumer", + observer.consume() instanceof ErrorMessageIterator); + assertEquals(null, observer.consume()); + } finally { + release.countDown(); + producer.join(2000); + observer.clear(); + } + } + + private static CommonKvStreamObserver, Integer> parsingObserver( + CountDownLatch parsing, CountDownLatch release) { + CommonKvStreamObserver, Integer> observer = new CommonKvStreamObserver<>(values -> { + parsing.countDown(); + await(release); + return values.iterator(); + }, ignored -> ResultState.FINISHED); + observer.setRequestSender(ignored -> { }); + observer.setTransferComplete(ignored -> { }); + return observer; + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(2, TimeUnit.SECONDS)); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new AssertionError(error); + } + } + + private static List drain(HgKvIterator stream) { + List values = new ArrayList<>(); + while (stream.hasNext()) { + values.add(stream.next()); + } + return values; + } + + private static MultiStreamIterator merged(List queues) throws Exception { + List> streams = new ArrayList<>(); + for (int value : Arrays.asList(1, 2, 3)) { + RecordingQueue queue = new RecordingQueue(); + queues.add(queue); + CommonKvStreamObserver, Integer> observer = observer(queue); + observer.onNext(Collections.singletonList(value)); + observer.onCompleted(); + streams.add(stream(observer)); + } + return new MultiStreamIterator<>(streams); + } + + private static StreamKvIterator stream(CommonKvStreamObserver, Integer> observer) { + return new StreamKvIterator<>(ignored -> observer.clear(), observer::consume); + } + + private static CommonKvStreamObserver, Integer> observer(RecordingQueue queue) throws Exception { + CommonKvStreamObserver, Integer> observer = + new CommonKvStreamObserver<>(List::iterator, ignored -> ResultState.FINISHED); + observer.setRequestSender(ignored -> { }); + observer.setTransferComplete(ignored -> { }); + Field field = CommonKvStreamObserver.class.getDeclaredField("queue"); + field.setAccessible(true); + field.set(observer, queue); + return observer; + } + + private static class RecordingQueue extends LinkedBlockingQueue> { + + private static final long serialVersionUID = 1L; + final AtomicInteger emptyTimedPolls = new AtomicInteger(); + + @Override + public Iterator poll(long timeout, TimeUnit unit) throws InterruptedException { + if (super.isEmpty()) { + this.emptyTimedPolls.incrementAndGet(); + } + return super.poll(timeout, unit); + } + } + + private static final class GatedQueue extends RecordingQueue { + + private static final long serialVersionUID = 1L; + private final CountDownLatch boundary; + private final CountDownLatch published; + private final AtomicBoolean armed = new AtomicBoolean(true); + + private GatedQueue(CountDownLatch boundary, CountDownLatch published) { + this.boundary = boundary; + this.published = published; + } + + private void publishAtBoundary() { + if (this.armed.compareAndSet(true, false)) { + this.boundary.countDown(); + try { + assertTrue("final batch must be published", this.published.await(2, TimeUnit.SECONDS)); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new AssertionError(error); + } + } + } + + @Override + public boolean isEmpty() { + boolean empty = super.isEmpty(); + this.publishAtBoundary(); + return empty; + } + + @Override + public Iterator poll(long timeout, TimeUnit unit) throws InterruptedException { + this.publishAtBoundary(); + return super.poll(timeout, unit); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIteratorTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIteratorTest.java new file mode 100644 index 0000000000..675c6cb9b9 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIteratorTest.java @@ -0,0 +1,78 @@ +/* + * 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.store.client.query; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; + +import java.util.List; + +import org.apache.hugegraph.store.HgKvIterator; +import org.junit.Test; + +public class MultiStreamIteratorTest { + + @Test + public void testCloseCancelsEveryChildBeforeIterationAndIsTerminal() { + HgKvIterator first = mock(HgKvIterator.class); + HgKvIterator second = mock(HgKvIterator.class); + MultiStreamIterator composite = new MultiStreamIterator<>(List.of(first, second, first)); + composite.close(); + composite.close(); + assertFalse(composite.hasNext()); + verify(first).close(); + verify(second).close(); + verifyNoMoreInteractions(first, second); + try { + composite.next(); + fail("Closed iterator must not resume a child stream"); + } catch (java.util.NoSuchElementException expected) { + // Closed is terminal. + } + } + + @Test + public void testCloseAttemptsAllChildrenAndPreservesPrimaryFailure() { + HgKvIterator first = mock(HgKvIterator.class); + HgKvIterator second = mock(HgKvIterator.class); + HgKvIterator third = mock(HgKvIterator.class); + RuntimeException primary = new IllegalStateException("first"); + RuntimeException later = new IllegalArgumentException("second"); + doThrow(primary).when(first).close(); + doThrow(later).when(second).close(); + MultiStreamIterator composite = new MultiStreamIterator<>(List.of(first, second, third)); + try { + composite.close(); + fail("Cleanup failure must be visible"); + } catch (RuntimeException error) { + assertSame(primary, error); + assertEquals(1, error.getSuppressed().length); + assertSame(later, error.getSuppressed()[0]); + } + verify(first).close(); + verify(second).close(); + verify(third).close(); + assertFalse(composite.hasNext()); + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/BatchScanHalfCloseTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/BatchScanHalfCloseTest.java new file mode 100644 index 0000000000..6d60ff32e2 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/BatchScanHalfCloseTest.java @@ -0,0 +1,360 @@ +/* + * 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.store.node.grpc; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Queue; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.hugegraph.rocksdb.access.RocksDBSession; +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.store.HgScanQuery; +import org.apache.hugegraph.store.buffer.KVByteBuffer; +import org.apache.hugegraph.store.client.grpc.KvBatchScanner; +import org.apache.hugegraph.store.client.grpc.KvBatchScannerMerger; +import org.apache.hugegraph.store.grpc.common.Header; +import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.stream.HgStoreStreamGrpc; +import org.apache.hugegraph.store.grpc.stream.KvPageRes; +import org.apache.hugegraph.store.grpc.stream.KvStream; +import org.apache.hugegraph.store.grpc.stream.ScanCancelRequest; +import org.apache.hugegraph.store.grpc.stream.ScanQueryRequest; +import org.apache.hugegraph.store.grpc.stream.ScanStreamBatchReq; +import org.junit.Test; + +import io.grpc.ManagedChannel; +import io.grpc.Server; +import io.grpc.Status; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.stub.StreamObserver; + +public class BatchScanHalfCloseTest { + + @Test(timeout = 5000) + public void testQueuedQueryHalfCloseDeliversFinalPartialBatch() throws Exception { + for (boolean encoded : new boolean[]{false, true}) { + DeferredExecutor executor = new DeferredExecutor(); + ScanIterator source = rows(3); + HgStoreWrapperEx wrapper = wrapper(source); + Output output = new Output(); + StreamObserver input = input(encoded, output, wrapper, executor); + try { + input.onNext(request(2)); + input.onCompleted(); + assertEquals("half-close must preserve the queued query", 0, output.completions); + executor.runAll(); + assertEquals(3, output.rows); + assertTrue(output.over); + assertEquals(1, output.completions); + assertEquals(null, output.error); + verify(source).close(); + input.onCompleted(); + assertEquals(1, output.completions); + } finally { + input.onError(Status.CANCELLED.asRuntimeException()); + executor.runAll(); + executor.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + public void testHalfCloseWithExhaustedCreditFailsInsteadOfTruncating() throws Exception { + for (boolean encoded : new boolean[]{false, true}) { + DeferredExecutor executor = new DeferredExecutor(); + Output output = new Output(); + ScanIterator iterator; + StreamObserver input; + if (encoded) { + ScanBatchResponse response = new ScanBatchResponse(output.streams(), mock(HgStoreWrapperEx.class), + executor); + iterator = batches(20); + set(response, "iterator", iterator); + set(response, "limit", Long.MAX_VALUE); + set(response, "queryReceived", true); + input = response; + } else { + iterator = rows(21); + input = input(false, output, wrapper(iterator), executor); + input.onNext(request(1)); + } + try { + input.onCompleted(); + executor.runAll(); + assertTrue("granted batches must precede the credit failure", output.rows > 0); + assertTrue("the full result exceeds granted credit", output.rows < 20); + assertFalse(output.over); + assertEquals(0, output.completions); + assertEquals(Status.Code.FAILED_PRECONDITION, Status.fromThrowable(output.error).getCode()); + verify(iterator).close(); + input.onCompleted(); + assertEquals(1, output.errors); + } finally { + input.onError(Status.CANCELLED.asRuntimeException()); + executor.runAll(); + executor.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + public void testExplicitAndTransportCancellationWinBeforeFinalDelivery() throws Exception { + for (boolean encoded : new boolean[]{false, true}) { + for (boolean transport : new boolean[]{false, true}) { + DeferredExecutor executor = new DeferredExecutor(); + Output output = new Output(); + AtomicReference> request = new AtomicReference<>(); + AtomicInteger checks = new AtomicInteger(); + ScanIterator source = rows(2); + when(source.hasNext()).thenAnswer(ignored -> { + if (checks.getAndIncrement() == 0) { + if (transport) { + request.get().onError(Status.CANCELLED.asRuntimeException()); + } else { + request.get().onNext(ScanStreamBatchReq.newBuilder().setCancelRequest( + ScanCancelRequest.getDefaultInstance()).build()); + } + } + return false; + }); + StreamObserver input; + if (encoded) { + ScanBatchResponse response = new ScanBatchResponse(output.streams(), mock(HgStoreWrapperEx.class), + executor); + set(response, "iterator", source); + set(response, "limit", Long.MAX_VALUE); + set(response, "queryReceived", true); + input = response; + } else { + input = input(false, output, wrapper(source), executor); + input.onNext(request(1)); + } + request.set(input); + try { + input.onCompleted(); + executor.runAll(); + assertEquals(0, output.rows); + assertFalse(output.over); + assertEquals(0, output.completions); + assertEquals(1, output.errors); + assertEquals(Status.Code.CANCELLED, Status.fromThrowable(output.error).getCode()); + input.onCompleted(); + input.onError(Status.CANCELLED.asRuntimeException()); + verify(source).close(); + assertEquals(1, output.errors); + } finally { + executor.shutdownNow(); + } + } + } + } + + @Test(timeout = 5000) + public void testClientEarlyCloseSendsCancelButNormalFinishOnlyHalfCloses() throws Exception { + List commands = new ArrayList<>(); + AtomicInteger completions = new AtomicInteger(); + String name = InProcessServerBuilder.generateName(); + Server server = InProcessServerBuilder.forName(name).directExecutor().addService( + new HgStoreStreamGrpc.HgStoreStreamImplBase() { + @Override + public StreamObserver scanBatch2(StreamObserver output) { + return new StreamObserver() { + @Override + public void onNext(ScanStreamBatchReq value) { + commands.add(value.getQueryCase()); + } + + @Override + public void onError(Throwable error) { + throw new AssertionError(error); + } + + @Override + public void onCompleted() { + completions.incrementAndGet(); + output.onCompleted(); + } + }; + } + }).build().start(); + ManagedChannel channel = InProcessChannelBuilder.forName(name).directExecutor().build(); + try { + KvBatchScanner early = new KvBatchScanner(HgStoreStreamGrpc.newStub(channel), "g", HgScanQuery.tableOf("t"), + mock(KvBatchScannerMerger.class)); + early.close(); + early.close(); + assertEquals(java.util.Arrays.asList(ScanStreamBatchReq.QueryCase.QUERY_REQUEST, + ScanStreamBatchReq.QueryCase.CANCEL_REQUEST), commands); + assertEquals(1, completions.get()); + commands.clear(); + KvBatchScanner normal = new KvBatchScanner(HgStoreStreamGrpc.newStub(channel), "g", HgScanQuery.tableOf("t"), + mock(KvBatchScannerMerger.class)); + normal.dataComplete(); + assertEquals(Collections.singletonList(ScanStreamBatchReq.QueryCase.QUERY_REQUEST), commands); + assertEquals(2, completions.get()); + } finally { + channel.shutdownNow(); + server.shutdownNow(); + } + } + + private static StreamObserver input(boolean encoded, Output output, + HgStoreWrapperEx wrapper, DeferredExecutor executor) { + return encoded ? new ScanBatchResponse(output.streams(), wrapper, executor) : + ScanBatchResponse3.of(output.pages(), wrapper, executor); + } + + private static ScanStreamBatchReq request(int pageSize) { + return ScanStreamBatchReq.newBuilder().setHeader(Header.newBuilder().setGraph("g")) + .setQueryRequest(ScanQueryRequest.newBuilder().setMethod(ScanMethod.ALL).setTable("t") + .setLimit(Long.MAX_VALUE).setPerKeyMax(Long.MAX_VALUE).setPageSize(pageSize)).build(); + } + + private static HgStoreWrapperEx wrapper(ScanIterator source) { + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(source); + return wrapper; + } + + private static ScanIterator rows(int size) { + ScanIterator source = mock(ScanIterator.class); + AtomicInteger read = new AtomicInteger(); + when(source.hasNext()).thenAnswer(ignored -> read.get() < size); + when(source.next()).thenAnswer(ignored -> { + read.incrementAndGet(); + return RocksDBSession.BackendColumn.of(new byte[4], new byte[4]); + }); + when(source.position()).thenReturn(new byte[4]); + return source; + } + + private static ScanIterator batches(int size) { + ScanIterator source = mock(ScanIterator.class); + AtomicInteger read = new AtomicInteger(); + when(source.hasNext()).thenAnswer(ignored -> read.get() < size); + when(source.next()).thenAnswer(ignored -> { + read.incrementAndGet(); + return Collections.singletonList(ParallelScanIterator.KV.of( + RocksDBSession.BackendColumn.of(new byte[4], new byte[4]))); + }); + return source; + } + + private static void set(Object target, String name, Object value) throws Exception { + Field field = target.getClass().getDeclaredField(name); + field.setAccessible(true); + field.set(target, value); + } + + private static final class DeferredExecutor extends ThreadPoolExecutor { + + private final Queue tasks = new ArrayDeque<>(); + + DeferredExecutor() { + super(1, 1, 0, TimeUnit.SECONDS, new LinkedBlockingQueue<>()); + } + + @Override + public void execute(Runnable task) { + this.tasks.add(task); + } + + void runAll() { + Runnable task; + while ((task = this.tasks.poll()) != null) { + task.run(); + } + } + } + + private static final class Output { + + int rows; + int completions; + int errors; + boolean over; + Throwable error; + + StreamObserver pages() { + return new StreamObserver() { + @Override + public void onNext(KvPageRes value) { + rows += value.getDataCount(); + over |= value.getOver(); + } + + @Override + public void onError(Throwable failure) { + error = failure; + errors++; + } + + @Override + public void onCompleted() { + completions++; + } + }; + } + + StreamObserver streams() { + return new StreamObserver() { + @Override + public void onNext(KvStream value) { + over |= value.getOver(); + if (value.getStream() != null) { + KVByteBuffer buffer = new KVByteBuffer(value.getStream().duplicate()); + while (buffer.hasRemaining()) { + buffer.getBytes(); + buffer.getBytes(); + rows++; + } + } + } + + @Override + public void onError(Throwable failure) { + error = failure; + errors++; + } + + @Override + public void onCompleted() { + completions++; + } + }; + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryLifecycleTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryLifecycleTest.java new file mode 100644 index 0000000000..0760ef10ce --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryLifecycleTest.java @@ -0,0 +1,940 @@ +/* + * 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.store.node.grpc.query; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.grpc.Status; +import io.grpc.stub.StreamObserver; +import java.util.Arrays; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.FutureTask; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.hugegraph.rocksdb.access.RocksDBSession; +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.store.business.MultiPartitionIterator; +import org.apache.hugegraph.store.grpc.query.QueryRequest; +import org.apache.hugegraph.store.grpc.query.QueryResponse; +import org.apache.hugegraph.store.node.grpc.query.model.PipelineResult; +import org.apache.hugegraph.store.node.grpc.query.model.QueryPlan; +import org.junit.Test; + +public class AggregativeQueryLifecycleTest extends AggregativeQueryTestSupport { + + @Test(timeout = 5000) + public void testPartitionTimeoutDoesNotClearPlanOrFinishBeforeChildFinally() throws Exception { + ThreadPoolExecutor pool = pool(2); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch releaseChild = new CountDownLatch(1); + CountDownLatch parentReturned = new CountDownLatch(1); + CountDownLatch childClosed = new CountDownLatch(1); + ScanIterator child = mock(ScanIterator.class); + when(child.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + awaitUninterruptibly(releaseChild); + return false; + }); + doAnswer(invocation -> { + childClosed.countDown(); + return null; + }).when(child).close(); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(child)); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + AggregativeQueryService service = new AggregativeQueryService(pool, 50, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + AggregativeQueryObserver observer = super.newObserver(sender); + return fixture(observer, iterator, plan, parentReturned); + } + }; + ResponseRecorder response = new ResponseRecorder(); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("timeout-query").build()); + assertTrue(reading.await(1, TimeUnit.SECONDS)); + assertTrue("parent timeout must not wait on its own executor's queued children", + parentReturned.await(1, TimeUnit.SECONDS)); + assertTrue("a live client must receive the timeout error", response.received.await(1, TimeUnit.SECONDS)); + assertError(response, "timeout-query"); + assertEquals("response must not complete before child finally", 0, response.completed.get()); + verify(plan, never()).clear(); + verify(iterator, never()).close(); + FutureTask closing = close(service); + awaitShutdownWait(start(closing)); + assertFalse(closing.isDone()); + releaseChild.countDown(); + closing.get(2, TimeUnit.SECONDS); + assertTrue(childClosed.await(1, TimeUnit.SECONDS)); + verify(plan).clear(); + verify(iterator).close(); + assertTrue(pool.awaitTermination(1, TimeUnit.SECONDS)); + assertEquals(0, response.completed.get()); + assertEquals(1, response.errors.get()); + } finally { + releaseChild.countDown(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testRejectedPartitionsReleaseAllOpenedIterators() throws Exception { + AtomicInteger submissions = new AtomicInteger(); + ThreadPoolExecutor pool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, + new LinkedBlockingQueue<>()) { + @Override + public void execute(Runnable task) { + if (submissions.incrementAndGet() > 1) { + throw new RejectedExecutionException("partition queue is full"); + } + super.execute(task); + } + }; + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(first, second)); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + CountDownLatch parentReturned = new CountDownLatch(1); + AggregativeQueryService service = new AggregativeQueryService(pool, 500, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + return fixture(super.newObserver(sender), iterator, plan, parentReturned); + } + }; + try { + service.query(sender()).onNext(QueryRequest.getDefaultInstance()); + assertTrue(parentReturned.await(1, TimeUnit.SECONDS)); + FutureTask closing = close(service); + start(closing); + closing.get(2, TimeUnit.SECONDS); + verify(first).close(); + verify(second).close(); + verify(iterator).close(); + verify(plan).clear(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testHasNextFailureSendsErrorBeforeSingleCompletion() throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + when(iterator.hasNext()).thenThrow(new IllegalStateException("iterator read failed")); + QueryPlan plan = mock(QueryPlan.class); + when(plan.isEmpty()).thenReturn(true); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, plan, 500); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("read-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "read-query"); + assertEquals(1, response.completed.get()); + assertEquals(0, response.errors.get()); + service.shutdownQueries(); + verify(iterator).close(); + verify(plan).clear(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testPartitionFailureSendsErrorBeforeSingleCompletion() throws Exception { + ThreadPoolExecutor pool = pool(2); + ScanIterator child = mock(ScanIterator.class); + when(child.hasNext()).thenThrow(new IllegalStateException("partition read failed")); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(child)); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, plan, 500); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("partition-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "partition-query"); + assertEquals(1, response.completed.get()); + service.shutdownQueries(); + verify(child).close(); + verify(iterator).close(); + verify(plan).clear(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testAcceptedChildAndRejectedRemainingChildrenRetainCleanupOwnership() throws Exception { + CountDownLatch childStarted = new CountDownLatch(1); + CountDownLatch releaseChild = new CountDownLatch(1); + CountDownLatch parentReturned = new CountDownLatch(1); + AtomicInteger submissions = new AtomicInteger(); + ThreadPoolExecutor pool = new ThreadPoolExecutor(2, 2, 0, TimeUnit.SECONDS, + new LinkedBlockingQueue<>()) { + @Override + public void execute(Runnable task) { + int submitted = submissions.incrementAndGet(); + if (submitted >= 3) { + throw new RejectedExecutionException("second partition rejected"); + } + super.execute(task); + if (submitted == 2) { + awaitUninterruptibly(childStarted); + } + } + }; + ScanIterator accepted = mock(ScanIterator.class); + when(accepted.hasNext()).thenAnswer(invocation -> { + childStarted.countDown(); + awaitUninterruptibly(releaseChild); + return false; + }); + ScanIterator rejected = mock(ScanIterator.class); + ScanIterator notSubmitted = mock(ScanIterator.class); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(accepted, rejected, notSubmitted)); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + AggregativeQueryService service = new AggregativeQueryService(pool, 500, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + return fixture(super.newObserver(sender), iterator, plan, parentReturned); + } + }; + ResponseRecorder response = new ResponseRecorder(); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("mixed-query").build()); + assertTrue(childStarted.await(1, TimeUnit.SECONDS)); + assertTrue(parentReturned.await(1, TimeUnit.SECONDS)); + assertError(response, "mixed-query"); + assertEquals(0, response.completed.get()); + verify(rejected).close(); + verify(notSubmitted).close(); + verify(accepted, never()).close(); + verify(iterator, never()).close(); + verify(plan, never()).clear(); + FutureTask closing = close(service); + awaitShutdownWait(start(closing)); + assertFalse("accepted child's finally still owns the query", closing.isDone()); + releaseChild.countDown(); + closing.get(2, TimeUnit.SECONDS); + assertEquals(0, response.completed.get()); + assertEquals(1, response.errors.get()); + verify(accepted).close(); + verify(iterator).close(); + verify(plan).clear(); + } finally { + releaseChild.countDown(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testShutdownReportsUnavailableWithoutReleasingActiveIteratorEarly() throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + when(iterator.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + awaitUninterruptibly(release); + return false; + }); + QueryPlan plan = new QueryPlan(); + ResponseRecorder response = new ResponseRecorder() { + @Override + public void onError(Throwable error) { + assertEquals(Status.Code.UNAVAILABLE, Status.fromThrowable(error).getCode()); + super.onError(error); + } + }; + AggregativeQueryService service = service(pool, iterator, plan, 5000); + try { + StreamObserver request = service.query(response); + request.onNext(QueryRequest.newBuilder().setQueryId("shutdown-query").build()); + assertTrue(reading.await(1, TimeUnit.SECONDS)); + FutureTask closing = close(service); + start(closing); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertFalse("termination must retain active iterator ownership", closing.isDone()); + verify(iterator, never()).close(); + release.countDown(); + closing.get(2, TimeUnit.SECONDS); + request.onCompleted(); + request.onError(Status.CANCELLED.asRuntimeException()); + assertTrue(response.responses.isEmpty()); + assertEquals(0, response.completed.get()); + assertEquals(1, response.errors.get()); + verify(iterator).close(); + assertTrue(pool.awaitTermination(1, TimeUnit.SECONDS)); + } finally { + release.countDown(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testShutdownReportsUnavailableForIdleQueryOnce() { + ThreadPoolExecutor pool = pool(1); + AggregativeQueryService service = service(pool, mock(ScanIterator.class), new QueryPlan(), 500); + ResponseRecorder response = new ResponseRecorder(); + try { + StreamObserver request = service.query(response); + service.shutdownQueries(); + request.onCompleted(); + service.shutdownQueries(); + assertEquals(0, response.completed.get()); + assertEquals(1, response.errors.get()); + assertTrue(response.responses.isEmpty()); + } finally { + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testNormalIdleHalfCloseCompletesOnceAndExternalErrorNeverCompletesNormally() { + ThreadPoolExecutor pool = pool(1); + AggregativeQueryService service = service(pool, mock(ScanIterator.class), new QueryPlan(), 500); + try { + ResponseRecorder normal = new ResponseRecorder(); + StreamObserver normalRequest = service.query(normal); + normalRequest.onCompleted(); + normalRequest.onCompleted(); + assertEquals(1, normal.completed.get()); + ResponseRecorder failed = new ResponseRecorder(); + StreamObserver failedRequest = service.query(failed); + failedRequest.onError(Status.CANCELLED.asRuntimeException()); + failedRequest.onCompleted(); + assertEquals(0, failed.completed.get()); + service.shutdownQueries(); + assertEquals(1, normal.completed.get()); + assertEquals(0, failed.completed.get()); + } finally { + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testActiveNormalHalfCloseCompletesAfterWorkerCleanupOnce() throws Exception { + ThreadPoolExecutor pool = pool(1); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + ScanIterator iterator = mock(ScanIterator.class); + AtomicInteger reads = new AtomicInteger(); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[]{1}, new byte[]{2})); + when(iterator.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + assertTrue(release.await(2, TimeUnit.SECONDS)); + return reads.incrementAndGet() == 1; + }); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, new QueryPlan(), 500); + try { + StreamObserver request = service.query(response); + request.onNext(QueryRequest.newBuilder().setQueryId("half-close-query").build()); + assertTrue(reading.await(1, TimeUnit.SECONDS)); + request.onCompleted(); + assertFalse(response.finished.await(50, TimeUnit.MILLISECONDS)); + release.countDown(); + assertTrue("normal half-close must complete after cleanup", response.finished.await(1, TimeUnit.SECONDS)); + request.onCompleted(); + service.shutdownQueries(); + assertEquals(1, response.completed.get()); + assertEquals(0, response.errors.get()); + assertEquals(1, response.responses.size()); + assertTrue(response.responses.get(0).getIsFinished()); + assertTrue(response.responses.get(0).getIsOk()); + assertEquals(1, response.responses.get(0).getDataCount()); + assertEquals(com.google.protobuf.ByteString.copyFrom(new byte[]{2}), + response.responses.get(0).getData(0).getValue()); + verify(iterator).close(); + } finally { + release.countDown(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testHalfCloseWithoutEnoughFeedbackReportsError() throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + when(iterator.hasNext()).thenReturn(true); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[]{1}, new byte[]{2})); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, new QueryPlan(), 5000); + try { + StreamObserver request = service.query(response); + request.onNext(QueryRequest.newBuilder().setQueryId("feedback-half-close").build()); + request.onCompleted(); + assertTrue(response.finished.await(2, TimeUnit.SECONDS)); + assertEquals(17, response.responses.size()); + QueryResponse error = response.responses.get(16); + assertFalse(error.getIsOk()); + assertTrue(error.getMessage().contains("without enough feedback")); + assertEquals(1, response.completed.get()); + verify(iterator).close(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 15000) + public void testTransportCancellationAfterHalfCloseReleasesWorker() throws Exception { + assertTransportCancellationAfterHalfClose(false); + } + + @Test(timeout = 15000) + public void testDeadlineAfterHalfCloseReleasesWorker() throws Exception { + assertTransportCancellationAfterHalfClose(true); + } + + private static void assertTransportCancellationAfterHalfClose(boolean deadline) throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator source = mock(ScanIterator.class); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + CountDownLatch released = new CountDownLatch(1); + CountDownLatch fallbackRelease = new CountDownLatch(1); + when(source.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + try { + fallbackRelease.await(); + } catch (InterruptedException expected) { + interrupted.countDown(); + Thread.currentThread().interrupt(); + } + return false; + }); + doAnswer(invocation -> { + released.countDown(); + return null; + }).when(source).close(); + AggregativeQueryService service = new AggregativeQueryService(pool, 5000, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + return fixture(super.newObserver(sender), source, new QueryPlan(), null); + } + }; + String name = io.grpc.inprocess.InProcessServerBuilder.generateName(); + io.grpc.Server server = io.grpc.inprocess.InProcessServerBuilder.forName(name) + .directExecutor().addService(service).build().start(); + io.grpc.ManagedChannel channel = io.grpc.inprocess.InProcessChannelBuilder.forName(name) + .directExecutor().build(); + io.grpc.CallOptions options = deadline ? io.grpc.CallOptions.DEFAULT.withDeadlineAfter(2, TimeUnit.SECONDS) : + io.grpc.CallOptions.DEFAULT; + io.grpc.ClientCall call = channel.newCall( + org.apache.hugegraph.store.grpc.query.QueryServiceGrpc.getQueryMethod(), options); + CountDownLatch terminated = new CountDownLatch(1); + java.util.concurrent.atomic.AtomicReference terminal = + new java.util.concurrent.atomic.AtomicReference<>(); + try { + call.start(new io.grpc.ClientCall.Listener() { + @Override + public void onClose(Status status, io.grpc.Metadata trailers) { + terminal.set(status); + terminated.countDown(); + } + }, new io.grpc.Metadata()); + call.request(1); + call.sendMessage(QueryRequest.newBuilder().setQueryId("half-close-cancel").build()); + assertTrue(reading.await(1, TimeUnit.SECONDS)); + call.halfClose(); + if (!deadline) { + call.cancel("cancel after half-close", null); + } + assertTrue(terminated.await(3, TimeUnit.SECONDS)); + assertEquals(deadline ? Status.Code.DEADLINE_EXCEEDED : Status.Code.CANCELLED, + terminal.get().getCode()); + assertTrue("transport cancellation must interrupt the worker after half-close", + interrupted.await(1, TimeUnit.SECONDS)); + assertTrue("transport cancellation must release the source without manual progress", + released.await(1, TimeUnit.SECONDS)); + verify(source).close(); + } finally { + fallbackRelease.countDown(); + call.cancel("test cleanup", null); + channel.shutdownNow(); + server.shutdownNow(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 10000) + public void testRealQueryClientCloseCancelsEarlyAndPreservesNormalCompletion() throws Exception { + assertClientIteratorClose(true); + assertClientIteratorClose(false); + } + + private static void assertClientIteratorClose(boolean early) throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator source = mock(ScanIterator.class); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch released = new CountDownLatch(1); + CountDownLatch completed = new CountDownLatch(1); + if (early) { + when(source.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + try { + new CountDownLatch(1).await(); + } catch (InterruptedException expected) { + Thread.currentThread().interrupt(); + } + return false; + }); + } + doAnswer(invocation -> { + released.countDown(); + return null; + }).when(source).close(); + java.util.concurrent.atomic.AtomicReference observed = + new java.util.concurrent.atomic.AtomicReference<>(); + AggregativeQueryService service = new AggregativeQueryService(pool, 5000, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + StreamObserver tracked = new StreamObserver() { + @Override + public void onNext(QueryResponse value) { + sender.onNext(value); + } + + @Override + public void onError(Throwable failure) { + sender.onError(failure); + } + + @Override + public void onCompleted() { + sender.onCompleted(); + completed.countDown(); + } + }; + AggregativeQueryObserver observer = fixture(super.newObserver(tracked), source, new QueryPlan(), null); + observed.set(observer); + return observer; + } + }; + String name = io.grpc.inprocess.InProcessServerBuilder.generateName(); + io.grpc.Server server = io.grpc.inprocess.InProcessServerBuilder.forName(name) + .directExecutor().addService(service).build().start(); + io.grpc.ManagedChannel channel = io.grpc.inprocess.InProcessChannelBuilder.forName(name) + .directExecutor().build(); + org.apache.hugegraph.store.client.query.QueryV2Client.setTestChannel(channel); + try { + org.apache.hugegraph.store.client.query.QueryExecutor executor = + new org.apache.hugegraph.store.client.query.QueryExecutor(null, null, 5000L); + java.lang.reflect.Method getIterator = executor.getClass().getDeclaredMethod( + "getIterator", String.class, QueryRequest.class); + getIterator.setAccessible(true); + org.apache.hugegraph.store.HgKvIterator iterator = + (org.apache.hugegraph.store.HgKvIterator) getIterator.invoke( + executor, "in-process-query-close", QueryRequest.newBuilder() + .setQueryId("client-close").setTable("test").build()); + if (early) { + assertTrue(reading.await(2, TimeUnit.SECONDS)); + FutureTask first = new FutureTask<>(() -> { + iterator.close(); + return null; + }); + Thread closer = start(first); + iterator.close(); + first.get(2, TimeUnit.SECONDS); + closer.join(1000); + } else { + assertFalse(iterator.hasNext()); + iterator.close(); + // A finished batch arrives before transport completion; don't cancel that RPC in teardown. + assertTrue("normal RPC must finish before channel teardown", completed.await(2, TimeUnit.SECONDS)); + } + assertTrue("iterator close must release resources before service shutdown", + released.await(2, TimeUnit.SECONDS)); + if (early) { + org.mockito.ArgumentCaptor failure = + org.mockito.ArgumentCaptor.forClass(Throwable.class); + verify(observed.get()).onError(failure.capture()); + verify(observed.get()).cancel(); + assertEquals(Status.Code.CANCELLED, Status.fromThrowable(failure.getValue()).getCode()); + } else { + verify(observed.get(), never()).onError(org.mockito.ArgumentMatchers.any()); + verify(observed.get(), never()).cancel(); + } + verify(observed.get()).onNext(org.mockito.ArgumentMatchers.any()); + verify(source).close(); + } finally { + channel.shutdownNow(); + server.shutdownNow(); + service.shutdownQueries(); + pool.shutdownNow(); + org.apache.hugegraph.store.client.query.QueryV2Client.setTestChannel(null); + } + } + + @Test(timeout = 5000) + public void testIteratorInitializationFailureKeepsOriginalThrowWithoutNormalCompletion() { + ThreadPoolExecutor pool = pool(1); + IllegalStateException failure = new IllegalStateException("iterator initialization failed"); + AggregativeQueryService service = new AggregativeQueryService(pool, 500, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + AggregativeQueryObserver observer = spy(super.newObserver(sender)); + doThrow(failure).when(observer).getIterator(org.mockito.ArgumentMatchers.any()); + return observer; + } + }; + ResponseRecorder response = new ResponseRecorder(); + try { + try { + service.query(response).onNext(QueryRequest.getDefaultInstance()); + fail("synchronous iterator initialization must fail"); + } catch (IllegalStateException expected) { + org.junit.Assert.assertSame(failure, expected); + } + assertEquals(0, response.completed.get()); + assertTrue(response.responses.isEmpty()); + service.shutdownQueries(); + } finally { + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testParentSubmissionRejectionKeepsOriginalThrowWithoutNormalCompletion() { + ThreadPoolExecutor pool = pool(1); + pool.shutdown(); + ScanIterator iterator = mock(ScanIterator.class); + QueryPlan plan = mock(QueryPlan.class); + AggregativeQueryService service = service(pool, iterator, plan, 500); + ResponseRecorder response = new ResponseRecorder(); + try { + service.query(response).onNext(QueryRequest.getDefaultInstance()); + fail("synchronous parent submission must be rejected"); + } catch (RejectedExecutionException expected) { + assertEquals(0, response.completed.get()); + assertTrue(response.responses.isEmpty()); + verify(iterator).close(); + verify(plan).clear(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testFinalSuccessIsSentOnlyAfterPlanAndIteratorCleanup() throws Exception { + ThreadPoolExecutor pool = pool(2); + ScanIterator child = mock(ScanIterator.class); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(child)); + ScanIterator result = mock(ScanIterator.class); + when(result.hasNext()).thenReturn(true, false); + when(result.next()).thenReturn(PipelineResult.EMPTY); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + when(plan.execute(PipelineResult.EMPTY)).thenReturn(result); + AtomicBoolean cleared = new AtomicBoolean(); + AtomicBoolean closed = new AtomicBoolean(); + doAnswer(invocation -> { + cleared.set(true); + return null; + }).when(plan).clear(); + doAnswer(invocation -> { + closed.set(true); + return null; + }).when(iterator).close(); + ResponseRecorder response = new ResponseRecorder() { + @Override + public void onNext(QueryResponse batch) { + assertTrue("success cannot precede plan cleanup", cleared.get()); + assertTrue("success cannot precede iterator cleanup", closed.get()); + super.onNext(batch); + } + }; + AggregativeQueryService service = service(pool, iterator, plan, 500); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("success-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertEquals(1, response.responses.size()); + assertTrue(response.responses.get(0).getIsOk()); + assertTrue(response.responses.get(0).getIsFinished()); + service.shutdownQueries(); + verify(child).close(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test + public void testPartitionInitializationFailureClosesEveryOpenedChild() { + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + when(first.hasNext()).thenReturn(true); + when(second.hasNext()).thenThrow(new IllegalStateException("child initialization failed")); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), + (id, key) -> id == 1 ? first : second); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::getIterators); + verify(first).close(); + verify(second).close(); + iterator.close(); + } + + @Test + public void testPartitionInitializationCleanupFailureIsRetainedByParent() { + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + when(first.hasNext()).thenReturn(true); + when(second.hasNext()).thenThrow(new IllegalStateException("child initialization failed")); + doThrow(new IllegalStateException("first cleanup failed")).when(first).close(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), + (id, key) -> id == 1 ? first : second); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::getIterators); + verify(first).close(); + verify(second).close(); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close); + verify(first).close(); + } + + @Test + public void testEmptyPartitionIteratorsAreReleasedBeforeOwnershipTransfer() { + ScanIterator empty = mock(ScanIterator.class); + ScanIterator nonempty = mock(ScanIterator.class); + when(nonempty.hasNext()).thenReturn(true); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2, 3), + (id, key) -> id == 1 ? empty : + id == 2 ? null : nonempty); + assertEquals(Arrays.asList(nonempty), iterator.getIterators()); + verify(empty).close(); + verify(nonempty, never()).close(); + nonempty.close(); + } + + @Test + public void testSequentialScanReleasesEveryEmptyPartitionBeforeAdvancing() { + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + ScanIterator nonempty = mock(ScanIterator.class); + when(nonempty.hasNext()).thenReturn(true); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2, 3, 4), + (id, key) -> id == 1 ? first : id == 2 ? second : id == 3 ? null : nonempty); + + assertTrue(iterator.hasNext()); + verify(first).close(); + verify(second).close(); + verify(nonempty, never()).close(); + iterator.close(); + iterator.close(); + verify(nonempty).close(); + } + + @Test + public void testClosedSequentialIteratorDoesNotOpenRemainingPartitions() { + for (boolean initialized : new boolean[]{false, true}) { + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + when(first.hasNext()).thenReturn(true); + when(second.hasNext()).thenReturn(true); + AtomicInteger opened = new AtomicInteger(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), (id, key) -> { + opened.incrementAndGet(); + return id == 1 ? first : second; + }); + if (initialized) { + assertTrue(iterator.hasNext()); + } + iterator.close(); + iterator.close(); + assertFalse(iterator.hasNext()); + assertFalse(iterator.isValid()); + assertEquals(0L, iterator.count()); + assertTrue(iterator.getIterators().isEmpty()); + org.junit.Assert.assertThrows(java.util.NoSuchElementException.class, iterator::next); + assertEquals(initialized ? 1 : 0, opened.get()); + if (initialized) { + verify(first).close(); + } else { + verify(first, never()).close(); + } + verify(second, never()).close(); + } + } + + @Test + public void testSequentialCloseFailureDoesNotReopenRemainingPartitions() { + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + when(first.hasNext()).thenReturn(true); + when(second.hasNext()).thenReturn(true); + IllegalStateException cleanup = new IllegalStateException("partition cleanup failed"); + doThrow(cleanup).when(first).close(); + AtomicInteger opened = new AtomicInteger(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), (id, key) -> { + opened.incrementAndGet(); + return id == 1 ? first : second; + }); + assertTrue(iterator.hasNext()); + IllegalStateException retained = org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close); + org.junit.Assert.assertSame(cleanup, retained.getCause()); + assertFalse(iterator.hasNext()); + assertFalse(iterator.isValid()); + assertEquals(0L, iterator.count()); + assertTrue(iterator.getIterators().isEmpty()); + org.junit.Assert.assertThrows(java.util.NoSuchElementException.class, iterator::next); + org.junit.Assert.assertSame(retained, + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close)); + assertEquals(1, opened.get()); + verify(first).close(); + verify(second, never()).close(); + } + + @Test + public void testSequentialInitializationFailureReleasesItsCreatedIterator() { + ScanIterator child = mock(ScanIterator.class); + IllegalStateException original = new IllegalStateException("partition initialization failed"); + when(child.hasNext()).thenThrow(original); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + + org.junit.Assert.assertSame(original, + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::hasNext)); + verify(child).close(); + iterator.close(); + } + + @Test + public void testSequentialInitializationCleanupFailureRetainsOriginalErrorAndBarrier() { + ScanIterator child = mock(ScanIterator.class); + IllegalStateException original = new IllegalStateException("partition initialization failed"); + IllegalStateException cleanup = new IllegalStateException("partition cleanup failed"); + when(child.hasNext()).thenThrow(original); + doThrow(cleanup).when(child).close(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + + org.junit.Assert.assertSame(original, + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::hasNext)); + assertEquals(1, original.getSuppressed().length); + org.junit.Assert.assertSame(cleanup, original.getSuppressed()[0]); + IllegalStateException retained = org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close); + org.junit.Assert.assertSame(cleanup, retained.getCause()); + verify(child).close(); + } + + @Test + public void testSequentialCountIncludesAlreadyOpenedPartitionAndClosesAllSources() { + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + when(first.hasNext()).thenReturn(true); + when(second.hasNext()).thenReturn(true); + when(first.count()).thenReturn(7L); + when(second.count()).thenReturn(11L); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), + (id, key) -> id == 1 ? first : second); + + assertTrue(iterator.hasNext()); + assertEquals(18L, iterator.count()); + assertFalse(iterator.hasNext()); + iterator.close(); + verify(first).close(); + verify(second).close(); + } + + @Test + public void testSequentialCountFailureReleasesCurrentSource() { + ScanIterator child = mock(ScanIterator.class); + IllegalStateException original = new IllegalStateException("partition count failed"); + when(child.hasNext()).thenReturn(true); + when(child.count()).thenThrow(original); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + + org.junit.Assert.assertSame(original, + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::count)); + iterator.close(); + verify(child).close(); + } + + @Test + public void testSequentialCountCleanupFailureRetainsOriginalErrorAndBarrier() { + ScanIterator child = mock(ScanIterator.class); + IllegalStateException original = new IllegalStateException("partition count failed"); + IllegalStateException cleanup = new IllegalStateException("partition cleanup failed"); + when(child.hasNext()).thenReturn(true); + when(child.count()).thenThrow(original); + doThrow(cleanup).when(child).close(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + + org.junit.Assert.assertSame(original, + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::count)); + assertEquals(1, original.getSuppressed().length); + org.junit.Assert.assertSame(cleanup, original.getSuppressed()[0]); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close); + verify(child).close(); + } + + @Test + public void testSequentialExhaustionCleanupFailureRemainsVisibleToParent() { + ScanIterator child = mock(ScanIterator.class); + when(child.hasNext()).thenReturn(true, false); + IllegalStateException cleanup = new IllegalStateException("exhausted partition cleanup failed"); + doThrow(cleanup).when(child).close(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + + assertTrue(iterator.hasNext()); + org.junit.Assert.assertSame(cleanup, + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::next)); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close); + verify(child).close(); + } + + @Test + public void testPlanCleanupContinuesAfterStageFailure() { + QueryStage failed = mock(QueryStage.class); + QueryStage remaining = mock(QueryStage.class); + doThrow(new IllegalStateException("stage cleanup failed")).when(failed).close(); + QueryPlan plan = new QueryPlan(); + plan.addStage(failed); + plan.addStage(remaining); + org.junit.Assert.assertThrows(IllegalStateException.class, plan::clear); + verify(failed).close(); + verify(remaining).close(); + assertTrue(plan.isEmpty()); + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryTestSupport.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryTestSupport.java new file mode 100644 index 0000000000..b8d7f40563 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryTestSupport.java @@ -0,0 +1,165 @@ +/* + * 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.store.node.grpc.query; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +import io.grpc.stub.StreamObserver; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.LockSupport; +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.store.grpc.query.QueryResponse; +import org.apache.hugegraph.store.node.grpc.query.model.QueryPlan; + +abstract class AggregativeQueryTestSupport { + + static void awaitShutdownWait(Thread closer) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (System.nanoTime() < deadline) { + if (closer.getState() == Thread.State.TIMED_WAITING && + Arrays.stream(closer.getStackTrace()).anyMatch(frame -> + frame.getClassName().equals(AggregativeQueryService.class.getName()) && + frame.getMethodName().equals("shutdownQueries"))) { + return; + } + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); + } + fail("shutdown must enter the aggregate-query cleanup wait"); + } + + static void assertError(ResponseRecorder response, String queryId) { + assertEquals(1, response.responses.size()); + QueryResponse error = response.responses.get(0); + assertEquals(queryId, error.getQueryId()); + assertFalse(error.getIsOk()); + assertFalse(error.getIsFinished()); + assertFalse(error.getMessage().isEmpty()); + } + + static ThreadPoolExecutor pool(int size) { + return (ThreadPoolExecutor) Executors.newFixedThreadPool(size); + } + + static AggregativeQueryService service(ThreadPoolExecutor pool, ScanIterator iterator, + QueryPlan plan, long timeout) { + return new AggregativeQueryService(pool, timeout, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + return fixture(super.newObserver(sender), iterator, plan, null); + } + }; + } + + static AggregativeQueryObserver fixture(AggregativeQueryObserver observer, + ScanIterator iterator, + QueryPlan plan, CountDownLatch parentReturned) { + AggregativeQueryObserver managed = spy(observer); + doReturn(iterator).when(managed).getIterator(org.mockito.ArgumentMatchers.any()); + doReturn(plan).when(managed).buildPlan(org.mockito.ArgumentMatchers.any()); + if (parentReturned != null) { + doAnswer(invocation -> { + try { + invocation.callRealMethod(); + return null; + } finally { + parentReturned.countDown(); + } + }).when(managed).sendData(); + } + return managed; + } + + @SuppressWarnings("unchecked") + static StreamObserver sender() { + return mock(StreamObserver.class); + } + + static FutureTask close(AggregativeQueryService service) { + return new FutureTask<>(() -> { + service.shutdownQueries(); + return null; + }); + } + + static Thread start(FutureTask task) { + Thread thread = new Thread(task, "test-aggregate-query-close"); + thread.setDaemon(true); + thread.start(); + return thread; + } + + static void awaitUninterruptibly(CountDownLatch latch) { + boolean interrupted = false; + try { + for (;;) { + try { + latch.await(); + return; + } catch (InterruptedException e) { + interrupted = true; + } + } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + } + + static class ResponseRecorder implements StreamObserver { + + final List responses = new CopyOnWriteArrayList<>(); + final AtomicInteger completed = new AtomicInteger(); + final AtomicInteger errors = new AtomicInteger(); + final CountDownLatch received = new CountDownLatch(1); + final CountDownLatch finished = new CountDownLatch(1); + + @Override + public void onNext(QueryResponse response) { + this.responses.add(response); + this.received.countDown(); + } + + @Override + public void onError(Throwable error) { + this.errors.incrementAndGet(); + this.finished.countDown(); + } + + @Override + public void onCompleted() { + this.completed.incrementAndGet(); + this.finished.countDown(); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/UnaryQueryLifecycleTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/UnaryQueryLifecycleTest.java new file mode 100644 index 0000000000..f48164563f --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/UnaryQueryLifecycleTest.java @@ -0,0 +1,544 @@ +/* + * 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.store.node.grpc.query; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.grpc.Context; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.grpc.stub.ServerCallStreamObserver; +import io.grpc.stub.StreamObserver; +import java.lang.reflect.Field; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.FutureTask; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.store.business.MultiPartitionIterator; +import org.apache.hugegraph.store.grpc.query.AggregateFunc; +import org.apache.hugegraph.store.grpc.query.QueryRequest; +import org.apache.hugegraph.store.grpc.query.QueryResponse; +import org.apache.hugegraph.store.query.KvSerializer; +import org.junit.Test; + +public class UnaryQueryLifecycleTest extends AggregativeQueryTestSupport { + + private static AggregativeQueryService unaryService(ThreadPoolExecutor pool, ScanIterator iterator, + RuntimeException initializationFailure) { + return new AggregativeQueryService(pool, 1000, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + AggregativeQueryObserver owner = spy(super.newObserver(sender)); + if (initializationFailure == null) { + doReturn(iterator).when(owner).getIterator(org.mockito.ArgumentMatchers.any()); + doReturn(iterator).when(owner).getCountIterator(org.mockito.ArgumentMatchers.any()); + } else { + doThrow(initializationFailure).when(owner).getIterator(org.mockito.ArgumentMatchers.any()); + doThrow(initializationFailure).when(owner).getCountIterator(org.mockito.ArgumentMatchers.any()); + } + return owner; + } + }; + } + + private static void call(AggregativeQueryService service, boolean count, + StreamObserver response) { + QueryRequest request = QueryRequest.newBuilder().setQueryId("unary-query").build(); + if (count) { + service.count(request, response); + } else { + service.query0(request, response); + } + } + + @Test(timeout = 5000) + public void testSuccessFollowsIteratorReleaseAndCompletesOnce() { + for (boolean count : new boolean[]{false, true}) { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + AtomicBoolean closed = new AtomicBoolean(); + doAnswer(invocation -> { closed.set(true); return null; }).when(iterator).close(); + ResponseRecorder response = new ResponseRecorder() { + @Override + public void onNext(QueryResponse value) { + assertTrue(closed.get()); + super.onNext(value); + } + }; + AggregativeQueryService service = unaryService(pool, iterator, null); + try { + call(service, count, response); + assertEquals(1, response.responses.size()); + assertTrue(response.responses.get(0).getIsOk()); + assertTrue(response.responses.get(0).getIsFinished()); + assertEquals(count ? 1 : 0, response.responses.get(0).getDataCount()); + assertEquals(1, response.completed.get()); + verify(iterator).close(); + service.shutdownQueries(); + } finally { + pool.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + public void testCountPreservesFunctionPayload() { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + when(iterator.hasNext()).thenReturn(true, true, false); + AggregativeQueryService service = unaryService(pool, iterator, null); + ResponseRecorder response = new ResponseRecorder(); + try { + QueryRequest request = QueryRequest.newBuilder().setQueryId("unary-query") + .addFunctions(AggregateFunc.getDefaultInstance()) + .addFunctions(AggregateFunc.getDefaultInstance()).build(); + service.count(request, response); + assertArrayEquals(KvSerializer.toBytes(List.of()), + response.responses.get(0).getData(0).getKey().toByteArray()); + assertArrayEquals(KvSerializer.toBytes(List.of(new AtomicLong(2), new AtomicLong(2))), + response.responses.get(0).getData(0).getValue().toByteArray()); + service.shutdownQueries(); + } finally { + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testCountPartitionsRunInParallelAndCloseBeforeResponse() throws Exception { + ThreadPoolExecutor pool = pool(2); + CountDownLatch entered = new CountDownLatch(2); + CountDownLatch release = new CountDownLatch(1); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + when(iterator.getIterators()).thenReturn(List.of(first, second)); + for (ScanIterator child : List.of(first, second)) { + when(child.hasNext()).thenAnswer(invocation -> { + entered.countDown(); + awaitUninterruptibly(release); + return false; + }); + } + AggregativeQueryService service = unaryService(pool, iterator, null); + ResponseRecorder response = new ResponseRecorder(); + FutureTask request = new FutureTask<>(() -> { call(service, true, response); return null; }); + try { + start(request); + assertTrue(entered.await(1, TimeUnit.SECONDS)); + assertTrue(response.responses.isEmpty()); + release.countDown(); + request.get(1, TimeUnit.SECONDS); + verify(first).close(); + verify(second).close(); + verify(iterator).close(); + assertTrue(response.responses.get(0).getIsOk()); + assertEquals(1, response.completed.get()); + service.shutdownQueries(); + } finally { + release.countDown(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testRejectedCountChildWaitsForAcceptedChildCleanup() throws Exception { + AtomicInteger submissions = new AtomicInteger(); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch rejected = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + ThreadPoolExecutor pool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, + new LinkedBlockingQueue<>()) { + @Override + public void execute(Runnable task) { + if (submissions.incrementAndGet() == 2) { + rejected.countDown(); + throw new RejectedExecutionException("rejected child"); + } + super.execute(task); + } + }; + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + when(iterator.getIterators()).thenReturn(List.of(first, second)); + when(first.hasNext()).thenAnswer(invocation -> { + entered.countDown(); + awaitUninterruptibly(release); + return false; + }); + AggregativeQueryService service = unaryService(pool, iterator, null); + ResponseRecorder response = new ResponseRecorder(); + FutureTask request = new FutureTask<>(() -> { call(service, true, response); return null; }); + try { + start(request); + assertTrue(entered.await(1, TimeUnit.SECONDS)); + assertTrue(rejected.await(1, TimeUnit.SECONDS)); + assertTrue(response.responses.isEmpty()); + assertFalse(request.isDone()); + release.countDown(); + request.get(1, TimeUnit.SECONDS); + verify(first).close(); + verify(second).close(); + verify(iterator).close(); + assertError(response, "unary-query"); + assertEquals("rejected child", response.responses.get(0).getMessage()); + assertEquals(1, response.completed.get()); + service.shutdownQueries(); + } finally { + release.countDown(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testInitializationFailureReportsSingleErrorAndReleasesOwner() { + for (boolean count : new boolean[]{false, true}) { + ThreadPoolExecutor pool = pool(1); + AggregativeQueryService service = unaryService(pool, null, new IllegalStateException("init")); + ResponseRecorder response = new ResponseRecorder(); + try { + call(service, count, response); + assertError(response, "unary-query"); + assertEquals("init", response.responses.get(0).getMessage()); + assertEquals(1, response.completed.get()); + service.shutdownQueries(); + } finally { + pool.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + public void testReadFailureClosesIteratorAndPreservesPrimaryError() { + for (boolean count : new boolean[]{false, true}) { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + when(iterator.hasNext()).thenThrow(new IllegalStateException("read")); + AggregativeQueryService service = unaryService(pool, iterator, null); + ResponseRecorder response = new ResponseRecorder(); + try { + call(service, count, response); + assertError(response, "unary-query"); + assertEquals("read", response.responses.get(0).getMessage()); + assertEquals(1, response.completed.get()); + verify(iterator).close(); + service.shutdownQueries(); + } finally { + pool.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + public void testCloseFailureRetainsOwnerAndSuppressesFailureOnReadError() throws Exception { + for (boolean count : new boolean[]{false, true}) { + for (boolean readFailure : new boolean[]{false, true}) { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + IllegalStateException primary = new IllegalStateException("read"); + IllegalStateException cleanup = new IllegalStateException("close"); + if (readFailure) { + when(iterator.hasNext()).thenThrow(primary); + } + doThrow(cleanup).when(iterator).close(); + AggregativeQueryService service = unaryService(pool, iterator, null); + ResponseRecorder response = new ResponseRecorder(); + try { + call(service, count, response); + assertError(response, "unary-query"); + assertEquals(readFailure ? "read" : "close", response.responses.get(0).getMessage()); + assertEquals(1, response.completed.get()); + if (readFailure) { + assertEquals(1, primary.getSuppressed().length); + assertSame(cleanup, primary.getSuppressed()[0]); + } + Field field = AggregativeQueryService.class.getDeclaredField("queries"); + field.setAccessible(true); + assertEquals(1, ((Set) field.get(service)).size()); + } finally { + pool.shutdownNow(); + } + } + } + } + + @Test(timeout = 5000) + public void testCancellationDuringReadWaitsForIteratorCleanup() throws Exception { + for (boolean count : new boolean[]{false, true}) { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + AtomicBoolean closed = new AtomicBoolean(); + when(iterator.hasNext()).thenAnswer(invocation -> { + entered.countDown(); + awaitUninterruptibly(release); + return false; + }); + doAnswer(invocation -> { closed.set(true); return null; }).when(iterator).close(); + AggregativeQueryService service = unaryService(pool, iterator, null); + ResponseRecorder response = new ResponseRecorder(); + FutureTask request = new FutureTask<>(() -> { call(service, count, response); return null; }); + FutureTask closing = close(service); + try { + start(request); + assertTrue(entered.await(1, TimeUnit.SECONDS)); + awaitShutdownWait(start(closing)); + assertFalse(closing.isDone()); + assertFalse(closed.get()); + release.countDown(); + request.get(1, TimeUnit.SECONDS); + closing.get(1, TimeUnit.SECONDS); + assertTrue(closed.get()); + assertTrue(response.responses.isEmpty()); + assertEquals(0, response.completed.get()); + assertEquals(1, response.errors.get()); + verify(iterator).close(); + } finally { + release.countDown(); + pool.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + @SuppressWarnings("unchecked") + public void testTransportCancellationClosesWithoutTerminalCallback() throws Exception { + for (boolean count : new boolean[]{false, true}) { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + when(iterator.hasNext()).thenAnswer(invocation -> { + entered.countDown(); + awaitUninterruptibly(release); + return false; + }); + AggregativeQueryService service = unaryService(pool, iterator, null); + ServerCallStreamObserver response = mock(ServerCallStreamObserver.class); + AtomicReference cancellation = new AtomicReference<>(); + doAnswer(invocation -> { cancellation.set(invocation.getArgument(0)); return null; }) + .when(response).setOnCancelHandler(org.mockito.ArgumentMatchers.any()); + FutureTask request = new FutureTask<>(() -> { call(service, count, response); return null; }); + try { + start(request); + assertTrue(entered.await(1, TimeUnit.SECONDS)); + cancellation.get().run(); + release.countDown(); + request.get(1, TimeUnit.SECONDS); + service.shutdownQueries(); + verify(iterator).close(); + verify(response, org.mockito.Mockito.never()).onCompleted(); + verify(response, org.mockito.Mockito.never()).onNext(org.mockito.ArgumentMatchers.any()); + verify(response, org.mockito.Mockito.never()).onError(org.mockito.ArgumentMatchers.any()); + } finally { + release.countDown(); + pool.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + public void testContextCancellationInterruptsBothUnaryReads() throws Exception { + for (boolean count : new boolean[]{false, true}) { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + CountDownLatch fallback = new CountDownLatch(1); + when(iterator.hasNext()).thenAnswer(invocation -> { + entered.countDown(); + try { + fallback.await(); + } catch (InterruptedException expected) { + interrupted.countDown(); + Thread.currentThread().interrupt(); + } + return false; + }); + AggregativeQueryService service = unaryService(pool, iterator, null); + ResponseRecorder response = new ResponseRecorder(); + Context.CancellableContext context = Context.current().withCancellation(); + FutureTask request = new FutureTask<>(() -> { + context.run(() -> call(service, count, response)); + return null; + }); + try { + start(request); + assertTrue(entered.await(1, TimeUnit.SECONDS)); + context.cancel(new IllegalStateException("client canceled")); + assertTrue("context must interrupt the unary read", interrupted.await(1, TimeUnit.SECONDS)); + request.get(1, TimeUnit.SECONDS); + verify(iterator).close(); + assertTrue(response.responses.isEmpty()); + assertEquals(0, response.completed.get()); + assertEquals(0, response.errors.get()); + assertNoUnaryOwners(service); + } finally { + fallback.countDown(); + context.cancel(null); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + public void testAlreadyCanceledContextNeverOpensUnaryIterator() throws Exception { + for (boolean count : new boolean[]{false, true}) { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + AggregativeQueryService service = unaryService(pool, iterator, null); + ResponseRecorder response = new ResponseRecorder(); + Context.CancellableContext context = Context.current().withCancellation(); + context.cancel(null); + try { + context.run(() -> call(service, count, response)); + verify(iterator, org.mockito.Mockito.never()).hasNext(); + verify(iterator, org.mockito.Mockito.never()).close(); + assertTrue(response.responses.isEmpty()); + assertEquals(0, response.completed.get()); + assertNoUnaryOwners(service); + } finally { + Thread.interrupted(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + } + + @Test(timeout = 15000) + public void testRealNettyCancellationInterruptsBothUnaryReads() throws Exception { + for (boolean count : new boolean[]{false, true}) { + ThreadPoolExecutor pool = pool(2); + ExecutorService callbacks = Executors.newFixedThreadPool(2); + ScanIterator iterator = mock(ScanIterator.class); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + CountDownLatch closed = new CountDownLatch(1); + CountDownLatch fallback = new CountDownLatch(1); + when(iterator.hasNext()).thenAnswer(invocation -> { + entered.countDown(); + try { + fallback.await(); + } catch (InterruptedException expected) { + interrupted.countDown(); + Thread.currentThread().interrupt(); + } + return false; + }); + doAnswer(invocation -> { closed.countDown(); return null; }).when(iterator).close(); + AggregativeQueryService service = unaryService(pool, iterator, null); + io.grpc.Server server = io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder.forPort(0) + .executor(callbacks).addService(service).build().start(); + io.grpc.ManagedChannel channel = io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder + .forAddress("localhost", server.getPort()).usePlaintext().build(); + io.grpc.ClientCall call = channel.newCall(count ? + org.apache.hugegraph.store.grpc.query.QueryServiceGrpc.getCountMethod() : + org.apache.hugegraph.store.grpc.query.QueryServiceGrpc.getQuery0Method(), + io.grpc.CallOptions.DEFAULT); + CountDownLatch terminated = new CountDownLatch(1); + AtomicReference terminal = new AtomicReference<>(); + AtomicInteger responses = new AtomicInteger(); + try { + call.start(new io.grpc.ClientCall.Listener() { + @Override + public void onMessage(QueryResponse response) { + responses.incrementAndGet(); + } + + @Override + public void onClose(Status status, io.grpc.Metadata trailers) { + terminal.set(status); + terminated.countDown(); + } + }, new io.grpc.Metadata()); + call.request(1); + call.sendMessage(QueryRequest.newBuilder().setQueryId("netty-unary").build()); + call.halfClose(); + assertTrue(entered.await(2, TimeUnit.SECONDS)); + call.cancel("cancel blocked unary read", null); + assertTrue(terminated.await(2, TimeUnit.SECONDS)); + assertEquals(Status.Code.CANCELLED, terminal.get().getCode()); + assertTrue("transport context must interrupt unary read", interrupted.await(2, TimeUnit.SECONDS)); + assertTrue("cancellation must release iterator", closed.await(2, TimeUnit.SECONDS)); + service.shutdownQueries(); + verify(iterator).close(); + assertEquals(0, responses.get()); + assertNoUnaryOwners(service); + } finally { + fallback.countDown(); + call.cancel("test cleanup", null); + channel.shutdownNow(); + server.shutdownNow(); + service.shutdownQueries(); + callbacks.shutdownNow(); + pool.shutdownNow(); + } + } + } + + private static void assertNoUnaryOwners(AggregativeQueryService service) throws Exception { + Field field = AggregativeQueryService.class.getDeclaredField("queries"); + field.setAccessible(true); + synchronized (service) { + assertTrue(((Set) field.get(service)).isEmpty()); + } + } + + @Test(timeout = 5000) + public void testBothUnaryEndpointsRejectRequestsAfterAdmissionCloses() { + ThreadPoolExecutor pool = pool(1); + AggregativeQueryService service = unaryService(pool, mock(ScanIterator.class), null); + try { + service.stopAcceptingQueries(); + for (boolean count : new boolean[]{false, true}) { + try { + call(service, count, new ResponseRecorder()); + org.junit.Assert.fail("closed service must reject unary RPCs"); + } catch (StatusRuntimeException e) { + assertEquals(Status.Code.UNAVAILABLE, e.getStatus().getCode()); + } + } + } finally { + pool.shutdownNow(); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GraphPartitionScanLifecycleTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GraphPartitionScanLifecycleTest.java new file mode 100644 index 0000000000..4ab3217d11 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GraphPartitionScanLifecycleTest.java @@ -0,0 +1,561 @@ +/* + * 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.store.service; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.FutureTask; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.store.business.BusinessHandler; +import org.apache.hugegraph.store.business.GraphStoreIterator; +import org.apache.hugegraph.store.grpc.Graphpb; +import org.apache.hugegraph.store.grpc.Graphpb.ScanPartitionRequest; +import org.apache.hugegraph.store.grpc.Graphpb.ScanResponse; +import org.apache.hugegraph.store.node.grpc.HgStoreStreamImpl; +import org.apache.hugegraph.store.node.grpc.scan.ScanResponseObserver; +import org.junit.Test; + +import io.grpc.Context; +import io.grpc.Status; +import io.grpc.stub.StreamObserver; + +public class GraphPartitionScanLifecycleTest extends GraphPartitionScanTestSupport { + + @Test + public void testNaturalEndAndLimitCloseIteratorBeforeCompleting() throws Exception { + for (int limit : new int[]{0, 1}) { + ThreadPoolExecutor executor = executor(); + try { + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + AtomicInteger consumed = new AtomicInteger(); + when(iterator.hasNext()).thenAnswer(call -> consumed.get() < 2); + when(iterator.next()).thenAnswer(call -> { + consumed.incrementAndGet(); + return Graphpb.Vertex.getDefaultInstance(); + }); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + AtomicInteger delivered = new AtomicInteger(); + doAnswer(call -> { + delivered.addAndGet(((ScanResponse) call.getArgument(0)).getVertexCount()); + return null; + }).when(response).onNext(any()); + StreamObserver request = service.scanGraphPartition(response, handler); + request.onNext(request(limit)); + verify(response, timeout(2000)).onCompleted(); + service.awaitScanCleanup(); + verify(iterator).close(); + assertEquals(limit == 0 ? 2 : 1, consumed.get()); + assertEquals(consumed.get(), delivered.get()); + assertTrue(registry(service).isEmpty()); + request.onError(new IllegalStateException("late cancellation")); + verify(iterator).close(); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + } + + @Test(timeout = 30000) + public void testFinalBatchPublishedAfterEmptyPollIsDelivered() throws Exception { + ThreadPoolExecutor executor = executor(); + CountDownLatch finalRead = new CountDownLatch(1); + CountDownLatch releaseRead = new CountDownLatch(1); + CountDownLatch emptyPoll = new CountDownLatch(1); + CountDownLatch releasePoll = new CountDownLatch(1); + AtomicInteger consumed = new AtomicInteger(); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + when(iterator.hasNext()).thenAnswer(call -> consumed.get() < 100001); + when(iterator.next()).thenAnswer(call -> { + if (consumed.get() == 100000) { + finalRead.countDown(); + assertTrue(releaseRead.await(10, TimeUnit.SECONDS)); + } + consumed.incrementAndGet(); + return Graphpb.Vertex.getDefaultInstance(); + }); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + AtomicInteger delivered = new AtomicInteger(); + doAnswer(call -> { + delivered.addAndGet(((ScanResponse) call.getArgument(0)).getVertexCount()); + return null; + }).when(response).onNext(any()); + ScanResponseObserver request = + new ScanResponseObserver<>(response, handler, executor); + LinkedBlockingQueue packages = new LinkedBlockingQueue(16) { + @Override + public ScanResponse poll(long timeout, TimeUnit unit) throws InterruptedException { + ScanResponse value = super.poll(timeout, unit); + if (value == null && emptyPoll.getCount() != 0) { + emptyPoll.countDown(); + assertTrue(releasePoll.await(10, TimeUnit.SECONDS)); + } + return value; + } + }; + Field field = ScanResponseObserver.class.getDeclaredField("packages"); + field.setAccessible(true); + field.set(request, packages); + try { + request.onNext(request(0)); + assertTrue(finalRead.await(10, TimeUnit.SECONDS)); + assertTrue("sender must have actually polled an empty queue", + emptyPoll.await(10, TimeUnit.SECONDS)); + assertEquals(100000, delivered.get()); + releaseRead.countDown(); + // The sender occupies one worker; this probe runs only after the real reader + // has published its final batch, closed its iterator and set readOver. + executor.submit(() -> { }).get(10, TimeUnit.SECONDS); + verify(iterator).close(); + assertEquals(1, packages.size()); + releasePoll.countDown(); + verify(response, timeout(5000)).onCompleted(); + assertEquals(100001, consumed.get()); + assertEquals(consumed.get(), delivered.get()); + assertTrue(packages.isEmpty()); + request.onCompleted(); + request.onError(new IllegalStateException("late cancellation")); + verify(response).onCompleted(); + verify(response, org.mockito.Mockito.never()).onError(any()); + verify(iterator).close(); + } finally { + releaseRead.countDown(); + releasePoll.countDown(); + executor.shutdown(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testEmptyScanClosesOnceAndReleasesBarrier() throws Exception { + ThreadPoolExecutor executor = executor(); + try { + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + service.scanGraphPartition(response, handler).onNext(request(0)); + verify(response, timeout(2000)).onCompleted(); + service.awaitScanCleanup(); + verify(iterator).close(); + assertTrue(registry(service).isEmpty()); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + + @Test + public void testRejectedWorkerClosesIteratorAndTerminatesResponse() throws Exception { + ThreadPoolExecutor executor = executor(); + executor.shutdown(); + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + service.scanGraphPartition(response, handler).onNext(request(0)); + verify(response).onError(any()); + verify(iterator).close(); + assertTrue(registry(service).isEmpty()); + } + + @Test + public void testReadFailureClosesIteratorAndPreservesPrimaryError() throws Exception { + ThreadPoolExecutor executor = executor(); + try { + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + RuntimeException failure = new IllegalStateException("native read failed"); + when(iterator.hasNext()).thenThrow(failure); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + service.scanGraphPartition(response, handler).onNext(request(0)); + verify(response, timeout(2000)).onError(failure); + service.awaitScanCleanup(); + verify(iterator).close(); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + + @Test + public void testClientCompletionBeforeInitializationReleasesBarrier() throws Exception { + ThreadPoolExecutor executor = executor(); + try { + HgStoreStreamImpl service = service(executor); + StreamObserver response = mock(StreamObserver.class); + StreamObserver request = + service.scanGraphPartition(response, mock(BusinessHandler.class)); + request.onCompleted(); + request.onCompleted(); + verify(response).onCompleted(); + assertTrue(registry(service).isEmpty()); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + + @Test + public void testReadFailureKeepsCloseFailureSuppressedAndSticky() throws Exception { + ThreadPoolExecutor executor = executor(); + try { + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + RuntimeException primary = new IllegalStateException("read"); + RuntimeException cleanup = new IllegalArgumentException("close"); + when(iterator.hasNext()).thenThrow(primary); + doThrow(cleanup).when(iterator).close(); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + service.scanGraphPartition(response, handler).onNext(request(0)); + verify(response, timeout(2000)).onError(primary); + assertEquals(1, primary.getSuppressed().length); + assertSame(cleanup, primary.getSuppressed()[0]); + verify(iterator).close(); + assertEquals(1, registry(service).size()); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + + @Test + public void testFullCreditWindowResumesWithOrderedPagesAndKeepsFinalPartialBatch() throws Exception { + checkCreditWindow(false, 7); + } + + @Test(timeout = 5000) + public void testHalfCloseBeforeWorkerStartsKeepsFinalPartialBatch() throws Exception { + ThreadPoolExecutor executor = executor(); + CountDownLatch workersBlocked = new CountDownLatch(2); + CountDownLatch release = new CountDownLatch(1); + List> blockers = new ArrayList<>(); + try { + // Finish service and mock initialization before establishing the worker gate. + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + AtomicInteger consumed = new AtomicInteger(); + AtomicInteger reads = new AtomicInteger(); + when(iterator.hasNext()).thenAnswer(call -> { + reads.incrementAndGet(); + return consumed.get() < 2; + }); + when(iterator.next()).thenAnswer(call -> { + consumed.incrementAndGet(); + return Graphpb.Vertex.getDefaultInstance(); + }); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + AtomicInteger delivered = new AtomicInteger(); + doAnswer(call -> { + delivered.addAndGet(((ScanResponse) call.getArgument(0)).getVertexCount()); + return null; + }).when(response).onNext(any()); + StreamObserver request = service.scanGraphPartition(response, handler); + for (int i = 0; i < 2; i++) { + blockers.add(executor.submit(() -> { + workersBlocked.countDown(); + try { + release.await(); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(failure); + } + })); + } + assertTrue(workersBlocked.await(2, TimeUnit.SECONDS)); + request.onNext(request(0)); + assertEquals("reading must not begin before half-close", 0, reads.get()); + request.onCompleted(); + assertEquals("queued reader must remain gated through half-close", 0, reads.get()); + release.countDown(); + for (Future blocker : blockers) { + blocker.get(2, TimeUnit.SECONDS); + } + verify(response, timeout(2000)).onCompleted(); + service.awaitScanCleanup(); + assertEquals(2, delivered.get()); + verify(iterator).close(); + request.onCompleted(); + verify(response).onCompleted(); + assertTrue(registry(service).isEmpty()); + } finally { + release.countDown(); + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + + @Test(timeout = 5000) + public void testContextCancellationInterruptsBlockedReaderBeforeCleanup() throws Exception { + checkBlockedReaderCancellation(false); + } + + @Test(timeout = 5000) + public void testContextDeadlineInterruptsBlockedReaderBeforeCleanup() throws Exception { + checkBlockedReaderCancellation(true); + } + + private static void checkBlockedReaderCancellation(boolean deadline) throws Exception { + ThreadPoolExecutor executor = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, + new LinkedBlockingQueue<>()); + ScheduledExecutorService timer = Executors.newSingleThreadScheduledExecutor(); + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch readExited = new CountDownLatch(1); + AtomicBoolean interrupted = new AtomicBoolean(); + when(iterator.hasNext()).thenAnswer(call -> { + reading.countDown(); + try { + new CountDownLatch(1).await(); + throw new AssertionError("blocked reader must exit through interruption"); + } catch (InterruptedException expected) { + interrupted.set(true); + Thread.currentThread().interrupt(); + return false; + } finally { + readExited.countDown(); + } + }); + AtomicBoolean closedAfterRead = new AtomicBoolean(); + doAnswer(call -> { + closedAfterRead.set(readExited.getCount() == 0); + return null; + }).when(iterator).close(); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + Context.CancellableContext context = deadline ? + Context.current().withDeadlineAfter(1, TimeUnit.SECONDS, timer) : + Context.current().withCancellation(); + Thread cancellation = null; + try { + StreamObserver request = context.call(() -> + service.scanGraphPartition(response, handler)); + request.onNext(request(0)); + assertTrue("real worker must enter its interruptible read", reading.await(500, TimeUnit.MILLISECONDS)); + if (!deadline) { + FutureTask cancel = new FutureTask<>(() -> { + context.cancel(null); + return null; + }); + cancellation = new Thread(cancel, "test-partition-context-cancel"); + cancellation.setDaemon(true); + cancellation.start(); + cancel.get(2, TimeUnit.SECONDS); + } + assertTrue("context cancellation must interrupt the owned reader", readExited.await(2, TimeUnit.SECONDS)); + assertTrue(interrupted.get()); + verify(response, timeout(2000)).onError(any()); + service.awaitScanCleanup(); + verify(iterator).close(); + assertTrue(closedAfterRead.get()); + assertTrue(registry(service).isEmpty()); + // A released worker must remain usable and free from a cancellation interrupt. + assertFalse(executor.submit(() -> Thread.currentThread().isInterrupted()).get(1, TimeUnit.SECONDS)); + } finally { + // Also release an old implementation's blocked reader when this regression fails. + executor.shutdownNow(); + context.cancel(null); + timer.shutdownNow(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + if (cancellation != null) { + cancellation.join(1000); + } + } + } + + @Test + public void testHalfCloseDeliversEntirePermittedWindowAndCompletes() throws Exception { + checkCreditWindow(true, 0); + } + + @Test + public void testHalfCloseReportsInsufficientCreditAfterDeliveringPermittedPages() throws Exception { + checkCreditWindow(true, 7); + } + + private void checkCreditWindow(boolean halfClose, int extraRows) throws Exception { + ThreadPoolExecutor executor = executor(); + HgStoreStreamImpl service = service(executor); + final int batchSize = 100000; + final int window = 8; + final int total = batchSize * window + extraRows; + AtomicInteger closes = new AtomicInteger(); + Field suppliers = org.apache.hugegraph.store.business.BusinessHandlerImpl.class + .getDeclaredField("GRAPH_SUPPLIER_CACHE"); + suppliers.setAccessible(true); + @SuppressWarnings("unchecked") + Map graphCache = + (Map) suppliers.get(null); + String graphName = "TEST/credit-window"; + org.apache.hugegraph.HugeGraphSupplier previous = + graphCache.put(graphName, mock(org.apache.hugegraph.HugeGraphSupplier.class)); + GraphStoreIterator iterator; + try { + iterator = new GraphStoreIterator( + mock(ScanIterator.class), request(0)) { + private int row; + + @Override + public boolean hasNext() { + return this.row < total; + } + + @Override + public Graphpb.Vertex next() { + return Graphpb.Vertex.newBuilder().setId(Graphpb.Variant.newBuilder() + .setType(Graphpb.VariantType.VT_LONG).setValueInt64(this.row++)).build(); + } + + @Override + public void close() { + closes.incrementAndGet(); + } + }; + } finally { + if (previous == null) { + graphCache.remove(graphName); + } else { + graphCache.put(graphName, previous); + } + } + LinkedBlockingQueue responses = new LinkedBlockingQueue<>(); + AtomicInteger terminals = new AtomicInteger(); + AtomicReference error = new AtomicReference<>(); + CountDownLatch finished = new CountDownLatch(1); + StreamObserver response = new StreamObserver() { + @Override + public void onNext(ScanResponse value) { + responses.add(value); + } + + @Override + public void onError(Throwable failure) { + error.set(failure); + terminals.incrementAndGet(); + finished.countDown(); + } + + @Override + public void onCompleted() { + terminals.incrementAndGet(); + finished.countDown(); + } + }; + try { + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver request = service.scanGraphPartition(response, handler); + request.onNext(request(0)); + if (halfClose) { + request.onCompleted(); + } + long expectedId = 0; + for (int page = 0; page < window; page++) { + ScanResponse value = responses.poll(20, TimeUnit.SECONDS); + assertTrue("Expected page " + page, value != null); + assertEquals(page, value.getSeqNo()); + assertEquals(batchSize, value.getVertexCount()); + for (Graphpb.Vertex vertex : value.getVertexList()) { + assertEquals(expectedId++, vertex.getId().getValueInt64()); + } + } + assertNull("No ninth page is allowed without credit", responses.poll(100, TimeUnit.MILLISECONDS)); + if (halfClose) { + assertTrue(finished.await(20, TimeUnit.SECONDS)); + if (extraRows == 0) { + assertNull(error.get()); + } else { + assertEquals(Status.Code.FAILED_PRECONDITION, Status.fromThrowable(error.get()).getCode()); + } + service.awaitScanCleanup(); + request.onCompleted(); + assertEquals(1, terminals.get()); + assertEquals(1, closes.get()); + assertTrue(registry(service).isEmpty()); + return; + } + assertEquals(0, terminals.get()); + for (int page = 0; page < window; page++) { + request.onNext(ScanPartitionRequest.newBuilder().setReplyRequest( + ScanPartitionRequest.Reply.newBuilder().setSeqNo(page)).build()); + } + ScanResponse last = responses.poll(20, TimeUnit.SECONDS); + assertTrue("The last partial batch must be delivered", last != null); + assertEquals(window, last.getSeqNo()); + assertEquals(7, last.getVertexCount()); + for (Graphpb.Vertex vertex : last.getVertexList()) { + assertEquals(expectedId++, vertex.getId().getValueInt64()); + } + assertEquals(total, expectedId); + assertTrue(finished.await(20, TimeUnit.SECONDS)); + assertNull(error.get()); + service.awaitScanCleanup(); + request.onCompleted(); + assertEquals(1, terminals.get()); + assertEquals(1, closes.get()); + assertTrue(responses.isEmpty()); + assertTrue(registry(service).isEmpty()); + } finally { + service.shutdownScans(); + executor.shutdown(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GraphPartitionScanTestSupport.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GraphPartitionScanTestSupport.java new file mode 100644 index 0000000000..1742ac64c5 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GraphPartitionScanTestSupport.java @@ -0,0 +1,55 @@ +/* + * 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.store.service; + +import java.lang.reflect.Field; +import java.util.Map; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; + +import org.apache.hugegraph.store.grpc.Graphpb; +import org.apache.hugegraph.store.grpc.Graphpb.ScanPartitionRequest; +import org.apache.hugegraph.store.node.grpc.HgStoreStreamImpl; + +abstract class GraphPartitionScanTestSupport { + + static ScanPartitionRequest request(long limit) { + return ScanPartitionRequest.newBuilder().setScanRequest(Graphpb.ScanPartitionRequest.Request.newBuilder() + .setGraphName("TEST/credit-window") + .setScanType(Graphpb.ScanPartitionRequest.ScanType.SCAN_VERTEX).setLimit(limit)).build(); + } + + static ThreadPoolExecutor executor() { + return new ThreadPoolExecutor(2, 2, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>()); + } + + static HgStoreStreamImpl service(ThreadPoolExecutor executor) throws Exception { + HgStoreStreamImpl service = new HgStoreStreamImpl(); + Field field = HgStoreStreamImpl.class.getDeclaredField("executor"); + field.setAccessible(true); + field.set(service, executor); + return service; + } + + static Map registry(HgStoreStreamImpl service) throws Exception { + Field field = HgStoreStreamImpl.class.getDeclaredField("scans"); + field.setAccessible(true); + return (Map) field.get(service); + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GrpcShutdownBarrierTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GrpcShutdownBarrierTest.java new file mode 100644 index 0000000000..756244a6f9 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GrpcShutdownBarrierTest.java @@ -0,0 +1,275 @@ +/* + * 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.store.service; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.io.ByteArrayInputStream; +import java.io.InputStream; +import java.lang.reflect.Field; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.apache.hugegraph.store.node.AppConfig; +import org.apache.hugegraph.store.node.grpc.GRpcServerConfig; +import org.apache.hugegraph.store.node.grpc.GrpcShutdownBarrier; +import org.apache.hugegraph.store.node.util.HgExecutorUtil; +import org.junit.Test; + +import io.grpc.CallOptions; +import io.grpc.ClientCall; +import io.grpc.Metadata; +import io.grpc.ManagedChannel; +import io.grpc.MethodDescriptor; +import io.grpc.Server; +import io.grpc.ServerCall; +import io.grpc.ServerServiceDefinition; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder; +import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder; +import io.grpc.stub.ClientCalls; +import io.grpc.stub.ServerCalls; +import io.grpc.stub.StreamObserver; + +public class GrpcShutdownBarrierTest { + + @Test(timeout = 10000) + public void testNettyCancellationSurvivesSaturatedSingleDispatchThread() throws Exception { + CountDownLatch callsStarted = new CountDownLatch(2); + CountDownLatch callbackStarted = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch cancelled = new CountDownLatch(2); + MethodDescriptor.Marshaller marshaller = new MethodDescriptor.Marshaller() { + @Override + public InputStream stream(byte[] value) { + return new ByteArrayInputStream(value); + } + + @Override + public byte[] parse(InputStream stream) { + return new byte[0]; + } + }; + MethodDescriptor method = MethodDescriptor.newBuilder() + .setType(MethodDescriptor.MethodType.BIDI_STREAMING).setFullMethodName("saturation/hold") + .setRequestMarshaller(marshaller).setResponseMarshaller(marshaller).build(); + GrpcShutdownBarrier barrier = new GrpcShutdownBarrier(); + NettyServerBuilder builder = NettyServerBuilder.forPort(0).intercept(barrier) + .addService(ServerServiceDefinition.builder("saturation").addMethod(method, (call, headers) -> { + callsStarted.countDown(); + call.request(1); + return new ServerCall.Listener() { + @Override + public void onMessage(byte[] message) { + callbackStarted.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + @Override + public void onCancel() { + cancelled.countDown(); + } + }; + }).build()); + AppConfig config = mock(AppConfig.class); + AppConfig.ThreadPoolGrpc grpc = mock(AppConfig.ThreadPoolGrpc.class); + when(config.getThreadPoolGrpc()).thenReturn(grpc); + when(grpc.getCore()).thenReturn(1); + when(grpc.getMax()).thenReturn(1); + when(grpc.getQueue()).thenReturn(0); + GRpcServerConfig serverConfig = new GRpcServerConfig(); + Field appConfig = GRpcServerConfig.class.getDeclaredField("appConfig"); + appConfig.setAccessible(true); + appConfig.set(serverConfig, config); + // Isolate the production named-executor cache for this fixture only. + Field poolsField = HgExecutorUtil.class.getDeclaredField("threadPoolMap"); + poolsField.setAccessible(true); + @SuppressWarnings("unchecked") + Map pools = (Map) poolsField.get(null); + ThreadPoolExecutor previous = pools.remove(GRpcServerConfig.EXECUTOR_NAME); + ThreadPoolExecutor executor = null; + Server server = null; + ManagedChannel channel = null; + try { + serverConfig.configure(builder); + executor = HgExecutorUtil.getThreadPoolExecutor(GRpcServerConfig.EXECUTOR_NAME); + assertEquals(1, executor.getCorePoolSize()); + assertEquals(1, executor.getMaximumPoolSize()); + server = builder.build().start(); + channel = NettyChannelBuilder.forAddress("localhost", server.getPort()).usePlaintext().build(); + ClientCall first = channel.newCall(method, CallOptions.DEFAULT); + ClientCall second = channel.newCall(method, CallOptions.DEFAULT); + first.start(new ClientCall.Listener() { }, new Metadata()); + second.start(new ClientCall.Listener() { }, new Metadata()); + first.request(1); + second.request(1); + assertTrue("both calls must register before saturation", callsStarted.await(2, TimeUnit.SECONDS)); + first.sendMessage(new byte[0]); + assertTrue(callbackStarted.await(2, TimeUnit.SECONDS)); + second.cancel("cancel idle call while dispatch is saturated", null); + first.cancel("cancel busy call", null); + long queuedDeadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2); + while (executor.getQueue().isEmpty() && System.nanoTime() < queuedDeadline) { + Thread.yield(); + } + assertFalse("idle call's real cancellation callback must queue behind the busy handler", + executor.getQueue().isEmpty()); + CountDownLatch waiting = new CountDownLatch(1); + FutureTask drained = new FutureTask<>(() -> { + waiting.countDown(); + barrier.awaitCallbacks(); + return null; + }); + Thread waiter = new Thread(drained, "test-netty-saturation-drain"); + waiter.setDaemon(true); + waiter.start(); + assertTrue(waiting.await(1, TimeUnit.SECONDS)); + try { + drained.get(100, TimeUnit.MILLISECONDS); + fail("active callback must still block database destruction"); + } catch (TimeoutException expected) { + // Cancelling the transport cannot discard the application's terminal callback. + } + release.countDown(); + assertTrue("both cancellation callbacks must be delivered", cancelled.await(2, TimeUnit.SECONDS)); + drained.get(2, TimeUnit.SECONDS); + executor.submit(() -> { }).get(1, TimeUnit.SECONDS); + barrier.stopAcceptingCalls(); + server.shutdownNow(); + assertTrue(server.awaitTermination(2, TimeUnit.SECONDS)); + barrier.awaitCallbacks(); + } finally { + release.countDown(); + if (channel != null) { + channel.shutdownNow(); + channel.awaitTermination(2, TimeUnit.SECONDS); + } + if (server != null) { + server.shutdownNow(); + server.awaitTermination(2, TimeUnit.SECONDS); + } + if (executor != null) { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + pools.remove(GRpcServerConfig.EXECUTOR_NAME, executor); + } + if (previous != null) { + pools.put(GRpcServerConfig.EXECUTOR_NAME, previous); + } + } + } + + @Test(timeout = 7000) + public void testTransportTerminationDoesNotReleaseActiveCallback() throws Exception { + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + AtomicBoolean finished = new AtomicBoolean(); + ExecutorService executor = Executors.newSingleThreadExecutor(); + MethodDescriptor.Marshaller marshaller = new MethodDescriptor.Marshaller() { + @Override + public InputStream stream(byte[] value) { + return new ByteArrayInputStream(value); + } + @Override + public byte[] parse(InputStream stream) { + return new byte[0]; + } + }; + MethodDescriptor method = MethodDescriptor.newBuilder() + .setType(MethodDescriptor.MethodType.UNARY).setFullMethodName("probe/hold") + .setRequestMarshaller(marshaller).setResponseMarshaller(marshaller).build(); + String name = InProcessServerBuilder.generateName(); + GrpcShutdownBarrier barrier = new GrpcShutdownBarrier(); + Server server = InProcessServerBuilder.forName(name).executor(executor).intercept(barrier) + .addService(ServerServiceDefinition.builder("probe").addMethod(method, + ServerCalls.asyncUnaryCall((request, response) -> { + started.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + finished.set(true); + } + response.onNext(new byte[0]); + response.onCompleted(); + })).build()).build().start(); + ManagedChannel channel = InProcessChannelBuilder.forName(name).directExecutor().build(); + try { + ClientCalls.asyncUnaryCall(channel.newCall(method, CallOptions.DEFAULT), new byte[0], + new StreamObserver() { + @Override + public void onNext(byte[] value) { + } + @Override + public void onError(Throwable error) { + } + @Override + public void onCompleted() { + } + }); + assertTrue("handler must start", started.await(2, TimeUnit.SECONDS)); + barrier.stopAcceptingCalls(); + server.shutdownNow(); + assertTrue(server.awaitTermination(1, TimeUnit.SECONDS)); + assertFalse("transport termination must not be mistaken for callback exit", finished.get()); + CountDownLatch waiting = new CountDownLatch(1); + FutureTask drained = new FutureTask<>(() -> { + waiting.countDown(); + barrier.awaitCallbacks(); + return null; + }); + Thread waiter = new Thread(drained, "test-rpc-drain"); + waiter.setDaemon(true); + waiter.start(); + assertTrue(waiting.await(1, TimeUnit.SECONDS)); + try { + drained.get(100, TimeUnit.MILLISECONDS); + fail("barrier passed while a handler still owns resources"); + } catch (TimeoutException expected) { + // The active handler must prevent database destruction. + } + release.countDown(); + drained.get(2, TimeUnit.SECONDS); + assertTrue(finished.get()); + } finally { + release.countDown(); + channel.shutdownNow(); + server.shutdownNow(); + executor.shutdown(); + assertTrue("cleanup must finish", executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/PartitionEngineShutdownTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/PartitionEngineShutdownTest.java new file mode 100644 index 0000000000..599d64b682 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/PartitionEngineShutdownTest.java @@ -0,0 +1,82 @@ +/* + * 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.store.service; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +import java.lang.reflect.Field; +import java.util.concurrent.atomic.AtomicInteger; + +import org.apache.hugegraph.store.HgStoreEngine; +import org.apache.hugegraph.store.PartitionEngine; +import org.apache.hugegraph.store.meta.ShardGroup; +import org.junit.Test; + +import com.alipay.sofa.jraft.RaftGroupService; + +public class PartitionEngineShutdownTest { + + @Test(timeout = 5000) + public void testInterruptedJoinRetriesBeforeMarkingPartitionStopped() throws Exception { + PartitionEngine engine = new PartitionEngine(mock(HgStoreEngine.class), + mock(ShardGroup.class)) { + @Override + public String toString() { + return "test-partition"; + } + }; + Field started = PartitionEngine.class.getDeclaredField("started"); + started.setAccessible(true); + started.setBoolean(engine, true); + RaftGroupService raft = mock(RaftGroupService.class); + Field service = PartitionEngine.class.getDeclaredField("raftGroupService"); + service.setAccessible(true); + service.set(engine, raft); + AtomicInteger joins = new AtomicInteger(); + doAnswer(invocation -> { + assertTrue("partition must remain started until Raft terminates", + started.getBoolean(engine)); + verify(raft).shutdown(); + if (joins.incrementAndGet() == 1) { + assertTrue("model an interrupt restored by the shutdown listener", + Thread.interrupted()); + throw new InterruptedException("shutdown thread interrupted"); + } + assertFalse("join retry must clear interruption while draining Raft", + Thread.currentThread().isInterrupted()); + return null; + }).when(raft).join(); + try { + Thread.currentThread().interrupt(); + engine.shutdown(); + assertEquals(2, joins.get()); + verify(raft, times(2)).join(); + assertFalse(started.getBoolean(engine)); + assertTrue("shutdown must restore caller interruption after Raft termination", + Thread.currentThread().isInterrupted()); + } finally { + Thread.interrupted(); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanContextCleanupTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanContextCleanupTest.java new file mode 100644 index 0000000000..4e33faeaed --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanContextCleanupTest.java @@ -0,0 +1,172 @@ +/* + * 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.store.service; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.apache.hugegraph.store.node.grpc.HgStoreStreamImpl; +import org.apache.hugegraph.store.node.grpc.GrpcShutdownBarrier; +import org.apache.hugegraph.store.node.grpc.query.AggregativeQueryService; +import org.apache.hugegraph.store.node.listener.ContextClosedListener; +import org.apache.hugegraph.store.node.task.TTLCleaner; +import org.junit.Test; +import org.lognet.springboot.grpc.context.GRpcServerInitializedEvent; + +import io.grpc.Server; +import org.springframework.context.annotation.AnnotationConfigApplicationContext; + + +public class ScanContextCleanupTest { + + @Test(timeout = 5000) + public void testScopedWorkersFinishBeforeBeanDestruction() throws Exception { + ThreadPoolExecutor scan = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + ThreadPoolExecutor ttl = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(); + CountDownLatch started = new CountDownLatch(3); + CountDownLatch interrupted = new CountDownLatch(3); + CountDownLatch cleanupStarted = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch cleanupAllowed = new CountDownLatch(1); + Runnable work = () -> { + started.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + interrupted.countDown(); + cleanupStarted.countDown(); + boolean done = false; + while (!done) { + try { + cleanupAllowed.await(); + done = true; + } catch (InterruptedException ignored) { + // Model resource cleanup which must finish before DB destruction. + } + } + Thread.currentThread().interrupt(); + } + }; + AtomicBoolean databaseClosed = new AtomicBoolean(); + AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(); + try { + HgStoreStreamImpl stream = mock(HgStoreStreamImpl.class); + TTLCleaner cleaner = mock(TTLCleaner.class); + when(stream.getRealExecutor()).thenReturn(scan); + doAnswer(invocation -> { + scan.shutdownNow(); + return null; + }).when(stream).shutdownScans(); + when(cleaner.getExecutor()).thenReturn(ttl); + when(cleaner.getScheduler()).thenReturn(scheduler); + context.getBeanFactory().registerSingleton("storeStream", stream); + context.getBeanFactory().registerSingleton("queryService", mock(AggregativeQueryService.class)); + context.getBeanFactory().registerSingleton("cleaner", cleaner); + context.getDefaultListableBeanFactory().registerDisposableBean( + "database", () -> databaseClosed.set(true)); + context.register(ContextClosedListener.class, GrpcShutdownBarrier.class); + context.refresh(); + scan.execute(work); + ttl.execute(work); + scheduler.execute(work); + assertTrue("all shutdown participants must be active", started.await(1, TimeUnit.SECONDS)); + + FutureTask closing = new FutureTask<>(() -> { + context.close(); + return null; + }); + Thread closeThread = new Thread(closing, "test-context-close"); + closeThread.setDaemon(true); + closeThread.start(); + // Scheduler shutdown must not let DB destruction overtake resource cleanup. + assertTrue(cleanupStarted.await(1, TimeUnit.SECONDS)); + assertFalse("context close must wait for worker cleanup", closing.isDone()); + assertFalse("database must remain open during worker cleanup", databaseClosed.get()); + cleanupAllowed.countDown(); + closing.get(2, TimeUnit.SECONDS); + assertTrue("close must interrupt active workers", interrupted.await(1, TimeUnit.SECONDS)); + assertTrue(scan.awaitTermination(1, TimeUnit.SECONDS)); + assertTrue(ttl.awaitTermination(1, TimeUnit.SECONDS)); + assertTrue(scheduler.awaitTermination(1, TimeUnit.SECONDS)); + assertTrue(databaseClosed.get()); + } finally { + cleanupAllowed.countDown(); + release.countDown(); + scan.shutdownNow(); + ttl.shutdownNow(); + scheduler.shutdownNow(); + context.close(); + } + } + + @Test(timeout = 5000) + public void testGrpcCallbacksFinishBeforeBeanDestructionEvenWhenInterrupted() throws Exception { + CountDownLatch awaitingCallbacks = new CountDownLatch(1); + CountDownLatch releaseCallbacks = new CountDownLatch(1); + AtomicBoolean terminated = new AtomicBoolean(); + AtomicBoolean databaseClosed = new AtomicBoolean(); + Server server = mock(Server.class); + when(server.shutdownNow()).thenReturn(server); + when(server.isTerminated()).thenAnswer(invocation -> terminated.get()); + when(server.awaitTermination(5, TimeUnit.SECONDS)).thenAnswer(invocation -> { + awaitingCallbacks.countDown(); + releaseCallbacks.await(); + terminated.set(true); + return true; + }); + AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(); + try { + context.getBeanFactory().registerSingleton("storeStream", mock(HgStoreStreamImpl.class)); + context.getBeanFactory().registerSingleton("queryService", mock(AggregativeQueryService.class)); + context.getBeanFactory().registerSingleton("cleaner", mock(TTLCleaner.class)); + context.getDefaultListableBeanFactory().registerDisposableBean("database", () -> databaseClosed.set(true)); + context.register(ContextClosedListener.class, GrpcShutdownBarrier.class); + context.refresh(); + context.publishEvent(new GRpcServerInitializedEvent(context, server)); + FutureTask closing = new FutureTask<>(() -> { + context.close(); + return Thread.currentThread().isInterrupted(); + }); + Thread closeThread = new Thread(closing, "test-grpc-context-close"); + closeThread.setDaemon(true); + closeThread.start(); + assertTrue(awaitingCallbacks.await(1, TimeUnit.SECONDS)); + closeThread.interrupt(); + assertFalse("DB must remain open while callbacks own its resources", databaseClosed.get()); + releaseCallbacks.countDown(); + assertTrue("shutdown must preserve interruption", closing.get(2, TimeUnit.SECONDS)); + assertTrue(terminated.get()); + assertTrue(databaseClosed.get()); + } finally { + releaseCallbacks.countDown(); + context.close(); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanLifecycleTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanLifecycleTest.java new file mode 100644 index 0000000000..443778b53c --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanLifecycleTest.java @@ -0,0 +1,1846 @@ +/* + * 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.store.service; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.SynchronousQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.ReentrantLock; + +import org.apache.hugegraph.pd.common.KVPair; +import org.apache.hugegraph.rocksdb.access.RocksDBSession; +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.store.grpc.common.Header; +import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.common.ScanOrderType; +import org.apache.hugegraph.store.grpc.stream.HgStoreStreamGrpc; +import org.apache.hugegraph.store.grpc.stream.KvPageRes; +import org.apache.hugegraph.store.grpc.stream.KvStream; +import org.apache.hugegraph.store.grpc.stream.ScanQueryRequest; +import org.apache.hugegraph.store.grpc.stream.ScanStreamBatchReq; +import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; +import org.apache.hugegraph.store.node.AppConfig; +import org.apache.hugegraph.store.node.grpc.HgStoreStreamImpl; +import org.apache.hugegraph.store.node.grpc.HgStoreWrapperEx; +import org.apache.hugegraph.store.node.grpc.ParallelScanIterator; +import org.apache.hugegraph.store.node.grpc.QueryCondition; +import org.apache.hugegraph.store.node.grpc.ScanBatchOneShotResponse; +import org.apache.hugegraph.store.node.grpc.ScanBatchResponse; +import org.apache.hugegraph.store.node.grpc.ScanBatchResponseFactory; +import org.apache.hugegraph.store.node.grpc.ScanBatchResponse3; +import org.apache.hugegraph.store.node.grpc.ScanOneShotResponse; +import org.apache.hugegraph.store.node.grpc.ScanStreamResponse; +import org.apache.hugegraph.store.node.util.HgChannel; +import org.junit.Test; +import org.mockito.ArgumentCaptor; + +import io.grpc.Context; +import io.grpc.ManagedChannel; +import io.grpc.Server; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.stub.StreamObserver; + +public class ScanLifecycleTest extends ScanTestSupport { + + @Test + public void testFailedSelectionCleanupBlocksLifecycleCompletion() throws Exception { + ScanIterator iterator = mock(ScanIterator.class); + IllegalStateException closing = new IllegalStateException("native release failed"); + doThrow(closing).when(iterator).close(); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + Class lifecycleType = Class.forName("org.apache.hugegraph.store.node.grpc.ScanLifecycle"); + java.lang.reflect.Constructor constructor = lifecycleType.getDeclaredConstructor(); + constructor.setAccessible(true); + Object lifecycle = constructor.newInstance(); + Method failedCleanup = lifecycleType.getDeclaredMethod("failedCleanup", Throwable.class); + failedCleanup.setAccessible(true); + AtomicBoolean finished = new AtomicBoolean(); + Method onFinished = lifecycleType.getDeclaredMethod("onFinished", Runnable.class); + onFinished.setAccessible(true); + onFinished.invoke(lifecycle, (Runnable) () -> finished.set(true)); + java.util.function.Consumer cleanup = failure -> { + try { + failedCleanup.invoke(lifecycle, failure); + } catch (ReflectiveOperationException error) { + throw new AssertionError(error); + } + }; + Class selector = Class.forName("org.apache.hugegraph.store.node.grpc.ScanUtil"); + Method getIterator = selector.getDeclaredMethod("getIterator", ScanStreamReq.class, + HgStoreWrapperEx.class, + java.util.function.Consumer.class); + getIterator.setAccessible(true); + ScanStreamReq request = ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setHeader(Header.newBuilder().setGraph("selection-cleanup")) + .setTable("unsupported-property-table") + .setSelects(org.apache.hugegraph.store.grpc.stream.SelectParam.newBuilder().addProperties(1)).build(); + try { + getIterator.invoke(null, request, wrapper, cleanup); + fail("Invalid property selection must fail"); + } catch (java.lang.reflect.InvocationTargetException error) { + assertTrue(error.getCause() instanceof IllegalArgumentException); + assertEquals(Collections.singletonList(closing), Arrays.asList(error.getCause().getSuppressed())); + } + Method finish = lifecycleType.getDeclaredMethod("finishWithoutResponse"); + finish.setAccessible(true); + finish.invoke(lifecycle); + assertFalse("Failed native release must remain a shutdown blocker", finished.get()); + verify(iterator).close(); + } + + @Test(timeout = 5000) + public void testInFlightReceiptDoesNotPublishOverAfterErrorCancellation() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + StreamObserver output = mock(StreamObserver.class); + ScanBatchResponse response = new ScanBatchResponse(output, mock(HgStoreWrapperEx.class), executor); + Field lockField = ScanBatchResponse.class.getDeclaredField("stateLock"); + lockField.setAccessible(true); + Field state = ScanBatchResponse.class.getDeclaredField("state"); + state.setAccessible(true); + Field cancelled = ScanBatchResponse.class.getDeclaredField("cancelled"); + cancelled.setAccessible(true); + AtomicBoolean cancelling = (AtomicBoolean) cancelled.get(response); + Object done = Arrays.stream(state.getType().getEnumConstants()) + .filter(value -> value.toString().equals("DONE")).findFirst().get(); + FutureTask receipt = new FutureTask<>(() -> { + response.onNext(ScanStreamBatchReq.newBuilder().setReceiptRequest( + org.apache.hugegraph.store.grpc.stream.ScanReceiptRequest.newBuilder().setTimes(1)).build()); + return null; + }); + FutureTask failing = new FutureTask<>(() -> { + response.onError(Status.INTERNAL.asRuntimeException()); + return null; + }); + Thread reader = new Thread(receipt, "receipt-at-terminal-transition"); + Thread closer = new Thread(failing, "error-at-terminal-transition"); + try { + synchronized (lockField.get(response)) { + reader.start(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (reader.getState() != Thread.State.BLOCKED && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals(Thread.State.BLOCKED, reader.getState()); + closer.start(); + deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (!cancelling.get() && System.nanoTime() < deadline) { + Thread.yield(); + } + assertTrue(cancelling.get()); + state.set(response, done); + } + receipt.get(2, TimeUnit.SECONDS); + failing.get(2, TimeUnit.SECONDS); + verify(output, never()).onNext(any(KvStream.class)); + verify(output).onError(any(Throwable.class)); + } finally { + executor.shutdownNow(); + reader.join(1000); + closer.join(1000); + } + } + + @Test(timeout = 10000) + public void testOrderedVertexRemainsConsecutiveAcrossBackpressureWithTwoScanners() throws Exception { + AtomicBoolean defer = new AtomicBoolean(true); + List initial = new ArrayList<>(); + ThreadPoolExecutor executor = new ThreadPoolExecutor(2, 2, 0, TimeUnit.SECONDS, + new java.util.concurrent.LinkedBlockingQueue<>()) { + @Override + public void execute(Runnable command) { + if (defer.get()) { + initial.add(command); + } else { + super.execute(command); + } + } + }; + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + AtomicInteger firstRows = new AtomicInteger(); + AtomicInteger secondRows = new AtomicInteger(); + CountDownLatch firstBackpressure = new CountDownLatch(1); + CountDownLatch secondReading = new CountDownLatch(1); + when(first.hasNext()).thenAnswer(ignored -> firstRows.get() < 8); + when(first.next()).thenAnswer(ignored -> { + if (firstRows.incrementAndGet() == 5) { + firstBackpressure.countDown(); + } + return RocksDBSession.BackendColumn.of(new byte[]{1, 0, 0, 0}, new byte[16]); + }); + when(second.hasNext()).thenAnswer(ignored -> secondRows.get() < 2); + when(second.next()).thenAnswer(ignored -> { + secondRows.incrementAndGet(); + secondReading.countDown(); + return RocksDBSession.BackendColumn.of(new byte[]{2, 0, 0, 0}, new byte[16]); + }); + AtomicInteger supplies = new AtomicInteger(); + Field bodySize = ParallelScanIterator.class.getDeclaredField("maxBodySize"); + bodySize.setAccessible(true); + int previous = bodySize.getInt(null); + bodySize.setInt(null, 16); + ParallelScanIterator scan = null; + try { + scan = ParallelScanIterator.of(() -> { + int index = supplies.incrementAndGet(); + return new KVPair<>(mock(QueryCondition.class), index == 1 ? first : index == 2 ? second : null); + }, () -> Long.MAX_VALUE, + ScanQueryRequest.newBuilder().setOrderType(ScanOrderType.ORDER_WITHIN_VERTEX).build(), executor); + // Force two owned scanners independently of the container's CPU-derived default. + Field registry = ParallelScanIterator.class.getDeclaredField("scanners"); + registry.setAccessible(true); + java.util.Queue scanners = (java.util.Queue) registry.get(scan); + Class scannerType = scanners.peek().getClass(); + java.lang.reflect.Constructor constructor = + scannerType.getDeclaredConstructor(ParallelScanIterator.class); + constructor.setAccessible(true); + Object other = constructor.newInstance(scan); + scanners.add(other); + Method run = scannerType.getDeclaredMethod("scanKV"); + run.setAccessible(true); + defer.set(false); + executor.execute(initial.get(0)); + assertTrue(firstBackpressure.await(2, TimeUnit.SECONDS)); + executor.execute(() -> { + try { + run.invoke(other); + } catch (ReflectiveOperationException e) { + throw new AssertionError(e); + } + }); + assertTrue(secondReading.await(2, TimeUnit.SECONDS)); + List vertices = new ArrayList<>(); + while (scan.hasNext()) { + for (Object row : scan.next()) { + Field key = row.getClass().getField("key"); + key.setAccessible(true); + vertices.add((int) ((byte[]) key.get(row))[0]); + } + } + assertEquals(Arrays.asList(1, 1, 1, 1, 1, 1, 1, 1, 2, 2), vertices); + verify(first).close(); + verify(second).close(); + } finally { + if (scan != null) { + scan.close(); + } + executor.shutdownNow(); + bodySize.setInt(null, previous); + } + } + + @Test(timeout = 5000) + public void testUnsupportedPausedBatchClosesOwnedIteratorAndRetainsFailure() throws Exception { + for (boolean fails : new boolean[]{false, true}) { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreStreamImpl service = scanService(executor, mock(HgStoreWrapperEx.class)); + StreamObserver output = mock(StreamObserver.class); + doThrow(new AssertionError("injected response callback failure")) + .when(output).onError(any(Throwable.class)); + StreamObserver input = service.scanBatch2(output); + ScanBatchResponse response = null; + for (Field field : input.getClass().getDeclaredFields()) { + field.setAccessible(true); + Object captured = field.get(input); + if (captured instanceof ScanBatchResponse) { + response = (ScanBatchResponse) captured; + } + } + org.junit.Assert.assertNotNull(response); + ScanIterator iterator = mock(ScanIterator.class); + if (fails) { + doThrow(new IllegalStateException("paused iterator close failed")).when(iterator).close(); + } + Field resource = ScanBatchResponse.class.getDeclaredField("iterator"); + resource.setAccessible(true); + resource.set(response, iterator); + try { + // An idle response can retain a native iterator while waiting for receipts. + input.onNext(ScanStreamBatchReq.getDefaultInstance()); + verify(iterator).close(); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onCompleted(); + assertEquals(fails ? 1 : 0, scanRegistry(service).size()); + service.shutdownScans(); + input.onCompleted(); + verify(iterator).close(); + } finally { + scanRegistry(service).clear(); + executor.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + public void testFailedCleanupAndResponseCallbackStillClearWorkerOwnership() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1, task -> { + Thread thread = new Thread(task, "scan-close-callback-failure"); + thread.setUncaughtExceptionHandler((ignored, failure) -> { }); + return thread; + }); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + StreamObserver output = mock(StreamObserver.class); + CountDownLatch sending = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + doAnswer(invocation -> { + // Exercise cleanup failure while the response callback is already in flight. + assertTrue(sending.await(2, TimeUnit.SECONDS)); + throw new IllegalStateException("native close failed"); + }).when(iterator).close(); + doAnswer(invocation -> { + sending.countDown(); + assertTrue(release.await(2, TimeUnit.SECONDS)); + return null; + }).when(output).onNext(any(KvPageRes.class)); + doThrow(new AssertionError("error callback failed")).when(output).onError(any(Throwable.class)); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(5); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return null; + }); + Thread caller = new Thread(request, "scan-close-callback-request"); + try { + caller.start(); + assertTrue(sending.await(1, TimeUnit.SECONDS)); + Field finished = ScanStreamResponse.class.getDeclaredField("finishFlag"); + finished.setAccessible(true); + AtomicBoolean terminal = (AtomicBoolean) finished.get(response); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (!terminal.get() && System.nanoTime() < deadline) { + Thread.yield(); + } + assertTrue("cleanup failure must claim error before normal completion", terminal.get()); + release.countDown(); + request.get(2, TimeUnit.SECONDS); + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + Field worker = ScanStreamResponse.class.getDeclaredField("worker"); + worker.setAccessible(true); + assertNull(worker.get(response)); + verify(iterator).close(); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onCompleted(); + } finally { + release.countDown(); + executor.shutdownNow(); + caller.join(1000); + } + } + + @Test(timeout = 5000) + public void testTerminalAdmissionRejectsRequestWaitingForLifecycleLock() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + HgStoreStreamImpl service = scanService(executor, wrapper); + StreamObserver input = service.scanBatch(mock(StreamObserver.class)); + Object lifecycle = scanRegistry(service).keySet().iterator().next(); + Method finished = lifecycle.getClass().getDeclaredMethod("finishWithoutResponse"); + finished.setAccessible(true); + FutureTask request = new FutureTask<>(() -> { + input.onNext(batchRequest()); + return null; + }); + Thread late = new Thread(request, "late-scan-admission"); + try { + synchronized (lifecycle) { + late.start(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (late.getState() != Thread.State.BLOCKED && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals(Thread.State.BLOCKED, late.getState()); + // Exercise the atomic admission invariant directly, without claiming a native fault. + finished.invoke(lifecycle); + assertTrue(scanRegistry(service).isEmpty()); + } + request.get(1, TimeUnit.SECONDS); + service.shutdownScans(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verifyNoInteractions(wrapper); + assertEquals(0, executor.getTaskCount()); + } finally { + executor.shutdownNow(); + late.join(1000); + } + } + + @Test(timeout = 10000) + public void testInProcessNormalCompletionDoesNotCancelCompletedScan() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(2); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + HgStoreStreamImpl service = scanService(executor, wrapper); + String name = InProcessServerBuilder.generateName(); + Server server = InProcessServerBuilder.forName(name) + .directExecutor().addService(service).build().start(); + ManagedChannel channel = InProcessChannelBuilder.forName(name) + .directExecutor().build(); + org.apache.logging.log4j.core.Logger logger = (org.apache.logging.log4j.core.Logger) + org.apache.logging.log4j.LogManager.getLogger(ScanStreamResponse.class); + List messages = new java.util.concurrent.CopyOnWriteArrayList<>(); + org.apache.logging.log4j.core.appender.AbstractAppender logs = + new org.apache.logging.log4j.core.appender.AbstractAppender( + "scan-normal-completion", null, null, false, + org.apache.logging.log4j.core.config.Property.EMPTY_ARRAY) { + @Override + public void append(org.apache.logging.log4j.core.LogEvent event) { + messages.add(event.getMessage().getFormattedMessage()); + } + }; + logs.start(); + logger.addAppender(logs); + CountDownLatch completed = new CountDownLatch(1); + AtomicInteger errors = new AtomicInteger(); + try { + StreamObserver request = + HgStoreStreamGrpc.newStub(channel) + .scan(new StreamObserver() { + @Override + public void onNext(KvPageRes value) { } + + @Override + public void onError(Throwable failure) { + errors.incrementAndGet(); + completed.countDown(); + } + + @Override + public void onCompleted() { + completed.countDown(); + } + }); + request.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + assertTrue(completed.await(2, TimeUnit.SECONDS)); + service.shutdownScans(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + service.awaitScanCleanup(); + assertEquals(0, errors.get()); + assertTrue(messages.stream().noneMatch(message -> message.contains("onError from client"))); + verify(iterator).close(); + assertTrue(scanRegistry(service).isEmpty()); + } finally { + logger.removeAppender(logs); + logs.stop(); + channel.shutdownNow(); + server.shutdownNow(); + executor.shutdownNow(); + } + } + + @Test(timeout = 10000) + public void testInProcessClientCancellationDrainsIterator() throws Exception { + assertInProcessCancellation(false); + assertInProcessCancellation(true); + } + + private static void assertInProcessCancellation(boolean completionRace) throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(2); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch releaseRead = new CountDownLatch(1); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + try { + releaseRead.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return false; + }); + HgStoreStreamImpl service = scanService(executor, wrapper); + // gRPC 1.39 dispatches context cancellation through the application executor too. + // Keep a slot for cancellation while scan.onNext waits for its worker's first page. + ExecutorService callbacks = Executors.newFixedThreadPool(2); + String name = InProcessServerBuilder.generateName(); + Server server = InProcessServerBuilder.forName(name) + .executor(callbacks).addService(service).build().start(); + ManagedChannel channel = InProcessChannelBuilder.forName(name) + .directExecutor().build(); + Context.CancellableContext context = Context.current().withCancellation(); + CountDownLatch ended = new CountDownLatch(1); + AtomicInteger terminals = new AtomicInteger(); + AtomicReference failure = new AtomicReference<>(); + try { + context.run(() -> { + StreamObserver request = + HgStoreStreamGrpc.newStub(channel) + .scan(new StreamObserver() { + @Override + public void onNext(KvPageRes value) { } + + @Override + public void onError(Throwable error) { + failure.set(error); + terminals.incrementAndGet(); + ended.countDown(); + } + + @Override + public void onCompleted() { + terminals.incrementAndGet(); + ended.countDown(); + } + }); + request.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + }); + assertTrue(reading.await(2, TimeUnit.SECONDS)); + if (completionRace) { + CountDownLatch start = new CountDownLatch(1); + Thread cancel = new Thread(() -> { + try { + start.await(); + } catch (InterruptedException e) { + throw new AssertionError(e); + } + context.cancel(null); + }, "scan-completion-cancel-race"); + cancel.start(); + start.countDown(); + releaseRead.countDown(); + cancel.join(2000); + assertFalse(cancel.isAlive()); + } else { + context.cancel(null); + } + assertTrue(ended.await(2, TimeUnit.SECONDS)); + verify(iterator, timeout(2000)).close(); + service.shutdownScans(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + service.awaitScanCleanup(); + assertEquals(1, terminals.get()); + if (!completionRace || failure.get() != null) { + assertEquals(Status.Code.CANCELLED, Status.fromThrowable(failure.get()).getCode()); + } + verify(iterator).close(); + assertTrue(scanRegistry(service).isEmpty()); + } finally { + releaseRead.countDown(); + context.cancel(null); + channel.shutdownNow(); + server.shutdownNow(); + executor.shutdownNow(); + callbacks.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testBatchRejectionReportsStatusAndClosesUnstartedIterator() throws Exception { + for (boolean stopped : new boolean[]{false, true}) { + ThreadPoolExecutor executor = new ThreadPoolExecutor(1, 1, 1, TimeUnit.SECONDS, + new SynchronousQueue<>()); + CountDownLatch release = new CountDownLatch(1); + if (stopped) { + executor.shutdown(); + } else { + executor.execute(() -> { + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + } + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + HgStoreStreamImpl service = scanService(executor, wrapper); + StreamObserver output = mock(StreamObserver.class); + try { + service.scanBatch(output).onNext(batchRequest()); + ArgumentCaptor error = ArgumentCaptor.forClass(Throwable.class); + verify(output).onError(error.capture()); + assertEquals(stopped ? Status.Code.UNAVAILABLE : Status.Code.RESOURCE_EXHAUSTED, + Status.fromThrowable(error.getValue()).getCode()); + verify(output, never()).onCompleted(); + assertTrue(scanRegistry(service).isEmpty()); + verifyNoInteractions(wrapper); + } finally { + release.countDown(); + executor.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + public void testBatchTimeoutCleanupDoesNotHoldAdmissionMonitor() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreStreamImpl service = scanService(executor, mock(HgStoreWrapperEx.class)); + ScanBatchResponse expired = mock(ScanBatchResponse.class); + CountDownLatch checking = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + ScanBatchResponseFactory factory = ScanBatchResponseFactory.getInstance(); + doAnswer(invocation -> { + assertFalse(Thread.holdsLock(service)); + checking.countDown(); + assertTrue(release.await(2, TimeUnit.SECONDS)); + return null; + }).when(expired).checkActiveTimeout(); + FutureTask> opening = new FutureTask<>(() -> + service.scanBatch2(mock(StreamObserver.class))); + Thread caller = new Thread(opening, "test-batch-timeout-cleanup"); + try { + factory.addStreamObserver(expired); + caller.start(); + assertTrue(checking.await(1, TimeUnit.SECONDS)); + FutureTask stopping = new FutureTask<>(() -> { + service.stopAcceptingScans(); + return null; + }); + Thread stopper = new Thread(stopping, "test-batch-stop-admission"); + stopper.start(); + stopping.get(1, TimeUnit.SECONDS); + assertUnavailable(() -> service.scan(mock(StreamObserver.class))); + assertFalse(opening.isDone()); + release.countDown(); + opening.get(1, TimeUnit.SECONDS).onError(Status.CANCELLED.asRuntimeException()); + assertTrue(scanRegistry(service).isEmpty()); + } finally { + release.countDown(); + factory.removeStreamObserver(expired); + caller.join(2000); + service.shutdownScans(); + executor.shutdownNow(); + } + } + + @Test(timeout = 10000) + public void testOneShotCancellationInterruptsCallerAndClosesIterator() throws Exception { + for (boolean batch : new boolean[]{false, true}) { + for (boolean transport : new boolean[]{false, true}) { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + StreamObserver output = mock(StreamObserver.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + when(iterator.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + try { + assertTrue(release.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException expected) { + interrupted.countDown(); + Thread.currentThread().interrupt(); + } + return false; + }); + HgStoreStreamImpl service = scanService(executor, wrapper); + Context.CancellableContext context = Context.current().withCancellation(); + FutureTask scanning = new FutureTask<>(() -> { + context.run(() -> { + if (batch) { + service.scanBatchOneShot(batchRequest(), output); + } else { + service.scanOneShot(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setLimit(10).build(), output); + } + }); + assertFalse("a cancelled call must not interrupt the next call", + Thread.currentThread().isInterrupted()); + doReturn(false).when(iterator).hasNext(); + StreamObserver nextOutput = mock(StreamObserver.class); + HgStoreStreamImpl nextService = scanService(executor, wrapper); + if (batch) { + nextService.scanBatchOneShot(batchRequest(), nextOutput); + } else { + nextService.scanOneShot(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setLimit(10).build(), nextOutput); + } + verify(nextOutput).onNext(any(KvPageRes.class)); + verify(nextOutput).onCompleted(); + verify(nextOutput, never()).onError(any(Throwable.class)); + return null; + }); + Thread caller = new Thread(scanning, "test-one-shot-blocked-read"); + try { + caller.start(); + assertTrue(reading.await(1, TimeUnit.SECONDS)); + if (transport) { + context.cancel(new IllegalStateException("client canceled")); + } else { + service.shutdownScans(); + } + assertTrue(interrupted.await(1, TimeUnit.SECONDS)); + scanning.get(1, TimeUnit.SECONDS); + assertCancelled(output); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verify(iterator, times(2)).close(); + assertTrue(scanRegistry(service).isEmpty()); + } finally { + release.countDown(); + context.cancel(null); + caller.join(2000); + service.shutdownScans(); + executor.shutdownNow(); + assertFalse(caller.isAlive()); + } + } + } + } + + private static void awaitCloseRelease(CountDownLatch release) { + boolean interrupted = false; + try { + for (;;) { + try { + assertTrue(release.await(5, TimeUnit.SECONDS)); + return; + } catch (InterruptedException ignored) { + interrupted = true; + } + } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + } + + @Test(timeout = 10000) + public void testInProcessOneShotCancellationDuringCloseReportsCancelled() throws Exception { + for (boolean batch : new boolean[]{false, true}) { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(false); + CountDownLatch closing = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + doAnswer(invocation -> { + closing.countDown(); + awaitCloseRelease(release); + return null; + }).when(iterator).close(); + HgStoreStreamImpl service = scanService(executor, wrapper); + String name = InProcessServerBuilder.generateName(); + Server server = InProcessServerBuilder.forName(name).directExecutor() + .addService(service).build().start(); + ManagedChannel channel = InProcessChannelBuilder.forName(name).directExecutor().build(); + FutureTask call = new FutureTask<>(() -> { + try { + HgStoreStreamGrpc.HgStoreStreamBlockingStub stub = + HgStoreStreamGrpc.newBlockingStub(channel).withDeadlineAfter(3, TimeUnit.SECONDS); + if (batch) { + stub.scanBatchOneShot(batchRequest()); + } else { + stub.scanOneShot(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setLimit(10).build()); + } + fail("cancelled unary scan reported success"); + return Status.Code.OK; + } catch (StatusRuntimeException error) { + return error.getStatus().getCode(); + } + }); + Thread caller = new Thread(call, "one-shot-cancellation-client"); + try { + caller.start(); + assertTrue(closing.await(2, TimeUnit.SECONDS)); + service.shutdownScans(); + assertEquals(1, scanRegistry(service).size()); + assertFalse(call.isDone()); + release.countDown(); + assertEquals(Status.Code.CANCELLED, call.get(2, TimeUnit.SECONDS)); + service.awaitScanCleanup(); + assertTrue(scanRegistry(service).isEmpty()); + verify(iterator).close(); + } finally { + release.countDown(); + channel.shutdownNow(); + server.shutdownNow(); + executor.shutdownNow(); + caller.join(2000); + assertFalse(caller.isAlive()); + assertTrue(channel.awaitTermination(2, TimeUnit.SECONDS)); + assertTrue(server.awaitTermination(2, TimeUnit.SECONDS)); + } + } + } + + @Test(timeout = 10000) + public void testOneShotCancelledDuringCloseDoesNotCompleteSuccessfully() throws Exception { + for (boolean batch : new boolean[]{false, true}) { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + StreamObserver output = mock(StreamObserver.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(false); + CountDownLatch closing = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + doAnswer(invocation -> { + closing.countDown(); + awaitCloseRelease(release); + return null; + }).when(iterator).close(); + HgStoreStreamImpl service = scanService(executor, wrapper); + FutureTask scanning = new FutureTask<>(() -> { + if (batch) { + service.scanBatchOneShot(batchRequest(), output); + } else { + service.scanOneShot(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setLimit(10).build(), output); + } + return null; + }); + Thread caller = new Thread(scanning, "one-shot-close-cancellation"); + try { + caller.start(); + assertTrue(closing.await(2, TimeUnit.SECONDS)); + service.shutdownScans(); + assertEquals(1, scanRegistry(service).size()); + assertFalse(scanning.isDone()); + verifyNoInteractions(output); + release.countDown(); + scanning.get(2, TimeUnit.SECONDS); + ArgumentCaptor error = ArgumentCaptor.forClass(Throwable.class); + verify(output).onError(error.capture()); + assertEquals(Status.Code.CANCELLED, Status.fromThrowable(error.getValue()).getCode()); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + assertTrue(scanRegistry(service).isEmpty()); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + } finally { + release.countDown(); + caller.join(2000); + executor.shutdownNow(); + assertFalse(caller.isAlive()); + } + } + } + + @Test(timeout = 5000) + public void testBatchShutdownInterruptsBlockedReaderBeforeIteratorClose() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + CountDownLatch readExited = new CountDownLatch(1); + AtomicBoolean closedAfterRead = new AtomicBoolean(); + when(iterator.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + try { + new CountDownLatch(1).await(); + return true; + } catch (InterruptedException expected) { + interrupted.countDown(); + return false; + } finally { + readExited.countDown(); + } + }); + doAnswer(invocation -> { + closedAfterRead.set(readExited.getCount() == 0); + return null; + }).when(iterator).close(); + + HgStoreStreamImpl service = scanService(executor, wrapper); + StreamObserver output = mock(StreamObserver.class); + StreamObserver input = service.scanBatch(output); + FutureTask shutdown = new FutureTask<>(() -> { + service.shutdownScans(); + return null; + }); + Thread shutdownThread = new Thread(shutdown, "batch-scan-shutdown-test"); + try { + input.onNext(batchRequest()); + assertTrue("batch reader must enter its blocking iterator read", + reading.await(1, TimeUnit.SECONDS)); + shutdownThread.start(); + shutdown.get(2, TimeUnit.SECONDS); + assertTrue(interrupted.await(1, TimeUnit.SECONDS)); + service.awaitScanCleanup(); + assertTrue(closedAfterRead.get()); + assertTrue(scanRegistry(service).isEmpty()); + assertCancelled(output); + verify(output, never()).onCompleted(); + verify(iterator).close(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + } finally { + executor.shutdownNow(); + shutdownThread.interrupt(); + shutdownThread.join(1000); + } + } + + @Test(timeout = 5000) + public void testEncodedBatchShutdownInterruptsBlockedScannerBeforeIteratorClose() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + CountDownLatch readExited = new CountDownLatch(1); + AtomicBoolean closedAfterRead = new AtomicBoolean(); + when(iterator.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + try { + new CountDownLatch(1).await(); + return true; + } catch (InterruptedException expected) { + interrupted.countDown(); + return false; + } finally { + readExited.countDown(); + } + }); + doAnswer(invocation -> { + closedAfterRead.set(readExited.getCount() == 0); + return null; + }).when(iterator).close(); + + HgStoreStreamImpl service = scanService(executor, wrapper); + StreamObserver output = mock(StreamObserver.class); + StreamObserver input = service.scanBatch2(output); + FutureTask shutdown = new FutureTask<>(() -> { + service.shutdownScans(); + return null; + }); + Thread shutdownThread = new Thread(shutdown, "encoded-batch-scan-shutdown-test"); + try { + input.onNext(batchRequest()); + assertTrue("parallel scanner must enter its blocking iterator read", + reading.await(1, TimeUnit.SECONDS)); + shutdownThread.start(); + shutdown.get(2, TimeUnit.SECONDS); + assertTrue(interrupted.await(1, TimeUnit.SECONDS)); + service.awaitScanCleanup(); + assertTrue(closedAfterRead.get()); + assertTrue(scanRegistry(service).isEmpty()); + ArgumentCaptor failure = ArgumentCaptor.forClass(Throwable.class); + verify(output).onError(failure.capture()); + assertEquals(Status.Code.CANCELLED, Status.fromThrowable(failure.getValue()).getCode()); + verify(output, never()).onCompleted(); + verify(iterator).close(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + } finally { + executor.shutdownNow(); + shutdownThread.interrupt(); + shutdownThread.join(1000); + } + } + + @Test(timeout = 5000) + public void testOneShotObservesContextCancellationAndClosesIterator() { + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + StreamObserver output = mock(StreamObserver.class); + Context.CancellableContext context = Context.current().withCancellation(); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenAnswer(invocation -> { + context.cancel(null); + return true; + }); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + try { + context.run(() -> ScanOneShotResponse.scanOneShot( + ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL).setLimit(100).build(), + output, wrapper)); + verify(iterator).close(); + verifyNoInteractions(output); + } finally { + context.cancel(null); + } + } + + @Test(timeout = 5000) + public void testBatchDoneRejectsSecondQueryBeforeAllocatingIterator() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator source = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(source); + StreamObserver output = mock(StreamObserver.class); + ScanBatchResponse response = new ScanBatchResponse(output, wrapper, executor); + Field state = ScanBatchResponse.class.getDeclaredField("state"); + state.setAccessible(true); + Field iterator = ScanBatchResponse.class.getDeclaredField("iterator"); + iterator.setAccessible(true); + try { + response.onNext(batchRequest()); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2); + while (!"DONE".equals(state.get(response).toString()) && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals("DONE", state.get(response).toString()); + verify(source).close(); + assertNull(iterator.get(response)); + response.onNext(batchRequest()); + // The single worker processes the queued request before this barrier. + executor.submit(() -> { }).get(2, TimeUnit.SECONDS); + assertNull("a completed query must reject another request before allocating", iterator.get(response)); + response.onCompleted(); + response.onCompleted(); + verify(wrapper).scanAll(anyString(), anyString(), any(byte[].class)); + verify(source).close(); + verify(output).onCompleted(); + verify(output, never()).onError(any(Throwable.class)); + } finally { + // Preserve cleanup even when the old implementation allocated a second producer. + Object remaining = iterator.get(response); + response.onCompleted(); + if (remaining instanceof ScanIterator) { + ((ScanIterator) remaining).close(); + } + executor.shutdownNow(); + } + } + + @Test + public void testPausedBatchClosesIteratorAndRejectsLaterQuery() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + StreamObserver output = mock(StreamObserver.class); + ScanBatchResponse response = new ScanBatchResponse(output, wrapper, executor); + Field field = ScanBatchResponse.class.getDeclaredField("iterator"); + field.setAccessible(true); + field.set(response, iterator); + try { + response.onCompleted(); + response.onNext(ScanStreamBatchReq.newBuilder() + .setQueryRequest(ScanQueryRequest.getDefaultInstance()) + .build()); + response.onCompleted(); + verify(iterator).close(); + verifyNoInteractions(wrapper); + assertNull(field.get(response)); + assertFalse(executor.isShutdown()); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testRejectedStreamReportsOnlyFailure() { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + executor.shutdown(); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + ScanStreamReq request = ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build(); + + response.onNext(request); + response.onNext(request); + response.onCompleted(); + + verify(output).onError(any(Throwable.class)); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verifyNoInteractions(wrapper); + } + + @Test(timeout = 5000) + public void testStreamIteratorFailureReportsOnlyFailure() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))) + .thenThrow(new IllegalStateException("iterator failed")); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return null; + }); + Thread caller = new Thread(request, "scan-failure-test"); + try { + caller.start(); + request.get(1, TimeUnit.SECONDS); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + response.onCompleted(); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + caller.join(1000); + } + } + + @Test(timeout = 5000) + public void testReceiveFailureDoesNotInterruptItsCallbackThread() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + CountDownLatch occupied = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + executor.execute(() -> { + occupied.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(0); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return Thread.currentThread().isInterrupted(); + }); + Thread caller = new Thread(request, "scan-receive-failure-test"); + try { + assertTrue(occupied.await(1, TimeUnit.SECONDS)); + caller.start(); + assertFalse(request.get(1, TimeUnit.SECONDS)); + release.countDown(); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verifyNoInteractions(wrapper); + } finally { + release.countDown(); + response.onCompleted(); + executor.shutdownNow(); + caller.join(1000); + } + } + + @Test(timeout = 10000) + public void testStreamFailureWaitsForPageCallback() throws Exception { + AtomicReference scanWorker = new AtomicReference<>(); + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1, task -> { + Thread thread = new Thread(task, "scan-page-failure-worker"); + scanWorker.set(thread); + return thread; + }); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + CountDownLatch pageEntered = new CountDownLatch(1); + CountDownLatch releasePage = new CountDownLatch(1); + CountDownLatch iteratorFailed = new CountDownLatch(1); + CountDownLatch errorReceived = new CountDownLatch(1); + AtomicInteger advances = new AtomicInteger(); + when(iterator.hasNext()).thenAnswer(invocation -> { + if (advances.incrementAndGet() <= 2) { + return true; + } + assertTrue(pageEntered.await(1, TimeUnit.SECONDS)); + iteratorFailed.countDown(); + throw new IllegalStateException("failure after first page"); + }); + StreamObserver output = mock(StreamObserver.class); + List signals = Collections.synchronizedList(new ArrayList<>()); + doAnswer(invocation -> { + signals.add("page-start"); + pageEntered.countDown(); + assertTrue(releasePage.await(2, TimeUnit.SECONDS)); + signals.add("page-end"); + return null; + }).when(output).onNext(any(KvPageRes.class)); + doAnswer(invocation -> { + signals.add("error"); + errorReceived.countDown(); + return null; + }).when(output).onError(any(Throwable.class)); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return null; + }); + Thread caller = new Thread(request, "scan-page-failure-test"); + try { + caller.start(); + assertTrue(pageEntered.await(1, TimeUnit.SECONDS)); + assertTrue(iteratorFailed.await(1, TimeUnit.SECONDS)); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (scanWorker.get().getState() != Thread.State.BLOCKED && + errorReceived.getCount() != 0 && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals("worker must reach the response serialization point", + Thread.State.BLOCKED, scanWorker.get().getState()); + assertEquals(1L, errorReceived.getCount()); + releasePage.countDown(); + request.get(1, TimeUnit.SECONDS); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + assertEquals(Arrays.asList("page-start", "page-end", "error"), signals); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + } finally { + releasePage.countDown(); + response.onCompleted(); + executor.shutdownNow(); + caller.join(1000); + } + } + + @Test(timeout = 10000) + public void testNormalStreamCancellationClaimsTerminalBeforeInterruptingWorker() throws Exception { + for (boolean closeRequest : new boolean[]{false, true}) { + CountDownLatch workerEntered = new CountDownLatch(1); + CountDownLatch iteratorClosed = new CountDownLatch(1); + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + AtomicInteger advances = new AtomicInteger(); + when(iterator.hasNext()).thenAnswer(invocation -> { + if (advances.incrementAndGet() <= 2) { + return true; + } + workerEntered.countDown(); + new CountDownLatch(1).await(); + return false; + }); + doAnswer(invocation -> { + iteratorClosed.countDown(); + return null; + }).when(iterator).close(); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + try { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + assertTrue(workerEntered.await(1, TimeUnit.SECONDS)); + if (closeRequest) { + response.onNext(ScanStreamReq.newBuilder().setCloseFlag(1).build()); + } else { + response.onCompleted(); + } + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(output, never()).onError(any(Throwable.class)); + verify(output).onCompleted(); + verify(iterator).close(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + } + + @Test + public void testRejectedBatchReportsFailureWithoutSuccessfulCompletion() { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + executor.shutdown(); + StreamObserver output = mock(StreamObserver.class); + ScanBatchResponse response = new ScanBatchResponse(output, mock(HgStoreWrapperEx.class), executor); + response.onNext(ScanStreamBatchReq.newBuilder() + .setQueryRequest(ScanQueryRequest.getDefaultInstance()).build()); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onCompleted(); + } + + @Test + public void testBatchIteratorFailureClosesResourcesWithoutSuccessfulCompletion() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + StreamObserver output = mock(StreamObserver.class); + ScanIterator iterator = mock(ScanIterator.class); + IllegalStateException failure = new IllegalStateException("iterator failed"); + when(iterator.hasNext()).thenThrow(failure); + ScanBatchResponse response = new ScanBatchResponse(output, mock(HgStoreWrapperEx.class), executor); + Field field = ScanBatchResponse.class.getDeclaredField("iterator"); + field.setAccessible(true); + field.set(response, iterator); + Method send = ScanBatchResponse.class.getDeclaredMethod("sendEntries"); + send.setAccessible(true); + try { + send.invoke(response); + verify(iterator).close(); + verify(output).onError(failure); + verify(output, never()).onCompleted(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testActiveStreamCancellationReleasesBlockedProducer() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(true); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + try { + // Consume just one page, then leave the producer waiting for its consumer. + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(1000).build()); + response.onError(Status.CANCELLED.asRuntimeException()); + executor.shutdown(); + assertTrue("cancel must release a producer without its 60-second timeout", + executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(iterator).close(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testParallelPausedScannerReleasesIterator() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + ScanIterator source = mock(ScanIterator.class); + when(source.hasNext()).thenReturn(true); + when(source.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + Field bodySize = ParallelScanIterator.class.getDeclaredField("maxBodySize"); + bodySize.setAccessible(true); + int originalBodySize = bodySize.getInt(null); + bodySize.setInt(null, 2); + ParallelScanIterator scan = null; + try { + scan = ParallelScanIterator.of( + () -> new KVPair<>(mock(QueryCondition.class), source), () -> Long.MAX_VALUE, + ScanQueryRequest.getDefaultInstance(), executor); + // One worker fills its output allowance and pauses while retaining its iterator. + executor.submit(() -> { }).get(2, TimeUnit.SECONDS); + verify(source, atLeastOnce()).next(); + verify(source, never()).close(); + scan.close(); + scan.close(); + verify(source).close(); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + } finally { + if (scan != null) { + scan.close(); + } + bodySize.setInt(null, originalBodySize); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testParallelOrderedScannerReleasesQueueLockOnEmptyIterator() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + ScanIterator first = mock(ScanIterator.class); + ScanIterator empty = mock(ScanIterator.class); + when(first.hasNext()).thenReturn(true); + when(first.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(empty.hasNext()).thenReturn(false); + + AtomicInteger supplies = new AtomicInteger(); + Field bodySize = ParallelScanIterator.class.getDeclaredField("maxBodySize"); + bodySize.setAccessible(true); + int originalBodySize = bodySize.getInt(null); + bodySize.setInt(null, 2); + ParallelScanIterator scan = null; + try { + scan = ParallelScanIterator.of( + () -> { + int supply = supplies.incrementAndGet(); + if (supply == 1) { + return new KVPair<>(mock(QueryCondition.class), first); + } + if (supply == 2) { + return new KVPair<>(mock(QueryCondition.class), empty); + } + return new KVPair<>(mock(QueryCondition.class), null); + }, + () -> 1L, + ScanQueryRequest.newBuilder().setOrderType(ScanOrderType.ORDER_WITHIN_VERTEX).build(), + executor); + + // Wait for the scanner task to finish before checking the lock owner. + executor.submit(() -> { }).get(2, TimeUnit.SECONDS); + Field queueLock = ParallelScanIterator.class.getDeclaredField("queueLock"); + queueLock.setAccessible(true); + assertFalse(((ReentrantLock) queueLock.get(scan)).isLocked()); + } finally { + if (scan != null) { + scan.close(); + } + bodySize.setInt(null, originalBodySize); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testParallelOrderedProducerCancelsWithFullOutputQueue() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + ScanIterator source = mock(ScanIterator.class); + Field bodySize = ParallelScanIterator.class.getDeclaredField("maxBodySize"); + bodySize.setAccessible(true); + byte[] value = new byte[bodySize.getInt(null)]; + CountDownLatch full = new CountDownLatch(1); + AtomicInteger rows = new AtomicInteger(); + when(source.hasNext()).thenReturn(true); + when(source.next()).thenAnswer(invocation -> { + if (rows.incrementAndGet() == 5) { + full.countDown(); + } + return RocksDBSession.BackendColumn.of(new byte[4], value); + }); + ParallelScanIterator scan = ParallelScanIterator.of( + () -> new KVPair<>(mock(QueryCondition.class), source), () -> Long.MAX_VALUE, + ScanQueryRequest.newBuilder().setOrderType(ScanOrderType.ORDER_WITHIN_VERTEX).build(), + executor); + try { + assertTrue("four queued batches must fill the single-scanner queue", + full.await(2, TimeUnit.SECONDS)); + scan.close(); + executor.shutdown(); + assertTrue("cancellation must release the blocked ordered producer", + executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(source).close(); + } finally { + scan.close(); + executor.shutdownNow(); + } + } + + @Test + public void testInterruptedOneShotDoesNotSendPartialSuccess() throws Exception { + for (boolean batch : new boolean[]{false, true}) { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + StreamObserver output = mock(StreamObserver.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(true); + when(iterator.position()).thenReturn(new byte[4]); + when(iterator.next()).thenAnswer(invocation -> { + Thread.currentThread().interrupt(); + return RocksDBSession.BackendColumn.of(new byte[4], new byte[0]); + }); + try { + HgStoreStreamImpl service = scanService(executor, wrapper); + if (batch) { + service.scanBatchOneShot(batchRequest(), output); + } else { + service.scanOneShot(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL).setLimit(10).build(), + output); + } + verify(iterator).next(); + assertTrue("scan must preserve an external interrupt", + Thread.currentThread().isInterrupted()); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + StreamObserver nextOutput = mock(StreamObserver.class); + if (batch) { + service.scanBatchOneShot(batchRequest(), nextOutput); + } else { + service.scanOneShot(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL).setLimit(10).build(), + nextOutput); + } + assertTrue("an interrupt present at entry must be preserved", Thread.currentThread().isInterrupted()); + verify(nextOutput).onError(any(Throwable.class)); + verify(nextOutput, never()).onNext(any(KvPageRes.class)); + verify(nextOutput, never()).onCompleted(); + } finally { + Thread.interrupted(); + executor.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + public void testInterruptedBatchWorkerDoesNotCompletePartialPage() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + StreamObserver output = mock(StreamObserver.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(true); + when(iterator.position()).thenReturn(new byte[4]); + when(iterator.next()).thenAnswer(invocation -> { + Thread.currentThread().interrupt(); + return RocksDBSession.BackendColumn.of(new byte[4], new byte[0]); + }); + StreamObserver response = ScanBatchResponse3.of(output, wrapper, executor); + try { + response.onNext(batchRequest()); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(iterator).next(); + assertCancelled(output); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + + @Test(timeout = 10000) + public void testCancellationDoesNotInterruptReceiverAfterHandoff() throws Exception { + assertReceiverInterruptAfterHandoff(false); + } + + @Test(timeout = 10000) + public void testCancellationPreservesExternalReceiverInterruptAfterHandoff() throws Exception { + assertReceiverInterruptAfterHandoff(true); + } + + @Test(timeout = 5000) + public void testReceivePreservesExternalInterrupt() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + Thread.currentThread().interrupt(); + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return Thread.currentThread().isInterrupted(); + }); + Thread caller = new Thread(request, "scan-external-interrupt-test"); + try { + caller.start(); + assertTrue(request.get(1, TimeUnit.SECONDS)); + assertCancelled(output); + verify(output, never()).onCompleted(); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + } finally { + response.onCompleted(); + executor.shutdownNow(); + caller.join(1000); + } + } + + private static void assertReceiverInterruptAfterHandoff(boolean externalInterrupt) throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + CountDownLatch beforePage = new CountDownLatch(1); + CountDownLatch releasePage = new CountDownLatch(1); + AtomicInteger advances = new AtomicInteger(); + when(iterator.hasNext()).thenAnswer(invocation -> { + int advance = advances.incrementAndGet(); + if (advance == 2) { + beforePage.countDown(); + assertTrue(releasePage.await(2, TimeUnit.SECONDS)); + } + return advance <= 2; + }); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + Field responseLock = ScanStreamResponse.class.getDeclaredField("responseLock"); + responseLock.setAccessible(true); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return Thread.currentThread().isInterrupted(); + }); + Thread caller = new Thread(request, "scan-late-cancellation-test"); + try { + caller.start(); + assertTrue(beforePage.await(1, TimeUnit.SECONDS)); + synchronized (responseLock.get(response)) { + // The real channel hands off a page before its callback reaches responseLock. + releasePage.countDown(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (caller.getState() != Thread.State.BLOCKED && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals(Thread.State.BLOCKED, caller.getState()); + if (externalInterrupt) { + caller.interrupt(); + } + response.onCompleted(); + } + assertEquals(externalInterrupt, request.get(1, TimeUnit.SECONDS)); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(output).onCompleted(); + verify(output, never()).onError(any(Throwable.class)); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(iterator).close(); + } finally { + releasePage.countDown(); + response.onCompleted(); + executor.shutdownNow(); + caller.join(1000); + } + } + + @Test(timeout = 5000) + public void testInterruptedStreamWorkerDoesNotCompletePartialPage() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(true); + when(iterator.position()).thenReturn(new byte[4]); + when(iterator.next()).thenAnswer(invocation -> { + Thread.currentThread().interrupt(); + return RocksDBSession.BackendColumn.of(new byte[4], new byte[0]); + }); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + try { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(iterator).next(); + assertCancelled(output); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testInterruptedStreamSenderReportsCancellation() throws Exception { + AtomicReference worker = new AtomicReference<>(); + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1, task -> { + Thread thread = new Thread(task, "scan-interrupted-sender-test"); + worker.set(thread); + return thread; + }); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(true); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + CountDownLatch errorReceived = new CountDownLatch(1); + doAnswer(invocation -> { + errorReceived.countDown(); + return null; + }).when(output).onError(any(Throwable.class)); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + try { + // Consume one page, then interrupt the actual worker waiting to hand off its next page. + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + awaitTimedWaiting(worker.get()); + worker.get().interrupt(); + assertTrue(errorReceived.await(1, TimeUnit.SECONDS)); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + assertCancelled(output); + verify(output).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testChannelCloseReleasesWaitingReceiverWithoutInterrupt() throws Exception { + HgChannel channel = HgChannel.of(60); + FutureTask receive = new FutureTask<>(() -> { + assertNull(channel.receive()); + return Thread.currentThread().isInterrupted(); + }); + Thread caller = new Thread(receive, "channel-close-receiver-test"); + try { + caller.start(); + awaitTimedWaiting(caller); + channel.close(); + assertFalse(receive.get(1, TimeUnit.SECONDS)); + assertNull(channel.receive()); + assertFalse(channel.send("after-close")); + } finally { + channel.close(); + caller.join(1000); + } + } + + @Test(timeout = 5000) + public void testChannelCloseReleasesWaitingProducer() throws Exception { + HgChannel channel = HgChannel.of(60); + FutureTask send = new FutureTask<>(() -> channel.send("page")); + Thread producer = new Thread(send, "channel-close-producer-test"); + try { + producer.start(); + awaitTimedWaiting(producer); + channel.close(); + assertFalse(send.get(1, TimeUnit.SECONDS)); + } finally { + channel.close(); + producer.join(1000); + } + } + + @Test(timeout = 5000) + public void testChannelRetainsTimeoutAndNormalHandoff() throws Exception { + HgChannel channel = HgChannel.of(1); + AtomicInteger timeouts = new AtomicInteger(); + long started = System.nanoTime(); + assertNull(channel.receive(timeout -> { + assertEquals(Long.valueOf(1L), timeout); + timeouts.incrementAndGet(); + })); + assertEquals(1, timeouts.get()); + assertTrue("short close polling must not shorten the configured timeout", + System.nanoTime() - started >= TimeUnit.MILLISECONDS.toNanos(900)); + FutureTask send = new FutureTask<>(() -> channel.send("page")); + Thread producer = new Thread(send, "channel-handoff-test"); + try { + producer.start(); + assertEquals("page", channel.receive()); + assertTrue(send.get(1, TimeUnit.SECONDS)); + } finally { + channel.close(); + producer.join(1000); + } + } + + @Test + public void testClosingPreventsLazyExecutorAndStateCreation() { + HgStoreStreamImpl service = new HgStoreStreamImpl(); + service.stopAcceptingScans(); + assertUnavailable(service::getExecutor); + assertUnavailable(service::getState); + assertUnavailable(() -> service.scan(mock(StreamObserver.class))); + service.shutdownScans(); + assertNull(service.getRealExecutor()); + } + + @Test(timeout = 5000) + public void testQueuedStreamCancellationDrainsWithoutOpeningIterator() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch occupied = new CountDownLatch(1); + executor.execute(() -> { + occupied.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return null; + }); + Thread caller = new Thread(request, "scan-shutdown-test"); + try { + assertTrue(occupied.await(1, TimeUnit.SECONDS)); + caller.start(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (executor.getQueue().isEmpty() && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals(1, executor.getQueue().size()); + response.onError(Status.CANCELLED.asRuntimeException()); + executor.shutdown(); + release.countDown(); + request.get(1, TimeUnit.SECONDS); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verifyNoInteractions(wrapper); + assertEquals(2L, executor.getCompletedTaskCount()); + } finally { + response.onCompleted(); + release.countDown(); + executor.shutdownNow(); + caller.join(1000); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanTestSupport.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanTestSupport.java new file mode 100644 index 0000000000..c3755a10d5 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanTestSupport.java @@ -0,0 +1,98 @@ +/* + * 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.store.service; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.util.Map; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; + +import org.apache.hugegraph.store.grpc.common.Header; +import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.stream.KvPageRes; +import org.apache.hugegraph.store.grpc.stream.ScanQueryRequest; +import org.apache.hugegraph.store.grpc.stream.ScanStreamBatchReq; +import org.apache.hugegraph.store.node.AppConfig; +import org.apache.hugegraph.store.node.grpc.HgStoreStreamImpl; +import org.apache.hugegraph.store.node.grpc.HgStoreWrapperEx; +import org.mockito.ArgumentCaptor; + +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.grpc.stub.StreamObserver; + +abstract class ScanTestSupport { + + static HgStoreStreamImpl scanService(ThreadPoolExecutor executor, + HgStoreWrapperEx wrapper) throws Exception { + HgStoreStreamImpl service = new HgStoreStreamImpl(); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(5); + for (String name : new String[]{"executor", "wrapper", "appConfig"}) { + Field field = HgStoreStreamImpl.class.getDeclaredField(name); + field.setAccessible(true); + field.set(service, name.equals("executor") ? executor : name.equals("wrapper") ? wrapper : config); + } + return service; + } + + static Map scanRegistry(HgStoreStreamImpl service) throws Exception { + Field scans = HgStoreStreamImpl.class.getDeclaredField("scans"); + scans.setAccessible(true); + return (Map) scans.get(service); + } + + static void awaitTimedWaiting(Thread thread) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + Thread.State observed = thread.getState(); + while (observed != Thread.State.TIMED_WAITING && System.nanoTime() < deadline) { + Thread.yield(); + observed = thread.getState(); + } + assertEquals(Thread.State.TIMED_WAITING, observed); + } + + static ScanStreamBatchReq batchRequest() { + return ScanStreamBatchReq.newBuilder().setHeader(Header.newBuilder().setGraph("g")) + .setQueryRequest(ScanQueryRequest.newBuilder().setMethod(ScanMethod.ALL) + .setTable("t").setLimit(10) + .setPerKeyMax(Long.MAX_VALUE) + .setPageSize(1)).build(); + } + + static void assertCancelled(StreamObserver output) { + ArgumentCaptor failure = ArgumentCaptor.forClass(Throwable.class); + verify(output).onError(failure.capture()); + assertEquals(Status.Code.CANCELLED, Status.fromThrowable(failure.getValue()).getCode()); + } + + static void assertUnavailable(Runnable action) { + try { + action.run(); + fail("Scan admission must be closed"); + } catch (StatusRuntimeException e) { + assertEquals(Status.Code.UNAVAILABLE, e.getStatus().getCode()); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ServerSuiteTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ServerSuiteTest.java index 2d115e78ed..9804cde7b0 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ServerSuiteTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ServerSuiteTest.java @@ -17,7 +17,11 @@ package org.apache.hugegraph.store.service; +import org.apache.hugegraph.rocksdb.access.RocksDBScanIteratorCloseTest; import org.apache.hugegraph.store.business.OrderedMultiPartitionIteratorTest; +import org.apache.hugegraph.store.node.grpc.query.AggregativeQueryLifecycleTest; +import org.apache.hugegraph.store.node.grpc.query.UnaryQueryLifecycleTest; +import org.apache.hugegraph.store.node.grpc.BatchScanHalfCloseTest; import org.junit.runner.RunWith; import org.junit.runners.Suite; @@ -25,7 +29,16 @@ @RunWith(Suite.class) @Suite.SuiteClasses({ - OrderedMultiPartitionIteratorTest.class + RocksDBScanIteratorCloseTest.class, + OrderedMultiPartitionIteratorTest.class, + ScanContextCleanupTest.class, + GrpcShutdownBarrierTest.class, + PartitionEngineShutdownTest.class, + AggregativeQueryLifecycleTest.class, + UnaryQueryLifecycleTest.class, + BatchScanHalfCloseTest.class, + ScanLifecycleTest.class, + GraphPartitionScanLifecycleTest.class }) @Slf4j