github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/clients/hadoopfs/src/main/java/io/lakefs/FileSystemTracer.java

package io.lakefs;

import io.lakefs.utils.ObjectLocation;
import io.lakefs.utils.StringUtils;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.*;
import org.apache.hadoop.fs.permission.FsPermission;
import org.apache.hadoop.fs.s3a.S3AFileSystem;
import org.apache.hadoop.util.Progressable;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.OutputStream;
import java.net.URI;

/**
 * This class implements a troubleshooting tool for the LakeFSFileSystem. You can use it to compare the output of file
 * system operations done by {@link LakeFSFileSystem} and by {@link S3AFileSystem}.
 *
 * How to use the FileSystemTracer:
 * To use the FileSystemTracer, you should set the value of the Hadoop configuration 'fs.lakefs.impl' to FileSystemTracer.
 * This configures the FileSystemTracer to be the file system that handles paths with the lakefs scheme. i.e. paths with
 * lakefs:// prefix.
 *
 * How does the FileSystemTracer work:
 * The FileSystemTracer holds instances of {@link LakeFSFileSystem} and {@link S3AFileSystem}. On a file system operation,
 * the FileSystemTracer invokes the operation on both file systems, logs the output of both calls, and returns the result
 * of one of the file systems based on a configuration (S3AFileSystem output by default).

 * Configuration:
 * - fs.lakefs.tracer.working_dir - the s3 location in which the tracer can operate. This should be an S3 bucket name or
 *   an absolute path of a directory on a bucket.
 * Optional -
 * - fs.lakefs.tracer.use_lakefs_output - tells the tracer whether it should return the response coming from the lakefs file
 *   system or return s3a's response. by default it is set to false and returns s3a's response.
 *
 * Assumptions:
 * - The content of lakefs://repository/branch/ and s3a://${fs.lakefs.tracer.working.dir}/ should be identical.
 * - The s3 credentials available for Spark allow access to fs.lakefs.tracer.working.dir on s3.
 */
public class FileSystemTracer extends FileSystem {

    private static final Logger LOG = LoggerFactory.getLogger(FileSystemTracer.class);
    private static final String TRACER_WORKING_DIR = "fs.lakefs.tracer.working_dir";
    private static final String USE_LAKEFS_OUTPUT = "fs.lakefs.tracer.use_lakefs_output";
    private static final Object S3_URI_SCHEME = "s3";
    private static final Object RESULTS_COMPARISON = "[RESULTS_COMPARISON]";


    /*
     A property that determines which file system's response the FileSystemTracer returns.
     */
    private boolean useLakeFSFileSystemResults;
    private LakeFSFileSystem lfsFileSystem;
    private FileSystem s3AFileSystem;
    private String s3aPathPrefix;
    private String lfsPathPrefix;

    /**
     * Transforms a lakefs path into an s3 path.
     * Example:
     *   in: lakefs://repository/branch/key=1/a.parquet
     *   out: s3://${fs.lakefs.tracer.working_dir}/key=1/a.parquet
     */
    private Path translateLakeFSPathToS3APath(Path path) {
        return replacePathPrefix(path, lfsPathPrefix, s3aPathPrefix);
    }

    /**
     * Transforms an S3A path into a lakefs path.
     * This method is used to transform the output of S3AFileSystem operations that include paths, so that Spark
     * continues to direct it's requests to the FileSystemTracer that handles paths with the "lakefs://" prefix.
     * Example:
     *   in: s3://${fs.lakefs.tracer.working_dir}/key=1/a.parquet
     *   out: lakefs://repository/branch/key=1/a.parquet
     */
    private Path translateS3APathToLakeFSPath(Path path) {
        return replacePathPrefix(path, s3aPathPrefix, lfsPathPrefix);
    }

    private Path replacePathPrefix(Path path, String curPrefix, String newPrefix) {
        String p = path.toString();
        boolean isValidPath = p.startsWith(curPrefix);
        if (isValidPath) {
            String objRelativePath = StringUtils.trimLeadingSlash(p.substring(curPrefix.length()));
            String newPath = String.format("%s/%s", newPrefix, objRelativePath);
            LOG.trace("Converted {} to {}", path, newPath);
            return new Path(newPath);
        }
        LOG.error("Invalid path {}", path);
        return null;
    }

