MB-48399: Prevent shard count change of initialized magma bucket

BenHuddleston · BenHuddleston · commit 584d3a1f864c · 2021-10-11T09:26:14.000Z
Prevent the shard count of a magma bucket from changing by writing a shard count to a file during first bucket initialization then on subsequent initialization read that file and use the previously persisted shard count. Change-Id: I1961c2664ee397540169bd5e3fb7bb940a3b091d Reviewed-on: http://review.couchbase.org/c/kv_engine/+/161959 Reviewed-by: Dave Rigby <daver@couchbase.com> Tested-by: Build Bot <build@couchbase.com>
diff --git a/engines/ep/src/ep_engine.cc b/engines/ep/src/ep_engine.cc
@@ -50,6 +50,8 @@
 #include <boost/algorithm/string/classification.hpp>
 #include <boost/algorithm/string/split.hpp>
 #include <boost/algorithm/string/trim.hpp>
+#include <boost/filesystem.hpp>
+#include <folly/portability/Unistd.h>
 #include <logger/logger.h>
 #include <memcached/audit_interface.h>
 #include <memcached/collections.h>
@@ -1933,6 +1935,96 @@ class EpEngineValueChangeListener : public ValueChangedListener {
     EventuallyPersistentEngine &engine;
 };
 
+size_t EventuallyPersistentEngine::getShardCount() {
+    auto configShardCount = configuration.getMaxNumShards();
+    if (configuration.getBackend() != "magma") {
+        return configuration.getMaxNumShards();
+    }
+
+    auto diskShardCount = getShardCountFromDisk();
+    if (!diskShardCount) {
+        return configShardCount;
+    }
+
+    return diskShardCount.value();
+}
+
+constexpr std::string_view magmaShardFile = "/magmaShardCount";
+
+std::optional<size_t> EventuallyPersistentEngine::getShardCountFromDisk() {
+    Expects(configuration.getBackend() == "magma");
+
+    // Look for the file
+    const auto shardFile = boost::filesystem::path(
+            configuration.getDbname().append(magmaShardFile));
+    if (boost::filesystem::exists(shardFile)) {
+        boost::filesystem::ifstream ifs(shardFile);
+        std::string data;
+        std::getline(ifs, data);
+        EP_LOG_INFO("Found shard file for magma with {} shards", data);
+        uint64_t shards;
+        if (safe_strtoull(data, shards)) {
+            return shards;
+        }
+
+        auto msg = "Couldn't read shard file or found invalid data";
+        EP_LOG_CRITICAL_RAW(msg);
+        throw std::logic_error(msg);
+    }
+
+    return {};
+}
+
+void EventuallyPersistentEngine::maybeSaveShardCount(WorkLoadPolicy& workload) {
+    if (configuration.getBackend() == "magma") {
+        // We should have created this directory already
+        Expects(boost::filesystem::exists(configuration.getDbname()));
+
+        const auto shardFilePath = boost::filesystem::path(
+                configuration.getDbname().append(magmaShardFile));
+
+        if (boost::filesystem::exists(shardFilePath)) {
+            // File already exists, don't overwrite it (we should have the same
+            // of shards, it's just pointless).
+            return;
+        }
+
+        auto* file = fopen(shardFilePath.string().c_str(), "w");
+        if (!file) {
+            throw std::runtime_error(
+                    "EventuallyPersistentEngine::maybeSaveShardCount: Could "
+                    "not load magma shard file");
+        }
+
+        auto shardStr = std::to_string(workload.getNumShards());
+
+        auto count = fwrite(shardStr.data(), shardStr.size(), 1, file);
+        if (!count) {
+            throw std::runtime_error(
+                    "EventuallyPersistentEngine::maybeSaveShardCount: Error "
+                    "writing shard count to file");
+        }
+
+        auto ret = fflush(file);
+        if (ret != 0) {
+            throw std::runtime_error(
+                    "EventuallyPersistentEngine::maybeSaveShardCount: Error "
+                    "flushing shard file: " +
+                    std::to_string(ret));
+        }
+
+        ret = fclose(file);
+        if (ret != 0) {
+            throw std::runtime_error(
+                    "EventuallyPersistentEngine::maybeSaveShardCount: Error "
+                    "closing shard file: " +
+                    std::to_string(ret));
+        }
+
+        Ensures(boost::filesystem::exists(shardFilePath));
+    }
+}
+
 cb::engine_errc EventuallyPersistentEngine::initialize(const char* config) {
     auto switchToEngine = acquireEngine(this);
     resetStats();
@@ -2020,9 +2112,16 @@ cb::engine_errc EventuallyPersistentEngine::initialize(const char* config) {
             "allow_sanitize_value_in_deletion",
             std::make_unique<EpEngineValueChangeListener>(*this));
 
-    auto numShards = configuration.getMaxNumShards();
+    // The number of shards for a magma bucket cannot be changed after the first
+    // bucket instantiation. This is because the number of shards determines
+    // the on disk structure of the data. To solve this problem we store a file
+    // to the data directory on first bucket creation that tells us how many
+    // shards are to be used. We read this file here if it exists and use that
+    // number, if not, this should be the first bucket creation.
     workload = std::make_unique<WorkLoadPolicy>(
-            configuration.getMaxNumWorkers(), numShards);
+            configuration.getMaxNumWorkers(), getShardCount());
+
+    maybeSaveShardCount(*workload);
 
     const auto& confResMode = configuration.getConflictResolutionType();
     if (!setConflictResolutionMode(confResMode)) {
@@ -6493,6 +6592,7 @@ cb::engine_errc EventuallyPersistentEngine::deleteVBucketInner(
         break;
 
     case cb::engine_errc::not_my_vbucket:
+
         EP_LOG_WARN(
                 "Deletion of {} failed because the vbucket doesn't exist!!!",
                 vbid);
diff --git a/engines/ep/src/ep_engine.h b/engines/ep/src/ep_engine.h
@@ -1273,6 +1273,24 @@ class EventuallyPersistentEngine : public EngineIface, public DcpIface {
                                                const cb::mcbp::Request& request,
                                                const AddResponseFn& response);
 
+    /**
+     * Get the configured shard count for the bucket
+     * @return shard count
+     */
+    size_t getShardCount();
+
+    /**
+     * Attempt to read the shard count from disk
+     * @return empty optional if it does not exist
+     */
+    std::optional<size_t> getShardCountFromDisk();
+
+    /**
+     * Save the shard count to disk if it does not exist
+     * @param workload
+     */
+    void maybeSaveShardCount(WorkLoadPolicy& workload);
+
 private:
     void doEngineStatsCouchDB(const BucketStatCollector& collector, const EPStats& epstats);
     void doEngineStatsMagma(const BucketStatCollector& collector);
diff --git a/engines/ep/tests/module_tests/evp_engine_test.cc b/engines/ep/tests/module_tests/evp_engine_test.cc
@@ -25,6 +25,7 @@
 #include <executor/cb3_taskqueue.h>
 
 #include <boost/algorithm/string/join.hpp>
+#include <boost/filesystem.hpp>
 #include <configuration_impl.h>
 #include <platform/dirutils.h>
 #include <chrono>
@@ -50,6 +51,10 @@ void EventuallyPersistentEngineTest::SetUp() {
         ExecutorPool::create();
     }
 
+    initializeEngine();
+}
+
+void EventuallyPersistentEngineTest::initializeEngine() {
     // Setup an engine with a single active vBucket.
     EXPECT_EQ(cb::engine_errc::success,
               create_ep_engine_instance(get_mock_server_api, &handle))
@@ -90,13 +95,23 @@ void EventuallyPersistentEngineTest::SetUp() {
 }
 
 void EventuallyPersistentEngineTest::TearDown() {
+    shutdownEngine();
+    ExecutorPool::shutdown();
+    // Cleanup any files we created - ignore if they don't exist
+    try {
+        cb::io::rmrf(test_dbname);
+    } catch (std::system_error& e) {
+        if (e.code() != std::error_code(ENOENT, std::system_category())) {
+            throw e;
+        }
+    }
+}
+
+void EventuallyPersistentEngineTest::shutdownEngine() {
     destroy_mock_cookie(cookie);
     // Need to force the destroy (i.e. pass true) because
     // NonIO threads may have been disabled (see DCPTest subclass).
     engine->destroy(true);
-    ExecutorPool::shutdown();
-    // Cleanup any files we created.
-    cb::io::rmrf(test_dbname);
 }
 
 queued_item EventuallyPersistentEngineTest::store_item(
@@ -367,6 +382,92 @@ TEST_P(DurabilityTest, DurabilityStateStats) {
     expectStatsForVB(vb);
 }
 
+TEST_P(EPEnginePersistentTest, ShardCountsOnSecondBucketInit) {
+    auto originalShardCount = engine->getWorkLoadPolicy().getNumShards();
+    auto newShardCount = originalShardCount + 1;
+
+    // We populate the config with shards from this value in the initialize fn
+    numShards = newShardCount;
+    shutdownEngine();
+    initializeEngine();
+
+    if (bucketType == "persistent_magma") {
+        EXPECT_EQ(originalShardCount,
+                  engine->getWorkLoadPolicy().getNumShards());
+    } else {
+        EXPECT_EQ(newShardCount, engine->getWorkLoadPolicy().getNumShards());
+    }
+}
+
+TEST_P(EPEnginePersistentTest, EngineInitReadOnlyDataDir) {
+    store_item(vbid, "key", "value");
+
+    shutdownEngine();
+
+    // As we're modifying the directory we still need execute permissions or we
+    // can't even look in it.
+    boost::filesystem::permissions(test_dbname,
+                                   boost::filesystem::others_read |
+                                           boost::filesystem::owner_read |
+                                           boost::filesystem::group_read |
+                                           boost::filesystem::others_exe |
+                                           boost::filesystem::owner_exe |
+                                           boost::filesystem::group_exe);
+
+    std::string config = config_string;
+    config += "dbname=" + test_dbname;
+    config += ";max_vbuckets=" + std::to_string(numVbuckets) +
+              ";max_num_shards=" + std::to_string(numShards) + ";";
+
+    // Set the bucketType
+    config += generateBucketTypeConfig(bucketType);
+
+    EXPECT_EQ(cb::engine_errc::success,
+              create_ep_engine_instance(get_mock_server_api, &handle))
+            << "Failed to create ep engine instance";
+    engine = reinterpret_cast<EventuallyPersistentEngine*>(handle);
+    ObjectRegistry::onSwitchThread(engine);
+
+    // Should come up fine, but in some sort of read only mode
+    EXPECT_EQ(cb::engine_errc::success, engine->initialize(config.c_str()));
+
+    // Set the filesystem permissions back for the next test
+    boost::filesystem::permissions(test_dbname, boost::filesystem::all_all);
+
+    // Reset our cookie to have a ptr to the new engine which is required when
+    // we destroy it in TearDown()
+    cookie = create_mock_cookie(engine);
+}
+
+// Tests that engine initializes fine even if the data dir doesn't exist
+TEST_P(EPEnginePersistentTest, EngineInitNoDataDir) {
+    shutdownEngine();
+
+    cb::io::rmrf(test_dbname);
+
+    std::string config = config_string;
+    config += "dbname=" + test_dbname + ";";
+
+    // Set the bucketType
+    config += generateBucketTypeConfig(bucketType);
+
+    EXPECT_EQ(cb::engine_errc::success,
+              create_ep_engine_instance(get_mock_server_api, &handle))
+            << "Failed to create ep engine instance";
+    engine = reinterpret_cast<EventuallyPersistentEngine*>(handle);
+    ObjectRegistry::onSwitchThread(engine);
+
+    EXPECT_EQ(cb::engine_errc::success, engine->initialize(config.c_str()));
+    cookie = create_mock_cookie(engine);
+}
+
+INSTANTIATE_TEST_SUITE_P(Persistent,
+                         EPEnginePersistentTest,
+                         EPEngineParamTest::persistentConfigValues(),
+                         [](const ::testing::TestParamInfo<std::string>& info) {
+                             return info.param;
+                         });
+
 INSTANTIATE_TEST_SUITE_P(EphemeralOrPersistent,
                          DurabilityTest,
                          EPEngineParamTest::allConfigValues(),
diff --git a/engines/ep/tests/module_tests/evp_engine_test.h b/engines/ep/tests/module_tests/evp_engine_test.h
@@ -37,6 +37,16 @@ class EventuallyPersistentEngineTest : virtual public ::testing::Test {
 
     void TearDown() override;
 
+    /**
+     * Set up the engine. Does not touch file or executor pool
+     */
+    void initializeEngine();
+
+    /**
+     * Shutdown the engine. Does not touch files or executor pool
+     */
+    void shutdownEngine();
+
     /* Helper methods for tests */
 
     /* Stores an item into the given vbucket. */
@@ -134,6 +144,8 @@ class EPEngineParamTest : virtual public EventuallyPersistentEngineTest,
     }
 };
 
+class EPEnginePersistentTest : public EPEngineParamTest {};
+
 /*
  * EPEngine-level test fixture for Durability.
  */
