somandal commented on code in PR #17028: URL: https://github.com/apache/pinot/pull/17028#discussion_r2467076449
########## pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/MultipleTreesBuilderCloseTest.java: ########## @@ -0,0 +1,397 @@ +/** + * 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.pinot.segment.local.startree.v2.builder; + +import java.io.File; +import java.io.IOException; +import java.lang.reflect.Field; +import java.util.Arrays; +import java.util.List; +import org.apache.commons.configuration2.PropertiesConfiguration; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.startree.StarTreeBuilderUtils; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.V1Constants; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.index.startree.StarTreeV2Constants; +import org.apache.pinot.spi.config.table.StarTreeIndexConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.testng.Assert.*; + +/** + * Unit test for MultipleTreesBuilder.close() method to verify exception handling + * when cleanup operations fail. + */ +public class MultipleTreesBuilderCloseTest { + private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "MultipleTreesBuilderCloseTest"); + private static final File INDEX_DIR = new File(TEMP_DIR, "testSegment"); + + @BeforeMethod + public void setUp() throws Exception { + FileUtils.deleteQuietly(TEMP_DIR); + FileUtils.forceMkdir(TEMP_DIR); + } + + @AfterMethod + public void tearDown() { + FileUtils.deleteQuietly(TEMP_DIR); + } + + @Test + public void testBuildFailureThenCloseFailureWithSuppressedException() throws Exception { + // This test verifies that when build() fails and close() also fails, + // the close() exception is added as a suppressed exception to the build() exception + + // Build a test segment with star-tree + buildTestSegment(); + + // Create a MultipleTreesBuilder with invalid config to force build() to fail + List<StarTreeV2BuilderConfig> builderConfigs = createInvalidBuilderConfigs(); + + MultipleTreesBuilder builder = null; + Exception finalException = null; + + try { + builder = new MultipleTreesBuilder(builderConfigs, INDEX_DIR, MultipleTreesBuilder.BuildMode.OFF_HEAP); + + // Manually set up the scenario for close() to fail + // First, create the separator temp directory and existing metadata + Field separatorTempDirField = MultipleTreesBuilder.class.getDeclaredField("_separatorTempDir"); + separatorTempDirField.setAccessible(true); + + File segmentDir = INDEX_DIR.listFiles()[0]; + File manualSeparatorTempDir = new File(segmentDir, StarTreeV2Constants.EXISTING_STAR_TREE_TEMP_DIR); + FileUtils.forceMkdir(manualSeparatorTempDir); + separatorTempDirField.set(builder, manualSeparatorTempDir); + + // Create temp files to simulate existing star-tree files that need restoration + File tempIndexFile = new File(manualSeparatorTempDir, StarTreeV2Constants.INDEX_FILE_NAME); + File tempMapFile = new File(manualSeparatorTempDir, StarTreeV2Constants.INDEX_MAP_FILE_NAME); + FileUtils.touch(tempIndexFile); + FileUtils.touch(tempMapFile); + + // Create existing star-tree metadata using reflection + Field existingStarTreeMetadataField = MultipleTreesBuilder.class.getDeclaredField("_existingStarTreeMetadata"); + existingStarTreeMetadataField.setAccessible(true); + PropertiesConfiguration existingMetadata = new PropertiesConfiguration(); + existingMetadata.setProperty("test.key", "test.value"); + existingStarTreeMetadataField.set(builder, existingMetadata); + + // Make the metadata file read-only to force close() restoration to fail + File metadataFile = new File(segmentDir, V1Constants.MetadataKeys.METADATA_FILE_NAME); + if (metadataFile.exists()) { + metadataFile.setReadOnly(); + } + + // Use try-with-resources just like SegmentPreProcessor does + // This will automatically call close() if build() throws an exception + try (MultipleTreesBuilder autoCloseBuilder = builder) { + // This should fail due to invalid config, setting _starTreeCreationFailed = true + autoCloseBuilder.build(); + fail("Expected build() to throw an exception due to invalid config"); + } catch (Exception e) { + finalException = e; + // Verify that _starTreeCreationFailed was set to true + Field starTreeCreationFailedField = MultipleTreesBuilder.class.getDeclaredField("_starTreeCreationFailed"); + starTreeCreationFailedField.setAccessible(true); + assertTrue((Boolean) starTreeCreationFailedField.get(builder), + "_starTreeCreationFailed should be true after build() fails"); + } finally { + // Restore write permissions for cleanup + if (metadataFile.exists()) { + metadataFile.setWritable(true); + } + } + } finally { + // Ensure builder is closed even if test fails + if (builder != null) { + try { + // Restore permissions first + File segmentDir = INDEX_DIR.listFiles()[0]; + File metadataFile = new File(segmentDir, V1Constants.MetadataKeys.METADATA_FILE_NAME); + if (metadataFile.exists()) { + metadataFile.setWritable(true); + } + builder.close(); + } catch (Exception e) { + // Ignore cleanup exceptions + } + } + } + + // Verify that the final exception occurred (this includes both build and close exceptions) + assertNotNull(finalException, "Expected an exception from try-with-resources (build + close)"); + + // The key test: verify that finalException has buildException as a suppressed exception + // OR that finalException is the same as buildException with suppressed exceptions added + Throwable[] suppressedExceptions = finalException.getSuppressed(); + + // The close() exception should have suppressed exceptions + assertTrue(suppressedExceptions.length > 0, + "Expected close() exception to have suppressed exceptions, but found none. " + + "Final exception: " + finalException.getMessage()); + + // Verify the suppressed exception is related to the cleanup failure + boolean foundCleanupException = false; + for (Throwable suppressed : suppressedExceptions) { + if (suppressed.getMessage().contains("Could not reset") || suppressed.getMessage().contains("Permission denied") + || suppressed.getMessage().contains("Unable to save") || suppressed instanceof IOException) { + foundCleanupException = true; + break; + } + } + + assertTrue(foundCleanupException, + "Expected to find cleanup-related suppressed exception, but found: " + + Arrays.toString(suppressedExceptions)); + + // This scenario should trigger the re-throw condition in SegmentPreProcessor.processStarTrees() + // because finalException.getSuppressed().length > 0 + System.out.println("SUCCESS: Exception with " + suppressedExceptions.length Review Comment: oops, yes had added for testing, removed -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