    @Override
    public void initialize(URI name, Configuration conf) throws IOException {
        lfsFileSystem = new LakeFSFileSystem();
        lfsFileSystem.initialize(name, conf);
        ObjectLocation loc = lfsFileSystem.pathToObjectLocation(new Path(name));
        lfsPathPrefix = ObjectLocation.formatPath(loc.getScheme(), loc.getRepository(), loc.getRef());

        String tracerWorkingDir = conf.get(TRACER_WORKING_DIR);
        if (tracerWorkingDir == null) {
            throw new IOException("tracerWorkingDir is null");
        }
        s3aPathPrefix = String.format("%s://%s", S3_URI_SCHEME, tracerWorkingDir);
        Path s3aPath = new Path(s3aPathPrefix);
        s3AFileSystem = s3aPath.getFileSystem(conf);

        useLakeFSFileSystemResults = conf.getBoolean(USE_LAKEFS_OUTPUT, false);
        LOG.trace("Initialization finished, fs.lakefs.tracer.use_lakefs_output: {}", useLakeFSFileSystemResults);
    }

    @Override
    public URI getUri() {
        LOG.trace("getUri");

        URI lakefsRes = lfsFileSystem.getUri();
        URI s3aRes = s3AFileSystem.getUri();
        LOG.trace("{}[getUri] lakefs: {}, s3a: {}", RESULTS_COMPARISON, lakefsRes, s3aRes);

        if (useLakeFSFileSystemResults) {
            return lakefsRes;
        }
        return s3aRes;
    }

    @Override
    public Path makeQualified(Path path) {
        LOG.trace("makeQualified");

        Path lakefsRes = lfsFileSystem.makeQualified(path);
        Path s3aRes = s3AFileSystem.makeQualified(translateLakeFSPathToS3APath(path));
        LOG.trace("{}[makeQualified] lakefs: {}, s3a: {}", RESULTS_COMPARISON, lakefsRes, s3aRes);

        if (useLakeFSFileSystemResults) {
            return lakefsRes;
        }
        return translateS3APathToLakeFSPath(s3aRes);
    }

    @Override
    public FSDataInputStream open(Path f, int bufferSize) throws IOException {
        LOG.trace("open(Path {}, bufferSize {})", f, bufferSize);
        FSDataInputStream lakefsRes = null;
        FSDataInputStream s3aRes = null;
        IOException lakeFSException = null;
        IOException s3aException = null;
        Path s3aPath = translateLakeFSPathToS3APath(f);
        try {
            lakefsRes = lfsFileSystem.open(f, bufferSize);
        } catch (IOException e) {
            lakeFSException = e;
            LOG.error("[open] Can't open {} with lakeFSFileSystem, exception {}", f, e.getMessage());
        }
        try {
            s3aRes = s3AFileSystem.open(translateLakeFSPathToS3APath(f), bufferSize);
        } catch (IOException e) {
            s3aException = e;
            LOG.error("[open] Can't open {} with S3AFileSystem, exception {}", s3aPath, e.getMessage());
        }

        if (useLakeFSFileSystemResults && lakeFSException != null) {
            LOG.trace("[open] exception by lakeFSFileSystem");
            throw lakeFSException;
        }
        if (!useLakeFSFileSystemResults && s3aException != null) {
            LOG.trace("[open] exception by S3AFileSystem");
            throw s3aException;
        }

        LOG.trace("{}[open] lakefs: {}, s3a: {}", RESULTS_COMPARISON, lakefsRes, s3aRes);
        if (useLakeFSFileSystemResults) {
            return lakefsRes;
        }
        return s3aRes;
    }

