[common] Release the BTree index file on every path that abandons the writer

paimon-common/src/main/java/org/apache/paimon/globalindex/btree/BTreeIndexWriter.java

@@ -30,11 +30,13 @@
 import org.apache.paimon.sst.BlockHandle;
 import org.apache.paimon.sst.BloomFilterHandle;
 import org.apache.paimon.sst.SstFileWriter;
+import org.apache.paimon.utils.IOUtils;
 import org.apache.paimon.utils.LazyField;
 import org.apache.paimon.utils.RoaringNavigableMap64;
 
 import javax.annotation.Nullable;
 
+import java.io.Closeable;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -68,10 +70,11 @@
  * <p>For efficiency, we combine entries with the same keys and store a compact list of row ids for
  * each key.
  */
-public class BTreeIndexWriter implements GlobalIndexSingleColumnWriter {
+public class BTreeIndexWriter implements GlobalIndexSingleColumnWriter, Closeable {
 
     private final String fileName;
     private final PositionOutputStream out;
+    private boolean closed;
 
     private final SstFileWriter writer;
     private final KeySerializer keySerializer;
@@ -92,10 +95,15 @@ public BTreeIndexWriter(
             throws IOException {
         this.fileName = indexFileWriter.newFileName(BTreeGlobalIndexerFactory.IDENTIFIER);
         this.out = indexFileWriter.newOutputStream(this.fileName);
-        this.keySerializer = keySerializer;
-        this.comparator = keySerializer.createComparator();
-        // todo: we may enable bf to accelerate equal and in predicate in the future
-        this.writer = new SstFileWriter(out, blockSize, null, compressionFactory);
+        try {
+            this.keySerializer = keySerializer;
+            this.comparator = keySerializer.createComparator();
+            // todo: we may enable bf to accelerate equal and in predicate in the future
+            this.writer = new SstFileWriter(out, blockSize, null, compressionFactory);
+        } catch (RuntimeException e) {
+            IOUtils.closeQuietly(out);
+            throw e;
+        }
     }
 
     @Override
@@ -164,7 +172,10 @@ public List<ResultEntry> finish() {
             writer.writeSlice(footerEncoding);
 
             out.close();
+            closed = true;
         } catch (IOException e) {
+            IOUtils.closeQuietly(out);
+            closed = true;
             throw new RuntimeException("Error in closing BTree index writer", e);
         }
 
@@ -202,4 +213,17 @@ private BlockHandle writeNullBitmap() throws IOException {
 
         return nullBitmapHandle;
     }
+
+    /**
+     * Releases the output stream for a build that is abandoned without {@link #finish()}. The owner
+     * cleanup paths reach a writer only through {@code instanceof AutoCloseable}, so without this
+     * the stream opened in the constructor stays open for the life of the process.
+     */
+    @Override
+    public void close() throws IOException {
+        if (!closed) {
+            closed = true;
+            out.close();
+        }
+    }
 }

paimon-common/src/test/java/org/apache/paimon/globalindex/btree/BTreeIndexWriterCloseTest.java

@@ -0,0 +1,161 @@
+/*
+ * 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.paimon.globalindex.btree;
+
+import org.apache.paimon.compression.BlockCompressionFactory;
+import org.apache.paimon.fs.PositionOutputStream;
+import org.apache.paimon.globalindex.KeySerializer;
+import org.apache.paimon.globalindex.io.GlobalIndexFileWriter;
+import org.apache.paimon.types.IntType;
+import org.apache.paimon.utils.IOUtils;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Tests that {@link BTreeIndexWriter} releases the file it opened, on every path that abandons it.
+ */
+public class BTreeIndexWriterCloseTest {
+
+    /**
+     * finish() writes the null bitmap, the bloom filter, the index block and the footer before
+     * closing. A write that fails part way through - a full disk is the ordinary cause - must not
+     * leave the file open.
+     */
+    @Test
+    public void testFinishReleasesTheFileWhenWritingFails() throws IOException {
+        AtomicInteger closed = new AtomicInteger();
+        BTreeIndexWriter writer =
+                new BTreeIndexWriter(
+                        failingWriter(closed),
+                        KeySerializer.create(new IntType()),
+                        1024,
+                        (BlockCompressionFactory) null);
+
+        for (int i = 0; i < 200; i++) {
+            writer.write(i, (long) i);
+        }
+
+        assertThatThrownBy(writer::finish)
+                .isInstanceOf(RuntimeException.class)
+                .hasMessageContaining("Error in closing BTree index writer");
+
+        assertThat(closed).hasValue(1);
+    }
+
+    /**
+     * A build can be abandoned before finish() ever runs - the row source fails, or a Flink or
+     * Spark index build task is cancelled. The owner cleanup paths, PkSortedIndexFile#build and
+     * SortedSingleColumnIndexWriter#close, release the writer only if they can see it as an
+     * AutoCloseable, and skip it silently otherwise. This is that exact idiom.
+     */
+    @Test
+    public void testTheOwnerCleanupPathReleasesAnAbandonedWriter() throws IOException {
+        AtomicInteger closed = new AtomicInteger();
+        BTreeIndexWriter writer =
+                new BTreeIndexWriter(
+                        failingWriter(closed),
+                        KeySerializer.create(new IntType()),
+                        1024,
+                        (BlockCompressionFactory) null);
+        writer.write(1, 1L);
+
+        if (writer instanceof AutoCloseable) {
+            IOUtils.closeQuietly((AutoCloseable) writer);
+        }
+
+        assertThat(closed).hasValue(1);
+    }
+
+    /**
+     * SortedGlobalIndexWriter holds the task writer in a try-with-resources, so close() can follow
+     * a successful finish(). It must not close the stream a second time.
+     */
+    @Test
+    public void testCloseIsIdempotent() throws IOException {
+        AtomicInteger closed = new AtomicInteger();
+        BTreeIndexWriter writer =
+                new BTreeIndexWriter(
+                        failingWriter(closed),
+                        KeySerializer.create(new IntType()),
+                        1024,
+                        (BlockCompressionFactory) null);
+        writer.write(1, 1L);
+
+        writer.close();
+        writer.close();
+
+        assertThat(closed).hasValue(1);
+    }
+
+    private static GlobalIndexFileWriter failingWriter(AtomicInteger closed) {
+        return new GlobalIndexFileWriter() {
+
+            @Override
+            public String newFileName(String prefix) {
+                return "test-btree" + prefix;
+            }
+
+            @Override
+            public PositionOutputStream newOutputStream(String fileName) {
+                return new PositionOutputStream() {
+
+                    private long pos;
+                    private final long capacity = 1500L;
+
+                    @Override
+                    public long getPos() {
+                        return pos;
+                    }
+
+                    @Override
+                    public void write(int b) throws IOException {
+                        write(new byte[] {(byte) b}, 0, 1);
+                    }
+
+                    @Override
+                    public void write(byte[] b) throws IOException {
+                        write(b, 0, b.length);
+                    }
+
+                    @Override
+                    public void write(byte[] b, int off, int len) throws IOException {
+                        if (pos + len > capacity) {
+                            throw new IOException("no space left on device");
+                        }
+                        pos += len;
+                    }
+
+                    @Override
+                    public void flush() {}
+
+                    @Override
+                    public void close() {
+                        closed.incrementAndGet();
+                    }
+                };
+            }
+        };
+    }
+}