From d3a58b5718433b205f9bb0e3fd6da2918f86d35c Mon Sep 17 00:00:00 2001 From: Mayank Shrivastava Date: Thu, 27 Aug 2026 20:16:49 -0700 Subject: [PATCH] Fix protobuf temp directory leak in ProtoBufUtils ProtoBufUtils.getFileCopiedToLocal() created temporary directories with prefix "pinot-protobuf" that were never cleaned up. Every call to ProtoBufMessageDecoder.init(), ProtoBufRecordReader.init(), or ProtoBufCodeGenMessageDecoder.init() leaked one directory under java.io.tmpdir. For descriptor file callers (ProtoBufMessageDecoder, ProtoBufRecordReader): replace the copy-to-temp-then-open pattern with readDescriptorFileBytes() which streams bytes directly via PinotFS.open() - no temp files at all. For the JAR caller (ProtoBufCodeGenMessageDecoder): resolve the JAR to a local File inline. Local JARs are used directly without copying. Remote JARs are copied to a temp directory that intentionally persists for the decoder's lifetime because the JVM may lazily resolve classes from the JAR at decode time via the URLClassLoader chain. Since StreamMessageDecoder does not extend Closeable, there is no lifecycle hook to clean up, but this is a one-time-per-consumer-init cost and consumers are long-lived. Removed dead methods: getFileCopiedToLocal(), getDescriptorFileInputStream(), createLocalFile(), withLocalFile(), and FileAction from ProtoBufUtils. Renamed loadClass(File) to createClassLoader(File) for clarity. Testing: - Added ProtoBufTempFileLeakTest with 5 tests covering all three code paths. Each test snapshots pinot-protobuf* directories in java.io.tmpdir before the operation, performs the operation with functional correctness assertions (decoding a message and checking field values), then asserts no new temp directories remain afterward. - testMessageDecoderInitDoesNotLeakTempDir: simple descriptor via ProtoBufMessageDecoder, verifies decode of sample record fields. - testMessageDecoderComplexDescriptorDoesNotLeakTempDir: complex nested descriptor via ProtoBufMessageDecoder. - testCodeGenDecoderInitDoesNotLeakTempDir: simple JAR via ProtoBufCodeGenMessageDecoder, verifies decode after init cleanup. - testCodeGenDecoderComplexJarDoesNotLeakTempDir: complex JAR with nested/repeated/map types via ProtoBufCodeGenMessageDecoder. - testRecordReaderLifecycleDoesNotLeakTempDir: full ProtoBufRecordReader lifecycle (init, read, close) with delimited protobuf data file. - All 5 tests confirmed to fail before the fix (each leaking 1 temp directory) and pass after. Full module suite: 172/172 tests pass. --- .../ProtoBufCodeGenMessageDecoder.java | 56 ++++- .../protobuf/ProtoBufMessageDecoder.java | 10 +- .../protobuf/ProtoBufRecordReader.java | 4 +- .../inputformat/protobuf/ProtoBufUtils.java | 43 +--- .../protobuf/ProtoBufTempFileLeakTest.java | 211 ++++++++++++++++++ .../protobuf/ProtoBufUtilsTest.java | 3 +- 6 files changed, 279 insertions(+), 48 deletions(-) create mode 100644 pinot-plugins/pinot-input-format/pinot-protobuf/src/test/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufTempFileLeakTest.java diff --git a/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufCodeGenMessageDecoder.java b/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufCodeGenMessageDecoder.java index b80977c9f796..69ccecbd8e3c 100644 --- a/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufCodeGenMessageDecoder.java +++ b/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufCodeGenMessageDecoder.java @@ -24,18 +24,33 @@ import java.io.File; import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; +import java.net.URI; import java.net.URL; import java.net.URLClassLoader; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.Arrays; import java.util.Map; import java.util.Set; +import org.apache.commons.io.FileUtils; import org.apache.pinot.plugin.inputformat.protobuf.codegen.MessageCodeGen; import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.filesystem.PinotFS; +import org.apache.pinot.spi.filesystem.PinotFSFactory; import org.apache.pinot.spi.stream.StreamMessageDecoder; import org.codehaus.janino.SimpleCompiler; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +/// Protobuf stream decoder that uses Janino-compiled code for extraction. The generated code is compiled at +/// init time from the protobuf descriptor found in the user-supplied JAR. The [URLClassLoader] and its backing +/// JAR file must remain available for the lifetime of this decoder because the JVM may lazily resolve classes +/// referenced by the generated code at decode time. Since [StreamMessageDecoder] does not extend +/// [java.io.Closeable], these resources are not explicitly released. public class ProtoBufCodeGenMessageDecoder implements StreamMessageDecoder { + private static final Logger LOGGER = LoggerFactory.getLogger(ProtoBufCodeGenMessageDecoder.class); + public static final String PROTOBUF_JAR_FILE_PATH = "jarFile"; public static final String PROTO_CLASS_NAME = "protoClassName"; private Method _decodeMethod; @@ -49,12 +64,18 @@ public void init(Map props, Set fieldsToRead, String top "Protocol Buffer Message class name must be provided"); String protoClassName = props.getOrDefault(PROTO_CLASS_NAME, ""); String jarPath = props.getOrDefault(PROTOBUF_JAR_FILE_PATH, ""); - ClassLoader protoMessageClsLoader = loadClass(jarPath); + File jarFile = resolveToLocalFile(jarPath); + ClassLoader protoMessageClsLoader = createClassLoader(jarFile); Descriptors.Descriptor descriptor = getDescriptorForProtoClass(protoMessageClsLoader, protoClassName); String codeGenCode = new MessageCodeGen().codegen(descriptor, fieldsToRead); Class recordExtractor = compileClass(protoMessageClsLoader, MessageCodeGen.EXTRACTOR_PACKAGE_NAME + "." + MessageCodeGen.EXTRACTOR_CLASS_NAME, codeGenCode); _decodeMethod = recordExtractor.getMethod(MessageCodeGen.EXTRACTOR_METHOD_NAME, byte[].class, GenericRow.class); + // NOTE: Do NOT close the URLClassLoader or delete the JAR file. The generated code may trigger + // lazy class resolution at decode time via the classloader chain (Janino -> URLClassLoader -> JAR). + // Closing prematurely would cause NoClassDefFoundError. For local JARs there is nothing to + // clean up. For remote JARs, the temp directory persists for the lifetime of this decoder + // (StreamMessageDecoder does not extend Closeable). } @Override @@ -75,12 +96,10 @@ public GenericRow decode(byte[] payload, int offset, int length, GenericRow dest return decode(payload, destination); } - public static ClassLoader loadClass(String jarFilePath) { + public static ClassLoader createClassLoader(File jarFile) { try { - File file = ProtoBufUtils.getFileCopiedToLocal(jarFilePath); - URL url = file.toURI().toURL(); - URL[] urls = new URL[]{url}; - return new URLClassLoader(urls); + URL url = jarFile.toURI().toURL(); + return new URLClassLoader(new URL[]{url}); } catch (Exception e) { throw new RuntimeException("Error loading protobuf class", e); } @@ -104,4 +123,29 @@ public static Descriptors.Descriptor getDescriptorForProtoClass(ClassLoader prot Class updateMessage = (Class) protoMessageClsLoader.loadClass(protoClassName); return (Descriptors.Descriptor) updateMessage.getMethod("getDescriptor").invoke(null); } + + /// Resolves a file path (URI string) to a local [File]. For local files (no scheme or `file://` scheme), + /// the original file is returned directly. For remote files, the file is copied to a local temporary + /// directory. The caller is responsible for the lifetime of the returned file - for remote files, the + /// backing temp directory is intentionally NOT cleaned up because the JAR must remain accessible for + /// lazy class loading at decode time. + private static File resolveToLocalFile(String filePath) + throws Exception { + URI fileURI = URI.create(filePath); + String scheme = fileURI.getScheme(); + if (scheme == null || PinotFSFactory.LOCAL_PINOT_FS_SCHEME.equals(scheme)) { + return new File(fileURI.getPath()); + } + PinotFS pinotFS = PinotFSFactory.create(scheme); + Path localTmpDir = Files.createTempDirectory(ProtoBufUtils.TMP_DIR_PREFIX); + try { + File localFile = new File(localTmpDir.toFile(), new File(fileURI.getPath()).getName()); + LOGGER.info("Copying protocol buffer JAR from {} to {}", filePath, localFile.getAbsolutePath()); + pinotFS.copyToLocalFile(fileURI, localFile); + return localFile; + } catch (Exception e) { + FileUtils.deleteDirectory(localTmpDir.toFile()); + throw e; + } + } } diff --git a/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufMessageDecoder.java b/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufMessageDecoder.java index cee7fac5c7ee..24b2b303d40f 100644 --- a/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufMessageDecoder.java +++ b/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufMessageDecoder.java @@ -24,7 +24,6 @@ import com.google.protobuf.DynamicMessage; import com.google.protobuf.Message; import java.io.IOException; -import java.io.InputStream; import java.util.Map; import java.util.Set; import org.apache.commons.lang3.StringUtils; @@ -48,19 +47,18 @@ public void init(Map props, Set fieldsToRead, String top "Protocol Buffer schema descriptor file must be provided"); _protoClassName = props.getOrDefault(PROTO_CLASS_NAME, ""); - InputStream descriptorFileInputStream = ProtoBufUtils.getDescriptorFileInputStream( - props.get(DESCRIPTOR_FILE_PATH)); - Descriptors.Descriptor descriptor = buildProtoBufDescriptor(descriptorFileInputStream); + byte[] descriptorBytes = ProtoBufUtils.readDescriptorFileBytes(props.get(DESCRIPTOR_FILE_PATH)); + Descriptors.Descriptor descriptor = buildProtoBufDescriptor(descriptorBytes); _recordExtractor = new ProtoBufRecordExtractor(); _recordExtractor.init(fieldsToRead, null); DynamicMessage dynamicMessage = DynamicMessage.getDefaultInstance(descriptor); _builder = dynamicMessage.newBuilderForType(); } - private Descriptors.Descriptor buildProtoBufDescriptor(InputStream fin) + private Descriptors.Descriptor buildProtoBufDescriptor(byte[] descriptorBytes) throws IOException { try { - DynamicSchema dynamicSchema = DynamicSchema.parseFrom(fin); + DynamicSchema dynamicSchema = DynamicSchema.parseFrom(descriptorBytes); if (!StringUtils.isEmpty(_protoClassName)) { return dynamicSchema.getMessageDescriptor(_protoClassName); diff --git a/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufRecordReader.java b/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufRecordReader.java index edd76987b7b2..546fb1bc334f 100644 --- a/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufRecordReader.java +++ b/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufRecordReader.java @@ -80,9 +80,9 @@ public void init(File dataFile, @Nullable Set fieldsToRead, @Nullable Re private Descriptors.Descriptor buildProtoBufDescriptor(ProtoBufRecordReaderConfig protoBufRecordReaderConfig) throws IOException { try { - InputStream fin = ProtoBufUtils.getDescriptorFileInputStream( + byte[] descriptorBytes = ProtoBufUtils.readDescriptorFileBytes( protoBufRecordReaderConfig.getDescriptorFile().toString()); - DescriptorProtos.FileDescriptorSet set = DescriptorProtos.FileDescriptorSet.parseFrom(fin); + DescriptorProtos.FileDescriptorSet set = DescriptorProtos.FileDescriptorSet.parseFrom(descriptorBytes); Descriptors.FileDescriptor fileDescriptor = Descriptors.FileDescriptor.buildFrom(set.getFile(0), new Descriptors.FileDescriptor[]{}); return fileDescriptor.getMessageTypes().get(0); diff --git a/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufUtils.java b/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufUtils.java index 6fd0e2866a91..fbb89f2644e8 100644 --- a/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufUtils.java +++ b/pinot-plugins/pinot-input-format/pinot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufUtils.java @@ -20,59 +20,36 @@ import com.google.protobuf.Descriptors; import com.google.protobuf.ProtobufInternalUtils; -import java.io.File; -import java.io.FileInputStream; import java.io.InputStream; import java.net.URI; -import java.nio.file.Files; -import java.nio.file.Path; import org.apache.pinot.spi.filesystem.PinotFS; import org.apache.pinot.spi.filesystem.PinotFSFactory; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; public class ProtoBufUtils { - private static final Logger LOGGER = LoggerFactory.getLogger(ProtoBufUtils.class); public static final String TMP_DIR_PREFIX = "pinot-protobuf"; public static final String PB_OUTER_CLASS_SUFFIX = "OuterClass"; private ProtoBufUtils() { } - public static File getFileCopiedToLocal(String filePath) + /// Reads the contents of a descriptor file (local or remote) into a byte array. The file is read via + /// [PinotFS#open] and the stream is closed before returning - no temporary files are created. + /// + /// @param descriptorFilePath URI string pointing to a `.desc` protobuf descriptor file + /// @return the raw bytes of the descriptor file + public static byte[] readDescriptorFileBytes(String descriptorFilePath) throws Exception { - URI fileURI = URI.create(filePath); + URI fileURI = URI.create(descriptorFilePath); String scheme = fileURI.getScheme(); if (scheme == null) { scheme = PinotFSFactory.LOCAL_PINOT_FS_SCHEME; } - if (PinotFSFactory.isSchemeSupported(scheme)) { - PinotFS pinotFS = PinotFSFactory.create(scheme); - Path localTmpDir = Files.createTempDirectory(TMP_DIR_PREFIX + System.currentTimeMillis()); - File localFile = createLocalFile(fileURI, localTmpDir.toFile()); - LOGGER.info("Copying protocol buffer jar/descriptor file from source: {} to dst: {}", filePath, - localFile.getAbsolutePath()); - pinotFS.copyToLocalFile(fileURI, localFile); - return localFile; - } else { - throw new RuntimeException(String.format("Scheme: %s not supported in PinotFSFactory" - + " for protocol buffer jar/descriptor file: %s.", scheme, filePath)); + PinotFS pinotFS = PinotFSFactory.create(scheme); + try (InputStream in = pinotFS.open(fileURI)) { + return in.readAllBytes(); } } - public static InputStream getDescriptorFileInputStream(String descriptorFilePath) - throws Exception { - return new FileInputStream(getFileCopiedToLocal(descriptorFilePath)); - } - - public static File createLocalFile(URI srcURI, File dstDir) { - String sourceURIPath = srcURI.getPath(); - File dstFile = new File(dstDir, new File(sourceURIPath).getName()); - LOGGER.debug("Created empty local temporary file {} to copy protocol " - + "buffer descriptor {}", dstFile.getAbsolutePath(), srcURI); - return dstFile; - } - public static String getFullJavaName(Descriptors.Descriptor descriptor) { String prefix; if (null != descriptor.getContainingType()) { diff --git a/pinot-plugins/pinot-input-format/pinot-protobuf/src/test/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufTempFileLeakTest.java b/pinot-plugins/pinot-input-format/pinot-protobuf/src/test/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufTempFileLeakTest.java new file mode 100644 index 000000000000..a4f9ddf1964d --- /dev/null +++ b/pinot-plugins/pinot-input-format/pinot-protobuf/src/test/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufTempFileLeakTest.java @@ -0,0 +1,211 @@ +/** + * 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.plugin.inputformat.protobuf; + +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.net.URL; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.testng.annotations.Test; + +import static org.apache.pinot.plugin.inputformat.protobuf.ProtoBufCodeGenMessageDecoder.PROTOBUF_JAR_FILE_PATH; +import static org.apache.pinot.plugin.inputformat.protobuf.ProtoBufCodeGenMessageDecoder.PROTO_CLASS_NAME; +import static org.apache.pinot.plugin.inputformat.protobuf.ProtoBufTestDataGenerator.createComplexTypeRecord; +import static org.apache.pinot.plugin.inputformat.protobuf.ProtoBufTestDataGenerator.getComplexTypeObject; +import static org.apache.pinot.plugin.inputformat.protobuf.ProtoBufTestDataGenerator.getFieldsInSampleRecord; +import static org.apache.pinot.plugin.inputformat.protobuf.ProtoBufTestDataGenerator.getSampleRecordMessage; +import static org.apache.pinot.plugin.inputformat.protobuf.ProtoBufTestDataGenerator.getSourceFieldsForComplexType; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/// Verifies that protobuf decoder and reader operations do not leak temporary directories. +/// +/// Each test snapshots the set of `pinot-protobuf*` directories in the system temp directory before the operation, +/// then asserts that no new ones remain afterward. The functional assertions confirm the operation itself still works +/// correctly. +public class ProtoBufTempFileLeakTest { + private static final Path TEMP_DIR = Path.of(System.getProperty("java.io.tmpdir")); + + /// [ProtoBufMessageDecoder#init] with a descriptor file should not leak a temp directory. + /// This is the streaming consumer path - called once per consumer init. + @Test + public void testMessageDecoderInitDoesNotLeakTempDir() + throws Exception { + Set before = listProtobufTempDirs(); + + Map decoderProps = new HashMap<>(); + URL descriptorFile = getClass().getClassLoader().getResource("sample.desc"); + decoderProps.put("descriptorFile", descriptorFile.toURI().toString()); + ProtoBufMessageDecoder decoder = new ProtoBufMessageDecoder(); + decoder.init(decoderProps, getFieldsInSampleRecord(), ""); + + // Verify functional correctness - decoding still works + Sample.SampleRecord sampleRecord = getSampleRecordMessage(); + GenericRow destination = new GenericRow(); + decoder.decode(sampleRecord.toByteArray(), destination); + assertEquals(destination.getValue("email"), "foobar@hello.com"); + assertEquals(destination.getValue("name"), "Alice"); + assertEquals(destination.getValue("id"), 18); + + assertNoNewTempDirs(before); + } + + /// [ProtoBufMessageDecoder#init] with a complex descriptor should not leak a temp directory. + @Test + public void testMessageDecoderComplexDescriptorDoesNotLeakTempDir() + throws Exception { + Set before = listProtobufTempDirs(); + + Map decoderProps = new HashMap<>(); + URL descriptorFile = getClass().getClassLoader().getResource("complex_types.desc"); + decoderProps.put("descriptorFile", descriptorFile.toURI().toString()); + ProtoBufMessageDecoder decoder = new ProtoBufMessageDecoder(); + decoder.init(decoderProps, getSourceFieldsForComplexType(), ""); + + // Verify functional correctness + Map inputRecord = createComplexTypeRecord(); + GenericRow destination = new GenericRow(); + decoder.decode(getComplexTypeObject(inputRecord).toByteArray(), destination); + assertNotNull(destination.getValue("string_field")); + assertEquals(destination.getValue("string_field"), "hello"); + + assertNoNewTempDirs(before); + } + + /// [ProtoBufCodeGenMessageDecoder#init] with a JAR file should not leak a temp directory. + /// This is the streaming consumer codegen path. + @Test + public void testCodeGenDecoderInitDoesNotLeakTempDir() + throws Exception { + Set before = listProtobufTempDirs(); + + Map decoderProps = new HashMap<>(); + URL jarFile = getClass().getClassLoader().getResource("sample.jar"); + decoderProps.put(PROTOBUF_JAR_FILE_PATH, jarFile.toURI().toString()); + decoderProps.put(PROTO_CLASS_NAME, + "org.apache.pinot.plugin.inputformat.protobuf.Sample$SampleRecord"); + ProtoBufCodeGenMessageDecoder decoder = new ProtoBufCodeGenMessageDecoder(); + decoder.init(decoderProps, getFieldsInSampleRecord(), ""); + + // Verify functional correctness - decoding still works after cleanup + Sample.SampleRecord sampleRecord = getSampleRecordMessage(); + GenericRow destination = new GenericRow(); + decoder.decode(sampleRecord.toByteArray(), destination); + assertEquals(destination.getValue("email"), "foobar@hello.com"); + assertEquals(destination.getValue("name"), "Alice"); + assertEquals(destination.getValue("id"), 18); + + assertNoNewTempDirs(before); + } + + /// [ProtoBufCodeGenMessageDecoder#init] with a complex JAR should not leak a temp directory, + /// and decoding complex/nested types must still work after cleanup. + @Test + public void testCodeGenDecoderComplexJarDoesNotLeakTempDir() + throws Exception { + Set before = listProtobufTempDirs(); + + Map decoderProps = new HashMap<>(); + URL jarFile = getClass().getClassLoader().getResource("complex_types.jar"); + decoderProps.put(PROTOBUF_JAR_FILE_PATH, jarFile.toURI().toString()); + decoderProps.put(PROTO_CLASS_NAME, + "org.apache.pinot.plugin.inputformat.protobuf.ComplexTypes$TestMessage"); + ProtoBufCodeGenMessageDecoder decoder = new ProtoBufCodeGenMessageDecoder(); + decoder.init(decoderProps, getSourceFieldsForComplexType(), ""); + + // Verify functional correctness + Map inputRecord = createComplexTypeRecord(); + GenericRow destination = new GenericRow(); + decoder.decode(getComplexTypeObject(inputRecord).toByteArray(), destination); + assertNotNull(destination.getValue("string_field")); + assertEquals(destination.getValue("string_field"), "hello"); + + assertNoNewTempDirs(before); + } + + /// [ProtoBufRecordReader] lifecycle (init, read, close) should not leak a temp directory. + /// This is the batch ingestion path. + @Test + public void testRecordReaderLifecycleDoesNotLeakTempDir() + throws Exception { + Set before = listProtobufTempDirs(); + + // Write a small protobuf data file + File tempDataDir = Files.createTempDirectory("protobuf-leak-test-data").toFile(); + File dataFile = new File(tempDataDir, "test.data"); + try { + try (FileOutputStream out = new FileOutputStream(dataFile)) { + getSampleRecordMessage().writeDelimitedTo(out); + } + + URL descriptorFile = getClass().getClassLoader().getResource("sample.desc"); + ProtoBufRecordReaderConfig config = new ProtoBufRecordReaderConfig(); + config.setDescriptorFile(descriptorFile.toURI()); + + try (ProtoBufRecordReader reader = new ProtoBufRecordReader()) { + reader.init(dataFile, getFieldsInSampleRecord(), config); + + // Verify functional correctness + assertTrue(reader.hasNext()); + GenericRow row = reader.next(new GenericRow()); + assertEquals(row.getValue("email"), "foobar@hello.com"); + assertEquals(row.getValue("name"), "Alice"); + } + } finally { + //noinspection ResultOfMethodCallIgnored + dataFile.delete(); + //noinspection ResultOfMethodCallIgnored + tempDataDir.delete(); + } + + assertNoNewTempDirs(before); + } + + private static void assertNoNewTempDirs(Set before) + throws IOException { + Set after = listProtobufTempDirs(); + Set leaked = after.stream() + .filter(p -> !before.contains(p)) + .collect(Collectors.toSet()); + assertTrue(leaked.isEmpty(), + "Leaked " + leaked.size() + " temp director" + (leaked.size() == 1 ? "y" : "ies") + ": " + leaked); + } + + private static Set listProtobufTempDirs() + throws IOException { + if (!Files.isDirectory(TEMP_DIR)) { + return Set.of(); + } + try (Stream entries = Files.list(TEMP_DIR)) { + return entries + .filter(p -> p.getFileName().toString().startsWith(ProtoBufUtils.TMP_DIR_PREFIX)) + .collect(Collectors.toSet()); + } + } +} diff --git a/pinot-plugins/pinot-input-format/pinot-protobuf/src/test/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufUtilsTest.java b/pinot-plugins/pinot-input-format/pinot-protobuf/src/test/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufUtilsTest.java index 109f1fa0deb3..359ab8066daa 100644 --- a/pinot-plugins/pinot-input-format/pinot-protobuf/src/test/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufUtilsTest.java +++ b/pinot-plugins/pinot-input-format/pinot-protobuf/src/test/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufUtilsTest.java @@ -19,6 +19,7 @@ package org.apache.pinot.plugin.inputformat.protobuf; import com.google.protobuf.Descriptors; +import java.io.File; import java.net.URL; import org.testng.Assert; import org.testng.annotations.DataProvider; @@ -54,7 +55,7 @@ public void testGetTypeStrFromProto(String fieldName, String javaType) { @Test public void testGetTypeStrFromProto() throws Exception { URL jarFile = getClass().getClassLoader().getResource("complex_types.jar"); - ClassLoader clsLoader = ProtoBufCodeGenMessageDecoder.loadClass(jarFile.getPath()); + ClassLoader clsLoader = ProtoBufCodeGenMessageDecoder.createClassLoader(new File(jarFile.toURI())); Descriptors.Descriptor desc = ProtoBufCodeGenMessageDecoder.getDescriptorForProtoClass(clsLoader, "org.apache.pinot.plugin.inputformat.protobuf.ComplexTypes$TestMessage"); Assert.assertEquals(ProtoBufUtils.getFullJavaName(desc),