storage: Fix null pointer crash in FileCache (#175)

Signed-off-by: Wish <breezewish@outlook.com>
pingcap · Aug 6, 2024 · 947945e · 947945e
1 parent 5b2725b
commit 947945e
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 3 deletions.
diff --git a/dbms/src/Common/FailPoint.cpp b/dbms/src/Common/FailPoint.cpp
@@ -107,6 +107,7 @@ namespace DB
     M(proactive_flush_force_set_type)                        \
     M(exception_when_fetch_disagg_pages)                     \
     M(cop_send_failure)                                      \
+    M(file_cache_fg_download_fail)                           \
     M(force_set_parallel_prehandle_threshold)                \
     M(force_raise_prehandle_exception)                       \
     M(force_agg_on_partial_block)                            \

diff --git a/dbms/src/Storages/DeltaMerge/tests/gtest_dm_vector_index.cpp b/dbms/src/Storages/DeltaMerge/tests/gtest_dm_vector_index.cpp
@@ -28,16 +28,16 @@
 #include <Storages/KVStore/KVStore.h>
 #include <Storages/PathPool.h>
 #include <Storages/S3/FileCache.h>
+#include <Storages/S3/FileCachePerf.h>
 #include <TestUtils/FunctionTestUtils.h>
 #include <TestUtils/InputStreamTestUtils.h>
 #include <TestUtils/TiFlashStorageTestBasic.h>
 #include <gtest/gtest.h>
 #include <tipb/executor.pb.h>
 
+#include <ext/scope_guard.h>
 #include <filesystem>
 
-#include "Storages/S3/FileCachePerf.h"
-
 namespace CurrentMetrics
 {
 extern const Metric DT_SnapshotOfRead;
@@ -46,6 +46,7 @@ extern const Metric DT_SnapshotOfRead;
 namespace DB::FailPoints
 {
 extern const char force_use_dmfile_format_v3[];
+extern const char file_cache_fg_download_fail[];
 } // namespace DB::FailPoints
 
 namespace DB::DM::tests
@@ -1878,5 +1879,30 @@ try
 }
 CATCH
 
+TEST_F(VectorIndexSegmentOnS3Test, S3Failure)
+try
+{
+    prepareWriteNodeStable();
+    DB::FailPointHelper::enableFailPoint(DB::FailPoints::file_cache_fg_download_fail);
+    SCOPE_EXIT({ DB::FailPointHelper::disableFailPoint(DB::FailPoints::file_cache_fg_download_fail); });
+
+    {
+        auto * file_cache = FileCache::instance();
+        ASSERT_EQ(0, file_cache->getAll().size());
+    }
+    {
+        auto scan_context = std::make_shared<ScanContext>();
+        auto stream = computeNodeANNQuery({5.0}, 1, scan_context);
+
+        ASSERT_THROW(
+            {
+                stream->readPrefix();
+                stream->read();
+            },
+            DB::Exception);
+    }
+}
+CATCH
+
 
 } // namespace DB::DM::tests
diff --git a/dbms/src/Storages/S3/FileCache.cpp b/dbms/src/Storages/S3/FileCache.cpp
@@ -14,6 +14,7 @@
 
 #include <Common/CurrentMetrics.h>
 #include <Common/Exception.h>
+#include <Common/FailPoint.h>
 #include <Common/ProfileEvents.h>
 #include <Common/Stopwatch.h>
 #include <Common/SyncPoint/SyncPoint.h>
@@ -54,6 +55,11 @@ extern const int S3_ERROR;
 extern const int FILE_DOESNT_EXIST;
 } // namespace DB::ErrorCodes
 
+namespace DB::FailPoints
+{
+extern const char file_cache_fg_download_fail[];
+} // namespace DB::FailPoints
+
 namespace DB
 {
 using FileType = FileSegment::FileType;
@@ -283,7 +289,7 @@ FileSegmentPtr FileCache::getOrWait(const S3::S3FilenameView & s3_fname, const s
 
     PerfContext::file_cache.fg_download_from_s3++;
     fgDownload(lock, s3_key, file_seg);
-    if (!file_seg->isReadyToRead())
+    if (!file_seg || !file_seg->isReadyToRead())
         throw Exception( //
             ErrorCodes::S3_ERROR,
             "Download object {} failed",
@@ -679,6 +685,7 @@ void FileCache::fgDownload(std::unique_lock<std::mutex> & cache_lock, const Stri
 
     try
     {
+        FAIL_POINT_TRIGGER_EXCEPTION(FailPoints::file_cache_fg_download_fail);
         GET_METRIC(tiflash_storage_remote_cache, type_dtfile_download).Increment();
         downloadImpl(s3_key, file_seg);
     }