From be111a9c72675c5aeaef99795aa087485c03f1b9 Mon Sep 17 00:00:00 2001 From: Xinyuan Lin Date: Fri, 24 Jul 2026 20:27:02 -0700 Subject: [PATCH] fix(amber): close Iceberg reader streams leaked by iterator probes IcebergDocument's read iterator opened a Parquet reader (and the S3InputStream beneath it) inside hasNext, so probe-shaped consumers - isEmpty/nonEmpty/hasNext-only checks, e.g. VirtualDocumentSpec's "clear the document" test - opened a stream and then abandoned the iterator with no close path. The GC finalizer eventually reclaimed each one, logging "[S3InputStream] Unclosed input stream created by ..." warnings all over CI. Bounded reads had the same hole: the limit close only ran on a subsequent hasNext call that bounded consumers never make. Restructure the iterator so hasNext acquires nothing: * hasNext claims the next data file from FileScanTask.recordCount metadata alone - the same field the whole-file skip already trusts; planFiles() never splits files in Iceberg 1.9.2 and the write path is strictly append-only, so the count is exact; * the Parquet reader opens lazily in next() via openPendingFile(); * next() closes the reader deterministically the moment a bounded read (getRange) has served its last record; * all closes go through an idempotent closeCurrentReader() that also resets the record iterator, so a closed reader is never polled. Record sequences, hasNext semantics, and the incremental-snapshot refresh behavior are unchanged - only where streams open and close moved. Also close the RESTCatalog in IcebergRestCatalogIntegrationSpec's afterAll: Iceberg 1.9.2 tracks per-table FileIO instances and closes them with the catalog, removing the sibling "Unclosed S3FileIO instance" finalizer warnings from the same CI logs. --- .../IcebergRestCatalogIntegrationSpec.scala | 11 ++ .../result/iceberg/IcebergDocument.scala | 116 +++++++++++++----- 2 files changed, 95 insertions(+), 32 deletions(-) diff --git a/amber/src/test/integration/org/apache/texera/amber/storage/iceberg/IcebergRestCatalogIntegrationSpec.scala b/amber/src/test/integration/org/apache/texera/amber/storage/iceberg/IcebergRestCatalogIntegrationSpec.scala index 807591dde59..581013ba217 100644 --- a/amber/src/test/integration/org/apache/texera/amber/storage/iceberg/IcebergRestCatalogIntegrationSpec.scala +++ b/amber/src/test/integration/org/apache/texera/amber/storage/iceberg/IcebergRestCatalogIntegrationSpec.scala @@ -47,6 +47,17 @@ class IcebergRestCatalogIntegrationSpec extends AnyFlatSpec with BeforeAndAfterA ) } + override def afterAll(): Unit = { + // RESTCatalog is Closeable: closing releases its HTTP client and the + // S3FileIO instances it created for table operations. Left unclosed, + // the finalizer reclaims them and logs "Unclosed S3FileIO instance" + // warnings with full stack traces after the spec finishes. + if (restCatalog != null) { + restCatalog.close() + } + super.afterAll() + } + behavior of "Iceberg REST catalog" it should "round-trip table metadata via the REST catalog" in { diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/result/iceberg/IcebergDocument.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/result/iceberg/IcebergDocument.scala index 182da2baac2..c9c00382ee4 100644 --- a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/result/iceberg/IcebergDocument.scala +++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/result/iceberg/IcebergDocument.scala @@ -186,6 +186,53 @@ private[storage] class IcebergDocument[T >: Null <: AnyRef]( private var currentRecordIterator: Iterator[Record] = Iterator.empty private var currentRecordIteratorCloser: AutoCloseable = () => () + // Next file to read, claimed from usableFileIterator by hasNext but + // not opened until next() actually needs a record. hasNext answers + // from the file's recordCount metadata alone, so probes such as + // isEmpty/nonEmpty (which abandon the iterator right after hasNext) + // never leave a Parquet reader / S3 stream open behind them. + private var pendingFile: Option[FileScanTask] = None + + // Idempotent release of the active file's reader (and the S3 stream + // beneath it). Also resets the record iterator so a closed reader is + // never polled again. + private def closeCurrentReader(): Unit = { + currentRecordIteratorCloser.close() + currentRecordIteratorCloser = () => () + currentRecordIterator = Iterator.empty + } + + // Open the file claimed by hasNext and point currentRecordIterator at + // its records, applying the one-off partial skip for the first file. + // Only next() calls this, keeping hasNext free of any Parquet/S3 + // resource acquisition. + private def openPendingFile(): Unit = { + val task = pendingFile.getOrElse( + throw new IllegalStateException("no pending file to open") + ) + pendingFile = None + val schemaToUse = columns match { + case Some(cols) => tableSchema.select(cols.asJava) + case None => tableSchema + } + // Release the prior file's reader before opening the next. + closeCurrentReader() + val nextIter = IcebergUtil.readDataFileAsIterator( + task.file(), + schemaToUse, + table.get + ) + currentRecordIteratorCloser = nextIter + currentRecordIterator = nextIter.asScala + + // Skip records within the file if necessary + val recordsToSkipInFile = from - numOfSkippedRecords + if (recordsToSkipInFile > 0) { + currentRecordIterator = currentRecordIterator.drop(recordsToSkipInFile) + numOfSkippedRecords += recordsToSkipInFile + } + } + // Util function to load the table's metadata private def loadTableMetadata(): Option[Table] = { IcebergUtil.loadTableMetadata( @@ -266,8 +313,7 @@ private[storage] class IcebergDocument[T >: Null <: AnyRef]( override def hasNext: Boolean = { if (numOfReturnedRecords >= totalRecordsToReturn) { // Caller-imposed limit reached; release the active file's reader. - currentRecordIteratorCloser.close() - currentRecordIteratorCloser = () => () + closeCurrentReader() return false } @@ -277,48 +323,54 @@ private[storage] class IcebergDocument[T >: Null <: AnyRef]( return true } - if (!usableFileIterator.hasNext) { - usableFileIterator = seekToUsableFile() - } + // The active file (if any) is exhausted; release its reader before + // deciding from metadata whether more records exist. + closeCurrentReader() - while (!currentRecordIterator.hasNext && usableFileIterator.hasNext) { - val nextFile = usableFileIterator.next() - val schemaToUse = columns match { - case Some(cols) => tableSchema.select(cols.asJava) - case None => tableSchema + if (pendingFile.isEmpty) { + if (!usableFileIterator.hasNext) { + usableFileIterator = seekToUsableFile() } - // Release the prior file's reader before opening the next. - currentRecordIteratorCloser.close() - val nextIter = IcebergUtil.readDataFileAsIterator( - nextFile.file(), - schemaToUse, - table.get - ) - currentRecordIteratorCloser = nextIter - currentRecordIterator = nextIter.asScala - - // Skip records within the file if necessary - val recordsToSkipInFile = from - numOfSkippedRecords - if (recordsToSkipInFile > 0) { - currentRecordIterator = currentRecordIterator.drop(recordsToSkipInFile) - numOfSkippedRecords += recordsToSkipInFile + // Claim the next file that still has usable records, judging by + // recordCount metadata alone (the same field the whole-file skip + // in seekToUsableFile already relies on). The partial skip + // `from - numOfSkippedRecords` can only be non-zero for the first + // claimed file: seekToUsableFile's dropWhile guarantees that file + // satisfies recordCount > partial skip, and openPendingFile zeroes + // the skip before any later file is judged here. + while (pendingFile.isEmpty && usableFileIterator.hasNext) { + val task = usableFileIterator.next() + if (task.file().recordCount() > from - numOfSkippedRecords) { + pendingFile = Some(task) + } } } - val hasMore = currentRecordIterator.hasNext - if (!hasMore) { - // All files exhausted; release the last file's reader. - currentRecordIteratorCloser.close() - currentRecordIteratorCloser = () => () - } - hasMore + pendingFile.nonEmpty } override def next(): T = { if (!hasNext) throw new NoSuchElementException("No more records available") + // hasNext only claims files by their metadata; the Parquet reader is + // opened lazily here. The loop is defensive: should a claimed file + // yield nothing after the partial skip, hasNext claims the next one + // (or reports exhaustion). + while (!currentRecordIterator.hasNext) { + openPendingFile() + if (!currentRecordIterator.hasNext && !hasNext) { + throw new NoSuchElementException("No more records available") + } + } + val record = currentRecordIterator.next() numOfReturnedRecords += 1 + if (numOfReturnedRecords >= totalRecordsToReturn) { + // Bounded read fully served: release the reader now instead of + // waiting for a further hasNext call that bounded consumers + // (e.g. getRange) rarely make. + closeCurrentReader() + } val schemaToUse = columns match { case Some(cols) => tableSchema.select(cols.asJava) case None => tableSchema