    @Override
    public FSDataOutputStream create(Path f, FsPermission permission, boolean overwrite, int bufferSize,
                                     short replication, long blockSize, Progressable progress) throws IOException {
        LOG.trace("create(Path {}, permission {}, overwrite {}, bufferSize {}, replication {}, blockSize {}, progress {})",
                f, permission, overwrite, bufferSize, replication, blockSize, progress);
        FSDataOutputStream lakeFSStream = null;
        FSDataOutputStream s3aStream = null;
        IOException lakeFSException = null;
        IOException s3aException = null;
        Path s3aPath = translateLakeFSPathToS3APath(f);
        try {
            lakeFSStream = lfsFileSystem.create(f, permission, overwrite, bufferSize, replication, blockSize, progress);
        } catch (IOException e) {
            lakeFSException = e;
            LOG.error("[create] Can't create {} with lakeFSFileSystem, exception {}", f, e.getMessage());
        }
        try {
            s3aStream = s3AFileSystem.create(s3aPath, permission, overwrite, bufferSize,
                replication, blockSize, progress);
        } catch (IOException e) {
            s3aException = e;
            LOG.error("[create] Can't create {} with S3AFileSystem, exception {}", s3aPath, e.getMessage());
        }

        if (useLakeFSFileSystemResults && lakeFSException != null) {
            LOG.trace("[create] exception by lakeFSFileSystem");
            if (s3aStream != null) {
                s3aStream.close();
            }
            throw lakeFSException;
        }
        if (!useLakeFSFileSystemResults && s3aException != null) {
            LOG.trace("[create] exception by S3AFileSystem");
            if (lakeFSStream != null) {
                lakeFSStream.close();
            }
            throw s3aException;
        }

        LOG.trace("{}[create] lakefs: {}, s3a: {}", RESULTS_COMPARISON, lakeFSStream, s3aStream);
        TracerOutputTStream tOutputStream = new TracerOutputTStream(lakeFSStream, s3aStream);
        return new FSDataOutputStream(tOutputStream, null);
    }

    @Override
    public FSDataOutputStream append(Path f, int bufferSize, Progressable progress) throws IOException {
        LOG.trace("append(f {}, bufferSize {}, progress {})", f, bufferSize, progress);

        FSDataOutputStream lakefsRes = lfsFileSystem.append(f, bufferSize, progress);
        FSDataOutputStream s3aRes = s3AFileSystem.append(translateLakeFSPathToS3APath(f), bufferSize, progress);
        LOG.trace("{}[append] lakefs: {}, s3a: {}", RESULTS_COMPARISON, lakefsRes, s3aRes);

        if (useLakeFSFileSystemResults) {
            return lakefsRes;
        }
        return s3aRes;
    }

    @Override
    public boolean rename(Path src, Path dst) throws IOException {
        LOG.trace("rename(src {}, dst {})", src, dst);
        boolean lakefsRes = false;
        boolean s3aRes = false;
        IOException lakeFSException = null;
        IOException s3aException = null;
        try {
            lakefsRes = lfsFileSystem.rename(src, dst);
        } catch (IOException e) {
            lakeFSException = e;
            LOG.error("[rename] Can't rename {} to {} with lakeFSFileSystem, exception {}", src, dst, e.getMessage());
        }
        try {
            s3aRes = s3AFileSystem.rename(translateLakeFSPathToS3APath(src), translateLakeFSPathToS3APath(dst));
        } catch (IOException e) {
            s3aException = e;
            LOG.error("[rename] Can't rename {} to {} with S3AFileSystem, exception {}", src, dst, e.getMessage());
        }

        if (useLakeFSFileSystemResults && lakeFSException != null) {
            LOG.trace("[rename] exception by lakeFSFileSystem");
            throw lakeFSException;
        }
        if (!useLakeFSFileSystemResults && s3aException != null) {
            LOG.trace("[rename] exception by S3AFileSystem");
            throw s3aException;
        }

        LOG.trace("{}[rename] lakefs: {}, s3a: {}", RESULTS_COMPARISON, lakefsRes, s3aRes);
        if (useLakeFSFileSystemResults) {
            return lakefsRes;
        }
        return s3aRes;
    }

