mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-12-16 19:32:07 +00:00
376 lines
11 KiB
C++
376 lines
11 KiB
C++
#include <zstd.h>
|
|
#include <sys/mman.h>
|
|
#if defined __APPLE__
|
|
#include <sys/mount.h>
|
|
#else
|
|
#include <sys/statfs.h>
|
|
#endif
|
|
#include <fcntl.h>
|
|
#include <sys/wait.h>
|
|
#include <unistd.h>
|
|
#include <cstdlib>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <iostream>
|
|
|
|
#include "types.h"
|
|
|
|
/// decompress part
|
|
int doDecompress(char * input, char * output, off_t & in_offset, off_t & out_offset,
|
|
off_t input_size, off_t output_size, ZSTD_DCtx* dctx)
|
|
{
|
|
size_t decompressed_size = ZSTD_decompressDCtx(dctx, output + out_offset, output_size, input + in_offset, input_size);
|
|
if (ZSTD_isError(decompressed_size))
|
|
{
|
|
std::cerr << "Error (ZSTD):" << decompressed_size << " " << ZSTD_getErrorName(decompressed_size) << std::endl;
|
|
return 1;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
/// decompress data from in_fd into out_fd
|
|
int decompress(char * input, char * output, off_t start, off_t end, size_t max_number_of_forks=10)
|
|
{
|
|
off_t in_pointer = start, out_pointer = 0;
|
|
off_t size = 0;
|
|
off_t max_block_size = 1ull<<27;
|
|
off_t decompressed_size = 0;
|
|
size_t number_of_forks = 0;
|
|
|
|
/// Create context
|
|
ZSTD_DCtx * dctx = ZSTD_createDCtx();
|
|
if (dctx == nullptr)
|
|
{
|
|
std::cerr << "Error (ZSTD): failed to create decompression context" << std::endl;
|
|
return 1;
|
|
}
|
|
pid_t pid;
|
|
bool error_happened = false;
|
|
|
|
/// Decompress data
|
|
while (in_pointer < end && !error_happened)
|
|
{
|
|
size = ZSTD_findFrameCompressedSize(input + in_pointer, max_block_size);
|
|
if (ZSTD_isError(size))
|
|
{
|
|
std::cerr << "Error (ZSTD): " << size << " " << ZSTD_getErrorName(size) << std::endl;
|
|
error_happened = true;
|
|
break;
|
|
}
|
|
|
|
decompressed_size = ZSTD_getFrameContentSize(input + in_pointer, max_block_size);
|
|
if (ZSTD_isError(decompressed_size))
|
|
{
|
|
std::cerr << "Error (ZSTD): " << decompressed_size << " " << ZSTD_getErrorName(decompressed_size) << std::endl;
|
|
error_happened = true;
|
|
break;
|
|
}
|
|
|
|
pid = fork();
|
|
if (-1 == pid)
|
|
{
|
|
perror(nullptr);
|
|
/// If fork failed just decompress data in main process.
|
|
if (0 != doDecompress(input, output, in_pointer, out_pointer, size, max_block_size, dctx))
|
|
break;
|
|
}
|
|
else if (pid == 0)
|
|
{
|
|
/// Decompress data in child process.
|
|
if (0 != doDecompress(input, output, in_pointer, out_pointer, size, max_block_size, dctx))
|
|
exit(1);
|
|
exit(0);
|
|
}
|
|
else
|
|
{
|
|
++number_of_forks;
|
|
while (number_of_forks >= max_number_of_forks)
|
|
{
|
|
/// Wait any fork
|
|
int status;
|
|
waitpid(0, &status, 0);
|
|
|
|
/// If error happened, stop processing
|
|
if (WEXITSTATUS(status) != 0)
|
|
{
|
|
error_happened = true;
|
|
break;
|
|
}
|
|
|
|
--number_of_forks;
|
|
}
|
|
in_pointer += size;
|
|
out_pointer += decompressed_size;
|
|
}
|
|
}
|
|
|
|
/// wait for all working decompressions
|
|
while (number_of_forks > 0)
|
|
{
|
|
/// Wait any fork
|
|
int status;
|
|
waitpid(0, &status, 0);
|
|
|
|
if (WEXITSTATUS(status) != 0)
|
|
error_happened = true;
|
|
|
|
--number_of_forks;
|
|
}
|
|
|
|
/// If error happen end of processed part will not reach end
|
|
if (in_pointer < end || error_happened)
|
|
return 1;
|
|
|
|
return 0;
|
|
}
|
|
|
|
|
|
/// Read data about files and decomrpess them.
|
|
int decompressFiles(int input_fd, char * path, char * name, bool & have_compressed_analoge)
|
|
{
|
|
/// Read data about output file.
|
|
/// Compressed data will replace data in file
|
|
struct stat info_in;
|
|
if (0 != fstat(input_fd, &info_in))
|
|
{
|
|
perror(nullptr);
|
|
return 1;
|
|
}
|
|
|
|
/// mmap input file
|
|
char * input = static_cast<char*>(mmap(nullptr, info_in.st_size, PROT_READ, MAP_PRIVATE, input_fd, 0));
|
|
if (input == MAP_FAILED)
|
|
{
|
|
perror(nullptr);
|
|
return 1;
|
|
}
|
|
|
|
/// Read metadata from end of file
|
|
MetaData metadata = *reinterpret_cast<MetaData*>(input + info_in.st_size - sizeof(MetaData));
|
|
|
|
/// Prepare to read information about files and decompress them
|
|
off_t files_pointer = metadata.start_of_files_data;
|
|
size_t decompressed_full_size = 0;
|
|
|
|
/// Read files metadata and check if decompression is possible
|
|
off_t check_pointer = metadata.start_of_files_data;
|
|
for (size_t i = 0; i < metadata.number_of_files; ++i)
|
|
{
|
|
FileData data = *reinterpret_cast<FileData*>(input + check_pointer);
|
|
decompressed_full_size += data.uncompressed_size;
|
|
check_pointer += sizeof(FileData) + data.name_length;
|
|
}
|
|
|
|
/// Check free space
|
|
struct statfs fs_info;
|
|
if (0 != fstatfs(input_fd, &fs_info))
|
|
{
|
|
perror(nullptr);
|
|
if (0 != munmap(input, info_in.st_size))
|
|
perror(nullptr);
|
|
return 1;
|
|
}
|
|
if (fs_info.f_blocks * info_in.st_blksize < decompressed_full_size)
|
|
{
|
|
std::cerr << "Not enough space for decompression. Have " << fs_info.f_blocks * info_in.st_blksize << ", need " << decompressed_full_size << std::endl;
|
|
return 1;
|
|
}
|
|
|
|
FileData file_info;
|
|
/// Decompress files with appropriate file names
|
|
for (size_t i = 0; i < metadata.number_of_files; ++i)
|
|
{
|
|
/// Read information about file
|
|
file_info = *reinterpret_cast<FileData*>(input + files_pointer);
|
|
files_pointer += sizeof(FileData);
|
|
|
|
size_t file_name_len =
|
|
(strcmp(input + files_pointer, name) ? file_info.name_length : file_info.name_length + 13);
|
|
|
|
size_t file_path_len = path ? strlen(path) + 1 + file_name_len : file_name_len;
|
|
|
|
char file_name[file_path_len];
|
|
memset(file_name, '\0', file_path_len);
|
|
if (path)
|
|
{
|
|
strcat(file_name, path);
|
|
strcat(file_name, "/");
|
|
}
|
|
strcat(file_name, input + files_pointer);
|
|
files_pointer += file_info.name_length;
|
|
if (file_name_len != file_info.name_length)
|
|
{
|
|
strcat(file_name, ".decompressed");
|
|
have_compressed_analoge = true;
|
|
}
|
|
|
|
int output_fd = open(file_name, O_RDWR | O_CREAT, file_info.umask);
|
|
|
|
if (output_fd == -1)
|
|
{
|
|
perror(nullptr);
|
|
if (0 != munmap(input, info_in.st_size))
|
|
perror(nullptr);
|
|
return 1;
|
|
}
|
|
|
|
/// Prepare output file
|
|
if (0 != ftruncate(output_fd, file_info.uncompressed_size))
|
|
{
|
|
perror(nullptr);
|
|
if (0 != munmap(input, info_in.st_size))
|
|
perror(nullptr);
|
|
return 1;
|
|
}
|
|
|
|
char * output = static_cast<char*>(
|
|
mmap(nullptr,
|
|
file_info.uncompressed_size,
|
|
PROT_READ | PROT_WRITE, MAP_SHARED,
|
|
output_fd,
|
|
0)
|
|
);
|
|
if (output == MAP_FAILED)
|
|
{
|
|
perror(nullptr);
|
|
if (0 != munmap(input, info_in.st_size))
|
|
perror(nullptr);
|
|
return 1;
|
|
}
|
|
|
|
/// Decompress data into file
|
|
if (0 != decompress(input, output, file_info.start, file_info.end))
|
|
{
|
|
if (0 != munmap(input, info_in.st_size))
|
|
perror(nullptr);
|
|
if (0 != munmap(output, file_info.uncompressed_size))
|
|
perror(nullptr);
|
|
return 1;
|
|
}
|
|
|
|
if (0 != fsync(output_fd))
|
|
perror(nullptr);
|
|
if (0 != close(output_fd))
|
|
perror(nullptr);
|
|
}
|
|
|
|
if (0 != munmap(input, info_in.st_size))
|
|
perror(nullptr);
|
|
return 0;
|
|
}
|
|
|
|
/// Copy particular part of command and update shift
|
|
void fill(char * dest, char * source, size_t length, size_t& shift)
|
|
{
|
|
memcpy(dest + shift, source, length);
|
|
shift += length;
|
|
}
|
|
|
|
/// Set command to `mv filename.decompressed filename && filename agrs...`
|
|
void fillCommand(char command[], int argc, char * argv[], size_t length)
|
|
{
|
|
memset(command, '\0', 3 + strlen(argv[0]) + 14 + strlen(argv[0]) + 4 + strlen(argv[0]) + length + argc);
|
|
|
|
/// position in command
|
|
size_t shift = 0;
|
|
|
|
/// Support variables to create command
|
|
char mv[] = "mv ";
|
|
char decompressed[] = ".decompressed ";
|
|
char add_command[] = " && ";
|
|
char space[] = " ";
|
|
|
|
fill(command, mv, 3, shift);
|
|
fill(command, argv[0], strlen(argv[0]), shift);
|
|
fill(command, decompressed, 14, shift);
|
|
fill(command, argv[0], strlen(argv[0]), shift);
|
|
fill(command, add_command, 4, shift);
|
|
fill(command, argv[0], strlen(argv[0]), shift);
|
|
fill(command, space, 1, shift);
|
|
|
|
/// forward all arguments
|
|
for (int i = 1; i < argc; ++i)
|
|
{
|
|
fill(command, argv[i], strlen(argv[i]), shift);
|
|
if (i != argc - 1)
|
|
fill(command, space, 1, shift);
|
|
}
|
|
}
|
|
|
|
int main(int argc, char* argv[])
|
|
{
|
|
char file_path[strlen(argv[0]) + 1];
|
|
memset(file_path, 0, sizeof(file_path));
|
|
strcpy(file_path, argv[0]);
|
|
|
|
char * path = nullptr;
|
|
char * name = strrchr(file_path, '/');
|
|
if (name)
|
|
{
|
|
path = file_path;
|
|
*name = 0;
|
|
++name;
|
|
}
|
|
else
|
|
name = file_path;
|
|
|
|
int input_fd = open(argv[0], O_RDONLY);
|
|
if (input_fd == -1)
|
|
{
|
|
perror(nullptr);
|
|
return 1;
|
|
}
|
|
|
|
bool have_compressed_analoge = false;
|
|
|
|
/// Decompress all files
|
|
if (0 != decompressFiles(input_fd, path, name, have_compressed_analoge))
|
|
{
|
|
printf("Error happened during decompression.\n");
|
|
if (0 != close(input_fd))
|
|
perror(nullptr);
|
|
return 1;
|
|
}
|
|
|
|
if (0 != close(input_fd))
|
|
perror(nullptr);
|
|
|
|
/// According to documentation `mv` will rename file if it
|
|
/// doesn't move to other directory.
|
|
/// Sometimes `rename` doesn't exist by default and
|
|
/// `rename.ul` is set instead. It will lead to errors
|
|
/// that can be easily avoided with help of `mv`
|
|
|
|
if (!have_compressed_analoge)
|
|
{
|
|
printf("No target executable - decompression only was performed.\n");
|
|
/// remove file
|
|
execlp("rm", "rm", argv[0], NULL);
|
|
perror(nullptr);
|
|
return 1;
|
|
}
|
|
else
|
|
{
|
|
/// move decompressed file instead of this binary and apply command
|
|
char bash[] = "sh";
|
|
char executable[] = "-c";
|
|
|
|
/// length of forwarded args
|
|
size_t length = 0;
|
|
for (int i = 1; i < argc; ++i)
|
|
length += strlen(argv[i]);
|
|
|
|
/// mv filename.decompressed filename && filename agrs...
|
|
char command[3 + strlen(argv[0]) + 14 + strlen(argv[0]) + 4 + strlen(argv[0]) + length + argc];
|
|
fillCommand(command, argc, argv, length);
|
|
|
|
/// replace file and call executable
|
|
char * newargv[] = { bash, executable, command, nullptr };
|
|
execvp(bash, newargv);
|
|
|
|
/// This part of code will be reached only if error happened
|
|
perror(nullptr);
|
|
return 1;
|
|
}
|
|
}
|