Add faulty injection in writer fuzzer

kewang1024 · kewang1024 · commit 7a347b493dcb · 2024-10-29T09:45:22.000-07:00
diff --git a/velox/common/file/tests/FaultyFileSystem.h b/velox/common/file/tests/FaultyFileSystem.h
@@ -55,11 +55,11 @@ class FaultyFileSystem : public FileSystem {
 
   std::unique_ptr<ReadFile> openFileForRead(
       std::string_view path,
-      const FileOptions& options) override;
+      const FileOptions& options = {}) override;
 
   std::unique_ptr<WriteFile> openFileForWrite(
       std::string_view path,
-      const FileOptions& options) override;
+      const FileOptions& options = {}) override;
 
   void remove(std::string_view path) override;
 
diff --git a/velox/dwio/common/tests/CMakeLists.txt b/velox/dwio/common/tests/CMakeLists.txt
@@ -13,6 +13,8 @@
 # limitations under the License.
 
 add_subdirectory(utils)
+velox_add_library(velox_dwio_faulty_file_sink FaultyFileSink.cpp)
+velox_link_libraries(velox_dwio_faulty_file_sink velox_file_test_utils)
 
 add_executable(
   velox_dwio_common_test
diff --git a/velox/dwio/common/tests/FaultyFileSink.cpp b/velox/dwio/common/tests/FaultyFileSink.cpp
@@ -0,0 +1,69 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed 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 "velox/common/base/Fs.h"
+#include "velox/common/file/FileSystems.h"
+#include "velox/dwio/common/FileSink.h"
+#include "velox/dwio/common/exception/Exception.h"
+#include "velox/dwio/common/tests/FaultyFileSink.h"
+#include "velox/common/file/tests/FaultyFileSystem.h"
+
+namespace facebook::velox::dwio::common {
+namespace {
+constexpr std::string_view kFileScheme("faulty:");
+constexpr std::string_view kFileSep("/");
+
+std::unique_ptr<FileSink> faultyFileSink(
+    const std::string& filePath,
+    const FileSink::Options& options) {
+  if (filePath.find(kFileScheme) == 0) {
+    return std::make_unique<FaultyFileSink>(filePath, options);
+  }
+  return nullptr;
+}
+} // namespace
+
+FaultyFileSink::FaultyFileSink(const std::string& name, const Options& options)
+    : FileSink{name, options}, writeFile_() {
+  const auto dir = fs::path(name.substr(7)).parent_path();
+  if (!fs::exists(dir)) {
+    VELOX_CHECK(velox::common::generateFileDirectory(dir.c_str()));
+  }
+  auto fs = std::dynamic_pointer_cast<tests::utils::FaultyFileSystem>(
+      filesystems::getFileSystem(name_, nullptr));
+  writeFile_ = fs->openFileForWrite(name_);
+}
+
+void FaultyFileSink::write(std::vector<DataBuffer<char>>& buffers) {
+  writeImpl(buffers, [&](auto& buffer) {
+    const uint64_t size = buffer.size();
+    writeFile_->append({buffer.data(), size});
+    return size;
+  });
+}
+
+void FaultyFileSink::doClose() {
+  if (writeFile_ != nullptr) {
+    writeFile_->close();
+  }
+}
+
+VELOX_REGISTER_DATA_SINK_METHOD_DEFINITION(FaultyFileSink, faultyFileSink);
+
+void registerFaultyFileSinks() {
+  dwio::common::FaultyFileSink::registerFactory();
+}
+} // namespace facebook::velox::dwio::common
diff --git a/velox/dwio/common/tests/FaultyFileSink.h b/velox/dwio/common/tests/FaultyFileSink.h
@@ -0,0 +1,56 @@
+/*
+* Copyright (c) Facebook, Inc. and its affiliates.
+*
+* Licensed 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.
+*/
+
+#pragma once
+
+#include <chrono>
+
+#include "velox/common/config/Config.h"
+#include "velox/common/file/File.h"
+#include "velox/common/io/IoStatistics.h"
+#include "velox/dwio/common/Closeable.h"
+#include "velox/dwio/common/DataBuffer.h"
+#include "velox/dwio/common/MetricsLog.h"
+#include "velox/dwio/common/FileSink.h"
+#include "velox/common/file/tests/FaultyFile.h"
+
+namespace facebook::velox::dwio::common {
+using namespace facebook::velox::io;
+
+class FaultyFileSink : public FileSink {
+public:
+ FaultyFileSink(const std::string& name, const Options& options);
+
+ ~FaultyFileSink() override {
+   destroy();
+ }
+
+ using FileSink::write;
+
+ void write(std::vector<DataBuffer<char>>& buffers) override;
+
+ static void registerFactory();
+
+protected:
+ void doClose() override;
+
+private:
+ std::unique_ptr<WriteFile> writeFile_;
+};
+
+void registerFaultyFileSinks();
+
+} // namespace facebook::velox::dwio::common
diff --git a/velox/exec/fuzzer/CMakeLists.txt b/velox/exec/fuzzer/CMakeLists.txt
@@ -29,7 +29,8 @@ target_link_libraries(
   velox_hive_connector
   velox_dwio_dwrf_reader
   velox_dwio_dwrf_writer
-  velox_dwio_catalog_fbhive)
+  velox_dwio_catalog_fbhive
+  velox_dwio_faulty_file_sink)
 
 add_library(velox_aggregation_fuzzer_base AggregationFuzzerBase.cpp)
 