    @Override
    public boolean delete(Path f, boolean recursive) throws IOException {
        LOG.trace("delete(f {}, recursive {})", f, recursive);
        boolean lakefsRes = false;
        boolean s3aRes = false;
        IOException lakeFSException = null;
        IOException s3aException = null;
        try {
            lakefsRes = delete(f, recursive);
        } catch (IOException e) {
            lakeFSException = e;
            LOG.error("[delete] Can't delete {} with lakeFSFileSystem, exception {}", f, e.getMessage());
        }
        try {
            s3aRes = s3AFileSystem.delete(translateLakeFSPathToS3APath(f), recursive);
        } catch (IOException e) {
            s3aException = e;
            LOG.error("[delete] Can't delete {} to {} with S3AFileSystem, exception {}", f, e.getMessage());
        }

        if (useLakeFSFileSystemResults && lakeFSException != null) {
            LOG.trace("[delete] exception by lakeFSFileSystem");
            throw lakeFSException;
        }
        if (!useLakeFSFileSystemResults && s3aException != null) {
            LOG.trace("[delete] exception by S3AFileSystem");
            throw s3aException;
        }

        LOG.trace("{}[delete] lakefs: {}, s3a: {}", RESULTS_COMPARISON, lakefsRes, s3aRes);
        if (useLakeFSFileSystemResults) {
            return lakefsRes;
        }
        return s3aRes;
    }

    @Override
    public FileStatus[] listStatus(Path f) throws FileNotFoundException, IOException {
        LOG.trace("listStatus(f {})", f);

        FileStatus[] lakefsRes = null;
        FileStatus[] s3aRes = null;
        IOException lakeFSException = null;
        IOException s3aException = null;
        Path s3aPath = translateLakeFSPathToS3APath(f);
        try {
            lakefsRes = lfsFileSystem.listStatus(f);
        } catch (IOException e) {
            lakeFSException = e;
            LOG.error("[listStatus] Can't list the status of {} with lakeFSFileSystem, exception {}", f, e.getMessage());
        }
        try {
            s3aRes = s3AFileSystem.listStatus(s3aPath);
        } catch (IOException e) {
            s3aException = e;
            LOG.error("[listStatus] Can't list the status of {} with S3AFileSystem, exception {}", s3aPath, e.getMessage());
        }

        if (useLakeFSFileSystemResults && lakeFSException != null) {
            LOG.trace("[listStatus] exception by lakeFSFileSystem");
            throw lakeFSException;
        }
        if (!useLakeFSFileSystemResults && s3aException != null) {
            LOG.trace("[listStatus] exception by S3AFileSystem");
            throw s3aException;
        }

        LOG.trace("{}[listStatus] lakefs: {}, s3a: {}", RESULTS_COMPARISON, lakefsRes, s3aRes);
        if (useLakeFSFileSystemResults) {
            return lakefsRes;
        }
        for (FileStatus stat : s3aRes) {
            Path filePath = stat.getPath();
            Path lfsPath = translateS3APathToLakeFSPath(filePath);
            stat.setPath(lfsPath);
        }
        return s3aRes;
    }

    @Override
    public void setWorkingDirectory(Path newDir) {
        LOG.trace("setWorkingDirectory(new_dir {})", newDir);

        lfsFileSystem.setWorkingDirectory(newDir);
        s3AFileSystem.setWorkingDirectory(translateLakeFSPathToS3APath(newDir));
    }

    @Override
    public Path getWorkingDirectory() {
        LOG.trace("getWorkingDirectory()");

        Path lakefsRes = lfsFileSystem.getWorkingDirectory();
        Path s3aRes = s3AFileSystem.getWorkingDirectory();
        LOG.trace("{}[getWorkingDirectory] lakefs: {}, s3a: {}", RESULTS_COMPARISON, lakefsRes, s3aRes);

        if (useLakeFSFileSystemResults) {
            return lakefsRes;
        }
        return s3aRes;
    }

