Repository navigation
fix: return errors instead of panicking on malformed JNI arguments - #7186
Conversation
Several JNI entry points converted their arguments with unwrap() or expect(), so a malformed argument surfaced as a CometNativeException carrying only the panic message. Return a CometError::Internal that names the bad argument instead: - convert_datatype_arrays reports which column's data type failed to deserialize. - writeSortedFileNative, traceBegin, traceEnd and logMemoryUsage report which string argument could not be read. - get_execution_context rejects a null handle. The exception class stays CometNativeException.
sunchao
left a comment
There was a problem hiding this comment.
Reviewed full head 35e2e2bf6461003c41de2dfa2951e32c8f0ec76b. No introduced P1/P2 issues found within this review.
The PR is non-draft. The full PR diff against supplied base 69111f68a5ff62904342910c445d191f53615aff, through merge-base 785fdbcb4cbf9057dfcfae3bd9dd45a77bb505bb, contains one commit changing native/core/src/execution/jni_api.rs. The base-only memory-management commit is not a PR-introduced reversal. There are no stacked prerequisites. The supplied snapshot and live discussion checks contained no reviews, issue comments, inline comments, or review threads.
Routed skills: review-comet-pr, review-comet-ffi-pr, and review-comet-shuffle-pr. Read AGENTS.md and the relevant contributor documentation.
Summary
- Prior state and problem: Malformed serialized types, unreadable strings, and null execution handles reached
unwrap()orexpect(). The JNI boundary caught these panics, but diagnostics lacked argument context and invoked panic logging. - Design approach: Replace those panic paths with fallible helpers and propagate
CometError::Internalthrough the existing JNI exception boundary. - Correctness: Checked every changed caller, exception mapping, and early-return cleanup. Valid inputs retain the same conversions. JNI array guards still drop on errors, and decoding fails before the affected writer or converter performs its work. No expression evaluation, partitioning, or row-format semantics change.
- Compatibility analysis: JNI signatures and protobuf encoding remain unchanged. Both old panic errors and new internal errors map to
CometNativeException. Checked Spark'sShuffleWriteProcessorsources at configured versions 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0: each stops the writer on an exception and rethrows it without depending on these diagnostic messages. The affected file on latest release branchbranch-1.2, at889856dc5af821de1f7c6ca2e9538bd663431d04, is byte-identical to the PR merge-base, so the release-relative changes are exactly this intended diagnostic change. - Key design decisions: Retaining the established exception class avoids changing JVM failure classification.
setShufflePartitionPusherpreserves its earlier explicitNullPointerExceptioncheck. Semantic validation insideto_arrow_datatyperemains unchanged, including its existing panic cases. - Implementation sketch:
deserialize_arrow_datatypeadds the column index to decode errors.string_argidentifies the export and argument.get_execution_contextreturns a result, and all three callers propagate it. Three Rust tests cover successful decoding, malformed bytes, and a null handle. - Performance: Error-message formatting occurs only on failure. Successful calls retain the existing decoding, string conversion, and pointer check without additional batch copies, buffering, or allocations for diagnostics. Avoiding unwinding and panic-backtrace logging improves the malformed-input path structurally, but no throughput claim was benchmarked.
- Design: The change is localized and preserves existing ownership, resource lifecycle, and shuffle contracts. The current FFI and shuffle documentation remains accurate.
- Abstraction & complexity: The two small conversion helpers centralize repeated diagnostics without introducing a broader framework. Returning
CometResultfrom the datatype-array helper makes error propagation straightforward. - Behavioral changes worth calling out: The covered malformed arguments now produce contextual error messages through ordinary error propagation instead of panic handling. Valid query results and shuffle output remain unchanged.
- Suggested improvements: No corrective change meeting the P1/P2 reporting bar was identified.
Exact-head CI: title checks, labeling, CodeQL, and Analyze Actions passed. Comet CI remained in Preflight. No failed check was reported, but there was no completed Rust build/test or Spark-suite verdict.
Validation: git diff --check passed. Two disposable harness tests passed using the unchanged production string and handle helpers, jni 0.22.4, and OpenJDK 21 with -Xcheck:jni. They exercised null-string diagnostics for all five call sites, empty/Unicode/embedded-NUL strings, null handles, and valid pointer identity. The harness used minimal error/context scaffolding and does not replace Comet integration tests. The focused repository test could not start because the offline registry lacked locked dependency aws-config 1.13.0. Consequently, the added repository tests and end-to-end Comet/Spark suites were not run locally. The project checkout remains unchanged.
|
Thanks @sunchao |
Which issue does this PR close?
Closes #7183.
Rationale for this change
Some JNI entry points convert their arguments with
unwrap()/expect(), so a malformed argument panics.try_unwrap_or_throwturns that panic into aCometNativeException, but the exception carries only the panic message, such ascalled `Result::unwrap()` on an `Err` value: ..., without saying which argument was bad. These conversions should return regular errors, the same way the other arguments of these entry points already do.What changes are included in this PR?
All changes are in
native/core/src/execution/jni_api.rs. Each conversion now returns aCometError::Internalthrough?, so the exception class staysCometNativeExceptionand only the message changes:convert_datatype_arrays, used bycolumnarToRowInit,writeSortedFileNativeandwriteSortedPartitionsNative, now decodes each type with a newdeserialize_arrow_datatypehelper. When decoding fails, the error says which column it was, for exampleCannot deserialize the data type of column 2: ....file_pathandcompression_codecarguments ofwriteSortedFileNative, and the name arguments oftraceBegin,traceEndandlogMemoryUsage, are now read with a newstring_arghelper. When a string cannot be read, for example because it is null, the error names the argument and the entry point, for exampleCannot read argument 'event' of traceBegin: ....get_execution_contextnow returns an error for a null handle (Native execution context handle is null). Its callers areexecutePlan,getShufflePartitionOffsetsandsetShufflePartitionPusher.setShufflePartitionPusheralready rejected a zero handle with aNullPointerExceptionbefore reaching this check, and it still does.A data type that decodes but is invalid, such as an unknown type id, can still panic in
serde::to_arrow_datatype. That function is shared with the planner and has many callers, so making it fallible is left out of this PR.How are these changes tested?
New Rust unit tests in
jni_api.rs, none of which need a JVM:deserialize_arrow_datatype_decodes_a_serialized_typedeserialize_arrow_datatype_rejects_malformed_bytesget_execution_context_rejects_a_null_handleThe two error tests also check that the error is
CometError::Internal, which maps toCometNativeException.The string conversions need a JVM to call, so no new tests cover them. They are covered by the existing tests of
CometNativeColumnarToRowSuiteandCometShuffleSuite.