be/test/runtime/tmp_file_mgr_test.cpp - doris - Git at Google

 // 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.

 #include "runtime/tmp_file_mgr.h"

 #include <gtest/gtest.h>

 #include <boost/filesystem.hpp>
 #include <boost/scoped_ptr.hpp>
 #include <cstdlib>

 #include "gen_cpp/Types_types.h" // for TUniqueId
 #include "util/disk_info.h"
 #include "util/filesystem_util.h"
 #include "util/logging.h"
 #include "util/metrics.h"

 using boost::filesystem::path;
 using std::string;
 using std::vector;
 using std::set;

 namespace doris {

 class TmpFileMgrTest : public ::testing::Test {
 protected:
     // Check that metric values are consistent with TmpFileMgr state.
     void check_metrics(TmpFileMgr* tmp_file_mgr) {
         std::vector<TmpFileMgr::DeviceId> active = tmp_file_mgr->active_tmp_devices();
         int64_t active_metric = DorisMetrics::instance()
                                         ->metric_registry()
                                         ->get_entity("server")
                                         ->get_metric("active_scratch_dirs")
                                         .value();
         EXPECT_EQ(active.size(), active_metric);
     }
 };

 // Regression test for IMPALA-2160. Verify that temporary file manager allocates blocks
 // at the expected file offsets and expands the temporary file to the correct size.
 TEST_F(TmpFileMgrTest, TestFileAllocation) {
     TmpFileMgr tmp_file_mgr;
     EXPECT_TRUE(tmp_file_mgr.init().ok());
     // Default configuration should give us one temporary device.
     EXPECT_EQ(1, tmp_file_mgr.num_active_tmp_devices());
     std::vector<TmpFileMgr::DeviceId> tmp_devices = tmp_file_mgr.active_tmp_devices();
     EXPECT_EQ(1, tmp_devices.size());
     TUniqueId id;
     TmpFileMgr::File* file;
     Status status = tmp_file_mgr.get_file(tmp_devices[0], id, &file);
     EXPECT_TRUE(status.ok());
     EXPECT_TRUE(file != NULL);
     // Apply writes of variable sizes and check space was allocated correctly.
     int64_t write_sizes[] = {1, 10, 1024, 4, 1024 * 1024 * 8, 1024 * 1024 * 8, 16, 10};
     int num_write_sizes = sizeof(write_sizes) / sizeof(write_sizes[0]);
     int64_t next_offset = 0;
     for (int i = 0; i < num_write_sizes; ++i) {
         int64_t offset;
         status = file->allocate_space(write_sizes[i], &offset);
         EXPECT_TRUE(status.ok());
         EXPECT_EQ(next_offset, offset);
         next_offset = offset + write_sizes[i];
         EXPECT_EQ(next_offset, boost::filesystem::file_size(file->path()));
     }
     // Check that cleanup is correct.
     status = file->remove();
     EXPECT_TRUE(status.ok());
     EXPECT_FALSE(boost::filesystem::exists(file->path()));
     // check_metrics(&tmp_file_mgr);
 }
 // Test that we can do initialization with two directories on same device and
 // that validations prevents duplication of directories.
 TEST_F(TmpFileMgrTest, TestOneDirPerDevice) {
     std::vector<string> tmp_dirs;
     tmp_dirs.push_back("/tmp/tmp-file-mgr-test.1");
     tmp_dirs.push_back("/tmp/tmp-file-mgr-test.2");
     for (int i = 0; i < tmp_dirs.size(); ++i) {
         EXPECT_TRUE(FileSystemUtil::create_directory(tmp_dirs[i]).ok());
     }
     TmpFileMgr tmp_file_mgr;
     tmp_file_mgr.init_custom(tmp_dirs, true);

     // Only the first directory should be used.
     EXPECT_EQ(1, tmp_file_mgr.num_active_tmp_devices());
     std::vector<TmpFileMgr::DeviceId> devices = tmp_file_mgr.active_tmp_devices();
     EXPECT_EQ(1, devices.size());
     TUniqueId id;
     TmpFileMgr::File* file;
     EXPECT_TRUE(tmp_file_mgr.get_file(devices[0], id, &file).ok());
     // Check the prefix is the expected temporary directory.
     EXPECT_EQ(0, file->path().find(tmp_dirs[0]));
     FileSystemUtil::remove_paths(tmp_dirs);
     // check_metrics(&tmp_file_mgr);
 }

 // Test that we can do custom initialization with two dirs on same device.
 TEST_F(TmpFileMgrTest, TestMultiDirsPerDevice) {
     std::vector<string> tmp_dirs;
     tmp_dirs.push_back("/tmp/tmp-file-mgr-test.1");
     tmp_dirs.push_back("/tmp/tmp-file-mgr-test.2");
     for (int i = 0; i < tmp_dirs.size(); ++i) {
         EXPECT_TRUE(FileSystemUtil::create_directory(tmp_dirs[i]).ok());
     }
     TmpFileMgr tmp_file_mgr;
     tmp_file_mgr.init_custom(tmp_dirs, false);

     // Both directories should be used.
     EXPECT_EQ(2, tmp_file_mgr.num_active_tmp_devices());
     std::vector<TmpFileMgr::DeviceId> devices = tmp_file_mgr.active_tmp_devices();
     EXPECT_EQ(2, devices.size());
     for (int i = 0; i < tmp_dirs.size(); ++i) {
         EXPECT_EQ(0, tmp_file_mgr.get_tmp_dir_path(devices[i]).find(tmp_dirs[i]));
         TUniqueId id;
         TmpFileMgr::File* file;
         EXPECT_TRUE(tmp_file_mgr.get_file(devices[i], id, &file).ok());
         // Check the prefix is the expected temporary directory.
         EXPECT_EQ(0, file->path().find(tmp_dirs[i]));
     }
     FileSystemUtil::remove_paths(tmp_dirs);
     // check_metrics(&tmp_file_mgr);
 }

 // Test that reporting a write error is possible but does not result in
 // blacklisting, which is disabled.
 TEST_F(TmpFileMgrTest, TestReportError) {
     std::vector<string> tmp_dirs;
     tmp_dirs.push_back("/tmp/tmp-file-mgr-test.1");
     tmp_dirs.push_back("/tmp/tmp-file-mgr-test.2");
     for (int i = 0; i < tmp_dirs.size(); ++i) {
         EXPECT_TRUE(FileSystemUtil::create_directory(tmp_dirs[i]).ok());
     }
     TmpFileMgr tmp_file_mgr;
     tmp_file_mgr.init_custom(tmp_dirs, false);

     // Both directories should be used.
     std::vector<TmpFileMgr::DeviceId> devices = tmp_file_mgr.active_tmp_devices();
     EXPECT_EQ(2, devices.size());
     // check_metrics(&tmp_file_mgr);

     // Inject an error on one device so that we can validate it is handled correctly.
     TUniqueId id;
     int good_device = 0;
     int bad_device = 1;
     TmpFileMgr::File* bad_file;
     EXPECT_TRUE(tmp_file_mgr.get_file(devices[bad_device], id, &bad_file).ok());
     // ErrorMsg errmsg(TErrorCode::GENERAL, "A fake error");
     // bad_file->ReportIOError(errmsg);
     bad_file->report_io_error("A fake error");

     // Blacklisting is disabled.
     EXPECT_FALSE(bad_file->is_blacklisted());
     // The second device should still be active.
     EXPECT_EQ(2, tmp_file_mgr.num_active_tmp_devices());
     std::vector<TmpFileMgr::DeviceId> devices_after = tmp_file_mgr.active_tmp_devices();
     EXPECT_EQ(2, devices_after.size());
     // check_metrics(&tmp_file_mgr);

     // Attempts to expand bad file should succeed.
     int64_t offset;
     EXPECT_TRUE(bad_file->allocate_space(128, &offset).ok());
     EXPECT_TRUE(bad_file->remove().ok());
     // The good device should still be usable.
     TmpFileMgr::File* good_file;
     EXPECT_TRUE(tmp_file_mgr.get_file(devices[good_device], id, &good_file).ok());
     EXPECT_TRUE(good_file != NULL);
     EXPECT_TRUE(good_file->allocate_space(128, &offset).ok());
     // Attempts to allocate new files on bad device should succeed.
     EXPECT_TRUE(tmp_file_mgr.get_file(devices[bad_device], id, &bad_file).ok());
     FileSystemUtil::remove_paths(tmp_dirs);
     // check_metrics(&tmp_file_mgr);
 }

 TEST_F(TmpFileMgrTest, TestAllocateFails) {
     string tmp_dir("/tmp/tmp-file-mgr-test.1");
     string scratch_subdir = tmp_dir + "/doris-scratch";
     std::vector<string> tmp_dirs(1, tmp_dir);
     EXPECT_TRUE(FileSystemUtil::create_directory(tmp_dir).ok());
     TmpFileMgr tmp_file_mgr;
     tmp_file_mgr.init_custom(tmp_dirs, false);

     TUniqueId id;
     TmpFileMgr::File* allocated_file1;
     TmpFileMgr::File* allocated_file2;
     int64_t offset;
     EXPECT_TRUE(tmp_file_mgr.get_file(0, id, &allocated_file1).ok());
     EXPECT_TRUE(tmp_file_mgr.get_file(0, id, &allocated_file2).ok());
     EXPECT_TRUE(allocated_file1->allocate_space(1, &offset).ok());

     // Make scratch non-writable and test for allocation errors at different stages:
     // new file creation, files with no allocated blocks. files with allocated space.
     chmod(scratch_subdir.c_str(), 0);
     // allocated_file1 already has space allocated.
     EXPECT_FALSE(allocated_file1->allocate_space(1, &offset).ok());
     // allocated_file2 has no space allocated.
     EXPECT_FALSE(allocated_file2->allocate_space(1, &offset).ok());
     // Creating a new File object can succeed because it is not immediately created on disk.
     TmpFileMgr::File* unallocated_file;
     EXPECT_TRUE(tmp_file_mgr.get_file(0, id, &unallocated_file).ok());

     chmod(scratch_subdir.c_str(), S_IRWXU);
     FileSystemUtil::remove_paths(tmp_dirs);
 }

 } // end namespace doris

 int main(int argc, char** argv) {
     // std::string conffile = std::string(getenv("DORIS_HOME")) + "/conf/be.conf";
     // if (!doris::config::init(conffile.c_str(), false)) {
     //     fprintf(stderr, "error read config file. \n");
     //     return -1;
     // }
     doris::config::query_scratch_dirs = "/tmp";
     doris::init_glog("be-test");
     ::testing::InitGoogleTest(&argc, argv);

     doris::DiskInfo::init();

     return RUN_ALL_TESTS();
 }
	// 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.

	#include "runtime/tmp_file_mgr.h"

	#include <gtest/gtest.h>

	#include <boost/filesystem.hpp>
	#include <boost/scoped_ptr.hpp>
	#include <cstdlib>

	#include "gen_cpp/Types_types.h" // for TUniqueId
	#include "util/disk_info.h"
	#include "util/filesystem_util.h"
	#include "util/logging.h"
	#include "util/metrics.h"

	using boost::filesystem::path;
	using std::string;
	using std::vector;
	using std::set;

	namespace doris {

	class TmpFileMgrTest : public ::testing::Test {
	protected:
	// Check that metric values are consistent with TmpFileMgr state.
	void check_metrics(TmpFileMgr* tmp_file_mgr) {
	std::vector<TmpFileMgr::DeviceId> active = tmp_file_mgr->active_tmp_devices();
	int64_t active_metric = DorisMetrics::instance()
	->metric_registry()
	->get_entity("server")
	->get_metric("active_scratch_dirs")
	.value();
	EXPECT_EQ(active.size(), active_metric);
	}
	};

	// Regression test for IMPALA-2160. Verify that temporary file manager allocates blocks
	// at the expected file offsets and expands the temporary file to the correct size.
	TEST_F(TmpFileMgrTest, TestFileAllocation) {
	TmpFileMgr tmp_file_mgr;
	EXPECT_TRUE(tmp_file_mgr.init().ok());
	// Default configuration should give us one temporary device.
	EXPECT_EQ(1, tmp_file_mgr.num_active_tmp_devices());
	std::vector<TmpFileMgr::DeviceId> tmp_devices = tmp_file_mgr.active_tmp_devices();
	EXPECT_EQ(1, tmp_devices.size());
	TUniqueId id;
	TmpFileMgr::File* file;
	Status status = tmp_file_mgr.get_file(tmp_devices[0], id, &file);
	EXPECT_TRUE(status.ok());
	EXPECT_TRUE(file != NULL);
	// Apply writes of variable sizes and check space was allocated correctly.
	int64_t write_sizes[] = {1, 10, 1024, 4, 1024 * 1024 * 8, 1024 * 1024 * 8, 16, 10};
	int num_write_sizes = sizeof(write_sizes) / sizeof(write_sizes[0]);
	int64_t next_offset = 0;
	for (int i = 0; i < num_write_sizes; ++i) {
	int64_t offset;
	status = file->allocate_space(write_sizes[i], &offset);
	EXPECT_TRUE(status.ok());
	EXPECT_EQ(next_offset, offset);
	next_offset = offset + write_sizes[i];
	EXPECT_EQ(next_offset, boost::filesystem::file_size(file->path()));
	}
	// Check that cleanup is correct.
	status = file->remove();
	EXPECT_TRUE(status.ok());
	EXPECT_FALSE(boost::filesystem::exists(file->path()));
	// check_metrics(&tmp_file_mgr);
	}
	// Test that we can do initialization with two directories on same device and
	// that validations prevents duplication of directories.
	TEST_F(TmpFileMgrTest, TestOneDirPerDevice) {
	std::vector<string> tmp_dirs;
	tmp_dirs.push_back("/tmp/tmp-file-mgr-test.1");
	tmp_dirs.push_back("/tmp/tmp-file-mgr-test.2");
	for (int i = 0; i < tmp_dirs.size(); ++i) {
	EXPECT_TRUE(FileSystemUtil::create_directory(tmp_dirs[i]).ok());
	}
	TmpFileMgr tmp_file_mgr;
	tmp_file_mgr.init_custom(tmp_dirs, true);

	// Only the first directory should be used.
	EXPECT_EQ(1, tmp_file_mgr.num_active_tmp_devices());
	std::vector<TmpFileMgr::DeviceId> devices = tmp_file_mgr.active_tmp_devices();
	EXPECT_EQ(1, devices.size());
	TUniqueId id;
	TmpFileMgr::File* file;
	EXPECT_TRUE(tmp_file_mgr.get_file(devices[0], id, &file).ok());
	// Check the prefix is the expected temporary directory.
	EXPECT_EQ(0, file->path().find(tmp_dirs[0]));
	FileSystemUtil::remove_paths(tmp_dirs);
	// check_metrics(&tmp_file_mgr);
	}

	// Test that we can do custom initialization with two dirs on same device.
	TEST_F(TmpFileMgrTest, TestMultiDirsPerDevice) {
	std::vector<string> tmp_dirs;
	tmp_dirs.push_back("/tmp/tmp-file-mgr-test.1");
	tmp_dirs.push_back("/tmp/tmp-file-mgr-test.2");
	for (int i = 0; i < tmp_dirs.size(); ++i) {
	EXPECT_TRUE(FileSystemUtil::create_directory(tmp_dirs[i]).ok());
	}
	TmpFileMgr tmp_file_mgr;
	tmp_file_mgr.init_custom(tmp_dirs, false);

	// Both directories should be used.
	EXPECT_EQ(2, tmp_file_mgr.num_active_tmp_devices());
	std::vector<TmpFileMgr::DeviceId> devices = tmp_file_mgr.active_tmp_devices();
	EXPECT_EQ(2, devices.size());
	for (int i = 0; i < tmp_dirs.size(); ++i) {
	EXPECT_EQ(0, tmp_file_mgr.get_tmp_dir_path(devices[i]).find(tmp_dirs[i]));
	TUniqueId id;
	TmpFileMgr::File* file;
	EXPECT_TRUE(tmp_file_mgr.get_file(devices[i], id, &file).ok());
	// Check the prefix is the expected temporary directory.
	EXPECT_EQ(0, file->path().find(tmp_dirs[i]));
	}
	FileSystemUtil::remove_paths(tmp_dirs);
	// check_metrics(&tmp_file_mgr);
	}

	// Test that reporting a write error is possible but does not result in
	// blacklisting, which is disabled.
	TEST_F(TmpFileMgrTest, TestReportError) {
	std::vector<string> tmp_dirs;
	tmp_dirs.push_back("/tmp/tmp-file-mgr-test.1");
	tmp_dirs.push_back("/tmp/tmp-file-mgr-test.2");
	for (int i = 0; i < tmp_dirs.size(); ++i) {
	EXPECT_TRUE(FileSystemUtil::create_directory(tmp_dirs[i]).ok());
	}
	TmpFileMgr tmp_file_mgr;
	tmp_file_mgr.init_custom(tmp_dirs, false);

	// Both directories should be used.
	std::vector<TmpFileMgr::DeviceId> devices = tmp_file_mgr.active_tmp_devices();
	EXPECT_EQ(2, devices.size());
	// check_metrics(&tmp_file_mgr);

	// Inject an error on one device so that we can validate it is handled correctly.
	TUniqueId id;
	int good_device = 0;
	int bad_device = 1;
	TmpFileMgr::File* bad_file;
	EXPECT_TRUE(tmp_file_mgr.get_file(devices[bad_device], id, &bad_file).ok());
	// ErrorMsg errmsg(TErrorCode::GENERAL, "A fake error");
	// bad_file->ReportIOError(errmsg);
	bad_file->report_io_error("A fake error");

	// Blacklisting is disabled.
	EXPECT_FALSE(bad_file->is_blacklisted());
	// The second device should still be active.
	EXPECT_EQ(2, tmp_file_mgr.num_active_tmp_devices());
	std::vector<TmpFileMgr::DeviceId> devices_after = tmp_file_mgr.active_tmp_devices();
	EXPECT_EQ(2, devices_after.size());
	// check_metrics(&tmp_file_mgr);

	// Attempts to expand bad file should succeed.
	int64_t offset;
	EXPECT_TRUE(bad_file->allocate_space(128, &offset).ok());
	EXPECT_TRUE(bad_file->remove().ok());
	// The good device should still be usable.
	TmpFileMgr::File* good_file;
	EXPECT_TRUE(tmp_file_mgr.get_file(devices[good_device], id, &good_file).ok());
	EXPECT_TRUE(good_file != NULL);
	EXPECT_TRUE(good_file->allocate_space(128, &offset).ok());
	// Attempts to allocate new files on bad device should succeed.
	EXPECT_TRUE(tmp_file_mgr.get_file(devices[bad_device], id, &bad_file).ok());
	FileSystemUtil::remove_paths(tmp_dirs);
	// check_metrics(&tmp_file_mgr);
	}

	TEST_F(TmpFileMgrTest, TestAllocateFails) {
	string tmp_dir("/tmp/tmp-file-mgr-test.1");
	string scratch_subdir = tmp_dir + "/doris-scratch";
	std::vector<string> tmp_dirs(1, tmp_dir);
	EXPECT_TRUE(FileSystemUtil::create_directory(tmp_dir).ok());
	TmpFileMgr tmp_file_mgr;
	tmp_file_mgr.init_custom(tmp_dirs, false);

	TUniqueId id;
	TmpFileMgr::File* allocated_file1;
	TmpFileMgr::File* allocated_file2;
	int64_t offset;
	EXPECT_TRUE(tmp_file_mgr.get_file(0, id, &allocated_file1).ok());
	EXPECT_TRUE(tmp_file_mgr.get_file(0, id, &allocated_file2).ok());
	EXPECT_TRUE(allocated_file1->allocate_space(1, &offset).ok());

	// Make scratch non-writable and test for allocation errors at different stages:
	// new file creation, files with no allocated blocks. files with allocated space.
	chmod(scratch_subdir.c_str(), 0);
	// allocated_file1 already has space allocated.
	EXPECT_FALSE(allocated_file1->allocate_space(1, &offset).ok());
	// allocated_file2 has no space allocated.
	EXPECT_FALSE(allocated_file2->allocate_space(1, &offset).ok());
	// Creating a new File object can succeed because it is not immediately created on disk.
	TmpFileMgr::File* unallocated_file;
	EXPECT_TRUE(tmp_file_mgr.get_file(0, id, &unallocated_file).ok());

	chmod(scratch_subdir.c_str(), S_IRWXU);
	FileSystemUtil::remove_paths(tmp_dirs);
	}

	} // end namespace doris

	int main(int argc, char** argv) {
	// std::string conffile = std::string(getenv("DORIS_HOME")) + "/conf/be.conf";
	// if (!doris::config::init(conffile.c_str(), false)) {
	// fprintf(stderr, "error read config file. \n");
	// return -1;
	// }
	doris::config::query_scratch_dirs = "/tmp";
	doris::init_glog("be-test");
	::testing::InitGoogleTest(&argc, argv);

	doris::DiskInfo::init();

	return RUN_ALL_TESTS();
	}