Skip to content

Fix protobuf temp directory leak in ProtoBufUtils #19378

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mayankshriv wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Fix protobuf temp directory leak in ProtoBufUtils #19378

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 50 additions & 6 deletions ...main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufCodeGenMessageDecoder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<byte[]> {
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;
Expand All @@ -49,12 +64,18 @@ public void init(Map<String, String> props, Set<String> 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
Expand All @@ -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);
}
Expand All @@ -104,4 +123,29 @@ public static Descriptors.Descriptor getDescriptorForProtoClass(ClassLoader prot
Class<? extends Message> updateMessage = (Class<Message>) 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;
}
}
}
10 changes: 4 additions & 6 deletions ...uf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufMessageDecoder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -48,19 +47,18 @@ public void init(Map<String, String> props, Set<String> 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);
Expand Down
4 changes: 2 additions & 2 deletions ...obuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufRecordReader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,9 +80,9 @@ public void init(File dataFile, @Nullable Set<String> 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);
Expand Down
43 changes: 10 additions & 33 deletions ...ot-protobuf/src/main/java/org/apache/pinot/plugin/inputformat/protobuf/ProtoBufUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()) {
Expand Down
Loading
Loading