diff --git a/velox/exec/fuzzer/WriterFuzzer.cpp b/velox/exec/fuzzer/WriterFuzzer.cpp
@@ -22,6 +22,7 @@
 #include "velox/common/base/Fs.h"
 #include "velox/common/encode/Base64.h"
 #include "velox/common/file/FileSystems.h"
+#include "velox/common/file/tests/FaultyFileSystem.h"
 #include "velox/connectors/hive/HiveConnector.h"
 #include "velox/connectors/hive/HiveConnectorSplit.h"
 #include "velox/connectors/hive/TableHandle.h"
@@ -123,7 +124,8 @@ class WriterFuzzer {
       const std::vector<std::string>& bucketColumns,
       int32_t sortColumnOffset,
       const std::vector<std::shared_ptr<const HiveSortingColumn>>& sortBy,
-      const std::string& outputDirectoryPath);
+      const std::string& outputDirectoryPath,
+      const bool writeErrorInjected);
 
   // Generates table column handles based on table column properties
   std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
@@ -256,7 +258,8 @@ void writerFuzzer(
 
 std::vector<std::string> listFolders(std::string_view path) {
   std::vector<std::string> folders;
-  auto fileSystem = filesystems::getFileSystem("/", nullptr);
+  auto fileSystem = std::dynamic_pointer_cast<tests::utils::FaultyFileSystem>(
+      filesystems::getFileSystem("/", nullptr));
   for (auto& p : std::filesystem::recursive_directory_iterator(
            fileSystem->extractPath(path))) {
     if (p.is_directory())
@@ -340,7 +343,26 @@ void WriterFuzzer::go() {
     }
     auto input = generateInputData(names, types, partitionOffset);
 
-    auto tempDirPath = exec::test::TempDirectoryPath::create();
+    auto tempDirPath = exec::test::TempDirectoryPath::create(true);
+
+    bool writeErrorInjected = false;
+    auto fs = std::dynamic_pointer_cast<tests::utils::FaultyFileSystem>(
+        filesystems::getFileSystem(tempDirPath->getPath(), {}));
+    if (partitionKeys.empty() && vectorFuzzer_.coinToss(1)) {
+      fs->setFileInjectionError(
+          std::make_exception_ptr(VeloxRuntimeError(
+              "file_name",
+              1,
+              "function_name()",
+              "operator()",
+              "test message",
+              "",
+              error_code::kArithmeticError,
+              false)),
+          {tests::utils::FaultFileOperation::Type::kWrite});
+      writeErrorInjected = true;
+    }
+
     verifyWriter(
         input,
         names,
@@ -351,7 +373,10 @@ void WriterFuzzer::go() {
         bucketColumns,
         sortColumnOffset,
         sortBy,
-        tempDirPath->getPath());
+        tempDirPath->getPath(),
+        writeErrorInjected);
+
+    fs->clearFileFaultInjections();
 
     LOG(INFO) << "==============================> Done with iteration "
               << iteration++;
@@ -423,7 +448,8 @@ void WriterFuzzer::verifyWriter(
     const std::vector<std::string>& bucketColumns,
     const int32_t sortColumnOffset,
     const std::vector<std::shared_ptr<const HiveSortingColumn>>& sortBy,
-    const std::string& outputDirectoryPath) {
+    const std::string& outputDirectoryPath,
+    const bool writeErrorInjected) {
   const auto plan = PlanBuilder()
                         .values(input)
                         .tableWrite(
@@ -436,7 +462,14 @@ void WriterFuzzer::verifyWriter(
 
   const auto maxDrivers =
       boost::random::uniform_int_distribution<int32_t>(1, 16)(rng_);
-  const auto result = veloxToPrestoResult(execute(plan, maxDrivers));
+  RowVectorPtr result;
+  try {
+    result = veloxToPrestoResult(execute(plan, maxDrivers));
+  } catch (VeloxRuntimeError&) {
+    VELOX_CHECK(
+        writeErrorInjected, "write plan failed with no writeErrorInjected");
+    return;
+  }
 
   const auto dropSql = "DROP TABLE IF EXISTS tmp_write";
   const auto sql = referenceQueryRunner_->toSql(plan).value();
@@ -661,7 +694,8 @@ void WriterFuzzer::comparePartitionAndBucket(
 // static
 std::map<std::string, int32_t> WriterFuzzer::getPartitionNameAndFilecount(
     const std::string& tableDirectoryPath) {
-  auto fileSystem = filesystems::getFileSystem("/", nullptr);
+  auto fileSystem = std::dynamic_pointer_cast<tests::utils::FaultyFileSystem>(
+      filesystems::getFileSystem("/", nullptr));
   auto directories = listFolders(tableDirectoryPath);
   std::map<std::string, int32_t> partitionNameAndFileCount;
 
diff --git a/velox/exec/fuzzer/WriterFuzzerRunner.h b/velox/exec/fuzzer/WriterFuzzerRunner.h
@@ -23,8 +23,10 @@
 #include <vector>
 
 #include "velox/common/file/FileSystems.h"
+#include "velox/common/file/tests/FaultyFileSystem.h"
 #include "velox/connectors/hive/HiveConnector.h"
 #include "velox/dwio/common/FileSink.h"
+#include "velox/dwio/common/tests/FaultyFileSink.h"
 #include "velox/dwio/dwrf/RegisterDwrfReader.h"
 #include "velox/dwio/dwrf/RegisterDwrfWriter.h"
 #include "velox/exec/fuzzer/FuzzerUtil.h"
@@ -74,6 +76,7 @@ class WriterFuzzerRunner {
       size_t seed,
       std::unique_ptr<ReferenceQueryRunner> referenceQueryRunner) {
     filesystems::registerLocalFileSystem();
+    tests::utils::registerFaultyFileSystem();
     connector::registerConnectorFactory(
         std::make_shared<connector::hive::HiveConnectorFactory>());
     auto hiveConnector =
@@ -87,6 +90,7 @@ class WriterFuzzerRunner {
     dwrf::registerDwrfReaderFactory();
     dwrf::registerDwrfWriterFactory();
     dwio::common::registerFileSinks();
+    dwio::common::registerFaultyFileSinks();
     facebook::velox::exec::test::writerFuzzer(
         seed, std::move(referenceQueryRunner));
     // Calling gtest here so that it can be recognized as tests in CI systems.
