Move FileIO close from RecordWriter to RecordWriterManager

[dejii]

Summary

• Remove FileIO field and close logic from RecordWriter - it no longer owns the FileIO lifecycle
• Move FileIO cleanup to RecordWriterManager.close(), which collects unique FileIO instances across all destinations and closes each once in a finally block
• Add tests covering the fix

Root Cause

• table.io() returns the catalog's shared FileIO instance. Catalogs like RESTSessionCatalog track and reuse a single FileIO across tables via FileIOTracker (ref). When RecordWriter.close() called io.close(), it shut down the shared connection pool.
• This manifests when using dynamic destinations: RecordWriterManager.close() iterates destinations and calls cache.invalidateAll() on each, which triggers RecordWriter.close() via the removal listener. The first destination's writers close successfully but kill the shared pool - subsequent destinations fail with:

IllegalStateException: Encountered 12 failed writer(s).
  Caused by: IOException: Failed to close PARQUET writer for table ...
    Caused by: IOException: Connection pool shut down
      Caused by: IllegalStateException: Connection pool shut down
  Suppressed: RuntimeException: Encountered an error when closing data writer for table ...
  Suppressed: RuntimeException: ...  

The fix moves FileIO close to RecordWriterManager.close(), which deduplicates by identity and closes each shared instance exactly once after all writers have flushed.

Related: #36438

I can confirm this works fix works great when running in Dataflow.

---

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

• Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
• Update CHANGES.md with noteworthy changes.
• If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refactors the management of FileIO instances within the Iceberg connector to resolve a critical issue where shared connection pools were being prematurely shut down. By centralizing the FileIO closing logic in RecordWriterManager, the system now correctly handles the lifecycle of these shared resources, ensuring stability and preventing Connection pool shut down errors, especially in scenarios involving dynamic destinations.

Highlights

• FileIO Lifecycle Management: The responsibility for closing FileIO instances has been shifted from RecordWriter to RecordWriterManager to prevent premature shutdown of shared connection pools.
• RecordWriter Simplification: The RecordWriter class no longer holds a FileIO field or contains logic for closing FileIO instances, simplifying its role.
• Centralized FileIO Closure: The RecordWriterManager.close() method now collects all unique FileIO instances across different destinations and ensures each is closed exactly once in a finally block, addressing issues with dynamic destinations.
• Enhanced Testing: New tests have been added to verify that RecordWriter does not close shared FileIO instances and that RecordWriterManager correctly handles FileIO closure for multiple writers sharing the same FileIO.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriter.java
  • Removed the FileIO field and its associated import.
  • Eliminated the logic responsible for closing FileIO within the close() method.
  • Updated comments to reflect that FileIO lifecycle is now managed externally.
• sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java
  • Added imports for Closeable, HashSet, and Set.
  • Introduced a finally block in the close() method to iterate through destinations, identify unique FileIO instances, and close them safely.
  • Moved the destinations.clear() call into the new finally block.
• sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/RecordWriterManagerTest.java
  • Added an import for org.apache.iceberg.data.Record.
  • Renamed an existing test method to testRecordWriterDoesNotCloseSharedFileIO and updated its assertion to confirm FileIO is not closed by RecordWriter.
  • Added a new test testMultipleWritersSharingFileIOSurviveBatchClose to verify shared FileIO instances remain open until RecordWriterManager closes them.
  • Added another new test testRecordWriterManagerCloseFlushesAllDestinations to ensure RecordWriterManager.close() correctly flushes data files for all destinations.

Activity

• The author confirmed that this fix works as expected when running in Dataflow.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[github-actions]

Assigning reviewers:

R: @Abacn for label java.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[ahmedabu98]

Did some digging to verify some assumptions. It looks like the catalog creates a new TableOperations instance per table:
BaseMetastoreCatalog#L48
RESTSessionCatalog#L575

Update: Ahh but RESTCatalog indeed uses the same FileIO instance across each table:
RESTSessionCatalog#L1167

I think your change is good, this is just some further nuance:
I'm seeing some other catalogs may have other implementations. Here is GlueCatalog creating a new FileIO per TableOps instance: https://github.com/apache/iceberg/blob/ed8a16bbeb549b0286d3c229beb5a0cf165f2f4b/aws/src/main/java/org/apache/iceberg/aws/glue/GlueTableOperations.java#L113-L118
And keeping track of each instance using a Catalog-level FileIOTracker: https://github.com/apache/iceberg/blob/ed8a16bbeb549b0286d3c229beb5a0cf165f2f4b/aws/src/main/java/org/apache/iceberg/aws/glue/GlueCatalog.java#L244
The tracker is in charge of closing all instances. But the tracker itself is only closed when the Catalog itself is closed. In this case, we'd need to close it once per table.

But your change still covers this case because it iterates over each table and checks for unique IOs to close

[ahmedabu98]

Thanks Deji! Had a few nits regarding the comments, and one request to add an extra check. Overall looks great

[ahmedabu98]

Edited some of my suggestions. It's hard to say with certainty whether it is one FileIO per table, or if a FileIO is shared across all tables. It appears the catalog has the freedom to implement it either way

[dejii]

@ahmedabu98 Thanks for the review! I've addressed your comments:

• Updated the comments to remove the "catalog's" qualifier - as you noted, whether FileIO is shared depends on the catalog implementation
• Added FileIO close verification to the manager test.