    @Override
    public boolean mkdirs(Path f, FsPermission permission) throws IOException {
        LOG.trace("mkdirs(f {}, permission {})", f, permission);
        boolean lakefsRes = false;
        boolean s3aRes = false;
        IOException lakeFSException = null;
        IOException s3aException = null;
        try {
            lakefsRes = lfsFileSystem.mkdirs(f, permission);
        } catch (IOException e) {
            lakeFSException = e;
            LOG.error("[mkdirs] Can't mkdir {} with lakeFSFileSystem, exception {}", f, e.getMessage());
        }
        try {
            s3aRes = s3AFileSystem.mkdirs(translateLakeFSPathToS3APath(f), permission);
        } catch (IOException e) {
            s3aException = e;
            LOG.error("[mkdirs] Can't mkdir {} to {} with S3AFileSystem, exception {}", f, e.getMessage());
        }

        if (useLakeFSFileSystemResults && lakeFSException != null) {
            LOG.trace("[mkdirs] exception by lakeFSFileSystem");
            throw lakeFSException;
        }
        if (!useLakeFSFileSystemResults && s3aException != null) {
            LOG.trace("[mkdirs] exception by S3AFileSystem");
            throw s3aException;
        }

        LOG.trace("{}[mkdirs] lakefs: {}, s3a: {}", RESULTS_COMPARISON, lakefsRes, s3aRes);
        if (useLakeFSFileSystemResults) {
            return lakefsRes;
        }
        return s3aRes;
    }

    @Override
    public FileStatus getFileStatus(Path f) throws IOException {
        LOG.trace("getFileStatus(f {})", f);

        FileStatus lakefsRes = null;
        FileStatus s3aRes = null;
        IOException lakeFSException = null;
        IOException s3aException = null;
        Path s3aPath = translateLakeFSPathToS3APath(f);
        try {
            lakefsRes = lfsFileSystem.getFileStatus(f);
        } catch (IOException e) {
            lakeFSException = e;
            LOG.error("[getFileStatus] Can't get {} file status with lakeFSFileSystem, exception {}", f, e.getMessage());
        }
        try {
            s3aRes = s3AFileSystem.getFileStatus(s3aPath);
        } catch (IOException e) {
            s3aException = e;
            LOG.error("[getFileStatus] Can't get {} file status with S3AFileSystem, exception {}", s3aPath, e.getMessage());
        }

        if (useLakeFSFileSystemResults && lakeFSException != null) {
            LOG.trace("[getFileStatus] exception by lakeFSFileSystem");
            throw lakeFSException;
        }
        if (!useLakeFSFileSystemResults && s3aException != null) {
            LOG.trace("[getFileStatus] exception by S3AFileSystem");
            throw s3aException;
        }

        LOG.trace("{}[getFileStatus] lakefs: {}, s3a: {}", RESULTS_COMPARISON, lakefsRes, s3aRes);
        if (useLakeFSFileSystemResults) {
            return lakefsRes;
        }
        Path lfsPath = translateS3APathToLakeFSPath(s3aPath);
        s3aRes.setPath(lfsPath);
        return s3aRes;
    }

    /**
     * An output stream encapsulating two output streams, one for each file system the tracer uses.
     */
    private class TracerOutputTStream extends OutputStream {

        private FSDataOutputStream lakeFSStream;
        private FSDataOutputStream s3aStream;

        public TracerOutputTStream(FSDataOutputStream lakeFSStream, FSDataOutputStream s3aStream) throws IOException {
            this.lakeFSStream = lakeFSStream;
            this.s3aStream = s3aStream;
        }

        @Override
        public void write(int b) throws IOException {
            if (lakeFSStream != null) {
                lakeFSStream.write(b);
            }
            if (s3aStream != null) {
                s3aStream.write(b);
            }
        }

        @Override
        public void write(@NotNull byte[] b) throws IOException {
            if (lakeFSStream != null) {
                lakeFSStream.write(b);
            }
            if (s3aStream != null) {
                s3aStream.write(b);
            }
        }

        @Override
        public void write(@NotNull byte[] b, int off, int len) throws IOException {
            if (lakeFSStream != null) {
                lakeFSStream.write(b, off, len);
            }
            if (s3aStream != null) {
                s3aStream.write(b, off, len);
            }
        }

        @Override
        public void flush() throws IOException {
            if (lakeFSStream != null) {
                lakeFSStream.flush();
            }
            if (s3aStream != null) {
                s3aStream.flush();
            }
        }

        @Override
        public void close() throws IOException {
            if (lakeFSStream != null) {
                lakeFSStream.close();
            }
            if (s3aStream != null) {
                s3aStream.close();
            }
        }
    }
}