From 713597a10dde25049db4ee6a9877369387814f86 Mon Sep 17 00:00:00 2001
From: Wenxuan <breezewish@outlook.com>
Date: Tue, 9 Apr 2024 11:25:37 +0800
Subject: [PATCH] storage: Fix null pointer crash in FileCache (#175)

Signed-off-by: Wish <breezewish@outlook.com>
---
 dbms/src/Common/FailPoint.cpp                 |  1 +
 .../tests/gtest_dm_vector_index.cpp           | 27 +++++++++++++++++++
 dbms/src/Storages/S3/FileCache.cpp            |  9 ++++++-
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/dbms/src/Common/FailPoint.cpp b/dbms/src/Common/FailPoint.cpp
index 1c909b9b36c..50019e4d2b1 100644
--- a/dbms/src/Common/FailPoint.cpp
+++ b/dbms/src/Common/FailPoint.cpp
@@ -106,6 +106,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
index 67c6c2dce42..4633b85c0d3 100644
--- a/dbms/src/Storages/DeltaMerge/tests/gtest_dm_vector_index.cpp
+++ b/dbms/src/Storages/DeltaMerge/tests/gtest_dm_vector_index.cpp
@@ -36,6 +36,7 @@
 #include <gtest/gtest.h>
 #include <tipb/executor.pb.h>
 
+#include <ext/scope_guard.h>
 #include <filesystem>
 
 namespace CurrentMetrics
@@ -46,6 +47,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
@@ -1879,5 +1881,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
index 712cdf59819..a83a2ad71be 100644
--- 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>
@@ -53,6 +54,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;
@@ -284,7 +290,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",
@@ -680,6 +686,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);
     }