[DFAE] Fix memory allocation failure in write request path to fail request instead of shard

[mgodwan]

Fix memory allocation failure in write request path to fail request instead of shard

Description

[Describe what this change achieves]

Related Issues

Resolves #[Issue number to be closed when this PR is merged]

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Resource Cleanup on Failure If writer.addDoc(doc) throws an unexpected exception (instead of returning a Failure), the test will exit before doc.close(), setPoolLimit reset, and writer.close() are called, leaking Arrow buffers and leaving the allocator pool constrained for subsequent tests. Consider wrapping the assertions in a try/finally that resets the pool limit and closes the writer/doc. nativeAllocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_INGEST, 1L); ParquetDocumentInput doc = new ParquetDocumentInput(); populateMetadataFields(doc); doc.addField(idField, 1); doc.addField(nameField, "alice"); doc.addField(scoreField, 100L); doc.setRowId(DocumentInput.ROW_ID_FIELD, 0); // addDoc must translate the OOM into a Failure result rather than propagating the throw. WriteResult result = writer.addDoc(doc); assertTrue("expected a Failure result but was: " + result, result instanceof WriteResult.Failure); WriteResult.Failure failure = (WriteResult.Failure) result; assertTrue( "expected an Arrow OutOfMemoryException cause but was: " + failure.cause(), failure.cause() instanceof OutOfMemoryException ); doc.close(); // Relieve the limit before teardown so writer/allocator cleanup can proceed cleanly. nativeAllocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_INGEST, Long.MAX_VALUE); writer.close();

[github-actions]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Roll back partial VSR state on OOM Catching OutOfMemoryException and returning a Failure is correct, but if vsrManager.addDocument partially wrote data into the VSR before the allocation failure, the VSR may be in an inconsistent state. Ensure the VSR is properly rolled back / row count is reverted for the partially-added document before transitioning to PENDING_ROLLBACK, otherwise a subsequent caller-driven rollback may operate on a corrupted VSR state. sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java [155-158] try { vsrManager.addDocument(d); } catch (MismatchedInputException | OutOfMemoryException e) { + vsrManager.rollbackLastDocument(); state = WriterState.PENDING_ROLLBACK; return new WriteResult.Failure(e, -1, -1, -1); } Suggestion importance[1-10]: 4 __ Why: The suggestion raises a valid concern about potential partial VSR state after an OOM, but the PR comment indicates the caller-driven rollback handles this, and the invented rollbackLastDocument() method may not exist. The concern is speculative without deeper knowledge of VSRManager behavior.	Low

[Bukhtawar]

Thanks lgtm

[github-actions]

❌ Gradle check result for c647ecc: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

✅ Gradle check result for c647ecc: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.43%. Comparing base (7fa9b56) to head (c647ecc).
⚠️ Report is 40 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22415      +/-   ##
============================================
- Coverage     73.45%   73.43%   -0.02%     
+ Complexity    76153    76121      -32     
============================================
  Files          6076     6076              
  Lines        345518   345518              
  Branches      49733    49733              
============================================
- Hits         253815   253747      -68     
- Misses        71486    71515      +29     
- Partials      20217    20256      +39     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[rayshrey]

Nit - can we also disable rebalancer to avoid flakiness ?

[rayshrey]

While going through the rollback flow, noticed that we might enter into a corrupt state, as we are just setting row counts while rolling back our VSR. Any stale values are not cleared and if the same VSR is re-used (which it is as we change the status back to ACTIVE from PENDING_ROLLBACK after the reset the row count), values from previous document might remain (in case those fields are empty in the new doc).

Trying to validate this scenario using an IT, but since it is an existing issue, not blocking this PR - will open an issue once able to validate using IT.

[rayshrey]

Created an issue here - #22417