Print storage info memory

Was testing a setup that wasn't used in production, it would
unnecessarily thrash small buffers.

.github/workflows/package_specific.yaml

@@ -0,0 +1,294 @@
+name: Package Specific
+
+# TODO(gitbuda): Cleanup docker container if GHA job was canceled.
+
+on:
+  workflow_dispatch:
+    inputs:
+      memgraph_version:
+        description: "Memgraph version to upload as. Leave this field empty if you don't want to upload binaries to S3. Format: 'X.Y.Z'"
+        required: false
+      build_type:
+        type: choice
+        description: "Memgraph Build type. Default value is Release."
+        default: 'Release'
+        options:
+          - Release
+          - RelWithDebInfo
+      target_os:
+        type: choice
+        description: "Target OS for which memgraph will be packaged. Default is Ubuntu 22.04"
+        default: 'ubuntu-22_04'
+        options:
+          - amzn-2
+          - centos-7
+          - centos-9
+          - debian-10
+          - debian-11
+          - debian-11-arm
+          - debian-11-platform
+          - docker
+          - fedora-36
+          - ubuntu-18_04
+          - ubuntu-20_04
+          - ubuntu-22_04
+          - ubuntu-22_04-arm
+
+jobs:
+  amzn-2:
+    if: ${{ github.event.inputs.target_os == 'amzn-2' }}
+    runs-on: [self-hosted, DockerMgBuild, X64]
+    timeout-minutes: 60
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package amzn-2 ${{ github.event.inputs.build_type }}
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: amzn-2
+          path: build/output/amzn-2/memgraph*.rpm
+
+  centos-7:
+    if: ${{ github.event.inputs.target_os == 'centos-7' }}
+    runs-on: [self-hosted, DockerMgBuild, X64]
+    timeout-minutes: 60
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package centos-7 ${{ github.event.inputs.build_type }}
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: centos-7
+          path: build/output/centos-7/memgraph*.rpm
+
+  centos-9:
+    if: ${{ github.event.inputs.target_os == 'centos-9' }}
+    runs-on: [self-hosted, DockerMgBuild, X64]
+    timeout-minutes: 60
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package centos-9 ${{ github.event.inputs.build_type }}
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: centos-9
+          path: build/output/centos-9/memgraph*.rpm
+
+  debian-10:
+    if: ${{ github.event.inputs.target_os == 'debian-10' }}
+    runs-on: [self-hosted, DockerMgBuild, X64]
+    timeout-minutes: 60
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package debian-10 ${{ github.event.inputs.build_type }}
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: debian-10
+          path: build/output/debian-10/memgraph*.deb
+
+  debian-11:
+    if: ${{ github.event.inputs.target_os == 'debian-11' }}
+    runs-on: [self-hosted, DockerMgBuild, X64]
+    timeout-minutes: 60
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package debian-11 ${{ github.event.inputs.build_type }}
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: debian-11
+          path: build/output/debian-11/memgraph*.deb
+
+  debian-11-arm:
+    if: ${{ github.event.inputs.target_os == 'debian-11-arm' }}
+    runs-on: [self-hosted, DockerMgBuild, ARM64, strange]
+    timeout-minutes: 120
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package debian-11-arm ${{ github.event.inputs.build_type }}
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: debian-11-aarch64
+          path: build/output/debian-11-arm/memgraph*.deb
+
+  debian-11-platform:
+    if: ${{ github.event.inputs.target_os == 'debian-11-platform' }}
+    runs-on: [self-hosted, DockerMgBuild, X64]
+    timeout-minutes: 60
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package debian-11 ${{ github.event.inputs.build_type }} --for-platform
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: debian-11-platform
+          path: build/output/debian-11/memgraph*.deb
+
+  docker:
+    if: ${{ github.event.inputs.target_os == 'docker' }}
+    runs-on: [self-hosted, DockerMgBuild, X64]
+    timeout-minutes: 60
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          cd release/package
+          ./run.sh package debian-11 ${{ github.event.inputs.build_type }} --for-docker
+          ./run.sh docker
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: docker
+          path: build/output/docker/memgraph*.tar.gz
+
+  fedora-36:
+    if: ${{ github.event.inputs.target_os == 'fedora-36' }}
+    runs-on: [self-hosted, DockerMgBuild, X64]
+    timeout-minutes: 60
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package fedora-36 ${{ github.event.inputs.build_type }}
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: fedora-36
+          path: build/output/fedora-36/memgraph*.rpm
+
+  ubuntu-18_04:
+    if: ${{ github.event.inputs.target_os == 'ubuntu-18_04' }}
+    runs-on: [self-hosted, DockerMgBuild, X64]
+    timeout-minutes: 60
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package ubuntu-18.04 ${{ github.event.inputs.build_type }}
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: ubuntu-18.04
+          path: build/output/ubuntu-18.04/memgraph*.deb
+
+  ubuntu-20_04:
+    if: ${{ github.event.inputs.target_os == 'ubuntu-20_04' }}
+    runs-on: [self-hosted, DockerMgBuild, X64]
+    timeout-minutes: 60
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package ubuntu-20.04 ${{ github.event.inputs.build_type }}
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: ubuntu-20.04
+          path: build/output/ubuntu-20.04/memgraph*.deb
+
+  ubuntu-22_04:
+    if: ${{ github.event.inputs.target_os == 'ubuntu-22_04' }}
+    runs-on: [self-hosted, DockerMgBuild, X64]
+    timeout-minutes: 60
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package ubuntu-22.04 ${{ github.event.inputs.build_type }}
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: ubuntu-22.04
+          path: build/output/ubuntu-22.04/memgraph*.deb
+
+  ubuntu-22_04-arm:
+    if: ${{ github.event.inputs.target_os == 'ubuntu-22_04-arm' }}
+    runs-on: [self-hosted, DockerMgBuild, ARM64, strange]
+    timeout-minutes: 120
+    steps:
+      - name: "Set up repository"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # Required because of release/get_version.py
+      - name: "Build package"
+        run: |
+          ./release/package/run.sh package ubuntu-22.04-arm ${{ github.event.inputs.build_type }}
+      - name: "Upload package"
+        uses: actions/upload-artifact@v3
+        with:
+          name: ubuntu-22.04-aarch64
+          path: build/output/ubuntu-22.04-arm/memgraph*.deb
+
+  upload-to-s3:
+    # only run upload if we specified version. Allows for runs without upload
+    if: "${{ github.event.inputs.memgraph_version != '' }}"
+    needs: [amzn-2, centos-7, centos-9, debian-10, debian-11, debian-11-arm, debian-11-platform, docker, fedora-36, ubuntu-18_04, ubuntu-20_04, ubuntu-22_04, ubuntu-22_04-arm]
+    runs-on: ubuntu-latest
+    steps:
+      - name: Download artifacts
+        uses: actions/download-artifact@v3
+        with:
+          # name: # if name input parameter is not provided, all artifacts are downloaded
+                  # and put in directories named after each one.
+          path: build/output/release
+      - name: Upload to S3
+        uses: jakejarvis/s3-sync-action@v0.5.1
+        env:
+          AWS_S3_BUCKET: "download.memgraph.com"
+          AWS_ACCESS_KEY_ID: ${{ secrets.S3_AWS_ACCESS_KEY_ID }}
+          AWS_SECRET_ACCESS_KEY: ${{ secrets.S3_AWS_SECRET_ACCESS_KEY }}
+          AWS_REGION: "eu-west-1"
+          SOURCE_DIR: "build/output/release"
+          DEST_DIR: "memgraph/v${{ github.event.inputs.memgraph_version }}/"

CMakeLists.txt

@@ -42,10 +42,6 @@ endif()
 
 project(memgraph LANGUAGES C CXX)
 
-# NOTE: once in a while this needs to be toggled to check headers are
-#       correct and PCH isn't masking any include issues
-set(CMAKE_DISABLE_PRECOMPILE_HEADERS OFF)
-
 #TODO: upgrade to cmake 3.24 + CheckIPOSupported
 #cmake_policy(SET CMP0138 NEW)
 #include(CheckIPOSupported)

query_modules/CMakeLists.txt

@@ -13,6 +13,7 @@ string(TOLOWER ${CMAKE_BUILD_TYPE} lower_build_type)
 add_library(example_c SHARED example.c)
 target_include_directories(example_c PRIVATE ${CMAKE_SOURCE_DIR}/include)
 target_compile_options(example_c PRIVATE -Wall)
+target_link_libraries(example_c PRIVATE -static-libgcc -static-libstdc++)
 # Strip C example in release build.
 if (lower_build_type STREQUAL "release")
   add_custom_command(TARGET example_c POST_BUILD
@@ -28,6 +29,7 @@ install(FILES example.c DESTINATION lib/memgraph/query_modules/src)
 add_library(example_cpp SHARED example.cpp)
 target_include_directories(example_cpp PRIVATE ${CMAKE_SOURCE_DIR}/include)
 target_compile_options(example_cpp PRIVATE -Wall)
+target_link_libraries(example_cpp PRIVATE -static-libgcc -static-libstdc++)
 # Strip C++ example in release build.
 if (lower_build_type STREQUAL "release")
   add_custom_command(TARGET example_cpp POST_BUILD
@@ -43,6 +45,7 @@ install(FILES example.cpp DESTINATION lib/memgraph/query_modules/src)
 add_library(schema SHARED schema.cpp)
 target_include_directories(schema PRIVATE ${CMAKE_SOURCE_DIR}/include)
 target_compile_options(schema PRIVATE -Wall)
+target_link_libraries(schema PRIVATE -static-libgcc -static-libstdc++)
 # Strip C++ example in release build.
 if (lower_build_type STREQUAL "release")
   add_custom_command(TARGET schema POST_BUILD

src/CMakeLists.txt

@@ -143,7 +143,7 @@ install(CODE "file(MAKE_DIRECTORY \$ENV{DESTDIR}/var/log/memgraph
 # Memgraph CSV Import Tool Executable
 # ----------------------------------------------------------------------------
 add_executable(mg_import_csv mg_import_csv.cpp)
-target_link_libraries(mg_import_csv mg-storage-v2)
+target_link_libraries(mg_import_csv mg-storage-v2 mg-dbms)
 
 # Strip the executable in release build.
 if(lower_build_type STREQUAL "release")

src/communication/CMakeLists.txt

@@ -25,4 +25,3 @@ target_link_libraries(mg-communication Boost::headers Threads::Threads mg-utils
 find_package(OpenSSL REQUIRED)
 target_link_libraries(mg-communication ${OPENSSL_LIBRARIES})
 target_include_directories(mg-communication SYSTEM PUBLIC ${OPENSSL_INCLUDE_DIR})
-target_precompile_headers(mg-communication INTERFACE http/server.hpp <boost/beast/websocket.hpp> bolt/v1/session.hpp)

src/communication/bolt/v1/states/handlers.hpp

@@ -209,7 +209,11 @@ State HandleRunV1(TSession &session, const State state, const Marker marker) {
 
   DMG_ASSERT(!session.encoder_buffer_.HasData(), "There should be no data to write in this state");
 
+#if MG_ENTERPRISE
   spdlog::debug("[Run - {}] '{}'", session.GetCurrentDB(), query.ValueString());
+#else
+  spdlog::debug("[Run] '{}'", query.ValueString());
+#endif
 
   // Increment number of queries in the metrics
   IncrementQueryMetrics(session);
@@ -276,7 +280,11 @@ State HandleRunV4(TSession &session, const State state, const Marker marker) {
     return HandleFailure(session, e);
   }
 
+#if MG_ENTERPRISE
   spdlog::debug("[Run - {}] '{}'", session.GetCurrentDB(), query.ValueString());
+#else
+  spdlog::debug("[Run] '{}'", query.ValueString());
+#endif
 
   // Increment number of queries in the metrics
   IncrementQueryMetrics(session);

src/dbms/CMakeLists.txt

@@ -1,3 +1,3 @@
 
-add_library(mg-dbms STATIC database.cpp)
+add_library(mg-dbms STATIC database.cpp replication_handler.cpp inmemory/replication_handlers.cpp)
 target_link_libraries(mg-dbms mg-utils mg-storage-v2 mg-query)

src/dbms/database.cpp

@@ -10,6 +10,8 @@
 // licenses/APL.txt.
 
 #include "dbms/database.hpp"
+#include "dbms/inmemory/storage_helper.hpp"
+#include "dbms/replication_handler.hpp"
 #include "flags/storage_mode.hpp"
 #include "storage/v2/disk/storage.hpp"
 #include "storage/v2/inmemory/storage.hpp"
@@ -19,14 +21,16 @@ template struct memgraph::utils::Gatekeeper<memgraph::dbms::Database>;
 
 namespace memgraph::dbms {
 
-Database::Database(const storage::Config &config)
+Database::Database(storage::Config config, const replication::ReplicationState &repl_state)
     : trigger_store_(config.durability.storage_directory / "triggers"),
-      streams_{config.durability.storage_directory / "streams"} {
+      streams_{config.durability.storage_directory / "streams"},
+      plan_cache_{FLAGS_query_plan_cache_max_size},
+      repl_state_(&repl_state) {
   if (config.storage_mode == memgraph::storage::StorageMode::ON_DISK_TRANSACTIONAL || config.force_on_disk ||
       utils::DirExists(config.disk.main_storage_directory)) {
-    storage_ = std::make_unique<storage::DiskStorage>(config);
+    storage_ = std::make_unique<storage::DiskStorage>(std::move(config));
   } else {
-    storage_ = std::make_unique<storage::InMemoryStorage>(config, config.storage_mode);
+    storage_ = dbms::CreateInMemoryStorage(std::move(config), repl_state);
   }
 }
 

src/dbms/database.hpp

@@ -24,6 +24,8 @@
 #include "query/trigger.hpp"
 #include "storage/v2/storage.hpp"
 #include "utils/gatekeeper.hpp"
+#include "utils/lru_cache.hpp"
+#include "utils/synchronized.hpp"
 
 namespace memgraph::dbms {
 
@@ -46,7 +48,7 @@ class Database {
    *
    * @param config storage configuration
    */
-  explicit Database(const storage::Config &config);
+  explicit Database(storage::Config config, const replication::ReplicationState &repl_state);
 
   /**
    * @brief Returns the raw storage pointer.
@@ -56,6 +58,7 @@ class Database {
    * @return storage::Storage*
    */
   storage::Storage *storage() { return storage_.get(); }
+  storage::Storage const *storage() const { return storage_.get(); }
 
   /**
    * @brief Storage's Accessor
@@ -65,12 +68,12 @@ class Database {
    */
   std::unique_ptr<storage::Storage::Accessor> Access(
       std::optional<storage::IsolationLevel> override_isolation_level = {}) {
-    return storage_->Access(override_isolation_level);
+    return storage_->Access(override_isolation_level, repl_state_->IsMain());
   }
 
   std::unique_ptr<storage::Storage::Accessor> UniqueAccess(
       std::optional<storage::IsolationLevel> override_isolation_level = {}) {
-    return storage_->UniqueAccess(override_isolation_level);
+    return storage_->UniqueAccess(override_isolation_level, repl_state_->IsMain());
   }
 
   /**
@@ -145,9 +148,9 @@ class Database {
   /**
    * @brief Returns the PlanCache vector raw pointer
    *
-   * @return utils::SkipList<query::PlanCacheEntry>*
+   * @return utils::Synchronized<utils::LRUCache<uint64_t, std::shared_ptr<PlanWrapper>>, utils::RWSpinLock>
    */
-  utils::SkipList<query::PlanCacheEntry> *plan_cache() { return &plan_cache_; }
+  query::PlanCacheLRU *plan_cache() { return &plan_cache_; }
 
  private:
   std::unique_ptr<storage::Storage> storage_;       //!< Underlying storage
@@ -156,7 +159,9 @@ class Database {
   query::stream::Streams streams_;                  //!< Streams associated with the storage
 
   // TODO: Move to a better place
-  utils::SkipList<query::PlanCacheEntry> plan_cache_;  //!< Plan cache associated with the storage
+  query::PlanCacheLRU plan_cache_;  //!< Plan cache associated with the storage
+
+  const replication::ReplicationState *repl_state_;
 };
 
 }  // namespace memgraph::dbms

src/dbms/database_handler.hpp

@@ -51,7 +51,8 @@ class DatabaseHandler : public Handler<Database> {
    * @param config Storage configuration
    * @return HandlerT::NewResult
    */
-  HandlerT::NewResult New(std::string_view name, storage::Config config) {
+  HandlerT::NewResult New(std::string_view name, storage::Config config,
+                          const replication::ReplicationState &repl_state) {
     // Control that no one is using the same data directory
     if (std::any_of(begin(), end(), [&](auto &elem) {
           auto db_acc = elem.second.access();
@@ -62,7 +63,7 @@ class DatabaseHandler : public Handler<Database> {
       return NewError::EXISTS;
     }
     config.name = name;  // Set storage id via config
-    return HandlerT::New(std::piecewise_construct, name, config);
+    return HandlerT::New(std::piecewise_construct, name, config, repl_state);
   }
 
   /**
@@ -93,5 +94,4 @@ class DatabaseHandler : public Handler<Database> {
 };
 
 }  // namespace memgraph::dbms
-
 #endif

src/dbms/dbms_handler.hpp

@@ -26,7 +26,9 @@
 #include "auth/auth.hpp"
 #include "constants.hpp"
 #include "dbms/database.hpp"
+#ifdef MG_ENTERPRISE
 #include "dbms/database_handler.hpp"
+#endif
 #include "global.hpp"
 #include "query/config.hpp"
 #include "query/interpreter_context.hpp"
@@ -81,32 +83,35 @@ static inline nlohmann::json ToJson(const Statistics &stats) {
   return res;
 }
 
-#ifdef MG_ENTERPRISE
-
-using DeleteResult = utils::BasicResult<DeleteError>;
-
 /**
  * @brief Multi-database session contexts handler.
  */
 class DbmsHandler {
  public:
   using LockT = utils::RWLock;
+#ifdef MG_ENTERPRISE
+
   using NewResultT = utils::BasicResult<NewError, DatabaseAccess>;
+  using DeleteResult = utils::BasicResult<DeleteError>;
 
   /**
    * @brief Initialize the handler.
    *
-   * @param configs storage and interpreter configurations
+   * @param configs storage configuration
    * @param auth pointer to the global authenticator
    * @param recovery_on_startup restore databases (and its content) and authentication data
    * @param delete_on_drop when dropping delete any associated directories on disk
    */
-  DbmsHandler(storage::Config config, auto *auth, bool recovery_on_startup, bool delete_on_drop)
-      : lock_{utils::RWLock::Priority::READ}, default_config_{std::move(config)}, delete_on_drop_(delete_on_drop) {
+  DbmsHandler(storage::Config config, const replication::ReplicationState &repl_state, auto *auth,
+              bool recovery_on_startup, bool delete_on_drop)
+      : lock_{utils::RWLock::Priority::READ},
+        default_config_{std::move(config)},
+        repl_state_(repl_state),
+        delete_on_drop_(delete_on_drop) {
     // TODO: Decouple storage config from dbms config
     // TODO: Save individual db configs inside the kvstore and restore from there
-    storage::UpdatePaths(*default_config_, default_config_->durability.storage_directory / "databases");
-    const auto &db_dir = default_config_->durability.storage_directory;
+    storage::UpdatePaths(default_config_, default_config_.durability.storage_directory / "databases");
+    const auto &db_dir = default_config_.durability.storage_directory;
     const auto durability_dir = db_dir / ".durability";
     utils::EnsureDirOrDie(db_dir);
     utils::EnsureDirOrDie(durability_dir);
@@ -114,7 +119,6 @@ class DbmsHandler {
 
     // Generate the default database
     MG_ASSERT(!NewDefault_().HasError(), "Failed while creating the default DB.");
-
     // Recover previous databases
     if (recovery_on_startup) {
       for (const auto &[name, _] : *durability_) {
@@ -132,7 +136,21 @@ class DbmsHandler {
       }
     }
   }
+#else
+  /**
+   * @brief Initialize the handler. A single database is supported in community edition.
+   *
+   * @param configs storage configuration
+   */
+  DbmsHandler(storage::Config config, const replication::ReplicationState &repl_state)
+      : db_gatekeeper_{[&] {
+                         config.name = kDefaultDB;
+                         return std::move(config);
+                       }(),
+                       repl_state} {}
+#endif
 
+#ifdef MG_ENTERPRISE
   /**
    * @brief Create a new Database associated with the "name" database
    *
@@ -151,11 +169,24 @@ class DbmsHandler {
    * @return DatabaseAccess
    * @throw UnknownDatabaseException if database not found
    */
-  DatabaseAccess Get(std::string_view name) {
+  DatabaseAccess Get(std::string_view name = kDefaultDB) {
     std::shared_lock<LockT> rd(lock_);
     return Get_(name);
   }
+#else
+  /**
+   * @brief Get the context associated with the default database
+   *
+   * @return DatabaseAccess
+   */
+  DatabaseAccess Get() {
+    auto acc = db_gatekeeper_.access();
+    MG_ASSERT(acc, "Failed to get default database!");
+    return *acc;
+  }
+#endif
 
+#ifdef MG_ENTERPRISE
   /**
    * @brief Delete database.
    *
@@ -201,6 +232,7 @@ class DbmsHandler {
 
     return {};  // Success
   }
+#endif
 
   /**
    * @brief Return all active databases.
@@ -208,8 +240,12 @@ class DbmsHandler {
    * @return std::vector<std::string>
    */
   std::vector<std::string> All() const {
+#ifdef MG_ENTERPRISE
     std::shared_lock<LockT> rd(lock_);
     return db_handler_.All();
+#else
+    return {db_gatekeeper_.access()->get()->id()};
+#endif
   }
 
   /**
@@ -220,24 +256,30 @@ class DbmsHandler {
   Statistics Stats() {
     Statistics stats{};
     // TODO: Handle overflow?
+#ifdef MG_ENTERPRISE
     std::shared_lock<LockT> rd(lock_);
     for (auto &[_, db_gk] : db_handler_) {
+#else
+    {
+      auto &db_gk = db_gatekeeper_;
+#endif
       auto db_acc_opt = db_gk.access();
-      if (!db_acc_opt) continue;
-      auto &db_acc = *db_acc_opt;
-      const auto &info = db_acc->GetInfo();
-      const auto &storage_info = info.storage_info;
-      stats.num_vertex += storage_info.vertex_count;
-      stats.num_edges += storage_info.edge_count;
-      stats.triggers += info.triggers;
-      stats.streams += info.streams;
-      ++stats.num_databases;
-      stats.indices += storage_info.label_indices + storage_info.label_property_indices;
-      stats.constraints += storage_info.existence_constraints + storage_info.unique_constraints;
-      ++stats.storage_modes[(int)storage_info.storage_mode];
-      ++stats.isolation_levels[(int)storage_info.isolation_level];
-      stats.snapshot_enabled += storage_info.durability_snapshot_enabled;
-      stats.wal_enabled += storage_info.durability_wal_enabled;
+      if (db_acc_opt) {
+        auto &db_acc = *db_acc_opt;
+        const auto &info = db_acc->GetInfo();
+        const auto &storage_info = info.storage_info;
+        stats.num_vertex += storage_info.vertex_count;
+        stats.num_edges += storage_info.edge_count;
+        stats.triggers += info.triggers;
+        stats.streams += info.streams;
+        ++stats.num_databases;
+        stats.indices += storage_info.label_indices + storage_info.label_property_indices;
+        stats.constraints += storage_info.existence_constraints + storage_info.unique_constraints;
+        ++stats.storage_modes[(int)storage_info.storage_mode];
+        ++stats.isolation_levels[(int)storage_info.isolation_level];
+        stats.snapshot_enabled += storage_info.durability_snapshot_enabled;
+        stats.wal_enabled += storage_info.durability_wal_enabled;
+      }
     }
     return stats;
   }
@@ -249,13 +291,19 @@ class DbmsHandler {
    */
   std::vector<DatabaseInfo> Info() {
     std::vector<DatabaseInfo> res;
-    res.reserve(std::distance(db_handler_.cbegin(), db_handler_.cend()));
+#ifdef MG_ENTERPRISE
     std::shared_lock<LockT> rd(lock_);
+    res.reserve(std::distance(db_handler_.cbegin(), db_handler_.cend()));
     for (auto &[_, db_gk] : db_handler_) {
+#else
+    {
+      auto &db_gk = db_gatekeeper_;
+#endif
       auto db_acc_opt = db_gk.access();
-      if (!db_acc_opt) continue;
-      auto &db_acc = *db_acc_opt;
-      res.push_back(db_acc->GetInfo());
+      if (db_acc_opt) {
+        auto &db_acc = *db_acc_opt;
+        res.push_back(db_acc->GetInfo());
+      }
     }
     return res;
   }
@@ -267,15 +315,21 @@ class DbmsHandler {
    * @param ic global InterpreterContext
    */
   void RestoreTriggers(query::InterpreterContext *ic) {
+#ifdef MG_ENTERPRISE
     std::lock_guard<LockT> wr(lock_);
     for (auto &[_, db_gk] : db_handler_) {
+#else
+    {
+      auto &db_gk = db_gatekeeper_;
+#endif
       auto db_acc_opt = db_gk.access();
-      if (!db_acc_opt) continue;
-      auto &db_acc = *db_acc_opt;
-      spdlog::debug("Restoring trigger for database \"{}\"", db_acc->id());
-      auto storage_accessor = db_acc->Access();
-      auto dba = memgraph::query::DbAccessor{storage_accessor.get()};
-      db_acc->trigger_store()->RestoreTriggers(&ic->ast_cache, &dba, ic->config.query, ic->auth_checker);
+      if (db_acc_opt) {
+        auto &db_acc = *db_acc_opt;
+        spdlog::debug("Restoring trigger for database \"{}\"", db_acc->id());
+        auto storage_accessor = db_acc->Access();
+        auto dba = memgraph::query::DbAccessor{storage_accessor.get()};
+        db_acc->trigger_store()->RestoreTriggers(&ic->ast_cache, &dba, ic->config.query, ic->auth_checker);
+      }
     }
   }
 
@@ -286,17 +340,67 @@ class DbmsHandler {
    * @param ic global InterpreterContext
    */
   void RestoreStreams(query::InterpreterContext *ic) {
+#ifdef MG_ENTERPRISE
     std::lock_guard<LockT> wr(lock_);
     for (auto &[_, db_gk] : db_handler_) {
+#else
+    {
+      auto &db_gk = db_gatekeeper_;
+#endif
       auto db_acc = db_gk.access();
-      if (!db_acc) continue;
-      auto *db = db_acc->get();
-      spdlog::debug("Restoring streams for database \"{}\"", db->id());
-      db->streams()->RestoreStreams(*db_acc, ic);
+      if (db_acc) {
+        auto *db = db_acc->get();
+        spdlog::debug("Restoring streams for database \"{}\"", db->id());
+        db->streams()->RestoreStreams(*db_acc, ic);
+      }
     }
   }
 
+  /**
+   * @brief todo
+   *
+   * @param f
+   */
+  void ForEach(auto f) {
+#ifdef MG_ENTERPRISE
+    std::shared_lock<LockT> rd(lock_);
+    for (auto &[_, db_gk] : db_handler_) {
+#else
+    {
+      auto &db_gk = db_gatekeeper_;
+#endif
+      auto db_acc = db_gk.access();
+      if (db_acc) {  // This isn't an error, just a defunct db
+        f(db_acc->get());
+      }
+    }
+  }
+
+  /**
+   * @brief todo
+   *
+   * @param f
+   */
+  void ForOne(auto f) {
+#ifdef MG_ENTERPRISE
+    std::shared_lock<LockT> rd(lock_);
+    for (auto &[_, db_gk] : db_handler_) {
+      auto db_acc = db_gk.access();
+      if (db_acc) {                   // This isn't an error, just a defunct db
+        if (f(db_acc->get())) break;  // Run until the first successful one
+      }
+    }
+#else
+    {
+      auto db_acc = db_gatekeeper_.access();
+      MG_ASSERT(db_acc, "Should always have the database");
+      f(db_acc->get());
+    }
+#endif
+  }
+
  private:
+#ifdef MG_ENTERPRISE
   /**
    * @brief return the storage directory of the associated database
    *
@@ -328,13 +432,9 @@ class DbmsHandler {
    * @return NewResultT context on success, error on failure
    */
   NewResultT New_(const std::string &name, std::filesystem::path storage_subdir) {
-    if (default_config_) {
-      auto config_copy = *default_config_;
-      storage::UpdatePaths(config_copy, default_config_->durability.storage_directory / storage_subdir);
-      return New_(name, config_copy);
-    }
-    spdlog::info("Trying to generate session context without any configurations.");
-    return NewError::NO_CONFIGS;
+    auto config_copy = default_config_;
+    storage::UpdatePaths(config_copy, default_config_.durability.storage_directory / storage_subdir);
+    return New_(name, config_copy);
   }
 
   /**
@@ -351,7 +451,7 @@ class DbmsHandler {
       return NewError::DEFUNCT;
     }
 
-    auto new_db = db_handler_.New(name, storage_config);
+    auto new_db = db_handler_.New(name, storage_config, repl_state_);
     if (new_db.HasValue()) {
       // Success
       if (durability_) durability_->Put(name, "ok");  // TODO: Serialize the configuration?
@@ -436,14 +536,16 @@ class DbmsHandler {
     throw UnknownDatabaseException("Tried to retrieve an unknown database \"{}\".", name);
   }
 
-  // Should storage objects ever be deleted?
-  mutable LockT lock_;                             //!< protective lock
-  DatabaseHandler db_handler_;                     //!< multi-tenancy storage handler
-  std::optional<storage::Config> default_config_;  //!< Storage configuration used when creating new databases
-  std::unique_ptr<kvstore::KVStore> durability_;   //!< list of active dbs (pointer so we can postpone its creation)
-  std::set<std::string> defunct_dbs_;              //!< Databases that are in an unknown state due to various failures
-  bool delete_on_drop_;                            //!< Flag defining if dropping storage also deletes its directory
-};
+  mutable LockT lock_;                               //!< protective lock
+  storage::Config default_config_;                   //!< Storage configuration used when creating new databases
+  const replication::ReplicationState &repl_state_;  //!< Global replication state
+  DatabaseHandler db_handler_;                       //!< multi-tenancy storage handler
+  std::unique_ptr<kvstore::KVStore> durability_;     //!< list of active dbs (pointer so we can postpone its creation)
+  bool delete_on_drop_;                              //!< Flag defining if dropping storage also deletes its directory
+  std::set<std::string> defunct_dbs_;                //!< Databases that are in an unknown state due to various failures
+#else
+  mutable utils::Gatekeeper<Database> db_gatekeeper_;  //!< Single databases gatekeeper
 #endif
+};
 
 }  // namespace memgraph::dbms

src/dbms/global.hpp

@@ -19,6 +19,7 @@
 
 namespace memgraph::dbms {
 
+#ifdef MG_ENTERPRISE
 enum class DeleteError : uint8_t {
   DEFAULT_DB,
   USING,
@@ -34,11 +35,7 @@ enum class NewError : uint8_t {
   GENERIC,
 };
 
-enum class SetForResult : uint8_t {
-  SUCCESS,
-  ALREADY_SET,
-  FAIL,
-};
+#endif
 
 /**
  * UnknownSession Exception

src/dbms/handler.hpp

@@ -49,7 +49,7 @@ class Handler {
    * @return NewResult
    */
   template <typename... Args>
-  NewResult New(std::piecewise_construct_t /* marker */, std::string_view name, Args... args) {
+  NewResult New(std::piecewise_construct_t /* marker */, std::string_view name, Args &&...args) {
     // Make sure the emplace will succeed, since we don't want to create temporary objects that could break something
     if (!Has(name)) {
       auto [itr, _] = items_.emplace(std::piecewise_construct, std::forward_as_tuple(name),

src/dbms/inmemory/replication_handlers.cpp

@@ -9,7 +9,11 @@
 // by the Apache License, Version 2.0, included in the file
 // licenses/APL.txt.
 
-#include "storage/v2/inmemory/replication/replication_server.hpp"
+#include "dbms/inmemory/replication_handlers.hpp"
+#include "dbms/constants.hpp"
+#include "dbms/dbms_handler.hpp"
+#include "replication/replication_server.hpp"
+#include "spdlog/spdlog.h"
 #include "storage/v2/durability/durability.hpp"
 #include "storage/v2/durability/snapshot.hpp"
 #include "storage/v2/durability/version.hpp"
@@ -17,9 +21,20 @@
 #include "storage/v2/inmemory/storage.hpp"
 #include "storage/v2/inmemory/unique_constraints.hpp"
 
-namespace memgraph::storage {
+using memgraph::storage::Delta;
+using memgraph::storage::EdgeAccessor;
+using memgraph::storage::EdgeRef;
+using memgraph::storage::EdgeTypeId;
+using memgraph::storage::LabelIndexStats;
+using memgraph::storage::LabelPropertyIndexStats;
+using memgraph::storage::PropertyId;
+using memgraph::storage::UniqueConstraints;
+using memgraph::storage::View;
+using memgraph::storage::durability::WalDeltaData;
+
+namespace memgraph::dbms {
 namespace {
-std::pair<uint64_t, durability::WalDeltaData> ReadDelta(durability::BaseDecoder *decoder) {
+std::pair<uint64_t, WalDeltaData> ReadDelta(storage::durability::BaseDecoder *decoder) {
   try {
     auto timestamp = ReadWalDeltaHeader(decoder);
     SPDLOG_INFO("       Timestamp {}", timestamp);
@@ -27,78 +42,114 @@ std::pair<uint64_t, durability::WalDeltaData> ReadDelta(durability::BaseDecoder
     return {timestamp, delta};
   } catch (const slk::SlkReaderException &) {
     throw utils::BasicException("Missing data!");
-  } catch (const durability::RecoveryFailure &) {
+  } catch (const storage::durability::RecoveryFailure &) {
     throw utils::BasicException("Invalid data!");
   }
 };
+
+std::optional<DatabaseAccess> GetDatabaseAccessor(dbms::DbmsHandler *dbms_handler, std::string_view db_name) {
+  try {
+#ifdef MG_ENTERPRISE
+    auto acc = dbms_handler->Get(db_name);
+#else
+    if (db_name != dbms::kDefaultDB) {
+      spdlog::warn("Trying to replicate a non-default database on a community replica.");
+      return std::nullopt;
+    }
+    auto acc = dbms_handler->Get();
+#endif
+    if (!acc) {
+      spdlog::error("Failed to get access to ", db_name);
+      return std::nullopt;
+    }
+    auto *inmem_storage = dynamic_cast<storage::InMemoryStorage *>(acc.get()->storage());
+    if (!inmem_storage || inmem_storage->storage_mode_ != storage::StorageMode::IN_MEMORY_TRANSACTIONAL) {
+      spdlog::error("Database \"{}\" is not IN_MEMORY_TRANSACTIONAL.", db_name);
+      return std::nullopt;
+    }
+    return std::optional{std::move(acc)};
+  } catch (const dbms::UnknownDatabaseException &e) {
+    spdlog::warn("No database \"{}\" on replica!", db_name);
+    return std::nullopt;
+  }
+}
 }  // namespace
 
-InMemoryReplicationServer::InMemoryReplicationServer(InMemoryStorage *storage,
-                                                     const memgraph::replication::ReplicationServerConfig &config,
-                                                     memgraph::replication::ReplicationEpoch *repl_epoch)
-    : ReplicationServer{config}, storage_(storage), repl_epoch_{repl_epoch} {
-  rpc_server_.Register<replication::HeartbeatRpc>([this](auto *req_reader, auto *res_builder) {
+void InMemoryReplicationHandlers::Register(dbms::DbmsHandler *dbms_handler, replication::ReplicationServer &server) {
+  server.rpc_server_.Register<storage::replication::HeartbeatRpc>([dbms_handler](auto *req_reader, auto *res_builder) {
     spdlog::debug("Received HeartbeatRpc");
-    this->HeartbeatHandler(req_reader, res_builder);
+    InMemoryReplicationHandlers::HeartbeatHandler(dbms_handler, req_reader, res_builder);
   });
-
-  rpc_server_.Register<replication::AppendDeltasRpc>([this](auto *req_reader, auto *res_builder) {
-    spdlog::debug("Received AppendDeltasRpc");
-    this->AppendDeltasHandler(req_reader, res_builder);
-  });
-  rpc_server_.Register<replication::SnapshotRpc>([this](auto *req_reader, auto *res_builder) {
+  server.rpc_server_.Register<storage::replication::AppendDeltasRpc>(
+      [dbms_handler](auto *req_reader, auto *res_builder) {
+        spdlog::debug("Received AppendDeltasRpc");
+        InMemoryReplicationHandlers::AppendDeltasHandler(dbms_handler, req_reader, res_builder);
+      });
+  server.rpc_server_.Register<storage::replication::SnapshotRpc>([dbms_handler](auto *req_reader, auto *res_builder) {
     spdlog::debug("Received SnapshotRpc");
-    this->SnapshotHandler(req_reader, res_builder);
+    InMemoryReplicationHandlers::SnapshotHandler(dbms_handler, req_reader, res_builder);
   });
-  rpc_server_.Register<replication::WalFilesRpc>([this](auto *req_reader, auto *res_builder) {
+  server.rpc_server_.Register<storage::replication::WalFilesRpc>([dbms_handler](auto *req_reader, auto *res_builder) {
     spdlog::debug("Received WalFilesRpc");
-    this->WalFilesHandler(req_reader, res_builder);
+    InMemoryReplicationHandlers::WalFilesHandler(dbms_handler, req_reader, res_builder);
   });
-  rpc_server_.Register<replication::CurrentWalRpc>([this](auto *req_reader, auto *res_builder) {
+  server.rpc_server_.Register<storage::replication::CurrentWalRpc>([dbms_handler](auto *req_reader, auto *res_builder) {
     spdlog::debug("Received CurrentWalRpc");
-    this->CurrentWalHandler(req_reader, res_builder);
+    InMemoryReplicationHandlers::CurrentWalHandler(dbms_handler, req_reader, res_builder);
   });
-  rpc_server_.Register<replication::TimestampRpc>([this](auto *req_reader, auto *res_builder) {
+  server.rpc_server_.Register<storage::replication::TimestampRpc>([dbms_handler](auto *req_reader, auto *res_builder) {
     spdlog::debug("Received TimestampRpc");
-    this->TimestampHandler(req_reader, res_builder);
+    InMemoryReplicationHandlers::TimestampHandler(dbms_handler, req_reader, res_builder);
   });
 }
 
-void InMemoryReplicationServer::HeartbeatHandler(slk::Reader *req_reader, slk::Builder *res_builder) {
-  replication::HeartbeatReq req;
+void InMemoryReplicationHandlers::HeartbeatHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader,
+                                                   slk::Builder *res_builder) {
+  storage::replication::HeartbeatReq req;
   slk::Load(&req, req_reader);
-  replication::HeartbeatRes res{true, storage_->repl_storage_state_.last_commit_timestamp_.load(),
-                                std::string{repl_epoch_->id()}};
+  auto const db_acc = GetDatabaseAccessor(dbms_handler, req.db_name);
+  if (!db_acc) return;
+
+  // TODO: this handler is agnostic of InMemory, move to be reused by on-disk
+  auto const *storage = db_acc->get()->storage();
+  storage::replication::HeartbeatRes res{storage->id(), true,
+                                         storage->repl_storage_state_.last_commit_timestamp_.load(),
+                                         std::string{storage->repl_storage_state_.epoch_.id()}};
   slk::Save(res, res_builder);
 }
 
-void InMemoryReplicationServer::AppendDeltasHandler(slk::Reader *req_reader, slk::Builder *res_builder) {
-  replication::AppendDeltasReq req;
+void InMemoryReplicationHandlers::AppendDeltasHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader,
+                                                      slk::Builder *res_builder) {
+  storage::replication::AppendDeltasReq req;
   slk::Load(&req, req_reader);
+  auto db_acc = GetDatabaseAccessor(dbms_handler, req.db_name);
+  if (!db_acc) return;
 
-  replication::Decoder decoder(req_reader);
+  storage::replication::Decoder decoder(req_reader);
 
   auto maybe_epoch_id = decoder.ReadString();
   MG_ASSERT(maybe_epoch_id, "Invalid replication message");
 
-  auto &repl_storage_state = storage_->repl_storage_state_;
-  if (*maybe_epoch_id != repl_epoch_->id()) {
-    auto prev_epoch = repl_epoch_->SetEpoch(*maybe_epoch_id);
+  auto *storage = static_cast<storage::InMemoryStorage *>(db_acc->get()->storage());
+  auto &repl_storage_state = storage->repl_storage_state_;
+  if (*maybe_epoch_id != storage->repl_storage_state_.epoch_.id()) {
+    auto prev_epoch = storage->repl_storage_state_.epoch_.SetEpoch(*maybe_epoch_id);
     repl_storage_state.AddEpochToHistoryForce(prev_epoch);
   }
 
-  if (storage_->wal_file_) {
-    if (req.seq_num > storage_->wal_file_->SequenceNumber() || *maybe_epoch_id != repl_epoch_->id()) {
-      storage_->wal_file_->FinalizeWal();
-      storage_->wal_file_.reset();
-      storage_->wal_seq_num_ = req.seq_num;
+  if (storage->wal_file_) {
+    if (req.seq_num > storage->wal_file_->SequenceNumber() ||
+        *maybe_epoch_id != storage->repl_storage_state_.epoch_.id()) {
+      storage->wal_file_->FinalizeWal();
+      storage->wal_file_.reset();
+      storage->wal_seq_num_ = req.seq_num;
       spdlog::trace("Finalized WAL file");
     } else {
-      MG_ASSERT(storage_->wal_file_->SequenceNumber() == req.seq_num, "Invalid sequence number of current wal file");
-      storage_->wal_seq_num_ = req.seq_num + 1;
+      MG_ASSERT(storage->wal_file_->SequenceNumber() == req.seq_num, "Invalid sequence number of current wal file");
+      storage->wal_seq_num_ = req.seq_num + 1;
     }
   } else {
-    storage_->wal_seq_num_ = req.seq_num;
+    storage->wal_seq_num_ = req.seq_num;
   }
 
   if (req.previous_commit_timestamp != repl_storage_state.last_commit_timestamp_.load()) {
@@ -107,144 +158,161 @@ void InMemoryReplicationServer::AppendDeltasHandler(slk::Reader *req_reader, slk
     while (!transaction_complete) {
       SPDLOG_INFO("Skipping delta");
       const auto [timestamp, delta] = ReadDelta(&decoder);
-      transaction_complete = durability::IsWalDeltaDataTypeTransactionEnd(
-          delta.type, durability::kVersion);  // TODO: Check if we are always using the latest version when replicating
+      transaction_complete = storage::durability::IsWalDeltaDataTypeTransactionEnd(
+          delta.type,
+          storage::durability::kVersion);  // TODO: Check if we are always using the latest version when replicating
     }
 
-    replication::AppendDeltasRes res{false, repl_storage_state.last_commit_timestamp_.load()};
+    storage::replication::AppendDeltasRes res{storage->id(), false, repl_storage_state.last_commit_timestamp_.load()};
     slk::Save(res, res_builder);
     return;
   }
 
-  ReadAndApplyDelta(storage_, &decoder,
-                    durability::kVersion);  // TODO: Check if we are always using the latest version when replicating
+  ReadAndApplyDelta(
+      storage, &decoder,
+      storage::durability::kVersion);  // TODO: Check if we are always using the latest version when replicating
 
-  replication::AppendDeltasRes res{true, repl_storage_state.last_commit_timestamp_.load()};
+  storage::replication::AppendDeltasRes res{storage->id(), true, repl_storage_state.last_commit_timestamp_.load()};
   slk::Save(res, res_builder);
   spdlog::debug("Replication recovery from append deltas finished, replica is now up to date!");
 }
 
-void InMemoryReplicationServer::SnapshotHandler(slk::Reader *req_reader, slk::Builder *res_builder) {
-  replication::SnapshotReq req;
+void InMemoryReplicationHandlers::SnapshotHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader,
+                                                  slk::Builder *res_builder) {
+  storage::replication::SnapshotReq req;
   slk::Load(&req, req_reader);
+  auto db_acc = GetDatabaseAccessor(dbms_handler, req.db_name);
+  if (!db_acc) return;
 
-  replication::Decoder decoder(req_reader);
+  storage::replication::Decoder decoder(req_reader);
 
-  utils::EnsureDirOrDie(storage_->snapshot_directory_);
+  auto *storage = static_cast<storage::InMemoryStorage *>(db_acc->get()->storage());
+  utils::EnsureDirOrDie(storage->snapshot_directory_);
 
-  const auto maybe_snapshot_path = decoder.ReadFile(storage_->snapshot_directory_);
+  const auto maybe_snapshot_path = decoder.ReadFile(storage->snapshot_directory_);
   MG_ASSERT(maybe_snapshot_path, "Failed to load snapshot!");
   spdlog::info("Received snapshot saved to {}", *maybe_snapshot_path);
 
-  auto storage_guard = std::unique_lock{storage_->main_lock_};
+  auto storage_guard = std::unique_lock{storage->main_lock_};
   spdlog::trace("Clearing database since recovering from snapshot.");
   // Clear the database
-  storage_->vertices_.clear();
-  storage_->edges_.clear();
+  storage->vertices_.clear();
+  storage->edges_.clear();
 
-  storage_->constraints_.existence_constraints_ = std::make_unique<ExistenceConstraints>();
-  storage_->constraints_.unique_constraints_ = std::make_unique<InMemoryUniqueConstraints>();
-  storage_->indices_.label_index_ = std::make_unique<InMemoryLabelIndex>();
-  storage_->indices_.label_property_index_ = std::make_unique<InMemoryLabelPropertyIndex>();
+  storage->constraints_.existence_constraints_ = std::make_unique<storage::ExistenceConstraints>();
+  storage->constraints_.unique_constraints_ = std::make_unique<storage::InMemoryUniqueConstraints>();
+  storage->indices_.label_index_ = std::make_unique<storage::InMemoryLabelIndex>();
+  storage->indices_.label_property_index_ = std::make_unique<storage::InMemoryLabelPropertyIndex>();
   try {
     spdlog::debug("Loading snapshot");
-    auto recovered_snapshot = durability::LoadSnapshot(
-        *maybe_snapshot_path, &storage_->vertices_, &storage_->edges_, &storage_->repl_storage_state_.history,
-        storage_->name_id_mapper_.get(), &storage_->edge_count_, storage_->config_);
+    auto recovered_snapshot = storage::durability::LoadSnapshot(
+        *maybe_snapshot_path, &storage->vertices_, &storage->edges_, &storage->repl_storage_state_.history,
+        storage->name_id_mapper_.get(), &storage->edge_count_, storage->config_);
     spdlog::debug("Snapshot loaded successfully");
     // If this step is present it should always be the first step of
     // the recovery so we use the UUID we read from snasphost
-    storage_->uuid_ = std::move(recovered_snapshot.snapshot_info.uuid);
-    repl_epoch_->SetEpoch(std::move(recovered_snapshot.snapshot_info.epoch_id));
+    storage->uuid_ = std::move(recovered_snapshot.snapshot_info.uuid);
+    storage->repl_storage_state_.epoch_.SetEpoch(std::move(recovered_snapshot.snapshot_info.epoch_id));
     const auto &recovery_info = recovered_snapshot.recovery_info;
-    storage_->vertex_id_ = recovery_info.next_vertex_id;
-    storage_->edge_id_ = recovery_info.next_edge_id;
-    storage_->timestamp_ = std::max(storage_->timestamp_, recovery_info.next_timestamp);
+    storage->vertex_id_ = recovery_info.next_vertex_id;
+    storage->edge_id_ = recovery_info.next_edge_id;
+    storage->timestamp_ = std::max(storage->timestamp_, recovery_info.next_timestamp);
 
     spdlog::trace("Recovering indices and constraints from snapshot.");
-    durability::RecoverIndicesAndConstraints(recovered_snapshot.indices_constraints, &storage_->indices_,
-                                             &storage_->constraints_, &storage_->vertices_);
-  } catch (const durability::RecoveryFailure &e) {
+    storage::durability::RecoverIndicesAndConstraints(recovered_snapshot.indices_constraints, &storage->indices_,
+                                                      &storage->constraints_, &storage->vertices_);
+  } catch (const storage::durability::RecoveryFailure &e) {
     LOG_FATAL("Couldn't load the snapshot because of: {}", e.what());
   }
   storage_guard.unlock();
 
-  replication::SnapshotRes res{true, storage_->repl_storage_state_.last_commit_timestamp_.load()};
+  storage::replication::SnapshotRes res{storage->id(), true,
+                                        storage->repl_storage_state_.last_commit_timestamp_.load()};
   slk::Save(res, res_builder);
 
   spdlog::trace("Deleting old snapshot files due to snapshot recovery.");
   // Delete other durability files
-  auto snapshot_files = durability::GetSnapshotFiles(storage_->snapshot_directory_, storage_->uuid_);
+  auto snapshot_files = storage::durability::GetSnapshotFiles(storage->snapshot_directory_, storage->uuid_);
   for (const auto &[path, uuid, _] : snapshot_files) {
     if (path != *maybe_snapshot_path) {
       spdlog::trace("Deleting snapshot file {}", path);
-      storage_->file_retainer_.DeleteFile(path);
+      storage->file_retainer_.DeleteFile(path);
     }
   }
 
   spdlog::trace("Deleting old WAL files due to snapshot recovery.");
-  auto wal_files = durability::GetWalFiles(storage_->wal_directory_, storage_->uuid_);
+  auto wal_files = storage::durability::GetWalFiles(storage->wal_directory_, storage->uuid_);
   if (wal_files) {
     for (const auto &wal_file : *wal_files) {
       spdlog::trace("Deleting WAL file {}", wal_file.path);
-      storage_->file_retainer_.DeleteFile(wal_file.path);
+      storage->file_retainer_.DeleteFile(wal_file.path);
     }
 
-    storage_->wal_file_.reset();
+    storage->wal_file_.reset();
   }
   spdlog::debug("Replication recovery from snapshot finished!");
 }
 
-void InMemoryReplicationServer::WalFilesHandler(slk::Reader *req_reader, slk::Builder *res_builder) {
-  replication::WalFilesReq req;
+void InMemoryReplicationHandlers::WalFilesHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader,
+                                                  slk::Builder *res_builder) {
+  storage::replication::WalFilesReq req;
   slk::Load(&req, req_reader);
+  auto db_acc = GetDatabaseAccessor(dbms_handler, req.db_name);
+  if (!db_acc) return;
 
   const auto wal_file_number = req.file_number;
   spdlog::debug("Received WAL files: {}", wal_file_number);
 
-  replication::Decoder decoder(req_reader);
+  storage::replication::Decoder decoder(req_reader);
 
-  utils::EnsureDirOrDie(storage_->wal_directory_);
+  auto *storage = static_cast<storage::InMemoryStorage *>(db_acc->get()->storage());
+  utils::EnsureDirOrDie(storage->wal_directory_);
 
   for (auto i = 0; i < wal_file_number; ++i) {
-    LoadWal(storage_, *repl_epoch_, &decoder);
+    LoadWal(storage, &decoder);
   }
 
-  replication::WalFilesRes res{true, storage_->repl_storage_state_.last_commit_timestamp_.load()};
+  storage::replication::WalFilesRes res{storage->id(), true,
+                                        storage->repl_storage_state_.last_commit_timestamp_.load()};
   slk::Save(res, res_builder);
   spdlog::debug("Replication recovery from WAL files ended successfully, replica is now up to date!");
 }
 
-void InMemoryReplicationServer::CurrentWalHandler(slk::Reader *req_reader, slk::Builder *res_builder) {
-  replication::CurrentWalReq req;
+void InMemoryReplicationHandlers::CurrentWalHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader,
+                                                    slk::Builder *res_builder) {
+  storage::replication::CurrentWalReq req;
   slk::Load(&req, req_reader);
+  auto db_acc = GetDatabaseAccessor(dbms_handler, req.db_name);
+  if (!db_acc) return;
 
-  replication::Decoder decoder(req_reader);
+  storage::replication::Decoder decoder(req_reader);
 
-  utils::EnsureDirOrDie(storage_->wal_directory_);
+  auto *storage = static_cast<storage::InMemoryStorage *>(db_acc->get()->storage());
+  utils::EnsureDirOrDie(storage->wal_directory_);
 
-  LoadWal(storage_, *repl_epoch_, &decoder);
+  LoadWal(storage, &decoder);
 
-  replication::CurrentWalRes res{true, storage_->repl_storage_state_.last_commit_timestamp_.load()};
+  storage::replication::CurrentWalRes res{storage->id(), true,
+                                          storage->repl_storage_state_.last_commit_timestamp_.load()};
   slk::Save(res, res_builder);
   spdlog::debug("Replication recovery from current WAL ended successfully, replica is now up to date!");
 }
 
-void InMemoryReplicationServer::LoadWal(InMemoryStorage *storage, memgraph::replication::ReplicationEpoch &epoch,
-                                        replication::Decoder *decoder) {
-  const auto temp_wal_directory = std::filesystem::temp_directory_path() / "memgraph" / durability::kWalDirectory;
+void InMemoryReplicationHandlers::LoadWal(storage::InMemoryStorage *storage, storage::replication::Decoder *decoder) {
+  const auto temp_wal_directory =
+      std::filesystem::temp_directory_path() / "memgraph" / storage::durability::kWalDirectory;
   utils::EnsureDir(temp_wal_directory);
   auto maybe_wal_path = decoder->ReadFile(temp_wal_directory);
   MG_ASSERT(maybe_wal_path, "Failed to load WAL!");
   spdlog::trace("Received WAL saved to {}", *maybe_wal_path);
   try {
-    auto wal_info = durability::ReadWalInfo(*maybe_wal_path);
+    auto wal_info = storage::durability::ReadWalInfo(*maybe_wal_path);
     if (wal_info.seq_num == 0) {
       storage->uuid_ = wal_info.uuid;
     }
-
-    if (wal_info.epoch_id != epoch.id()) {
-      auto prev_epoch = epoch.SetEpoch(wal_info.epoch_id);
+    auto &replica_epoch = storage->repl_storage_state_.epoch_;
+    if (wal_info.epoch_id != replica_epoch.id()) {
+      auto prev_epoch = replica_epoch.SetEpoch(wal_info.epoch_id);
       storage->repl_storage_state_.AddEpochToHistoryForce(prev_epoch);
     }
 
@@ -259,11 +327,12 @@ void InMemoryReplicationServer::LoadWal(InMemoryStorage *storage, memgraph::repl
       storage->wal_seq_num_ = wal_info.seq_num;
     }
     spdlog::trace("Loading WAL deltas from {}", *maybe_wal_path);
-    durability::Decoder wal;
-    const auto version = wal.Initialize(*maybe_wal_path, durability::kWalMagic);
+    storage::durability::Decoder wal;
+    const auto version = wal.Initialize(*maybe_wal_path, storage::durability::kWalMagic);
     spdlog::debug("WAL file {} loaded successfully", *maybe_wal_path);
-    if (!version) throw durability::RecoveryFailure("Couldn't read WAL magic and/or version!");
-    if (!durability::IsVersionSupported(*version)) throw durability::RecoveryFailure("Invalid WAL version!");
+    if (!version) throw storage::durability::RecoveryFailure("Couldn't read WAL magic and/or version!");
+    if (!storage::durability::IsVersionSupported(*version))
+      throw storage::durability::RecoveryFailure("Invalid WAL version!");
     wal.SetPosition(wal_info.offset_deltas);
 
     for (size_t i = 0; i < wal_info.num_deltas;) {
@@ -271,38 +340,46 @@ void InMemoryReplicationServer::LoadWal(InMemoryStorage *storage, memgraph::repl
     }
 
     spdlog::debug("Replication from current WAL successful!");
-  } catch (const durability::RecoveryFailure &e) {
+  } catch (const storage::durability::RecoveryFailure &e) {
     LOG_FATAL("Couldn't recover WAL deltas from {} because of: {}", *maybe_wal_path, e.what());
   }
 }
 
-void InMemoryReplicationServer::TimestampHandler(slk::Reader *req_reader, slk::Builder *res_builder) {
-  replication::TimestampReq req;
+void InMemoryReplicationHandlers::TimestampHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader,
+                                                   slk::Builder *res_builder) {
+  storage::replication::TimestampReq req;
   slk::Load(&req, req_reader);
+  auto const db_acc = GetDatabaseAccessor(dbms_handler, req.db_name);
+  if (!db_acc) return;
 
-  replication::TimestampRes res{true, storage_->repl_storage_state_.last_commit_timestamp_.load()};
+  // TODO: this handler is agnostic of InMemory, move to be reused by on-disk
+  auto const *storage = db_acc->get()->storage();
+  storage::replication::TimestampRes res{storage->id(), true,
+                                         storage->repl_storage_state_.last_commit_timestamp_.load()};
   slk::Save(res, res_builder);
 }
 
-uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage, durability::BaseDecoder *decoder,
-                                                      const uint64_t version) {
+uint64_t InMemoryReplicationHandlers::ReadAndApplyDelta(storage::InMemoryStorage *storage,
+                                                        storage::durability::BaseDecoder *decoder,
+                                                        const uint64_t version) {
   auto edge_acc = storage->edges_.access();
   auto vertex_acc = storage->vertices_.access();
 
   constexpr bool kUniqueAccess = true;
+  constexpr bool kSharedAccess = false;
 
-  std::optional<std::pair<uint64_t, InMemoryStorage::ReplicationAccessor>> commit_timestamp_and_accessor;
+  std::optional<std::pair<uint64_t, storage::InMemoryStorage::ReplicationAccessor>> commit_timestamp_and_accessor;
   auto get_transaction = [storage, &commit_timestamp_and_accessor](uint64_t commit_timestamp,
-                                                                   bool unique = !kUniqueAccess) {
+                                                                   bool unique = kSharedAccess) {
     if (!commit_timestamp_and_accessor) {
-      std::unique_ptr<Storage::Accessor> acc = nullptr;
+      std::unique_ptr<storage::Storage::Accessor> acc = nullptr;
       if (unique) {
-        acc = storage->UniqueAccess(std::nullopt);
+        acc = storage->UniqueAccess(std::nullopt, false /*not main*/);
       } else {
-        acc = storage->Access(std::nullopt);
+        acc = storage->Access(std::nullopt, false /*not main*/);
       }
-      auto inmem_acc = std::unique_ptr<InMemoryStorage::InMemoryAccessor>(
-          static_cast<InMemoryStorage::InMemoryAccessor *>(acc.release()));
+      auto inmem_acc = std::unique_ptr<storage::InMemoryStorage::InMemoryAccessor>(
+          static_cast<storage::InMemoryStorage::InMemoryAccessor *>(acc.release()));
       commit_timestamp_and_accessor.emplace(commit_timestamp, std::move(*inmem_acc));
     } else if (commit_timestamp_and_accessor->first != commit_timestamp) {
       throw utils::BasicException("Received more than one transaction!");
@@ -319,7 +396,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
       max_commit_timestamp = timestamp;
     }
 
-    transaction_complete = durability::IsWalDeltaDataTypeTransactionEnd(delta.type, version);
+    transaction_complete = storage::durability::IsWalDeltaDataTypeTransactionEnd(delta.type, version);
 
     if (timestamp < storage->timestamp_) {
       continue;
@@ -327,13 +404,13 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
 
     SPDLOG_INFO("  Delta {}", applied_deltas);
     switch (delta.type) {
-      case durability::WalDeltaData::Type::VERTEX_CREATE: {
+      case WalDeltaData::Type::VERTEX_CREATE: {
         spdlog::trace("       Create vertex {}", delta.vertex_create_delete.gid.AsUint());
         auto *transaction = get_transaction(timestamp);
         transaction->CreateVertexEx(delta.vertex_create_delete.gid);
         break;
       }
-      case durability::WalDeltaData::Type::VERTEX_DELETE: {
+      case WalDeltaData::Type::VERTEX_DELETE: {
         spdlog::trace("       Delete vertex {}", delta.vertex_create_delete.gid.AsUint());
         auto *transaction = get_transaction(timestamp);
         auto vertex = transaction->FindVertex(delta.vertex_create_delete.gid, View::NEW);
@@ -342,7 +419,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         if (ret.HasError() || !ret.GetValue()) throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::VERTEX_ADD_LABEL: {
+      case WalDeltaData::Type::VERTEX_ADD_LABEL: {
         spdlog::trace("       Vertex {} add label {}", delta.vertex_add_remove_label.gid.AsUint(),
                       delta.vertex_add_remove_label.label);
         auto *transaction = get_transaction(timestamp);
@@ -352,7 +429,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         if (ret.HasError() || !ret.GetValue()) throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::VERTEX_REMOVE_LABEL: {
+      case WalDeltaData::Type::VERTEX_REMOVE_LABEL: {
         spdlog::trace("       Vertex {} remove label {}", delta.vertex_add_remove_label.gid.AsUint(),
                       delta.vertex_add_remove_label.label);
         auto *transaction = get_transaction(timestamp);
@@ -362,7 +439,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         if (ret.HasError() || !ret.GetValue()) throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::VERTEX_SET_PROPERTY: {
+      case WalDeltaData::Type::VERTEX_SET_PROPERTY: {
         spdlog::trace("       Vertex {} set property {} to {}", delta.vertex_edge_set_property.gid.AsUint(),
                       delta.vertex_edge_set_property.property, delta.vertex_edge_set_property.value);
         auto *transaction = get_transaction(timestamp);
@@ -373,7 +450,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         if (ret.HasError()) throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::EDGE_CREATE: {
+      case WalDeltaData::Type::EDGE_CREATE: {
         spdlog::trace("       Create edge {} of type {} from vertex {} to vertex {}",
                       delta.edge_create_delete.gid.AsUint(), delta.edge_create_delete.edge_type,
                       delta.edge_create_delete.from_vertex.AsUint(), delta.edge_create_delete.to_vertex.AsUint());
@@ -388,7 +465,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         if (edge.HasError()) throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::EDGE_DELETE: {
+      case WalDeltaData::Type::EDGE_DELETE: {
         spdlog::trace("       Delete edge {} of type {} from vertex {} to vertex {}",
                       delta.edge_create_delete.gid.AsUint(), delta.edge_create_delete.edge_type,
                       delta.edge_create_delete.from_vertex.AsUint(), delta.edge_create_delete.to_vertex.AsUint());
@@ -406,7 +483,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         if (ret.HasError()) throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::EDGE_SET_PROPERTY: {
+      case WalDeltaData::Type::EDGE_SET_PROPERTY: {
         spdlog::trace("       Edge {} set property {} to {}", delta.vertex_edge_set_property.gid.AsUint(),
                       delta.vertex_edge_set_property.property, delta.vertex_edge_set_property.value);
         if (!storage->config_.items.properties_on_edges)
@@ -469,17 +546,18 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         break;
       }
 
-      case durability::WalDeltaData::Type::TRANSACTION_END: {
+      case WalDeltaData::Type::TRANSACTION_END: {
         spdlog::trace("       Transaction end");
         if (!commit_timestamp_and_accessor || commit_timestamp_and_accessor->first != timestamp)
           throw utils::BasicException("Invalid commit data!");
-        auto ret = commit_timestamp_and_accessor->second.Commit(commit_timestamp_and_accessor->first);
+        auto ret =
+            commit_timestamp_and_accessor->second.Commit(commit_timestamp_and_accessor->first, false /* not main */);
         if (ret.HasError()) throw utils::BasicException("Invalid transaction!");
         commit_timestamp_and_accessor = std::nullopt;
         break;
       }
 
-      case durability::WalDeltaData::Type::LABEL_INDEX_CREATE: {
+      case WalDeltaData::Type::LABEL_INDEX_CREATE: {
         spdlog::trace("       Create label index on :{}", delta.operation_label.label);
         // Need to send the timestamp
         auto *transaction = get_transaction(timestamp, kUniqueAccess);
@@ -487,14 +565,14 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
           throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::LABEL_INDEX_DROP: {
+      case WalDeltaData::Type::LABEL_INDEX_DROP: {
         spdlog::trace("       Drop label index on :{}", delta.operation_label.label);
         auto *transaction = get_transaction(timestamp, kUniqueAccess);
         if (transaction->DropIndex(storage->NameToLabel(delta.operation_label.label)).HasError())
           throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::LABEL_INDEX_STATS_SET: {
+      case WalDeltaData::Type::LABEL_INDEX_STATS_SET: {
         spdlog::trace("       Set label index statistics on :{}", delta.operation_label_stats.label);
         // Need to send the timestamp
         auto *transaction = get_transaction(timestamp);
@@ -506,7 +584,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         transaction->SetIndexStats(label, stats);
         break;
       }
-      case durability::WalDeltaData::Type::LABEL_INDEX_STATS_CLEAR: {
+      case WalDeltaData::Type::LABEL_INDEX_STATS_CLEAR: {
         const auto &info = delta.operation_label;
         spdlog::trace("       Clear label index statistics on :{}", info.label);
         // Need to send the timestamp
@@ -514,7 +592,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         transaction->DeleteLabelIndexStats(storage->NameToLabel(info.label));
         break;
       }
-      case durability::WalDeltaData::Type::LABEL_PROPERTY_INDEX_CREATE: {
+      case WalDeltaData::Type::LABEL_PROPERTY_INDEX_CREATE: {
         spdlog::trace("       Create label+property index on :{} ({})", delta.operation_label_property.label,
                       delta.operation_label_property.property);
         auto *transaction = get_transaction(timestamp, kUniqueAccess);
@@ -525,7 +603,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
           throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::LABEL_PROPERTY_INDEX_DROP: {
+      case WalDeltaData::Type::LABEL_PROPERTY_INDEX_DROP: {
         spdlog::trace("       Drop label+property index on :{} ({})", delta.operation_label_property.label,
                       delta.operation_label_property.property);
         auto *transaction = get_transaction(timestamp, kUniqueAccess);
@@ -536,7 +614,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
           throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::LABEL_PROPERTY_INDEX_STATS_SET: {
+      case WalDeltaData::Type::LABEL_PROPERTY_INDEX_STATS_SET: {
         const auto &info = delta.operation_label_property_stats;
         spdlog::trace("       Set label-property index statistics on :{}", info.label);
         // Need to send the timestamp
@@ -550,7 +628,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         transaction->SetIndexStats(label, property, stats);
         break;
       }
-      case durability::WalDeltaData::Type::LABEL_PROPERTY_INDEX_STATS_CLEAR: {
+      case WalDeltaData::Type::LABEL_PROPERTY_INDEX_STATS_CLEAR: {
         const auto &info = delta.operation_label;
         spdlog::trace("       Clear label-property index statistics on :{}", info.label);
         // Need to send the timestamp
@@ -558,7 +636,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         transaction->DeleteLabelPropertyIndexStats(storage->NameToLabel(info.label));
         break;
       }
-      case durability::WalDeltaData::Type::EXISTENCE_CONSTRAINT_CREATE: {
+      case WalDeltaData::Type::EXISTENCE_CONSTRAINT_CREATE: {
         spdlog::trace("       Create existence constraint on :{} ({})", delta.operation_label_property.label,
                       delta.operation_label_property.property);
         auto *transaction = get_transaction(timestamp, kUniqueAccess);
@@ -568,7 +646,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
         if (ret.HasError()) throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::EXISTENCE_CONSTRAINT_DROP: {
+      case WalDeltaData::Type::EXISTENCE_CONSTRAINT_DROP: {
         spdlog::trace("       Drop existence constraint on :{} ({})", delta.operation_label_property.label,
                       delta.operation_label_property.property);
         auto *transaction = get_transaction(timestamp, kUniqueAccess);
@@ -579,7 +657,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
           throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::UNIQUE_CONSTRAINT_CREATE: {
+      case WalDeltaData::Type::UNIQUE_CONSTRAINT_CREATE: {
         std::stringstream ss;
         utils::PrintIterable(ss, delta.operation_label_properties.properties);
         spdlog::trace("       Create unique constraint on :{} ({})", delta.operation_label_properties.label, ss.str());
@@ -594,7 +672,7 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
           throw utils::BasicException("Invalid transaction!");
         break;
       }
-      case durability::WalDeltaData::Type::UNIQUE_CONSTRAINT_DROP: {
+      case WalDeltaData::Type::UNIQUE_CONSTRAINT_DROP: {
         std::stringstream ss;
         utils::PrintIterable(ss, delta.operation_label_properties.properties);
         spdlog::trace("       Drop unique constraint on :{} ({})", delta.operation_label_properties.label, ss.str());
@@ -621,4 +699,4 @@ uint64_t InMemoryReplicationServer::ReadAndApplyDelta(InMemoryStorage *storage,
   return applied_deltas;
 }
 
-}  // namespace memgraph::storage
+}  // namespace memgraph::dbms

src/dbms/inmemory/replication_handlers.hpp

@@ -0,0 +1,49 @@
+// Copyright 2023 Memgraph Ltd.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+// License, and you may not use this file except in compliance with the Business Source License.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+#pragma once
+
+#include "replication/replication_server.hpp"
+#include "replication/state.hpp"
+#include "storage/v2/replication/serialization.hpp"
+
+namespace memgraph::storage {
+class InMemoryStorage;
+}
+namespace memgraph::dbms {
+
+class DbmsHandler;
+
+class InMemoryReplicationHandlers {
+ public:
+  static void Register(dbms::DbmsHandler *dbms_handler, replication::ReplicationServer &server);
+
+ private:
+  // RPC handlers
+  static void HeartbeatHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader, slk::Builder *res_builder);
+
+  static void AppendDeltasHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader, slk::Builder *res_builder);
+
+  static void SnapshotHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader, slk::Builder *res_builder);
+
+  static void WalFilesHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader, slk::Builder *res_builder);
+
+  static void CurrentWalHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader, slk::Builder *res_builder);
+
+  static void TimestampHandler(dbms::DbmsHandler *dbms_handler, slk::Reader *req_reader, slk::Builder *res_builder);
+
+  static void LoadWal(storage::InMemoryStorage *storage, storage::replication::Decoder *decoder);
+
+  static uint64_t ReadAndApplyDelta(storage::InMemoryStorage *storage, storage::durability::BaseDecoder *decoder,
+                                    uint64_t version);
+};
+
+}  // namespace memgraph::dbms

src/dbms/inmemory/storage_helper.hpp

@@ -0,0 +1,67 @@
+// Copyright 2023 Memgraph Ltd.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+// License, and you may not use this file except in compliance with the Business Source License.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+#pragma once
+
+#include <variant>
+
+#include "dbms/constants.hpp"
+#include "dbms/replication_handler.hpp"
+#include "replication/state.hpp"
+#include "storage/v2/config.hpp"
+#include "storage/v2/inmemory/storage.hpp"
+#include "storage/v2/storage.hpp"
+
+namespace memgraph::dbms {
+
+#ifdef MG_EXPERIMENTAL_REPLICATION_MULTITENANCY
+constexpr bool allow_mt_repl = true;
+#else
+constexpr bool allow_mt_repl = false;
+#endif
+
+inline std::unique_ptr<storage::Storage> CreateInMemoryStorage(
+    storage::Config config, const ::memgraph::replication::ReplicationState &repl_state) {
+  const auto wal_mode = config.durability.snapshot_wal_mode;
+  const auto name = config.name;
+  auto storage = std::make_unique<storage::InMemoryStorage>(std::move(config));
+
+  // Connect replication state and storage
+  storage->CreateSnapshotHandler(
+      [storage = storage.get(), &repl_state]() -> utils::BasicResult<storage::InMemoryStorage::CreateSnapshotError> {
+        if (repl_state.IsReplica()) {
+          return storage::InMemoryStorage::CreateSnapshotError::DisabledForReplica;
+        }
+        return storage->CreateSnapshot();
+      });
+
+  if (allow_mt_repl || name == dbms::kDefaultDB) {
+    // Handle global replication state
+    spdlog::info("Replication configuration will be stored and will be automatically restored in case of a crash.");
+    // RECOVER REPLICA CONNECTIONS
+    memgraph::dbms::RestoreReplication(repl_state, *storage);
+  } else if (const ::memgraph::replication::RoleMainData *data =
+                 std::get_if<::memgraph::replication::RoleMainData>(&repl_state.ReplicationData());
+             data && !data->registered_replicas_.empty()) {
+    spdlog::warn("Multi-tenant replication is currently not supported!");
+  }
+
+  if (wal_mode == storage::Config::Durability::SnapshotWalMode::DISABLED && repl_state.IsMain()) {
+    spdlog::warn(
+        "The instance has the MAIN replication role, but durability logs and snapshots are disabled. Please consider "
+        "enabling durability by using --storage-snapshot-interval-sec and --storage-wal-enabled flags because "
+        "without write-ahead logs this instance is not replicating any data.");
+  }
+
+  return std::move(storage);
+}
+
+}  // namespace memgraph::dbms

src/dbms/replication_handler.cpp

@@ -0,0 +1,234 @@
+// Copyright 2023 Memgraph Ltd.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+// License, and you may not use this file except in compliance with the Business Source License.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+#include "dbms/replication_handler.hpp"
+
+#include "dbms/constants.hpp"
+#include "dbms/dbms_handler.hpp"
+#include "dbms/inmemory/replication_handlers.hpp"
+#include "dbms/inmemory/storage_helper.hpp"
+#include "replication/state.hpp"
+
+using memgraph::replication::ReplicationClientConfig;
+using memgraph::replication::ReplicationState;
+using memgraph::replication::RoleMainData;
+using memgraph::replication::RoleReplicaData;
+
+namespace memgraph::dbms {
+
+namespace {
+
+std::string RegisterReplicaErrorToString(RegisterReplicaError error) {
+  switch (error) {
+    using enum RegisterReplicaError;
+    case NAME_EXISTS:
+      return "NAME_EXISTS";
+    case END_POINT_EXISTS:
+      return "END_POINT_EXISTS";
+    case CONNECTION_FAILED:
+      return "CONNECTION_FAILED";
+    case COULD_NOT_BE_PERSISTED:
+      return "COULD_NOT_BE_PERSISTED";
+  }
+}
+}  // namespace
+
+bool ReplicationHandler::SetReplicationRoleMain() {
+  auto const main_handler = [](RoleMainData const &) {
+    // If we are already MAIN, we don't want to change anything
+    return false;
+  };
+  auto const replica_handler = [this](RoleReplicaData const &) {
+    // STEP 1) bring down all REPLICA servers
+    dbms_handler_.ForEach([](Database *db) {
+      auto *storage = db->storage();
+      // Remember old epoch + storage timestamp association
+      storage->PrepareForNewEpoch();
+    });
+
+    // STEP 2) Change to MAIN
+    // TODO: restore replication servers if false?
+    if (!repl_state_.SetReplicationRoleMain()) {
+      // TODO: Handle recovery on failure???
+      return false;
+    }
+
+    // STEP 3) We are now MAIN, update storage local epoch
+    dbms_handler_.ForEach([&](Database *db) {
+      auto *storage = db->storage();
+      storage->repl_storage_state_.epoch_ = std::get<RoleMainData>(std::as_const(repl_state_).ReplicationData()).epoch_;
+    });
+
+    return true;
+  };
+
+  // TODO: under lock
+  return std::visit(utils::Overloaded{main_handler, replica_handler}, repl_state_.ReplicationData());
+}
+
+bool ReplicationHandler::SetReplicationRoleReplica(const memgraph::replication::ReplicationServerConfig &config) {
+  // We don't want to restart the server if we're already a REPLICA
+  if (repl_state_.IsReplica()) {
+    return false;
+  }
+
+  // Remove registered replicas
+  dbms_handler_.ForEach([&](Database *db) {
+    auto *storage = db->storage();
+    storage->repl_storage_state_.replication_clients_.WithLock([](auto &clients) { clients.clear(); });
+  });
+
+  // Creates the server
+  repl_state_.SetReplicationRoleReplica(config);
+
+  // Start
+  const auto success =
+      std::visit(utils::Overloaded{[](auto) {
+                                     // ASSERT
+                                     return false;
+                                   },
+                                   [this](RoleReplicaData const &data) {
+                                     // Register handlers
+                                     InMemoryReplicationHandlers::Register(&dbms_handler_, *data.server);
+                                     if (!data.server->Start()) {
+                                       spdlog::error("Unable to start the replication server.");
+                                       return false;
+                                     }
+                                     return true;
+                                   }},
+                 repl_state_.ReplicationData());
+  // TODO Handle error (restore to main?)
+  return success;
+}
+
+auto ReplicationHandler::RegisterReplica(const memgraph::replication::ReplicationClientConfig &config)
+    -> memgraph::utils::BasicResult<RegisterReplicaError> {
+  MG_ASSERT(repl_state_.IsMain(), "Only main instance can register a replica!");
+
+  auto res = repl_state_.RegisterReplica(config);
+  switch (res) {
+    case memgraph::replication::RegisterReplicaError::NOT_MAIN:
+      MG_ASSERT(false, "Only main instance can register a replica!");
+      return {};
+    case memgraph::replication::RegisterReplicaError::NAME_EXISTS:
+      return memgraph::dbms::RegisterReplicaError::NAME_EXISTS;
+    case memgraph::replication::RegisterReplicaError::END_POINT_EXISTS:
+      return memgraph::dbms::RegisterReplicaError::END_POINT_EXISTS;
+    case memgraph::replication::RegisterReplicaError::COULD_NOT_BE_PERSISTED:
+      return memgraph::dbms::RegisterReplicaError::COULD_NOT_BE_PERSISTED;
+    case memgraph::replication::RegisterReplicaError::SUCCESS:
+      break;
+  }
+
+  bool all_clients_good = true;
+
+  if (!allow_mt_repl && dbms_handler_.All().size() > 1) {
+    spdlog::warn("Multi-tenant replication is currently not supported!");
+  }
+
+  dbms_handler_.ForEach([&](Database *db) {
+    auto *storage = db->storage();
+    if (!allow_mt_repl && storage->id() != kDefaultDB) {
+      return;
+    }
+    // TODO: ATM only IN_MEMORY_TRANSACTIONAL, fix other modes
+    if (storage->storage_mode_ != storage::StorageMode::IN_MEMORY_TRANSACTIONAL) return;
+
+    all_clients_good &=
+        storage->repl_storage_state_.replication_clients_.WithLock([storage, &config](auto &clients) -> bool {
+          auto client = storage->CreateReplicationClient(config, &storage->repl_storage_state_.epoch_);
+          client->Start();
+
+          if (client->State() == storage::replication::ReplicaState::INVALID) {
+            return false;
+          }
+          clients.push_back(std::move(client));
+          return true;
+        });
+  });
+  if (!all_clients_good) return RegisterReplicaError::CONNECTION_FAILED;  // TODO: this happen to 1 or many...what to do
+  return {};
+}
+
+auto ReplicationHandler::UnregisterReplica(std::string_view name) -> UnregisterReplicaResult {
+  auto const replica_handler = [](RoleReplicaData const &) -> UnregisterReplicaResult {
+    return UnregisterReplicaResult::NOT_MAIN;
+  };
+  auto const main_handler = [this, name](RoleMainData &mainData) -> UnregisterReplicaResult {
+    if (!repl_state_.TryPersistUnregisterReplica(name)) {
+      return UnregisterReplicaResult::COULD_NOT_BE_PERSISTED;
+    }
+    auto const n_unregistered =
+        std::erase_if(mainData.registered_replicas_,
+                      [&](ReplicationClientConfig const &registered_config) { return registered_config.name == name; });
+
+    dbms_handler_.ForEach([&](Database *db) {
+      db->storage()->repl_storage_state_.replication_clients_.WithLock(
+          [&](auto &clients) { std::erase_if(clients, [&](const auto &client) { return client->Name() == name; }); });
+    });
+
+    return n_unregistered != 0 ? UnregisterReplicaResult::SUCCESS : UnregisterReplicaResult::CAN_NOT_UNREGISTER;
+  };
+
+  return std::visit(utils::Overloaded{main_handler, replica_handler}, repl_state_.ReplicationData());
+}
+
+auto ReplicationHandler::GetRole() const -> memgraph::replication::ReplicationRole { return repl_state_.GetRole(); }
+
+bool ReplicationHandler::IsMain() const { return repl_state_.IsMain(); }
+
+bool ReplicationHandler::IsReplica() const { return repl_state_.IsReplica(); }
+
+void RestoreReplication(const replication::ReplicationState &repl_state, storage::Storage &storage) {
+  spdlog::info("Restoring replication role.");
+
+  /// MAIN
+  auto const recover_main = [&storage](RoleMainData const &mainData) {
+    for (const auto &config : mainData.registered_replicas_) {
+      spdlog::info("Replica {} restoration started for {}.", config.name, storage.id());
+
+      auto register_replica = [&storage](const memgraph::replication::ReplicationClientConfig &config)
+          -> memgraph::utils::BasicResult<RegisterReplicaError> {
+        return storage.repl_storage_state_.replication_clients_.WithLock(
+            [&storage, &config](auto &clients) -> utils::BasicResult<RegisterReplicaError> {
+              auto client = storage.CreateReplicationClient(config, &storage.repl_storage_state_.epoch_);
+              client->Start();
+
+              if (client->State() == storage::replication::ReplicaState::INVALID) {
+                spdlog::warn("Connection failed when registering replica {}. Replica will still be registered.",
+                             client->Name());
+              }
+              clients.push_back(std::move(client));
+              return {};
+            });
+      };
+
+      auto ret = register_replica(config);
+      if (ret.HasError()) {
+        MG_ASSERT(RegisterReplicaError::CONNECTION_FAILED != ret.GetError());
+        LOG_FATAL("Failure when restoring replica {}: {}.", config.name, RegisterReplicaErrorToString(ret.GetError()));
+      }
+      spdlog::info("Replica {} restored for {}.", config.name, storage.id());
+    }
+    spdlog::info("Replication role restored to MAIN.");
+  };
+
+  /// REPLICA
+  auto const recover_replica = [](RoleReplicaData const &data) { /*nothing to do*/ };
+
+  std::visit(
+      utils::Overloaded{
+          recover_main,
+          recover_replica,
+      },
+      std::as_const(repl_state).ReplicationData());
+}
+}  // namespace memgraph::dbms

src/dbms/replication_handler.hpp

@@ -12,6 +12,7 @@
 #pragma once
 
 #include "replication/role.hpp"
+#include "storage/v2/storage.hpp"
 #include "utils/result.hpp"
 
 // BEGIN fwd declares
@@ -20,14 +21,10 @@ struct ReplicationState;
 struct ReplicationServerConfig;
 struct ReplicationClientConfig;
 }  // namespace memgraph::replication
-namespace memgraph::storage {
-class Storage;
-}
-// END fwd declares
 
-namespace memgraph::storage {
+namespace memgraph::dbms {
+class DbmsHandler;
 
-enum class RegistrationMode : std::uint8_t { MUST_BE_INSTANTLY_VALID, RESTORE };
 enum class RegisterReplicaError : uint8_t { NAME_EXISTS, END_POINT_EXISTS, CONNECTION_FAILED, COULD_NOT_BE_PERSISTED };
 enum class UnregisterReplicaResult : uint8_t {
   NOT_MAIN,
@@ -39,8 +36,8 @@ enum class UnregisterReplicaResult : uint8_t {
 /// A handler type that keep in sync current ReplicationState and the MAIN/REPLICA-ness of Storage
 /// TODO: extend to do multiple storages
 struct ReplicationHandler {
-  ReplicationHandler(memgraph::replication::ReplicationState &replState, Storage &storage)
-      : repl_state_(replState), storage_(storage) {}
+  ReplicationHandler(memgraph::replication::ReplicationState &replState, DbmsHandler &dbms_handler)
+      : repl_state_(replState), dbms_handler_(dbms_handler) {}
 
   // as REPLICA, become MAIN
   bool SetReplicationRoleMain();
@@ -49,16 +46,12 @@ struct ReplicationHandler {
   bool SetReplicationRoleReplica(const memgraph::replication::ReplicationServerConfig &config);
 
   // as MAIN, define and connect to REPLICAs
-  auto RegisterReplica(RegistrationMode registration_mode, const memgraph::replication::ReplicationClientConfig &config)
+  auto RegisterReplica(const memgraph::replication::ReplicationClientConfig &config)
       -> utils::BasicResult<RegisterReplicaError>;
 
   // as MAIN, remove a REPLICA connection
   auto UnregisterReplica(std::string_view name) -> UnregisterReplicaResult;
 
-  // Generic restoration
-  // TODO: decouple storage restoration from epoch restoration
-  void RestoreReplication();
-
   // Helper pass-through (TODO: remove)
   auto GetRole() const -> memgraph::replication::ReplicationRole;
   bool IsMain() const;
@@ -66,6 +59,11 @@ struct ReplicationHandler {
 
  private:
   memgraph::replication::ReplicationState &repl_state_;
-  Storage &storage_;
+  DbmsHandler &dbms_handler_;
 };
-}  // namespace memgraph::storage
+
+/// A handler type that keep in sync current ReplicationState and the MAIN/REPLICA-ness of Storage
+/// TODO: extend to do multiple storages
+void RestoreReplication(const replication::ReplicationState &repl_state, storage::Storage &storage);
+
+}  // namespace memgraph::dbms

src/glue/CMakeLists.txt

@@ -8,4 +8,3 @@ target_sources(mg-glue PRIVATE auth.cpp
                                MonitoringServerT.cpp
                                run_id.cpp)
 target_link_libraries(mg-glue mg-query mg-auth mg-audit mg-flags)
-target_precompile_headers(mg-glue INTERFACE auth_checker.hpp auth_handler.hpp)

src/glue/run_id.hpp

@@ -11,6 +11,8 @@
 
 #pragma once
 
+#include <string>
+
 namespace memgraph::glue {
 extern const std::string run_id_;
 }  // namespace memgraph::glue

src/http_handlers/metrics.hpp

@@ -65,7 +65,7 @@ class MetricsService {
     return MetricsResponse{.vertex_count = info.vertex_count,
                            .edge_count = info.edge_count,
                            .average_degree = info.average_degree,
-                           .memory_usage = info.memory_usage,
+                           .memory_usage = info.memory_res,
                            .disk_usage = info.disk_usage,
                            .event_counters = GetEventCounters(),
                            .event_gauges = GetEventGauges(),

src/integrations/kafka/consumer.cpp

@@ -1,4 +1,4 @@
-// Copyright 2022 Memgraph Ltd.
+// Copyright 2023 Memgraph Ltd.
 //
 // Use of this software is governed by the Business Source License
 // included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source

src/io/network/socket.cpp

@@ -216,7 +216,7 @@ bool Socket::Write(const uint8_t *data, size_t len, bool have_more) {
   return true;
 }
 
-bool Socket::Write(const std::string &s, bool have_more) {
+bool Socket::Write(std::string_view s, bool have_more) {
   return Write(reinterpret_cast<const uint8_t *>(s.data()), s.size(), have_more);
 }
 

src/io/network/socket.hpp

@@ -1,4 +1,4 @@
-// Copyright 2022 Memgraph Ltd.
+// Copyright 2023 Memgraph Ltd.
 //
 // Use of this software is governed by the Business Source License
 // included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
@@ -149,7 +149,7 @@ class Socket {
    *             false if write failed
    */
   bool Write(const uint8_t *data, size_t len, bool have_more = false);
-  bool Write(const std::string &s, bool have_more = false);
+  bool Write(std::string_view s, bool have_more = false);
 
   /**
    * Read data from the socket.

src/license/license_sender.cpp

@@ -42,7 +42,7 @@ void LicenseInfoSender::SendData() {
   license_info_.WithLock([&data, this](const auto &license_info) mutable {
     if (license_info && !license_info->organization_name.empty()) {
       const auto memory_info = utils::GetMemoryInfo();
-      const auto memory_usage = utils::GetMemoryUsage();
+      const auto memory_res = utils::GetMemoryRES();
       data = {{"run_id", uuid_},
               {"machine_id", machine_id_},
               {"type", "license-check"},
@@ -52,7 +52,7 @@ void LicenseInfoSender::SendData() {
               {"valid", license_info->is_valid},
               {"physical_memory_size", memory_info.memory},
               {"swap_memory_size", memory_info.swap},
-              {"memory_used", memory_usage},
+              {"memory_used", memory_res},
               {"runtime_memory_limit", memory_limit_},
               {"license_memory_limit", license_info->license.memory_limit},
               {"timestamp", utils::Timestamp::Now().SecWithNsecSinceTheEpoch()}};

src/memgraph.cpp

@@ -9,11 +9,13 @@
 // by the Apache License, Version 2.0, included in the file
 // licenses/APL.txt.
 
+#include <cstdint>
 #include "audit/log.hpp"
 #include "communication/metrics.hpp"
 #include "communication/websocket/auth.hpp"
 #include "communication/websocket/server.hpp"
 #include "dbms/constants.hpp"
+#include "dbms/inmemory/replication_handlers.hpp"
 #include "flags/all.hpp"
 #include "flags/run_time_configurable.hpp"
 #include "glue/MonitoringServerT.hpp"
@@ -43,9 +45,11 @@
 #include "query/auth_query_handler.hpp"
 #include "query/interpreter_context.hpp"
 
+namespace {
 constexpr const char *kMgUser = "MEMGRAPH_USER";
 constexpr const char *kMgPassword = "MEMGRAPH_PASSWORD";
 constexpr const char *kMgPassfile = "MEMGRAPH_PASSFILE";
+constexpr uint64_t kMgVmMaxMapCount = 262144;
 
 // TODO: move elsewhere so that we can remove need of interpreter.hpp
 void InitFromCypherlFile(memgraph::query::InterpreterContext &ctx, memgraph::dbms::DatabaseAccess &db_acc,
@@ -107,6 +111,7 @@ void InitSignalHandlers(const std::function<void()> &shutdown_fun) {
                                                             block_shutdown_signals),
             "Unable to register SIGINT handler!");
 }
+}  // namespace
 
 int main(int argc, char **argv) {
   memgraph::memory::SetHooks();
@@ -204,6 +209,17 @@ int main(int argc, char **argv) {
   std::cout << "You are running Memgraph v" << gflags::VersionString() << std::endl;
   std::cout << "To get started with Memgraph, visit https://memgr.ph/start" << std::endl;
 
+  const auto vm_max_map_count = memgraph::utils::GetVmMaxMapCount();
+  if (vm_max_map_count.has_value()) {
+    if (vm_max_map_count.value() < kMgVmMaxMapCount) {
+      std::cout << "Max virtual memory areas vm.max_map_count " << vm_max_map_count.value()
+                << " is too low, increase to at least " << kMgVmMaxMapCount << std::endl;
+    }
+  } else {
+    std::cout << "Can't get info on vm.max_map_count, check whether it is too low, vm.max_map_count is at least "
+              << kMgVmMaxMapCount << std::endl;
+  }
+
   auto data_directory = std::filesystem::path(FLAGS_data_directory);
 
   memgraph::utils::EnsureDirOrDie(data_directory);
@@ -352,22 +368,35 @@ int main(int argc, char **argv) {
   std::unique_ptr<memgraph::query::AuthChecker> auth_checker;
   auth_glue(&auth_, auth_handler, auth_checker);
 
+  memgraph::replication::ReplicationState repl_state(ReplicationStateRootPath(db_config));
+
+  memgraph::dbms::DbmsHandler dbms_handler(db_config, repl_state
 #ifdef MG_ENTERPRISE
-  memgraph::dbms::DbmsHandler new_handler(db_config, &auth_, FLAGS_data_recovery_on_startup,
-                                          FLAGS_storage_delete_on_drop);
-  auto db_acc = new_handler.Get(memgraph::dbms::kDefaultDB);
-  memgraph::query::InterpreterContext interpreter_context_(interp_config, &new_handler, auth_handler.get(),
-                                                           auth_checker.get());
-#else
-  memgraph::utils::Gatekeeper<memgraph::dbms::Database> db_gatekeeper{db_config};
-  auto db_acc_opt = db_gatekeeper.access();
-  MG_ASSERT(db_acc_opt, "Failed to access the main database");
-  auto &db_acc = *db_acc_opt;
-  memgraph::query::InterpreterContext interpreter_context_(interp_config, &db_gatekeeper, auth_handler.get(),
-                                                           auth_checker.get());
+                                           ,
+                                           &auth_, FLAGS_data_recovery_on_startup, FLAGS_storage_delete_on_drop
 #endif
+  );
+  auto db_acc = dbms_handler.Get();
+  memgraph::query::InterpreterContext interpreter_context_(interp_config, &dbms_handler, &repl_state,
+                                                           auth_handler.get(), auth_checker.get());
   MG_ASSERT(db_acc, "Failed to access the main database");
 
+  // TODO: Move it somewhere better
+  // Startup replication state (if recovered at startup)
+  MG_ASSERT(std::visit(memgraph::utils::Overloaded{[](memgraph::replication::RoleMainData const &) { return true; },
+                                                   [&](memgraph::replication::RoleReplicaData const &data) {
+                                                     // Register handlers
+                                                     memgraph::dbms::InMemoryReplicationHandlers::Register(
+                                                         &dbms_handler, *data.server);
+                                                     if (!data.server->Start()) {
+                                                       spdlog::error("Unable to start the replication server.");
+                                                       return false;
+                                                     }
+                                                     return true;
+                                                   }},
+                       repl_state.ReplicationData()),
+            "Replica recovery failure!");
+
   memgraph::query::procedure::gModuleRegistry.SetModulesDirectory(memgraph::flags::ParseQueryModulesDirectory(),
                                                                   FLAGS_data_directory);
   memgraph::query::procedure::gModuleRegistry.UnloadAndLoadModulesFromDirectories();
@@ -388,8 +417,8 @@ int main(int argc, char **argv) {
   }
 
 #ifdef MG_ENTERPRISE
-  new_handler.RestoreTriggers(&interpreter_context_);
-  new_handler.RestoreStreams(&interpreter_context_);
+  dbms_handler.RestoreTriggers(&interpreter_context_);
+  dbms_handler.RestoreStreams(&interpreter_context_);
 #else
   {
     // Triggers can execute query procedures, so we need to reload the modules first and then
@@ -432,11 +461,10 @@ int main(int argc, char **argv) {
   if (FLAGS_telemetry_enabled) {
     telemetry.emplace(telemetry_server, data_directory / "telemetry", memgraph::glue::run_id_, machine_id,
                       service_name == "BoltS", FLAGS_data_directory, std::chrono::minutes(10));
+    telemetry->AddStorageCollector(dbms_handler, auth_);
 #ifdef MG_ENTERPRISE
-    telemetry->AddStorageCollector(new_handler, auth_);
-    telemetry->AddDatabaseCollector(new_handler);
+    telemetry->AddDatabaseCollector(dbms_handler);
 #else
-    telemetry->AddStorageCollector(db_gatekeeper, auth_);
     telemetry->AddDatabaseCollector();
 #endif
     telemetry->AddClientCollector();
@@ -496,17 +524,16 @@ int main(int argc, char **argv) {
 
   if (!FLAGS_init_data_file.empty()) {
     spdlog::info("Running init data file.");
+    auto db_acc = dbms_handler.Get();
+    MG_ASSERT(db_acc, "Failed to gain access to the main database");
 #ifdef MG_ENTERPRISE
-    auto db_acc = new_handler.Get(memgraph::dbms::kDefaultDB);
     if (memgraph::license::global_license_checker.IsEnterpriseValidFast()) {
       InitFromCypherlFile(interpreter_context_, db_acc, FLAGS_init_data_file, &audit_log);
     } else {
       InitFromCypherlFile(interpreter_context_, db_acc, FLAGS_init_data_file);
     }
 #else
-    auto db_acc_2 = db_gatekeeper.access();
-    MG_ASSERT(db_acc_2, "Failed to gain access to the main database");
-    InitFromCypherlFile(interpreter_context_, *db_acc_2, FLAGS_init_data_file);
+    InitFromCypherlFile(interpreter_context_, db_acc, FLAGS_init_data_file);
 #endif
   }
 

src/memory/global_memory_control.cpp

@@ -61,11 +61,8 @@ void *my_alloc(extent_hooks_t *extent_hooks, void *new_addr, size_t size, size_t
   // This needs to be before, to throw exception in case of too big alloc
   if (*commit) [[likely]] {
     memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(size));
-    if (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
-      auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
-      if (memory_tracker != nullptr) [[likely]] {
-        memory_tracker->Alloc(static_cast<int64_t>(size));
-      }
+    if (GetQueriesMemoryControl().IsThreadTracked()) [[unlikely]] {
+      GetQueriesMemoryControl().TrackAllocOnCurrentThread(size);
     }
   }
 
@@ -73,11 +70,8 @@ void *my_alloc(extent_hooks_t *extent_hooks, void *new_addr, size_t size, size_t
   if (ptr == nullptr) [[unlikely]] {
     if (*commit) {
       memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
-      if (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
-        auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
-        if (memory_tracker != nullptr) [[likely]] {
-          memory_tracker->Free(static_cast<int64_t>(size));
-        }
+      if (GetQueriesMemoryControl().IsThreadTracked()) [[unlikely]] {
+        GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
       }
     }
     return ptr;
@@ -96,11 +90,8 @@ static bool my_dalloc(extent_hooks_t *extent_hooks, void *addr, size_t size, boo
   if (committed) [[likely]] {
     memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
 
-    if (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
-      auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
-      if (memory_tracker != nullptr) [[likely]] {
-        memory_tracker->Free(static_cast<int64_t>(size));
-      }
+    if (GetQueriesMemoryControl().IsThreadTracked()) [[unlikely]] {
+      GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
     }
   }
 
@@ -110,11 +101,8 @@ static bool my_dalloc(extent_hooks_t *extent_hooks, void *addr, size_t size, boo
 static void my_destroy(extent_hooks_t *extent_hooks, void *addr, size_t size, bool committed, unsigned arena_ind) {
   if (committed) [[likely]] {
     memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
-    if (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
-      auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
-      if (memory_tracker != nullptr) [[likely]] {
-        memory_tracker->Free(static_cast<int64_t>(size));
-      }
+    if (GetQueriesMemoryControl().IsThreadTracked()) [[unlikely]] {
+      GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
     }
   }
 
@@ -130,11 +118,8 @@ static bool my_commit(extent_hooks_t *extent_hooks, void *addr, size_t size, siz
   }
 
   memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(length));
-  if (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
-    auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
-    if (memory_tracker != nullptr) [[likely]] {
-      memory_tracker->Alloc(static_cast<int64_t>(size));
-    }
+  if (GetQueriesMemoryControl().IsThreadTracked()) [[unlikely]] {
+    GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
   }
 
   return false;
@@ -150,11 +135,8 @@ static bool my_decommit(extent_hooks_t *extent_hooks, void *addr, size_t size, s
   }
 
   memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(length));
-  if (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
-    auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
-    if (memory_tracker != nullptr) [[likely]] {
-      memory_tracker->Free(static_cast<int64_t>(size));
-    }
+  if (GetQueriesMemoryControl().IsThreadTracked()) [[unlikely]] {
+    GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
   }
 
   return false;
@@ -170,11 +152,8 @@ static bool my_purge_forced(extent_hooks_t *extent_hooks, void *addr, size_t siz
   }
   memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(length));
 
-  if (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
-    auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
-    if (memory_tracker != nullptr) [[likely]] {
-      memory_tracker->Alloc(static_cast<int64_t>(size));
-    }
+  if (GetQueriesMemoryControl().IsThreadTracked()) [[unlikely]] {
+    GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
   }
 
   return false;
@@ -200,8 +179,11 @@ void SetHooks() {
     return;
   }
 
+  // Needs init due to error we might encounter otherwise
+  // https://gcc.gnu.org/bugzilla/show_bug.cgi?id=13684
+  [[maybe_unused]] const auto &queries_memory_control = GetQueriesMemoryControl();
+
   for (int i = 0; i < n_arenas; i++) {
-    GetQueriesMemoryControl().InitializeArenaCounter(i);
     std::string func_name = "arena." + std::to_string(i) + ".extent_hooks";
 
     size_t hooks_len = sizeof(old_hooks);
@@ -261,7 +243,6 @@ void UnsetHooks() {
   }
 
   for (int i = 0; i < n_arenas; i++) {
-    GetQueriesMemoryControl().InitializeArenaCounter(i);
     std::string func_name = "arena." + std::to_string(i) + ".extent_hooks";
 
     MG_ASSERT(old_hooks);

src/memory/query_memory_control.cpp

@@ -10,6 +10,7 @@
 // licenses/APL.txt.
 
 #include <atomic>
+#include <cassert>
 #include <cstdint>
 #include <iostream>
 #include <optional>
@@ -21,6 +22,7 @@
 #include "query_memory_control.hpp"
 #include "utils/exceptions.hpp"
 #include "utils/logging.hpp"
+#include "utils/memory.hpp"
 #include "utils/memory_tracker.hpp"
 #include "utils/rw_spin_lock.hpp"
 
@@ -32,20 +34,6 @@ namespace memgraph::memory {
 
 #if USE_JEMALLOC
 
-unsigned QueriesMemoryControl::GetArenaForThread() {
-  unsigned thread_arena{0};
-  size_t size_thread_arena = sizeof(thread_arena);
-  int err = mallctl("thread.arena", &thread_arena, &size_thread_arena, nullptr, 0);
-  if (err) {
-    LOG_FATAL("Can't get arena for thread.");
-  }
-  return thread_arena;
-}
-
-void QueriesMemoryControl::AddTrackingOnArena(unsigned arena_id) { arena_tracking[arena_id].fetch_add(1); }
-
-void QueriesMemoryControl::RemoveTrackingOnArena(unsigned arena_id) { arena_tracking[arena_id].fetch_sub(1); }
-
 void QueriesMemoryControl::UpdateThreadToTransactionId(const std::thread::id &thread_id, uint64_t transaction_id) {
   auto accessor = thread_id_to_transaction_id.access();
   accessor.insert({thread_id, transaction_id});
@@ -58,34 +46,58 @@ void QueriesMemoryControl::EraseThreadToTransactionId(const std::thread::id &thr
   accessor.remove(thread_id);
 }
 
-utils::MemoryTracker *QueriesMemoryControl::GetTrackerCurrentThread() {
+void QueriesMemoryControl::TrackAllocOnCurrentThread(size_t size) {
   auto thread_id_to_transaction_id_accessor = thread_id_to_transaction_id.access();
 
   // we might be just constructing mapping between thread id and transaction id
   // so we miss this allocation
   auto thread_id_to_transaction_id_elem = thread_id_to_transaction_id_accessor.find(std::this_thread::get_id());
   if (thread_id_to_transaction_id_elem == thread_id_to_transaction_id_accessor.end()) {
-    return nullptr;
+    return;
   }
 
   auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
   auto transaction_id_to_tracker =
       transaction_id_to_tracker_accessor.find(thread_id_to_transaction_id_elem->transaction_id);
-  return &transaction_id_to_tracker->tracker;
+
+  // It can happen that some allocation happens between mapping thread to
+  // transaction id, so we miss this allocation
+  if (transaction_id_to_tracker == transaction_id_to_tracker_accessor.end()) [[unlikely]] {
+    return;
+  }
+  auto &query_tracker = transaction_id_to_tracker->tracker;
+  query_tracker.TrackAlloc(size);
+}
+
+void QueriesMemoryControl::TrackFreeOnCurrentThread(size_t size) {
+  auto thread_id_to_transaction_id_accessor = thread_id_to_transaction_id.access();
+
+  // we might be just constructing mapping between thread id and transaction id
+  // so we miss this allocation
+  auto thread_id_to_transaction_id_elem = thread_id_to_transaction_id_accessor.find(std::this_thread::get_id());
+  if (thread_id_to_transaction_id_elem == thread_id_to_transaction_id_accessor.end()) {
+    return;
+  }
+
+  auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
+  auto transaction_id_to_tracker =
+      transaction_id_to_tracker_accessor.find(thread_id_to_transaction_id_elem->transaction_id);
+
+  // It can happen that some allocation happens between mapping thread to
+  // transaction id, so we miss this allocation
+  if (transaction_id_to_tracker == transaction_id_to_tracker_accessor.end()) [[unlikely]] {
+    return;
+  }
+  auto &query_tracker = transaction_id_to_tracker->tracker;
+  query_tracker.TrackFree(size);
 }
 
 void QueriesMemoryControl::CreateTransactionIdTracker(uint64_t transaction_id, size_t inital_limit) {
   auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
 
-  auto [elem, result] = transaction_id_to_tracker_accessor.insert({transaction_id, utils::MemoryTracker{}});
+  auto [elem, result] = transaction_id_to_tracker_accessor.insert({transaction_id, utils::QueryMemoryTracker{}});
 
-  elem->tracker.SetMaximumHardLimit(inital_limit);
-  elem->tracker.SetHardLimit(inital_limit);
-}
-
-bool QueriesMemoryControl::CheckTransactionIdTrackerExists(uint64_t transaction_id) {
-  auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
-  return transaction_id_to_tracker_accessor.contains(transaction_id);
+  elem->tracker.SetQueryLimit(inital_limit);
 }
 
 bool QueriesMemoryControl::EraseTransactionIdTracker(uint64_t transaction_id) {
@@ -94,27 +106,68 @@ bool QueriesMemoryControl::EraseTransactionIdTracker(uint64_t transaction_id) {
   return removed;
 }
 
-bool QueriesMemoryControl::IsArenaTracked(unsigned arena_ind) {
-  return arena_tracking[arena_ind].load(std::memory_order_acquire) != 0;
+bool QueriesMemoryControl::CheckTransactionIdTrackerExists(uint64_t transaction_id) {
+  auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
+  return transaction_id_to_tracker_accessor.contains(transaction_id);
 }
 
-void QueriesMemoryControl::InitializeArenaCounter(unsigned arena_ind) {
-  arena_tracking[arena_ind].store(0, std::memory_order_relaxed);
+void QueriesMemoryControl::TryCreateTransactionProcTracker(uint64_t transaction_id, int64_t procedure_id,
+                                                           size_t limit) {
+  auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
+  auto query_tracker = transaction_id_to_tracker_accessor.find(transaction_id);
+
+  if (query_tracker == transaction_id_to_tracker_accessor.end()) {
+    return;
+  }
+
+  query_tracker->tracker.TryCreateProcTracker(procedure_id, limit);
 }
 
+void QueriesMemoryControl::SetActiveProcIdTracker(uint64_t transaction_id, int64_t procedure_id) {
+  auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
+  auto query_tracker = transaction_id_to_tracker_accessor.find(transaction_id);
+
+  if (query_tracker == transaction_id_to_tracker_accessor.end()) {
+    return;
+  }
+
+  query_tracker->tracker.SetActiveProc(procedure_id);
+}
+
+void QueriesMemoryControl::PauseProcedureTracking(uint64_t transaction_id) {
+  auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
+  auto query_tracker = transaction_id_to_tracker_accessor.find(transaction_id);
+
+  if (query_tracker == transaction_id_to_tracker_accessor.end()) {
+    return;
+  }
+
+  query_tracker->tracker.StopProcTracking();
+}
+
+inline int &Get_Thread_Tracker() {
+  // store variable in bss segment for each thread
+  // https://cs-fundamentals.com/c-programming/memory-layout-of-c-program-code-data-segments#size-of-code-data-bss-segments
+  static thread_local int is_thread_tracked{0};
+  return is_thread_tracked;
+}
+
+bool QueriesMemoryControl::IsThreadTracked() { return Get_Thread_Tracker() == 1; }
+
 #endif
 
 void StartTrackingCurrentThreadTransaction(uint64_t transaction_id) {
 #if USE_JEMALLOC
+  Get_Thread_Tracker() = 0;
   GetQueriesMemoryControl().UpdateThreadToTransactionId(std::this_thread::get_id(), transaction_id);
-  GetQueriesMemoryControl().AddTrackingOnArena(QueriesMemoryControl::GetArenaForThread());
+  Get_Thread_Tracker() = 1;
 #endif
 }
 
 void StopTrackingCurrentThreadTransaction(uint64_t transaction_id) {
 #if USE_JEMALLOC
+  Get_Thread_Tracker() = 0;
   GetQueriesMemoryControl().EraseThreadToTransactionId(std::this_thread::get_id(), transaction_id);
-  GetQueriesMemoryControl().RemoveTrackingOnArena(QueriesMemoryControl::GetArenaForThread());
 #endif
 }
 
@@ -137,4 +190,29 @@ void TryStopTrackingOnTransaction(uint64_t transaction_id) {
 #endif
 }
 
+#if USE_JEMALLOC
+bool IsTransactionTracked(uint64_t transaction_id) {
+  return GetQueriesMemoryControl().CheckTransactionIdTrackerExists(transaction_id);
+}
+#else
+bool IsTransactionTracked(uint64_t /*transaction_id*/) { return false; }
+#endif
+
+void CreateOrContinueProcedureTracking(uint64_t transaction_id, int64_t procedure_id, size_t limit) {
+#if USE_JEMALLOC
+  if (!GetQueriesMemoryControl().CheckTransactionIdTrackerExists(transaction_id)) {
+    LOG_FATAL("Memory tracker for transaction was not set");
+  }
+
+  GetQueriesMemoryControl().TryCreateTransactionProcTracker(transaction_id, procedure_id, limit);
+  GetQueriesMemoryControl().SetActiveProcIdTracker(transaction_id, procedure_id);
+#endif
+}
+
+void PauseProcedureTracking(uint64_t transaction_id) {
+#if USE_JEMALLOC
+  GetQueriesMemoryControl().PauseProcedureTracking(transaction_id);
+#endif
+}
+
 }  // namespace memgraph::memory

src/memory/query_memory_control.hpp

@@ -16,10 +16,13 @@
 #include <unordered_map>
 
 #include "utils/memory_tracker.hpp"
+#include "utils/query_memory_tracker.hpp"
 #include "utils/skip_list.hpp"
 
 namespace memgraph::memory {
 
+static constexpr int64_t UNLIMITED_MEMORY{0};
+
 #if USE_JEMALLOC
 
 // Track memory allocations per query.
@@ -30,25 +33,6 @@ namespace memgraph::memory {
 // it is necessary to restart tracking at the beginning of new query for that transaction.
 class QueriesMemoryControl {
  public:
-  /*
-  Arena stats
-  */
-
-  static unsigned GetArenaForThread();
-
-  // Add counter on threads allocating inside arena
-  void AddTrackingOnArena(unsigned);
-
-  // Remove counter on threads allocating in arena
-  void RemoveTrackingOnArena(unsigned);
-
-  // Are any threads using current arena for allocations
-  // Multiple threads can allocate inside one arena
-  bool IsArenaTracked(unsigned);
-
-  // Initialize arena counter
-  void InitializeArenaCounter(unsigned);
-
   /*
     Transaction id <-> tracker
   */
@@ -75,20 +59,25 @@ class QueriesMemoryControl {
   // Important to reset if one thread gets reused for different transaction
   void EraseThreadToTransactionId(const std::thread::id &, uint64_t);
 
-  // C-API functionality for thread to transaction mapping
-  void UpdateThreadToTransactionId(const char *, uint64_t);
+  // Find tracker for current thread if exists, track
+  // query allocation and procedure allocation if
+  // necessary
+  void TrackAllocOnCurrentThread(size_t size);
 
-  // C-API functionality for thread to transaction unmapping
-  void EraseThreadToTransactionId(const char *, uint64_t);
+  // Find tracker for current thread if exists, track
+  // query allocation and procedure allocation if
+  // necessary
+  void TrackFreeOnCurrentThread(size_t size);
 
-  // Get tracker to current thread if exists, otherwise return
-  // nullptr. This can happen only if tracker is still
-  // being constructed.
-  utils::MemoryTracker *GetTrackerCurrentThread();
+  void TryCreateTransactionProcTracker(uint64_t, int64_t, size_t);
+
+  void SetActiveProcIdTracker(uint64_t, int64_t);
+
+  void PauseProcedureTracking(uint64_t);
+
+  bool IsThreadTracked();
 
  private:
-  std::unordered_map<unsigned, std::atomic<int>> arena_tracking;
-
   struct ThreadIdToTransactionId {
     std::thread::id thread_id;
     uint64_t transaction_id;
@@ -102,7 +91,7 @@ class QueriesMemoryControl {
 
   struct TransactionIdToTracker {
     uint64_t transaction_id;
-    utils::MemoryTracker tracker;
+    utils::QueryMemoryTracker tracker;
 
     bool operator<(const TransactionIdToTracker &other) const { return transaction_id < other.transaction_id; }
     bool operator==(const TransactionIdToTracker &other) const { return transaction_id == other.transaction_id; }
@@ -138,4 +127,15 @@ void TryStartTrackingOnTransaction(uint64_t transaction_id, size_t limit);
 // Does nothing if jemalloc is not enabled. Does nothing if tracker doesn't exist
 void TryStopTrackingOnTransaction(uint64_t transaction_id);
 
+// Is transaction with given id tracked in memory tracker
+bool IsTransactionTracked(uint64_t transaction_id);
+
+// Creates tracker on procedure if doesn't exist. Sets query tracker
+// to track procedure with id.
+void CreateOrContinueProcedureTracking(uint64_t transaction_id, int64_t procedure_id, size_t limit);
+
+// Pauses procedure tracking. This enables to continue
+// tracking on procedure once procedure execution resumes.
+void PauseProcedureTracking(uint64_t transaction_id);
+
 }  // namespace memgraph::memory

src/mg_import_csv.cpp

@@ -19,7 +19,9 @@
 #include <regex>
 #include <unordered_map>
 
+#include "dbms/inmemory/storage_helper.hpp"
 #include "helpers.hpp"
+#include "replication/state.hpp"
 #include "storage/v2/config.hpp"
 #include "storage/v2/edge_accessor.hpp"
 #include "storage/v2/inmemory/storage.hpp"
@@ -702,14 +704,16 @@ int main(int argc, char *argv[]) {
   }
 
   std::unordered_map<NodeId, memgraph::storage::Gid> node_id_map;
-  auto store = std::make_unique<memgraph::storage::InMemoryStorage>(memgraph::storage::Config{
+  memgraph::storage::Config config{
 
       .items = {.properties_on_edges = FLAGS_storage_properties_on_edges},
       .durability = {.storage_directory = FLAGS_data_directory,
                      .recover_on_startup = false,
                      .snapshot_wal_mode = memgraph::storage::Config::Durability::SnapshotWalMode::DISABLED,
                      .snapshot_on_exit = true},
-  });
+  };
+  memgraph::replication::ReplicationState repl_state{memgraph::storage::ReplicationStateRootPath(config)};
+  auto store = memgraph::dbms::CreateInMemoryStorage(config, repl_state);
 
   memgraph::utils::Timer load_timer;
 

src/query/cypher_query_interpreter.cpp

@@ -12,15 +12,16 @@
 #include "query/cypher_query_interpreter.hpp"
 #include "query/frontend/ast/cypher_main_visitor.hpp"
 #include "query/frontend/opencypher/parser.hpp"
+#include "utils/synchronized.hpp"
 
 // NOLINTNEXTLINE (cppcoreguidelines-avoid-non-const-global-variables)
 DEFINE_bool(query_cost_planner, true, "Use the cost-estimating query planner.");
 // NOLINTNEXTLINE (cppcoreguidelines-avoid-non-const-global-variables)
-DEFINE_VALIDATED_int32(query_plan_cache_ttl, 60, "Time to live for cached query plans, in seconds.",
+DEFINE_VALIDATED_int32(query_plan_cache_max_size, 1000, "Maximum number of query plans to cache.",
                        FLAG_IN_RANGE(0, std::numeric_limits<int32_t>::max()));
 
 namespace memgraph::query {
-CachedPlan::CachedPlan(std::unique_ptr<LogicalPlan> plan) : plan_(std::move(plan)) {}
+PlanWrapper::PlanWrapper(std::unique_ptr<LogicalPlan> plan) : plan_(std::move(plan)) {}
 
 ParsedQuery ParseQuery(const std::string &query_string, const std::map<std::string, storage::PropertyValue> &params,
                        utils::SkipList<QueryCacheEntry> *cache, const InterpreterConfig::Query &query_config) {
@@ -127,28 +128,24 @@ std::unique_ptr<LogicalPlan> MakeLogicalPlan(AstStorage ast_storage, CypherQuery
                                                  std::move(symbol_table));
 }
 
-std::shared_ptr<CachedPlan> CypherQueryToPlan(uint64_t hash, AstStorage ast_storage, CypherQuery *query,
-                                              const Parameters &parameters, utils::SkipList<PlanCacheEntry> *plan_cache,
-                                              DbAccessor *db_accessor,
-                                              const std::vector<Identifier *> &predefined_identifiers) {
-  std::optional<utils::SkipList<PlanCacheEntry>::Accessor> plan_cache_access;
+std::shared_ptr<PlanWrapper> CypherQueryToPlan(uint64_t hash, AstStorage ast_storage, CypherQuery *query,
+                                               const Parameters &parameters, PlanCacheLRU *plan_cache,
+                                               DbAccessor *db_accessor,
+                                               const std::vector<Identifier *> &predefined_identifiers) {
   if (plan_cache) {
-    plan_cache_access.emplace(plan_cache->access());
-    auto it = plan_cache_access->find(hash);
-    if (it != plan_cache_access->end()) {
-      if (it->second->IsExpired()) {
-        plan_cache_access->remove(hash);
-      } else {
-        return it->second;
-      }
+    auto existing_plan = plan_cache->WithLock([&](auto &cache) { return cache.get(hash); });
+    if (existing_plan.has_value()) {
+      return existing_plan.value();
     }
   }
 
-  auto plan = std::make_shared<CachedPlan>(
+  auto plan = std::make_shared<PlanWrapper>(
       MakeLogicalPlan(std::move(ast_storage), query, parameters, db_accessor, predefined_identifiers));
-  if (plan_cache_access) {
-    plan_cache_access->insert({hash, plan});
+
+  if (plan_cache) {
+    plan_cache->WithLock([&](auto &cache) { cache.put(hash, plan); });
   }
+
   return plan;
 }
 }  // namespace memgraph::query

src/query/cypher_query_interpreter.hpp

@@ -17,12 +17,14 @@
 #include "query/frontend/stripped.hpp"
 #include "query/plan/planner.hpp"
 #include "utils/flag_validation.hpp"
+#include "utils/lru_cache.hpp"
+#include "utils/synchronized.hpp"
 #include "utils/timer.hpp"
 
 // NOLINTNEXTLINE (cppcoreguidelines-avoid-non-const-global-variables)
 DECLARE_bool(query_cost_planner);
 // NOLINTNEXTLINE (cppcoreguidelines-avoid-non-const-global-variables)
-DECLARE_int32(query_plan_cache_ttl);
+DECLARE_int32(query_plan_cache_max_size);
 
 namespace memgraph::query {
 
@@ -45,23 +47,17 @@ class LogicalPlan {
   virtual const AstStorage &GetAstStorage() const = 0;
 };
 
-class CachedPlan {
+class PlanWrapper {
  public:
-  explicit CachedPlan(std::unique_ptr<LogicalPlan> plan);
+  explicit PlanWrapper(std::unique_ptr<LogicalPlan> plan);
 
   const auto &plan() const { return plan_->GetRoot(); }
   double cost() const { return plan_->GetCost(); }
   const auto &symbol_table() const { return plan_->GetSymbolTable(); }
   const auto &ast_storage() const { return plan_->GetAstStorage(); }
 
-  bool IsExpired() const {
-    // NOLINTNEXTLINE (modernize-use-nullptr)
-    return cache_timer_.Elapsed() > std::chrono::seconds(FLAGS_query_plan_cache_ttl);
-  };
-
  private:
   std::unique_ptr<LogicalPlan> plan_;
-  utils::Timer cache_timer_;
 };
 
 struct CachedQuery {
@@ -82,18 +78,6 @@ struct QueryCacheEntry {
   CachedQuery second;
 };
 
-struct PlanCacheEntry {
-  bool operator==(const PlanCacheEntry &other) const { return first == other.first; }
-  bool operator<(const PlanCacheEntry &other) const { return first < other.first; }
-  bool operator==(const uint64_t &other) const { return first == other; }
-  bool operator<(const uint64_t &other) const { return first < other; }
-
-  uint64_t first;
-  // TODO: Maybe store the query string here and use it as a key with the hash
-  // so that we eliminate the risk of hash collisions.
-  std::shared_ptr<CachedPlan> second;
-};
-
 /**
  * A container for data related to the parsing of a query.
  */
@@ -129,6 +113,9 @@ class SingleNodeLogicalPlan final : public LogicalPlan {
   SymbolTable symbol_table_;
 };
 
+using PlanCacheLRU =
+    utils::Synchronized<utils::LRUCache<uint64_t, std::shared_ptr<query::PlanWrapper>>, utils::RWSpinLock>;
+
 std::unique_ptr<LogicalPlan> MakeLogicalPlan(AstStorage ast_storage, CypherQuery *query, const Parameters &parameters,
                                              DbAccessor *db_accessor,
                                              const std::vector<Identifier *> &predefined_identifiers);
@@ -141,9 +128,9 @@ std::unique_ptr<LogicalPlan> MakeLogicalPlan(AstStorage ast_storage, CypherQuery
  * If an identifier is contained there, we inject it at that place and remove it,
  * because a predefined identifier can be used only in one scope.
  */
-std::shared_ptr<CachedPlan> CypherQueryToPlan(uint64_t hash, AstStorage ast_storage, CypherQuery *query,
-                                              const Parameters &parameters, utils::SkipList<PlanCacheEntry> *plan_cache,
-                                              DbAccessor *db_accessor,
-                                              const std::vector<Identifier *> &predefined_identifiers = {});
+std::shared_ptr<PlanWrapper> CypherQueryToPlan(uint64_t hash, AstStorage ast_storage, CypherQuery *query,
+                                               const Parameters &parameters, PlanCacheLRU *plan_cache,
+                                               DbAccessor *db_accessor,
+                                               const std::vector<Identifier *> &predefined_identifiers = {});
 
 }  // namespace memgraph::query

src/query/frame_change.hpp

@@ -43,23 +43,23 @@ struct CachedValue {
     return cache_.get_allocator().GetMemoryResource();
   }
 
-  // Func to check if cache_ contains value
-  bool CacheValue(const TypedValue &value) {
-    if (!value.IsList()) {
+  bool CacheValue(const TypedValue &maybe_list) {
+    if (!maybe_list.IsList()) {
       return false;
     }
-    const auto &list = value.ValueList();
+    auto &list = maybe_list.ValueList();
     TypedValue::Hash hash{};
-    for (const TypedValue &element : list) {
+    for (auto &element : list) {
       const auto key = hash(element);
       auto &vector_values = cache_[key];
       if (!IsValueInVec(vector_values, element)) {
-        vector_values.push_back(element);
+        vector_values.emplace_back(element);
       }
     }
     return true;
   }
-  // Func to cache_value inside cache_
+
+  // Func to check if cache_ contains value
   bool ContainsValue(const TypedValue &value) const {
     TypedValue::Hash hash{};
     const auto key = hash(value);
@@ -70,7 +70,7 @@ struct CachedValue {
   }
 
  private:
-  bool IsValueInVec(const std::vector<TypedValue> &vec_values, const TypedValue &value) const {
+  static bool IsValueInVec(const std::vector<TypedValue> &vec_values, const TypedValue &value) {
     return std::any_of(vec_values.begin(), vec_values.end(), [&value](auto &vec_value) {
       const auto is_value_equal = vec_value == value;
       if (is_value_equal.IsNull()) return false;
@@ -80,43 +80,37 @@ struct CachedValue {
 };
 
 // Class tracks keys for which user can cache values which help with faster search or faster retrieval
-// in the future.
+// in the future. Used for IN LIST operator.
 class FrameChangeCollector {
  public:
-  explicit FrameChangeCollector(utils::MemoryResource *mem) : tracked_values_(mem){};
+  explicit FrameChangeCollector() : tracked_values_(&memory_resource_){};
 
-  // Add tracking key to cache later value
   CachedValue &AddTrackingKey(const std::string &key) {
     const auto &[it, _] = tracked_values_.emplace(key, tracked_values_.get_allocator().GetMemoryResource());
     return it->second;
   }
 
-  // Is key tracked
   bool IsKeyTracked(const std::string &key) const { return tracked_values_.contains(key); }
 
-  // Is value for given key cached
   bool IsKeyValueCached(const std::string &key) const {
-    return tracked_values_.contains(key) && !tracked_values_.at(key).cache_.empty();
+    return IsKeyTracked(key) && !tracked_values_.at(key).cache_.empty();
   }
 
-  // Reset value for tracking key
   bool ResetTrackingValue(const std::string &key) {
-    if (tracked_values_.contains(key)) {
-      tracked_values_.erase(key);
-      AddTrackingKey(key);
+    if (!tracked_values_.contains(key)) {
+      return false;
     }
-
+    tracked_values_.erase(key);
+    AddTrackingKey(key);
     return true;
   }
 
-  // Get value cached for tracking key, throws if key is not in tracked
   CachedValue &GetCachedValue(const std::string &key) { return tracked_values_.at(key); }
 
-  // Checks for keys tracked
   bool IsTrackingValues() const { return !tracked_values_.empty(); }
 
  private:
-  // Key is output of utils::GetFrameChangeId, value is utils::pmr::unordered_map
+  utils::MonotonicBufferResource memory_resource_{0};
   memgraph::utils::pmr::unordered_map<std::string, CachedValue> tracked_values_;
 };
 }  // namespace memgraph::query

src/query/frontend/semantic/symbol_generator.cpp

@@ -480,12 +480,20 @@ bool SymbolGenerator::PreVisit(None &none) {
 }
 
 bool SymbolGenerator::PreVisit(Reduce &reduce) {
+  auto &scope = scopes_.back();
+  scope.in_reduce = true;
   reduce.initializer_->Accept(*this);
   reduce.list_->Accept(*this);
   VisitWithIdentifiers(reduce.expression_, {reduce.accumulator_, reduce.identifier_});
   return false;
 }
 
+bool SymbolGenerator::PostVisit(Reduce & /*reduce*/) {
+  auto &scope = scopes_.back();
+  scope.in_reduce = false;
+  return true;
+}
+
 bool SymbolGenerator::PreVisit(Extract &extract) {
   extract.list_->Accept(*this);
   VisitWithIdentifiers(extract.expression_, {extract.identifier_});
@@ -496,15 +504,23 @@ bool SymbolGenerator::PreVisit(Exists &exists) {
   auto &scope = scopes_.back();
 
   if (scope.in_set_property) {
-    throw utils::NotYetImplemented("Set property can not be used with exists, but only during matching!");
+    throw utils::NotYetImplemented("Exists cannot be used within SET clause.!");
   }
 
   if (scope.in_with) {
-    throw utils::NotYetImplemented("WITH can not be used with exists, but only during matching!");
+    throw utils::NotYetImplemented("Exists cannot be used within WITH!");
   }
 
   if (scope.in_return) {
-    throw utils::NotYetImplemented("RETURN can not be used with exists, but only during matching!");
+    throw utils::NotYetImplemented("Exists cannot be used within RETURN!");
+  }
+
+  if (scope.in_reduce) {
+    throw utils::NotYetImplemented("Exists cannot be used within REDUCE!");
+  }
+
+  if (scope.num_if_operators) {
+    throw utils::NotYetImplemented("IF operator cannot be used with exists, but only during matching!");
   }
 
   scope.in_exists = true;

src/query/frontend/semantic/symbol_generator.hpp

@@ -84,6 +84,7 @@ class SymbolGenerator : public HierarchicalTreeVisitor {
   bool PreVisit(Any &) override;
   bool PreVisit(None &) override;
   bool PreVisit(Reduce &) override;
+  bool PostVisit(Reduce &) override;
   bool PreVisit(Extract &) override;
   bool PreVisit(Exists & /*exists*/) override;
   bool PostVisit(Exists & /*exists*/) override;
@@ -123,6 +124,7 @@ class SymbolGenerator : public HierarchicalTreeVisitor {
     bool in_match{false};
     bool in_foreach{false};
     bool in_exists{false};
+    bool in_reduce{false};
     bool in_set_property{false};
     bool in_call_subquery{false};
     bool has_return{false};

src/query/interpret/awesome_memgraph_functions.cpp

@@ -1156,11 +1156,16 @@ void MapNumericParameters(auto &parameter_mappings, const auto &input_parameters
 }
 
 TypedValue Date(const TypedValue *args, int64_t nargs, const FunctionContext &ctx) {
-  FType<Optional<Or<String, Map>>>("date", args, nargs);
+  FType<Optional<Or<String, Map, LocalDateTime>>>("date", args, nargs);
   if (nargs == 0) {
     return TypedValue(utils::LocalDateTime(ctx.timestamp).date, ctx.memory);
   }
 
+  if (args[0].IsLocalDateTime()) {
+    utils::Date date{args[0].ValueLocalDateTime().date};
+    return TypedValue(date, ctx.memory);
+  }
+
   if (args[0].IsString()) {
     const auto &[date_parameters, is_extended] = utils::ParseDateParameters(args[0].ValueString());
     return TypedValue(utils::Date(date_parameters), ctx.memory);
@@ -1178,12 +1183,17 @@ TypedValue Date(const TypedValue *args, int64_t nargs, const FunctionContext &ct
 }
 
 TypedValue LocalTime(const TypedValue *args, int64_t nargs, const FunctionContext &ctx) {
-  FType<Optional<Or<String, Map>>>("localtime", args, nargs);
+  FType<Optional<Or<String, Map, LocalDateTime>>>("localtime", args, nargs);
 
   if (nargs == 0) {
     return TypedValue(utils::LocalDateTime(ctx.timestamp).local_time, ctx.memory);
   }
 
+  if (args[0].IsLocalDateTime()) {
+    utils::LocalTime local_time{args[0].ValueLocalDateTime().local_time};
+    return TypedValue(local_time, ctx.memory);
+  }
+
   if (args[0].IsString()) {
     const auto &[local_time_parameters, is_extended] = utils::ParseLocalTimeParameters(args[0].ValueString());
     return TypedValue(utils::LocalTime(local_time_parameters), ctx.memory);

src/query/interpret/eval.hpp

@@ -270,36 +270,33 @@ class ExpressionEvaluator : public ExpressionVisitor<TypedValue> {
   }
 
   TypedValue Visit(InListOperator &in_list) override {
-    TypedValue *_list_ptr = nullptr;
-    TypedValue _list;
     auto literal = in_list.expression1_->Accept(*this);
 
-    auto get_list_literal = [this, &in_list, &_list, &_list_ptr]() -> void {
+    auto get_list_literal = [this, &in_list]() -> TypedValue {
       ReferenceExpressionEvaluator reference_expression_evaluator{frame_, symbol_table_, ctx_};
-      _list_ptr = in_list.expression2_->Accept(reference_expression_evaluator);
-      if (nullptr == _list_ptr) {
-        _list = in_list.expression2_->Accept(*this);
-        _list_ptr = &_list;
+      auto *list_ptr = in_list.expression2_->Accept(reference_expression_evaluator);
+      if (nullptr == list_ptr) {
+        return in_list.expression2_->Accept(*this);
       }
+      return *list_ptr;
     };
 
-    auto do_list_literal_checks = [this, &literal, &_list_ptr]() -> std::optional<TypedValue> {
-      MG_ASSERT(_list_ptr, "List literal should have been defined");
-      if (_list_ptr->IsNull()) {
+    auto do_list_literal_checks = [this, &literal](const TypedValue &list) -> std::optional<TypedValue> {
+      if (list.IsNull()) {
         return TypedValue(ctx_->memory);
       }
       // Exceptions have higher priority than returning nulls when list expression
       // is not null.
-      if (_list_ptr->type() != TypedValue::Type::List) {
-        throw QueryRuntimeException("IN expected a list, got {}.", _list_ptr->type());
+      if (list.type() != TypedValue::Type::List) {
+        throw QueryRuntimeException("IN expected a list, got {}.", list.type());
       }
-      const auto &list = _list_ptr->ValueList();
+      const auto &list_value = list.ValueList();
 
       // If literal is NULL there is no need to try to compare it with every
       // element in the list since result of every comparison will be NULL. There
       // is one special case that we must test explicitly: if list is empty then
       // result is false since no comparison will be performed.
-      if (list.empty()) return TypedValue(false, ctx_->memory);
+      if (list_value.empty()) return TypedValue(false, ctx_->memory);
       if (literal.IsNull()) return TypedValue(ctx_->memory);
       return {};
     };
@@ -312,14 +309,14 @@ class ExpressionEvaluator : public ExpressionVisitor<TypedValue> {
       if (!frame_change_collector_->IsKeyValueCached(*cached_id)) {
         // Check only first time if everything is okay, later when we use
         // cache there is no need to check again as we did check first time
-        get_list_literal();
-        auto preoperational_checks = do_list_literal_checks();
+        auto list = get_list_literal();
+        auto preoperational_checks = do_list_literal_checks(list);
         if (preoperational_checks) {
           return std::move(*preoperational_checks);
         }
         auto &cached_value = frame_change_collector_->GetCachedValue(*cached_id);
-        cached_value.CacheValue(*_list_ptr);
-        spdlog::trace("Value cached {}", *cached_id);
+        // Don't move here because we don't want to remove the element from the frame
+        cached_value.CacheValue(list);
       }
       const auto &cached_value = frame_change_collector_->GetCachedValue(*cached_id);
 
@@ -334,16 +331,15 @@ class ExpressionEvaluator : public ExpressionVisitor<TypedValue> {
     }
     // When caching is not an option, we need to evaluate list literal every time
     // and do the checks
-    get_list_literal();
-    auto preoperational_checks = do_list_literal_checks();
+    const auto list = get_list_literal();
+    auto preoperational_checks = do_list_literal_checks(list);
     if (preoperational_checks) {
       return std::move(*preoperational_checks);
     }
 
-    const auto &list = _list.ValueList();
-    spdlog::trace("Not using cache on IN LIST operator");
+    const auto &list_value = list.ValueList();
     auto has_null = false;
-    for (const auto &element : list) {
+    for (const auto &element : list_value) {
       auto result = literal == element;
       if (result.IsNull()) {
         has_null = true;

src/query/interpreter.cpp

@@ -37,6 +37,7 @@
 #include "dbms/database.hpp"
 #include "dbms/dbms_handler.hpp"
 #include "dbms/global.hpp"
+#include "dbms/inmemory/storage_helper.hpp"
 #include "flags/run_time_configurable.hpp"
 #include "glue/communication.hpp"
 #include "license/license.hpp"
@@ -94,16 +95,17 @@
 #include "utils/on_scope_exit.hpp"
 #include "utils/readable_size.hpp"
 #include "utils/settings.hpp"
+#include "utils/stat.hpp"
 #include "utils/string.hpp"
 #include "utils/tsc.hpp"
 #include "utils/typeinfo.hpp"
 #include "utils/variant_helpers.hpp"
 
 #include "dbms/dbms_handler.hpp"
+#include "dbms/replication_handler.hpp"
 #include "query/auth_query_handler.hpp"
 #include "query/interpreter_context.hpp"
 #include "replication/state.hpp"
-#include "storage/v2/replication/replication_handler.hpp"
 
 namespace memgraph::metrics {
 extern Event ReadQuery;
@@ -270,7 +272,8 @@ inline auto convertToReplicationMode(const ReplicationQuery::SyncMode &sync_mode
 
 class ReplQueryHandler final : public query::ReplicationQueryHandler {
  public:
-  explicit ReplQueryHandler(storage::Storage *db) : db_(db), handler_{db_->repl_state_, *db_} {}
+  explicit ReplQueryHandler(dbms::DbmsHandler *dbms_handler, memgraph::replication::ReplicationState *repl_state)
+      : dbms_handler_(dbms_handler), handler_{*repl_state, *dbms_handler} {}
 
   /// @throw QueryRuntimeException if an error ocurred.
   void SetReplicationRole(ReplicationQuery::ReplicationRole replication_role, std::optional<int64_t> port) override {
@@ -314,10 +317,6 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
       throw QueryRuntimeException("Replica can't register another replica!");
     }
 
-    if (name == memgraph::replication::kReservedReplicationRoleName) {
-      throw QueryRuntimeException("This replica name is reserved and can not be used as replica name!");
-    }
-
     auto repl_mode = convertToReplicationMode(sync_mode);
 
     auto maybe_ip_and_port =
@@ -330,8 +329,7 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
                                                          .port = port,
                                                          .replica_check_frequency = replica_check_frequency,
                                                          .ssl = std::nullopt};
-      using storage::RegistrationMode;
-      auto ret = handler_.RegisterReplica(RegistrationMode::MUST_BE_INSTANTLY_VALID, config);
+      auto ret = handler_.RegisterReplica(config);
       if (ret.HasError()) {
         throw QueryRuntimeException(fmt::format("Couldn't register replica '{}'!", name));
       }
@@ -344,7 +342,7 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
   void DropReplica(std::string_view replica_name) override {
     auto const result = handler_.UnregisterReplica(replica_name);
     switch (result) {
-      using enum memgraph::storage::UnregisterReplicaResult;
+      using enum memgraph::dbms::UnregisterReplicaResult;
       case NOT_MAIN:
         throw QueryRuntimeException("Replica can't unregister a replica!");
       case COULD_NOT_BE_PERSISTED:
@@ -358,13 +356,22 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
 
   using Replica = ReplicationQueryHandler::Replica;
   std::vector<Replica> ShowReplicas() const override {
-    auto const &replState = db_->repl_state_;
-    if (replState.IsReplica()) {
+    if (handler_.IsReplica()) {
       // replica can't show registered replicas (it shouldn't have any)
       throw QueryRuntimeException("Replica can't show registered replicas (it shouldn't have any)!");
     }
 
-    auto repl_infos = db_->ReplicasInfo();
+    // TODO: Combine results? Have a single place with clients???
+    //       Also authentication checks (replica + database visibility)
+    std::vector<storage::ReplicaInfo> repl_infos{};
+    dbms_handler_->ForOne([&repl_infos](dbms::Database *db) -> bool {
+      auto infos = db->storage()->ReplicasInfo();
+      if (!infos.empty()) {
+        repl_infos = std::move(infos);
+        return true;
+      }
+      return false;
+    });
     std::vector<Replica> replicas;
     replicas.reserve(repl_infos.size());
 
@@ -408,8 +415,8 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
   }
 
  private:
-  storage::Storage *db_;
-  storage::ReplicationHandler handler_;
+  dbms::DbmsHandler *dbms_handler_;
+  dbms::ReplicationHandler handler_;
 };
 
 /// returns false if the replication role can't be set
@@ -418,7 +425,7 @@ class ReplQueryHandler final : public query::ReplicationQueryHandler {
 Callback HandleAuthQuery(AuthQuery *auth_query, InterpreterContext *interpreter_context, const Parameters &parameters) {
   AuthQueryHandler *auth = interpreter_context->auth;
 #ifdef MG_ENTERPRISE
-  auto *db_handler = interpreter_context->db_handler;
+  auto *db_handler = interpreter_context->dbms_handler;
 #endif
   // TODO: MemoryResource for EvaluationContext, it should probably be passed as
   // the argument to Callback.
@@ -702,8 +709,10 @@ Callback HandleAuthQuery(AuthQuery *auth_query, InterpreterContext *interpreter_
   }
 }  // namespace
 
-Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &parameters, storage::Storage *storage,
-                                const query::InterpreterConfig &config, std::vector<Notification> *notifications) {
+Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &parameters,
+                                dbms::DbmsHandler *dbms_handler, const query::InterpreterConfig &config,
+                                std::vector<Notification> *notifications,
+                                memgraph::replication::ReplicationState *repl_state) {
   // TODO: MemoryResource for EvaluationContext, it should probably be passed as
   // the argument to Callback.
   EvaluationContext evaluation_context;
@@ -723,7 +732,8 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
         notifications->emplace_back(SeverityLevel::WARNING, NotificationCode::REPLICA_PORT_WARNING,
                                     "Be careful the replication port must be different from the memgraph port!");
       }
-      callback.fn = [handler = ReplQueryHandler{storage}, role = repl_query->role_, maybe_port]() mutable {
+      callback.fn = [handler = ReplQueryHandler{dbms_handler, repl_state}, role = repl_query->role_,
+                     maybe_port]() mutable {
         handler.SetReplicationRole(role, maybe_port);
         return std::vector<std::vector<TypedValue>>();
       };
@@ -735,7 +745,7 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
     }
     case ReplicationQuery::Action::SHOW_REPLICATION_ROLE: {
       callback.header = {"replication role"};
-      callback.fn = [handler = ReplQueryHandler{storage}] {
+      callback.fn = [handler = ReplQueryHandler{dbms_handler, repl_state}] {
         auto mode = handler.ShowReplicationRole();
         switch (mode) {
           case ReplicationQuery::ReplicationRole::MAIN: {
@@ -754,7 +764,7 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
       auto socket_address = repl_query->socket_address_->Accept(evaluator);
       const auto replica_check_frequency = config.replication_replica_check_frequency;
 
-      callback.fn = [handler = ReplQueryHandler{storage}, name, socket_address, sync_mode,
+      callback.fn = [handler = ReplQueryHandler{dbms_handler, repl_state}, name, socket_address, sync_mode,
                      replica_check_frequency]() mutable {
         handler.RegisterReplica(name, std::string(socket_address.ValueString()), sync_mode, replica_check_frequency);
         return std::vector<std::vector<TypedValue>>();
@@ -765,7 +775,7 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
     }
     case ReplicationQuery::Action::DROP_REPLICA: {
       const auto &name = repl_query->replica_name_;
-      callback.fn = [handler = ReplQueryHandler{storage}, name]() mutable {
+      callback.fn = [handler = ReplQueryHandler{dbms_handler, repl_state}, name]() mutable {
         handler.DropReplica(name);
         return std::vector<std::vector<TypedValue>>();
       };
@@ -777,7 +787,7 @@ Callback HandleReplicationQuery(ReplicationQuery *repl_query, const Parameters &
       callback.header = {
           "name", "socket_address", "sync_mode", "current_timestamp_of_replica", "number_of_timestamp_behind_master",
           "state"};
-      callback.fn = [handler = ReplQueryHandler{storage}, replica_nfields = callback.header.size()] {
+      callback.fn = [handler = ReplQueryHandler{dbms_handler, repl_state}, replica_nfields = callback.header.size()] {
         const auto &replicas = handler.ShowReplicas();
         auto typed_replicas = std::vector<std::vector<TypedValue>>{};
         typed_replicas.reserve(replicas.size());
@@ -1203,7 +1213,7 @@ struct TxTimeout {
 };
 
 struct PullPlan {
-  explicit PullPlan(std::shared_ptr<CachedPlan> plan, const Parameters &parameters, bool is_profile_query,
+  explicit PullPlan(std::shared_ptr<PlanWrapper> plan, const Parameters &parameters, bool is_profile_query,
                     DbAccessor *dba, InterpreterContext *interpreter_context, utils::MemoryResource *execution_memory,
                     std::optional<std::string> username, std::atomic<TransactionStatus> *transaction_status,
                     std::shared_ptr<utils::AsyncTimer> tx_timer,
@@ -1216,7 +1226,7 @@ struct PullPlan {
                                                         std::map<std::string, TypedValue> *summary);
 
  private:
-  std::shared_ptr<CachedPlan> plan_ = nullptr;
+  std::shared_ptr<PlanWrapper> plan_ = nullptr;
   plan::UniqueCursorPtr cursor_ = nullptr;
   Frame frame_;
   ExecutionContext ctx_;
@@ -1243,7 +1253,7 @@ struct PullPlan {
   bool use_monotonic_memory_;
 };
 
-PullPlan::PullPlan(const std::shared_ptr<CachedPlan> plan, const Parameters &parameters, const bool is_profile_query,
+PullPlan::PullPlan(const std::shared_ptr<PlanWrapper> plan, const Parameters &parameters, const bool is_profile_query,
                    DbAccessor *dba, InterpreterContext *interpreter_context, utils::MemoryResource *execution_memory,
                    std::optional<std::string> username, std::atomic<TransactionStatus> *transaction_status,
                    std::shared_ptr<utils::AsyncTimer> tx_timer, TriggerContextCollector *trigger_context_collector,
@@ -1399,42 +1409,16 @@ std::optional<plan::ProfilingStatsWithTotalTime> PullPlan::Pull(AnyStream *strea
 
 using RWType = plan::ReadWriteTypeChecker::RWType;
 
-bool IsWriteQueryOnMainMemoryReplica(storage::Storage *storage,
-                                     const query::plan::ReadWriteTypeChecker::RWType query_type) {
-  if (auto storage_mode = storage->GetStorageMode(); storage_mode == storage::StorageMode::IN_MEMORY_ANALYTICAL ||
-                                                     storage_mode == storage::StorageMode::IN_MEMORY_TRANSACTIONAL) {
-    auto const &replState = storage->repl_state_;
-    return replState.IsReplica() && (query_type == RWType::W || query_type == RWType::RW);
-  }
-  return false;
-}
-
-bool IsReplica(storage::Storage *storage) {
-  if (auto storage_mode = storage->GetStorageMode(); storage_mode == storage::StorageMode::IN_MEMORY_ANALYTICAL ||
-                                                     storage_mode == storage::StorageMode::IN_MEMORY_TRANSACTIONAL) {
-    auto const &replState = storage->repl_state_;
-    return replState.IsReplica();
-  }
-  return false;
+bool IsQueryWrite(const query::plan::ReadWriteTypeChecker::RWType query_type) {
+  return query_type == RWType::W || query_type == RWType::RW;
 }
 
 }  // namespace
 
-#ifdef MG_ENTERPRISE
-InterpreterContext::InterpreterContext(InterpreterConfig interpreter_config, memgraph::dbms::DbmsHandler *handler,
-                                       query::AuthQueryHandler *ah, query::AuthChecker *ac)
-    : db_handler(handler), config(interpreter_config), auth(ah), auth_checker(ac) {}
-#else
-InterpreterContext::InterpreterContext(InterpreterConfig interpreter_config,
-                                       memgraph::utils::Gatekeeper<memgraph::dbms::Database> *db_gatekeeper,
-                                       query::AuthQueryHandler *ah, query::AuthChecker *ac)
-    : db_gatekeeper(db_gatekeeper), config(interpreter_config), auth(ah), auth_checker(ac) {}
-#endif
-
 Interpreter::Interpreter(InterpreterContext *interpreter_context) : interpreter_context_(interpreter_context) {
   MG_ASSERT(interpreter_context_, "Interpreter context must not be NULL");
 #ifndef MG_ENTERPRISE
-  auto db_acc = interpreter_context_->db_gatekeeper->access();
+  auto db_acc = interpreter_context_->dbms_handler->Get();
   MG_ASSERT(db_acc, "Database accessor needs to be valid");
   current_db_.db_acc_ = std::move(db_acc);
 #endif
@@ -1551,7 +1535,6 @@ inline static void TryCaching(const AstStorage &ast_storage, FrameChangeCollecto
       continue;
     }
     frame_change_collector->AddTrackingKey(*cached_id);
-    spdlog::trace("Tracking {} operator, by id: {}", InListOperator::kType.name, *cached_id);
   }
 }
 
@@ -1612,8 +1595,6 @@ PreparedQuery PrepareCypherQuery(ParsedQuery parsed_query, std::map<std::string,
   // If this is LOAD CSV query, use PoolResource without MonotonicMemoryResource as we want to reuse allocated memory
   auto use_monotonic_memory =
       !contains_csv && !IsCallBatchedProcedureQuery(clauses) && !IsAllShortestPathsQuery(clauses);
-  spdlog::trace("PrepareCypher has {} encountered all shortest paths and will {} use of monotonic memory",
-                IsAllShortestPathsQuery(clauses) ? "" : "not", use_monotonic_memory ? "" : "not");
 
   MG_ASSERT(current_db.execution_db_accessor_, "Cypher query expects a current DB transaction");
   auto *dba =
@@ -2124,10 +2105,7 @@ PreparedQuery PrepareAnalyzeGraphQuery(ParsedQuery parsed_query, bool in_explici
 
   // Creating an index influences computed plan costs.
   auto invalidate_plan_cache = [plan_cache = current_db.db_acc_->get()->plan_cache()] {
-    auto access = plan_cache->access();
-    for (auto &kv : access) {
-      access.remove(kv.first);
-    }
+    plan_cache->WithLock([&](auto &cache) { cache.reset(); });
   };
   utils::OnScopeExit cache_invalidator(invalidate_plan_cache);
 
@@ -2172,10 +2150,7 @@ PreparedQuery PrepareIndexQuery(ParsedQuery parsed_query, bool in_explicit_trans
 
   // Creating an index influences computed plan costs.
   auto invalidate_plan_cache = [plan_cache = db_acc->plan_cache()] {
-    auto access = plan_cache->access();
-    for (auto &kv : access) {
-      access.remove(kv.first);
-    }
+    plan_cache->WithLock([&](auto &cache) { cache.reset(); });
   };
 
   auto *storage = db_acc->storage();
@@ -2286,21 +2261,16 @@ PreparedQuery PrepareAuthQuery(ParsedQuery parsed_query, bool in_explicit_transa
 }
 
 PreparedQuery PrepareReplicationQuery(ParsedQuery parsed_query, bool in_explicit_transaction,
-                                      std::vector<Notification> *notifications, CurrentDB &current_db,
-                                      const InterpreterConfig &config) {
+                                      std::vector<Notification> *notifications, dbms::DbmsHandler &dbms_handler,
+                                      const InterpreterConfig &config,
+                                      memgraph::replication::ReplicationState *repl_state) {
   if (in_explicit_transaction) {
     throw ReplicationModificationInMulticommandTxException();
   }
 
-  MG_ASSERT(current_db.db_acc_, "Replication query expects a current DB");
-  storage::Storage *storage = current_db.db_acc_->get()->storage();
-
-  if (storage->GetStorageMode() == storage::StorageMode::ON_DISK_TRANSACTIONAL) {
-    throw ReplicationDisabledOnDiskStorage();
-  }
-
   auto *replication_query = utils::Downcast<ReplicationQuery>(parsed_query.query);
-  auto callback = HandleReplicationQuery(replication_query, parsed_query.parameters, storage, config, notifications);
+  auto callback = HandleReplicationQuery(replication_query, parsed_query.parameters, &dbms_handler, config,
+                                         notifications, repl_state);
 
   return PreparedQuery{callback.header, std::move(parsed_query.required_privileges),
                        [callback_fn = std::move(callback.fn), pull_plan = std::shared_ptr<PullPlanVector>{nullptr}](
@@ -2773,7 +2743,8 @@ PreparedQuery PrepareStorageModeQuery(ParsedQuery parsed_query, const bool in_ex
           "transactions using 'SHOW TRANSACTIONS' query and ensure no other transactions are active.");
     }
 
-    callback = [requested_mode, storage = db_acc->storage()]() -> std::function<void()> {
+    callback = [requested_mode,
+                storage = static_cast<storage::InMemoryStorage *>(db_acc->storage())]() -> std::function<void()> {
       // SetStorageMode will probably be handled at the Database level
       return [storage, requested_mode] { storage->SetStorageMode(requested_mode); };
     }();
@@ -2836,15 +2807,11 @@ PreparedQuery PrepareCreateSnapshotQuery(ParsedQuery parsed_query, bool in_expli
       std::move(parsed_query.required_privileges),
       [storage](AnyStream * /*stream*/, std::optional<int> /*n*/) -> std::optional<QueryHandlerResult> {
         auto *mem_storage = static_cast<storage::InMemoryStorage *>(storage);
-        if (auto maybe_error = mem_storage->CreateSnapshot(storage->repl_state_, {}); maybe_error.HasError()) {
+        if (auto maybe_error = mem_storage->CreateSnapshot(); maybe_error.HasError()) {
           switch (maybe_error.GetError()) {
             case storage::InMemoryStorage::CreateSnapshotError::DisabledForReplica:
               throw utils::BasicException(
                   "Failed to create a snapshot. Replica instances are not allowed to create them.");
-            case storage::InMemoryStorage::CreateSnapshotError::DisabledForAnalyticsPeriodicCommit:
-              spdlog::warn(utils::MessageWithLink("Periodic snapshots are disabled for analytical mode.",
-                                                  "https://memgr.ph/replication"));
-              break;
             case storage::InMemoryStorage::CreateSnapshotError::ReachedMaxNumTries:
               spdlog::warn("Failed to create snapshot. Reached max number of tries. Please contact support");
               break;
@@ -3116,14 +3083,22 @@ PreparedQuery PrepareSystemInfoQuery(ParsedQuery parsed_query, bool in_explicit_
       handler = [storage = current_db.db_acc_->get()->storage(), interpreter_isolation_level,
                  next_transaction_isolation_level] {
         auto info = storage->GetBaseInfo();
+        const auto vm_max_map_count = utils::GetVmMaxMapCount();
+        const int64_t vm_max_map_count_storage_info =
+            vm_max_map_count.has_value() ? vm_max_map_count.value() : memgraph::utils::VM_MAX_MAP_COUNT_DEFAULT;
+        auto memory_res = utils::GetReadableSize(static_cast<double>(info.memory_res));
+        auto memory_tracked = utils::GetReadableSize(static_cast<double>(utils::total_memory_tracker.Amount()));
+
+        spdlog::trace("Memory res: {}, memory tracked: {}", memory_res, memory_tracked);
         std::vector<std::vector<TypedValue>> results{
             {TypedValue("name"), TypedValue(storage->id())},
             {TypedValue("vertex_count"), TypedValue(static_cast<int64_t>(info.vertex_count))},
             {TypedValue("edge_count"), TypedValue(static_cast<int64_t>(info.edge_count))},
             {TypedValue("average_degree"), TypedValue(info.average_degree)},
-            {TypedValue("memory_usage"), TypedValue(utils::GetReadableSize(static_cast<double>(info.memory_usage)))},
+            {TypedValue("vm_max_map_count"), TypedValue(vm_max_map_count_storage_info)},
+            {TypedValue("memory_res"), TypedValue(utils::GetReadableSize(static_cast<double>(info.memory_res)))},
             {TypedValue("disk_usage"), TypedValue(utils::GetReadableSize(static_cast<double>(info.disk_usage)))},
-            {TypedValue("memory_allocated"),
+            {TypedValue("memory_tracked"),
              TypedValue(utils::GetReadableSize(static_cast<double>(utils::total_memory_tracker.Amount())))},
             {TypedValue("allocation_limit"),
              TypedValue(utils::GetReadableSize(static_cast<double>(utils::total_memory_tracker.HardLimit())))},
@@ -3377,22 +3352,23 @@ PreparedQuery PrepareConstraintQuery(ParsedQuery parsed_query, bool in_explicit_
 
 PreparedQuery PrepareMultiDatabaseQuery(ParsedQuery parsed_query, CurrentDB &current_db,
                                         InterpreterContext *interpreter_context,
-                                        std::optional<std::function<void(std::string_view)>> on_change_cb) {
+                                        std::optional<std::function<void(std::string_view)>> on_change_cb,
+                                        memgraph::replication::ReplicationState *repl_state) {
 #ifdef MG_ENTERPRISE
   if (!license::global_license_checker.IsEnterpriseValidFast()) {
     throw QueryException("Trying to use enterprise feature without a valid license.");
   }
   // TODO: Remove once replicas support multi-tenant replication
   if (!current_db.db_acc_) throw DatabaseContextRequiredException("Multi database queries require a defined database.");
-  if (IsReplica(current_db.db_acc_->get()->storage())) {
-    throw QueryException("Query forbidden on the replica!");
-  }
 
   auto *query = utils::Downcast<MultiDatabaseQuery>(parsed_query.query);
-  auto *db_handler = interpreter_context->db_handler;
+  auto *db_handler = interpreter_context->dbms_handler;
 
   switch (query->action_) {
     case MultiDatabaseQuery::Action::CREATE:
+      if (repl_state->IsReplica()) {
+        throw QueryException("Query forbidden on the replica!");
+      }
       return PreparedQuery{
           {"STATUS"},
           std::move(parsed_query.required_privileges),
@@ -3435,6 +3411,9 @@ PreparedQuery PrepareMultiDatabaseQuery(ParsedQuery parsed_query, CurrentDB &cur
       if (current_db.in_explicit_db_) {
         throw QueryException("Database switching is prohibited if session explicitly defines the used database");
       }
+      if (!dbms::allow_mt_repl && repl_state->IsReplica()) {
+        throw QueryException("Query forbidden on the replica!");
+      }
       return PreparedQuery{{"STATUS"},
                            std::move(parsed_query.required_privileges),
                            [db_name = query->db_name_, db_handler, &current_db, on_change_cb](
@@ -3466,6 +3445,9 @@ PreparedQuery PrepareMultiDatabaseQuery(ParsedQuery parsed_query, CurrentDB &cur
                            query->db_name_};
 
     case MultiDatabaseQuery::Action::DROP:
+      if (repl_state->IsReplica()) {
+        throw QueryException("Query forbidden on the replica!");
+      }
       return PreparedQuery{
           {"STATUS"},
           std::move(parsed_query.required_privileges),
@@ -3473,11 +3455,9 @@ PreparedQuery PrepareMultiDatabaseQuery(ParsedQuery parsed_query, CurrentDB &cur
               AnyStream *stream, std::optional<int> n) -> std::optional<QueryHandlerResult> {
             std::vector<std::vector<TypedValue>> status;
 
-            memgraph::dbms::DeleteResult success{};
-
             try {
               // Remove database
-              success = db_handler->Delete(db_name);
+              auto success = db_handler->Delete(db_name);
               if (!success.HasError()) {
                 // Remove from auth
                 auth->DeleteDatabase(db_name);
@@ -3526,14 +3506,9 @@ PreparedQuery PrepareShowDatabasesQuery(ParsedQuery parsed_query, CurrentDB &cur
   if (!license::global_license_checker.IsEnterpriseValidFast()) {
     throw QueryException("Trying to use enterprise feature without a valid license.");
   }
-  // TODO: Remove once replicas support multi-tenant replication
-  auto &replState = storage->repl_state_;
-  if (replState.IsReplica()) {
-    throw QueryException("SHOW DATABASES forbidden on the replica!");
-  }
 
   // TODO pick directly from ic
-  auto *db_handler = interpreter_context->db_handler;
+  auto *db_handler = interpreter_context->dbms_handler;
   AuthQueryHandler *auth = interpreter_context->auth;
 
   Callback callback;
@@ -3637,7 +3612,7 @@ void Interpreter::RollbackTransaction() {
 void Interpreter::SetCurrentDB(std::string_view db_name, bool in_explicit_db) {
   // Can throw
   // do we lock here?
-  current_db_.SetCurrentDB(interpreter_context_->db_handler->Get(db_name), in_explicit_db);
+  current_db_.SetCurrentDB(interpreter_context_->dbms_handler->Get(db_name), in_explicit_db);
 }
 #endif
 
@@ -3714,8 +3689,6 @@ Interpreter::PrepareResult Interpreter::Prepare(const std::string &query_string,
     // Setup QueryExecution
     // its MemoryResource is mostly used for allocations done on Frame and storing `row`s
     if (usePool) {
-      spdlog::trace("PrepareCypher has {} encountered all shortest paths, QueryExecution will use PoolResource",
-                    hasAllShortestPaths ? "" : "not");
       query_executions_.emplace_back(QueryExecution::Create(utils::PoolResource(128, kExecutionPoolMaxBlockSize)));
     } else {
       query_executions_.emplace_back(QueryExecution::Create(utils::MonotonicBufferResource(kExecutionMemoryBlockSize)));
@@ -3761,7 +3734,7 @@ Interpreter::PrepareResult Interpreter::Prepare(const std::string &query_string,
         std::visit([](auto &execution_memory) -> utils::MemoryResource * { return &execution_memory; },
                    query_execution->execution_memory);
     frame_change_collector_.reset();
-    frame_change_collector_.emplace(memory_resource);
+    frame_change_collector_.emplace();
     if (utils::Downcast<CypherQuery>(parsed_query.query)) {
       prepared_query = PrepareCypherQuery(std::move(parsed_query), &query_execution->summary, interpreter_context_,
                                           current_db_, memory_resource, &query_execution->notifications, username_,
@@ -3794,9 +3767,9 @@ Interpreter::PrepareResult Interpreter::Prepare(const std::string &query_string,
                                               &query_execution->notifications, current_db_);
     } else if (utils::Downcast<ReplicationQuery>(parsed_query.query)) {
       /// TODO: make replication DB agnostic
-      prepared_query =
-          PrepareReplicationQuery(std::move(parsed_query), in_explicit_transaction_, &query_execution->notifications,
-                                  current_db_, interpreter_context_->config);
+      prepared_query = PrepareReplicationQuery(std::move(parsed_query), in_explicit_transaction_,
+                                               &query_execution->notifications, *interpreter_context_->dbms_handler,
+                                               interpreter_context_->config, interpreter_context_->repl_state);
     } else if (utils::Downcast<LockPathQuery>(parsed_query.query)) {
       prepared_query = PrepareLockPathQuery(std::move(parsed_query), in_explicit_transaction_, current_db_);
     } else if (utils::Downcast<FreeMemoryQuery>(parsed_query.query)) {
@@ -3836,8 +3809,8 @@ Interpreter::PrepareResult Interpreter::Prepare(const std::string &query_string,
         throw MultiDatabaseQueryInMulticommandTxException();
       }
       /// SYSTEM (Replication) + INTERPRETER
-      prepared_query =
-          PrepareMultiDatabaseQuery(std::move(parsed_query), current_db_, interpreter_context_, on_change_);
+      prepared_query = PrepareMultiDatabaseQuery(std::move(parsed_query), current_db_, interpreter_context_, on_change_,
+                                                 interpreter_context_->repl_state);
     } else if (utils::Downcast<ShowDatabasesQuery>(parsed_query.query)) {
       /// SYSTEM PURE ("SHOW DATABASES")
       /// INTERPRETER (TODO: "SHOW DATABASE")
@@ -3859,7 +3832,7 @@ Interpreter::PrepareResult Interpreter::Prepare(const std::string &query_string,
 
     UpdateTypeCount(rw_type);
 
-    if (IsWriteQueryOnMainMemoryReplica(current_db_.db_acc_->get()->storage(), rw_type)) {
+    if (interpreter_context_->repl_state->IsReplica() && IsQueryWrite(rw_type)) {
       query_execution = nullptr;
       throw QueryException("Write query forbidden on the replica!");
     }
@@ -4088,7 +4061,8 @@ void Interpreter::Commit() {
 
   auto commit_confirmed_by_all_sync_repplicas = true;
 
-  auto maybe_commit_error = current_db_.db_transactional_accessor_->Commit();
+  auto maybe_commit_error =
+      current_db_.db_transactional_accessor_->Commit(std::nullopt, interpreter_context_->repl_state->IsMain());
   if (maybe_commit_error.HasError()) {
     const auto &error = maybe_commit_error.GetError();
 

src/query/interpreter_context.cpp

@@ -13,6 +13,12 @@
 
 #include "query/interpreter.hpp"
 namespace memgraph::query {
+
+InterpreterContext::InterpreterContext(InterpreterConfig interpreter_config, dbms::DbmsHandler *dbms_handler,
+                                       replication::ReplicationState *rs, query::AuthQueryHandler *ah,
+                                       query::AuthChecker *ac)
+    : dbms_handler(dbms_handler), config(interpreter_config), repl_state(rs), auth(ah), auth_checker(ac) {}
+
 std::vector<std::vector<TypedValue>> InterpreterContext::TerminateTransactions(
     std::vector<std::string> maybe_kill_transaction_ids, const std::optional<std::string> &username,
     std::function<bool(std::string const &)> privilege_checker) {

src/query/interpreter_context.hpp

@@ -21,17 +21,14 @@
 #include "query/config.hpp"
 #include "query/cypher_query_interpreter.hpp"
 #include "query/typed_value.hpp"
+#include "replication/state.hpp"
 #include "utils/gatekeeper.hpp"
 #include "utils/skip_list.hpp"
 #include "utils/spin_lock.hpp"
 #include "utils/synchronized.hpp"
 
 namespace memgraph::dbms {
-#ifdef MG_ENTERPRISE
 class DbmsHandler;
-#else
-class Database;
-#endif
 }  // namespace memgraph::dbms
 
 namespace memgraph::query {
@@ -48,20 +45,10 @@ class Interpreter;
  *
  */
 struct InterpreterContext {
-#ifdef MG_ENTERPRISE
-  InterpreterContext(InterpreterConfig interpreter_config, memgraph::dbms::DbmsHandler *db_handler,
-                     AuthQueryHandler *ah = nullptr, AuthChecker *ac = nullptr);
-#else
-  InterpreterContext(InterpreterConfig interpreter_config,
-                     memgraph::utils::Gatekeeper<memgraph::dbms::Database> *db_gatekeeper,
-                     query::AuthQueryHandler *ah = nullptr, query::AuthChecker *ac = nullptr);
-#endif
+  InterpreterContext(InterpreterConfig interpreter_config, dbms::DbmsHandler *dbms_handler,
+                     replication::ReplicationState *rs, AuthQueryHandler *ah = nullptr, AuthChecker *ac = nullptr);
 
-#ifdef MG_ENTERPRISE
-  memgraph::dbms::DbmsHandler *db_handler;
-#else
-  memgraph::utils::Gatekeeper<memgraph::dbms::Database> *db_gatekeeper;
-#endif
+  memgraph::dbms::DbmsHandler *dbms_handler;
 
   // Internal
   const InterpreterConfig config;
@@ -69,6 +56,7 @@ struct InterpreterContext {
   memgraph::utils::SkipList<QueryCacheEntry> ast_cache;
 
   // GLOBAL
+  memgraph::replication::ReplicationState *repl_state;
   AuthQueryHandler *auth;
   AuthChecker *auth_checker;
 

src/query/plan/operator.cpp

@@ -27,6 +27,7 @@
 
 #include <cppitertools/chain.hpp>
 #include <cppitertools/imap.hpp>
+#include "memory/query_memory_control.hpp"
 #include "query/common.hpp"
 #include "spdlog/spdlog.h"
 
@@ -57,6 +58,7 @@
 #include "utils/logging.hpp"
 #include "utils/memory.hpp"
 #include "utils/message.hpp"
+#include "utils/on_scope_exit.hpp"
 #include "utils/pmr/deque.hpp"
 #include "utils/pmr/list.hpp"
 #include "utils/pmr/unordered_map.hpp"
@@ -2588,39 +2590,68 @@ void Delete::DeleteCursor::UpdateDeleteBuffer(Frame &frame, ExecutionContext &co
     expression_results.emplace_back(expression->Accept(evaluator));
   }
 
+  auto vertex_auth_checker = [&context](const VertexAccessor &va) -> bool {
+#ifdef MG_ENTERPRISE
+    return !(license::global_license_checker.IsEnterpriseValidFast() && context.auth_checker &&
+             !context.auth_checker->Has(va, storage::View::NEW, query::AuthQuery::FineGrainedPrivilege::CREATE_DELETE));
+#else
+    return true;
+#endif
+  };
+
+  auto edge_auth_checker = [&context](const EdgeAccessor &ea) -> bool {
+#ifdef MG_ENTERPRISE
+    return !(
+        license::global_license_checker.IsEnterpriseValidFast() && context.auth_checker &&
+        !(context.auth_checker->Has(ea, query::AuthQuery::FineGrainedPrivilege::CREATE_DELETE) &&
+          context.auth_checker->Has(ea.To(), storage::View::NEW, query::AuthQuery::FineGrainedPrivilege::UPDATE) &&
+          context.auth_checker->Has(ea.From(), storage::View::NEW, query::AuthQuery::FineGrainedPrivilege::UPDATE)));
+#else
+    return true;
+#endif
+  };
+
   for (TypedValue &expression_result : expression_results) {
     AbortCheck(context);
     switch (expression_result.type()) {
       case TypedValue::Type::Vertex: {
         auto va = expression_result.ValueVertex();
-#ifdef MG_ENTERPRISE
-        if (license::global_license_checker.IsEnterpriseValidFast() && context.auth_checker &&
-            !context.auth_checker->Has(va, storage::View::NEW, query::AuthQuery::FineGrainedPrivilege::CREATE_DELETE)) {
+        if (vertex_auth_checker(va)) {
+          buffer_.nodes.push_back(va);
+        } else {
           throw QueryRuntimeException("Vertex not deleted due to not having enough permission!");
         }
-#endif
-        buffer_.nodes.push_back(va);
         break;
       }
       case TypedValue::Type::Edge: {
         auto ea = expression_result.ValueEdge();
-#ifdef MG_ENTERPRISE
-        if (license::global_license_checker.IsEnterpriseValidFast() && context.auth_checker &&
-            !(context.auth_checker->Has(ea, query::AuthQuery::FineGrainedPrivilege::CREATE_DELETE) &&
-              context.auth_checker->Has(ea.To(), storage::View::NEW, query::AuthQuery::FineGrainedPrivilege::UPDATE) &&
-              context.auth_checker->Has(ea.From(), storage::View::NEW,
-                                        query::AuthQuery::FineGrainedPrivilege::UPDATE))) {
+        if (edge_auth_checker(ea)) {
+          buffer_.edges.push_back(ea);
+        } else {
           throw QueryRuntimeException("Edge not deleted due to not having enough permission!");
         }
-#endif
-        buffer_.edges.push_back(ea);
         break;
       }
+      case TypedValue::Type::Path: {
+        auto path = expression_result.ValuePath();
+#ifdef MG_ENTERPRISE
+        auto edges_res = std::any_of(path.edges().cbegin(), path.edges().cend(),
+                                     [&edge_auth_checker](const auto &ea) { return !edge_auth_checker(ea); });
+        auto vertices_res = std::any_of(path.vertices().cbegin(), path.vertices().cend(),
+                                        [&vertex_auth_checker](const auto &va) { return !vertex_auth_checker(va); });
+
+        if (edges_res || vertices_res) {
+          throw QueryRuntimeException(
+              "Path not deleted due to not having enough permission on all edges and vertices on the path!");
+        }
+#endif
+        buffer_.nodes.insert(buffer_.nodes.begin(), path.vertices().begin(), path.vertices().end());
+        buffer_.edges.insert(buffer_.edges.begin(), path.edges().begin(), path.edges().end());
+      }
       case TypedValue::Type::Null:
         break;
-      // check we're not trying to delete anything except vertices and edges
       default:
-        throw QueryRuntimeException("Only edges and vertices can be deleted.");
+        throw QueryRuntimeException("Edges, vertices and paths can be deleted.");
     }
   }
 }
@@ -2789,15 +2820,13 @@ SetProperties::SetPropertiesCursor::SetPropertiesCursor(const SetProperties &sel
 namespace {
 
 template <typename T>
-concept AccessorWithProperties =
-    requires(T value, storage::PropertyId property_id, storage::PropertyValue property_value,
-             std::map<storage::PropertyId, storage::PropertyValue> properties) {
-      {
-        value.ClearProperties()
-      } -> std::same_as<storage::Result<std::map<storage::PropertyId, storage::PropertyValue>>>;
-      { value.SetProperty(property_id, property_value) };
-      { value.UpdateProperties(properties) };
-    };
+concept AccessorWithProperties = requires(T value, storage::PropertyId property_id,
+                                          storage::PropertyValue property_value,
+                                          std::map<storage::PropertyId, storage::PropertyValue> properties) {
+  { value.ClearProperties() } -> std::same_as<storage::Result<std::map<storage::PropertyId, storage::PropertyValue>>>;
+  {value.SetProperty(property_id, property_value)};
+  {value.UpdateProperties(properties)};
+};
 
 /// Helper function that sets the given values on either a Vertex or an Edge.
 ///
@@ -2869,6 +2898,8 @@ void SetPropertiesOnRecord(TRecordAccessor *record, const TypedValue &rhs, SetPr
 
   auto update_props = [&, record](PropertiesMap &new_properties) {
     auto updated_properties = UpdatePropertiesChecked(record, new_properties);
+    // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+    context->execution_stats[ExecutionStats::Key::UPDATED_PROPERTIES] += new_properties.size();
 
     if (should_register_change) {
       for (const auto &[id, old_value, new_value] : updated_properties) {
@@ -3432,7 +3463,7 @@ class AggregateCursor : public Cursor {
     SCOPED_PROFILE_OP_BY_REF(self_);
 
     if (!pulled_all_input_) {
-      ProcessAll(&frame, &context);
+      if (!ProcessAll(&frame, &context) && self_.AreAllAggregationsForCollecting()) return false;
       pulled_all_input_ = true;
       aggregation_it_ = aggregation_.begin();
 
@@ -3456,7 +3487,6 @@ class AggregateCursor : public Cursor {
         return true;
       }
     }
-
     if (aggregation_it_ == aggregation_.end()) return false;
 
     // place aggregation values on the frame
@@ -3536,12 +3566,16 @@ class AggregateCursor : public Cursor {
    * cache cardinality depends on number of
    * aggregation results, and not on the number of inputs.
    */
-  void ProcessAll(Frame *frame, ExecutionContext *context) {
+  bool ProcessAll(Frame *frame, ExecutionContext *context) {
     ExpressionEvaluator evaluator(frame, context->symbol_table, context->evaluation_context, context->db_accessor,
                                   storage::View::NEW);
+
+    bool pulled = false;
     while (input_cursor_->Pull(*frame, *context)) {
       ProcessOne(*frame, &evaluator);
+      pulled = true;
     }
+    if (!pulled) return false;
 
     // post processing
     for (size_t pos = 0; pos < self_.aggregations_.size(); ++pos) {
@@ -3575,6 +3609,7 @@ class AggregateCursor : public Cursor {
           break;
       }
     }
+    return true;
   }
 
   /**
@@ -3788,6 +3823,12 @@ UniqueCursorPtr Aggregate::MakeCursor(utils::MemoryResource *mem) const {
   return MakeUniqueCursorPtr<AggregateCursor>(mem, *this, mem);
 }
 
+auto Aggregate::AreAllAggregationsForCollecting() const -> bool {
+  return std::all_of(aggregations_.begin(), aggregations_.end(), [](const auto &agg) {
+    return agg.op == Aggregation::Op::COLLECT_LIST || agg.op == Aggregation::Op::COLLECT_MAP;
+  });
+}
+
 Skip::Skip(const std::shared_ptr<LogicalOperator> &input, Expression *expression)
     : input_(input), expression_(expression) {}
 
@@ -4588,7 +4629,7 @@ UniqueCursorPtr OutputTableStream::MakeCursor(utils::MemoryResource *mem) const
 
 CallProcedure::CallProcedure(std::shared_ptr<LogicalOperator> input, std::string name, std::vector<Expression *> args,
                              std::vector<std::string> fields, std::vector<Symbol> symbols, Expression *memory_limit,
-                             size_t memory_scale, bool is_write, bool void_procedure)
+                             size_t memory_scale, bool is_write, int64_t procedure_id, bool void_procedure)
     : input_(input ? input : std::make_shared<Once>()),
       procedure_name_(name),
       arguments_(args),
@@ -4597,6 +4638,7 @@ CallProcedure::CallProcedure(std::shared_ptr<LogicalOperator> input, std::string
       memory_limit_(memory_limit),
       memory_scale_(memory_scale),
       is_write_(is_write),
+      procedure_id_(procedure_id),
       void_procedure_(void_procedure) {}
 
 ACCEPT_WITH_INPUT(CallProcedure);
@@ -4625,7 +4667,7 @@ namespace {
 void CallCustomProcedure(const std::string_view fully_qualified_procedure_name, const mgp_proc &proc,
                          const std::vector<Expression *> &args, mgp_graph &graph, ExpressionEvaluator *evaluator,
                          utils::MemoryResource *memory, std::optional<size_t> memory_limit, mgp_result *result,
-                         const bool call_initializer = false) {
+                         int64_t procedure_id, uint64_t transaction_id, const bool call_initializer = false) {
   static_assert(std::uses_allocator_v<mgp_value, utils::Allocator<mgp_value>>,
                 "Expected mgp_value to use custom allocator and makes STL "
                 "containers aware of that");
@@ -4657,12 +4699,46 @@ void CallCustomProcedure(const std::string_view fully_qualified_procedure_name,
   if (memory_limit) {
     SPDLOG_INFO("Running '{}' with memory limit of {}", fully_qualified_procedure_name,
                 utils::GetReadableSize(*memory_limit));
-    utils::LimitedMemoryResource limited_mem(memory, *memory_limit);
-    mgp_memory proc_memory{&limited_mem};
+    // Only allocations which can leak memory are
+    // our own mgp object allocations. Jemalloc can track
+    // memory correctly, but some memory may not be released
+    // immediately, so we want to give user info on leak still
+    // considering our allocations
+    utils::MemoryTrackingResource memory_tracking_resource{memory, *memory_limit};
+    // if we are already tracking, no harm no faul
+    // if we are not tracking, we need to start now, with unlimited memory
+    // for query, but limited for procedure
+
+    // check if transaction is tracked currently, so we
+    // can disable tracking on that arena if it is not
+    // once we are done with procedure tracking
+
+    bool is_transaction_tracked = memgraph::memory::IsTransactionTracked(transaction_id);
+
+    if (!is_transaction_tracked) {
+      // start tracking with unlimited limit on query
+      // which is same as not being tracked at all
+      memgraph::memory::TryStartTrackingOnTransaction(transaction_id, memgraph::memory::UNLIMITED_MEMORY);
+    }
+    memgraph::memory::StartTrackingCurrentThreadTransaction(transaction_id);
+
+    // due to mgp_batch_read_proc and mgp_batch_write_proc
+    // we can return to execution without exhausting whole
+    // memory. Here we need to update tracking
+    memgraph::memory::CreateOrContinueProcedureTracking(transaction_id, procedure_id, *memory_limit);
+
+    mgp_memory proc_memory{&memory_tracking_resource};
     MG_ASSERT(result->signature == &proc.results);
+
+    utils::OnScopeExit on_scope_exit{[transaction_id = transaction_id]() {
+      memgraph::memory::StopTrackingCurrentThreadTransaction(transaction_id);
+      memgraph::memory::PauseProcedureTracking(transaction_id);
+    }};
+
     // TODO: What about cross library boundary exceptions? OMG C++?!
     proc.cb(&proc_args, &graph, result, &proc_memory);
-    size_t leaked_bytes = limited_mem.GetAllocatedBytes();
+
+    auto leaked_bytes = memory_tracking_resource.GetAllocatedBytes();
     if (leaked_bytes > 0U) {
       spdlog::warn("Query procedure '{}' leaked {} *tracked* bytes", fully_qualified_procedure_name, leaked_bytes);
     }
@@ -4770,8 +4846,10 @@ class CallProcedureCursor : public Cursor {
       auto *memory = self_->memory_resource;
       auto memory_limit = EvaluateMemoryLimit(evaluator, self_->memory_limit_, self_->memory_scale_);
       auto graph = mgp_graph::WritableGraph(*context.db_accessor, graph_view, context);
+      const auto transaction_id = context.db_accessor->GetTransactionId();
+      MG_ASSERT(transaction_id.has_value());
       CallCustomProcedure(self_->procedure_name_, *proc, self_->arguments_, graph, &evaluator, memory, memory_limit,
-                          result_, call_initializer);
+                          result_, self_->procedure_id_, transaction_id.value(), call_initializer);
 
       if (call_initializer) call_initializer = false;
 
@@ -5336,7 +5414,7 @@ class HashJoinCursor : public Cursor {
         // Check if the join value from the pulled frame is shared with any left frames
         ExpressionEvaluator evaluator(&frame, context.symbol_table, context.evaluation_context, context.db_accessor,
                                       storage::View::OLD);
-        auto right_value = self_.hash_join_condition_->expression1_->Accept(evaluator);
+        auto right_value = self_.hash_join_condition_->expression2_->Accept(evaluator);
         if (hashtable_.contains(right_value)) {
           // If so, finish pulling for now and proceed to joining the pulled frame
           right_op_frame_.assign(frame.elems().begin(), frame.elems().end());
@@ -5354,7 +5432,8 @@ class HashJoinCursor : public Cursor {
     restore_frame(self_.left_symbols_, *left_op_frame_it_);
 
     left_op_frame_it_++;
-    // When all left frames with the common value have been joined, move on to pulling and joining the next right frame
+    // When all left frames with the common value have been joined, move on to pulling and joining the next right
+    // frame
     if (common_value_found_ && left_op_frame_it_ == hashtable_[common_value].end()) {
       common_value_found_ = false;
     }
@@ -5383,7 +5462,7 @@ class HashJoinCursor : public Cursor {
     while (left_op_cursor_->Pull(frame, context)) {
       ExpressionEvaluator evaluator(&frame, context.symbol_table, context.evaluation_context, context.db_accessor,
                                     storage::View::OLD);
-      auto left_value = self_.hash_join_condition_->expression2_->Accept(evaluator);
+      auto left_value = self_.hash_join_condition_->expression1_->Accept(evaluator);
       if (left_value.type() != TypedValue::Type::Null) {
         hashtable_[left_value].emplace_back(frame.elems().begin(), frame.elems().end());
       }

src/query/plan/operator.hpp

@@ -1758,6 +1758,9 @@ class Aggregate : public memgraph::query::plan::LogicalOperator {
   Aggregate() = default;
   Aggregate(const std::shared_ptr<LogicalOperator> &input, const std::vector<Element> &aggregations,
             const std::vector<Expression *> &group_by, const std::vector<Symbol> &remember);
+
+  auto AreAllAggregationsForCollecting() const -> bool;
+
   bool Accept(HierarchicalLogicalOperatorVisitor &visitor) override;
   UniqueCursorPtr MakeCursor(utils::MemoryResource *) const override;
   std::vector<Symbol> ModifiedSymbols(const SymbolTable &) const override;
@@ -2361,7 +2364,7 @@ class CallProcedure : public memgraph::query::plan::LogicalOperator {
   CallProcedure() = default;
   CallProcedure(std::shared_ptr<LogicalOperator> input, std::string name, std::vector<Expression *> arguments,
                 std::vector<std::string> fields, std::vector<Symbol> symbols, Expression *memory_limit,
-                size_t memory_scale, bool is_write, bool void_procedure = false);
+                size_t memory_scale, bool is_write, int64_t procedure_id, bool void_procedure = false);
 
   bool Accept(HierarchicalLogicalOperatorVisitor &visitor) override;
   UniqueCursorPtr MakeCursor(utils::MemoryResource *) const override;
@@ -2383,6 +2386,7 @@ class CallProcedure : public memgraph::query::plan::LogicalOperator {
   Expression *memory_limit_{nullptr};
   size_t memory_scale_{1024U};
   bool is_write_;
+  int64_t procedure_id_;
   bool void_procedure_;
   mutable utils::MonotonicBufferResource monotonic_memory{1024UL * 1024UL};
   utils::MemoryResource *memory_resource = &monotonic_memory;
@@ -2405,6 +2409,7 @@ class CallProcedure : public memgraph::query::plan::LogicalOperator {
     object->memory_limit_ = memory_limit_ ? memory_limit_->Clone(storage) : nullptr;
     object->memory_scale_ = memory_scale_;
     object->is_write_ = is_write_;
+    object->procedure_id_ = procedure_id_;
     object->void_procedure_ = void_procedure_;
     return object;
   }

src/query/plan/rewrite/join.hpp

@@ -27,6 +27,7 @@
 
 #include "query/plan/operator.hpp"
 #include "query/plan/preprocess.hpp"
+#include "utils/algorithm.hpp"
 
 namespace memgraph::query::plan {
 
@@ -145,8 +146,11 @@ class JoinRewriter final : public HierarchicalLogicalOperatorVisitor {
 
   bool PreVisit(IndexedJoin &op) override {
     prev_ops_.push_back(&op);
-    return true;
+    RewriteBranch(&op.main_branch_);
+    RewriteBranch(&op.sub_branch_);
+    return false;
   }
+
   bool PostVisit(IndexedJoin &) override {
     prev_ops_.pop_back();
     return true;
@@ -523,6 +527,13 @@ class JoinRewriter final : public HierarchicalLogicalOperatorVisitor {
       auto rhs_property = rhs_lookup->property_;
       filter_exprs_for_removal_.insert(filter.expression);
       filters_.EraseFilter(filter);
+
+      if (utils::Contains(right_symbols, lhs_symbol) && utils::Contains(left_symbols, rhs_symbol)) {
+        // We need to duplicate this because expressions are shared between plans
+        join_condition = join_condition->Clone(ast_storage_);
+        std::swap(join_condition->expression1_, join_condition->expression2_);
+      }
+
       return std::make_unique<HashJoin>(left_op, left_symbols, right_op, right_symbols, join_condition);
     }
 

src/query/plan/rule_based_planner.hpp

@@ -12,6 +12,7 @@
 /// @file
 #pragma once
 
+#include <cstdint>
 #include <optional>
 #include <variant>
 
@@ -185,6 +186,10 @@ class RuleBasedPlanner {
 
         uint64_t merge_id = 0;
         uint64_t subquery_id = 0;
+        // procedures need to start from 1
+        // due to swapping mechanism of procedure
+        // tracking
+        uint64_t procedure_id = 1;
 
         for (const auto &clause : single_query_part.remaining_clauses) {
           MG_ASSERT(!utils::IsSubtype(*clause, Match::kType), "Unexpected Match in remaining clauses");
@@ -224,7 +229,7 @@ class RuleBasedPlanner {
             input_op = std::make_unique<plan::CallProcedure>(
                 std::move(input_op), call_proc->procedure_name_, call_proc->arguments_, call_proc->result_fields_,
                 result_symbols, call_proc->memory_limit_, call_proc->memory_scale_, call_proc->is_write_,
-                call_proc->void_procedure_);
+                procedure_id++, call_proc->void_procedure_);
           } else if (auto *load_csv = utils::Downcast<query::LoadCsv>(clause)) {
             const auto &row_sym = context.symbol_table->at(*load_csv->row_var_);
             context.bound_symbols.insert(row_sym);

src/query/trigger.cpp

@@ -169,7 +169,7 @@ Trigger::TriggerPlan::TriggerPlan(std::unique_ptr<LogicalPlan> logical_plan, std
 std::shared_ptr<Trigger::TriggerPlan> Trigger::GetPlan(DbAccessor *db_accessor,
                                                        const query::AuthChecker *auth_checker) const {
   std::lock_guard plan_guard{plan_lock_};
-  if (!parsed_statements_.is_cacheable || !trigger_plan_ || trigger_plan_->cached_plan.IsExpired()) {
+  if (!parsed_statements_.is_cacheable || !trigger_plan_) {
     auto identifiers = GetPredefinedIdentifiers(event_type_);
 
     AstStorage ast_storage;

src/query/trigger.hpp

@@ -62,7 +62,7 @@ struct Trigger {
 
     explicit TriggerPlan(std::unique_ptr<LogicalPlan> logical_plan, std::vector<IdentifierInfo> identifiers);
 
-    CachedPlan cached_plan;
+    PlanWrapper cached_plan;
     std::vector<IdentifierInfo> identifiers;
   };
   std::shared_ptr<TriggerPlan> GetPlan(DbAccessor *db_accessor, const query::AuthChecker *auth_checker) const;

src/replication/CMakeLists.txt

@@ -8,17 +8,19 @@ target_sources(mg-replication
         include/replication/mode.hpp
         include/replication/role.hpp
         include/replication/status.hpp
+        include/replication/replication_server.hpp
 
         PRIVATE
         state.cpp
         epoch.cpp
         config.cpp
         status.cpp
+        replication_server.cpp
 )
 target_include_directories(mg-replication PUBLIC include)
 
 find_package(fmt REQUIRED)
 target_link_libraries(mg-replication
-    PUBLIC mg::utils mg::kvstore lib::json
+    PUBLIC mg::utils mg::kvstore lib::json mg::rpc mg::slk
     PRIVATE fmt::fmt
 )

src/replication/include/replication/config.hpp

@@ -21,13 +21,12 @@ namespace memgraph::replication {
 
 inline constexpr uint16_t kDefaultReplicationPort = 10000;
 inline constexpr auto *kDefaultReplicationServerIp = "0.0.0.0";
-inline constexpr auto *kReservedReplicationRoleName{"__replication_role"};
 
 struct ReplicationClientConfig {
   std::string name;
-  ReplicationMode mode;
+  ReplicationMode mode{};
   std::string ip_address;
-  uint16_t port;
+  uint16_t port{};
 
   // The default delay between main checking/pinging replicas is 1s because
   // that seems like a reasonable timeframe in which main should notice a
@@ -42,18 +41,23 @@ struct ReplicationClientConfig {
   };
 
   std::optional<SSL> ssl;
+
+  friend bool operator==(ReplicationClientConfig const &, ReplicationClientConfig const &) = default;
 };
 
 struct ReplicationServerConfig {
   std::string ip_address;
-  uint16_t port;
+  uint16_t port{};
   struct SSL {
     std::string key_file;
     std::string cert_file;
     std::string ca_file;
-    bool verify_peer;
+    bool verify_peer{};
+    friend bool operator==(SSL const &, SSL const &) = default;
   };
 
   std::optional<SSL> ssl;
+
+  friend bool operator==(ReplicationServerConfig const &, ReplicationServerConfig const &) = default;
 };
 }  // namespace memgraph::replication

src/replication/include/replication/epoch.hpp

@@ -19,16 +19,21 @@ namespace memgraph::replication {
 
 struct ReplicationEpoch {
   ReplicationEpoch() : id_(memgraph::utils::GenerateUUID()) {}
-  ReplicationEpoch(ReplicationEpoch const &) = delete;
-  ReplicationEpoch(ReplicationEpoch &&) = delete;
-  ReplicationEpoch &operator=(ReplicationEpoch const &) = delete;
-  ReplicationEpoch &operator=(ReplicationEpoch &&) = delete;
+  explicit ReplicationEpoch(std::string explicit_id) : id_(std::move(explicit_id)) {}
+  ReplicationEpoch(ReplicationEpoch const &) = default;  // TODO: passkey idiom
+  ReplicationEpoch(ReplicationEpoch &&) = default;
+  ReplicationEpoch &operator=(ReplicationEpoch const &) = default;  // TODO: passkey idiom
+  ReplicationEpoch &operator=(ReplicationEpoch &&) = default;
 
   auto id() const -> std::string_view { return id_; }
 
-  auto NewEpoch() -> std::string { return std::exchange(id_, memgraph::utils::GenerateUUID()); }
+  // TODO: passkey idiom
+  friend struct ReplicationState;
+
   auto SetEpoch(std::string new_epoch) -> std::string { return std::exchange(id_, std::move(new_epoch)); }
 
+  friend bool operator==(ReplicationEpoch const &, ReplicationEpoch const &) = default;
+
  private:
   // UUID to distinguish different main instance runs for replication process
   // on SAME storage.

src/replication/include/replication/replication_server.hpp

@@ -14,9 +14,32 @@
 #include "replication/config.hpp"
 #include "rpc/server.hpp"
 #include "slk/streams.hpp"
-#include "storage/v2/replication/global.hpp"
 
-namespace memgraph::storage {
+namespace memgraph::replication {
+
+struct FrequentHeartbeatReq {
+  static const utils::TypeInfo kType;                            // TODO: make constexpr?
+  static const utils::TypeInfo &GetTypeInfo() { return kType; }  // WHAT?
+
+  static void Load(FrequentHeartbeatReq *self, memgraph::slk::Reader *reader);
+  static void Save(const FrequentHeartbeatReq &self, memgraph::slk::Builder *builder);
+  FrequentHeartbeatReq() {}
+};
+
+struct FrequentHeartbeatRes {
+  static const utils::TypeInfo kType;
+  static const utils::TypeInfo &GetTypeInfo() { return kType; }
+
+  static void Load(FrequentHeartbeatRes *self, memgraph::slk::Reader *reader);
+  static void Save(const FrequentHeartbeatRes &self, memgraph::slk::Builder *builder);
+  FrequentHeartbeatRes() {}
+  explicit FrequentHeartbeatRes(bool success) : success(success) {}
+
+  bool success;
+};
+
+// TODO: move to own header
+using FrequentHeartbeatRpc = rpc::RequestResponse<FrequentHeartbeatReq, FrequentHeartbeatRes>;
 
 class ReplicationServer {
  public:
@@ -31,10 +54,10 @@ class ReplicationServer {
   bool Start();
 
  protected:
-  static void FrequentHeartbeatHandler(slk::Reader *req_reader, slk::Builder *res_builder);
-
   communication::ServerContext rpc_server_context_;
-  rpc::Server rpc_server_;
+
+ public:
+  rpc::Server rpc_server_;  // TODO: Interface or something
 };
 
-}  // namespace memgraph::storage
+}  // namespace memgraph::replication

src/replication/include/replication/state.hpp

@@ -21,48 +21,70 @@
 #include "replication/epoch.hpp"
 #include "replication/mode.hpp"
 #include "replication/role.hpp"
+#include "replication_server.hpp"
+#include "status.hpp"
 #include "utils/result.hpp"
 
 namespace memgraph::replication {
 
 enum class RolePersisted : uint8_t { UNKNOWN_OR_NO, YES };
 
+enum class RegisterReplicaError : uint8_t { NAME_EXISTS, END_POINT_EXISTS, COULD_NOT_BE_PERSISTED, NOT_MAIN, SUCCESS };
+
+struct RoleMainData {
+  ReplicationEpoch epoch_;
+  std::vector<ReplicationClientConfig> registered_replicas_;
+};
+
+struct RoleReplicaData {
+  ReplicationServerConfig config;
+  std::unique_ptr<ReplicationServer> server;
+};
+
 struct ReplicationState {
-  ReplicationState(std::optional<std::filesystem::path> durability_dir);
+  explicit ReplicationState(std::optional<std::filesystem::path> durability_dir);
 
   ReplicationState(ReplicationState const &) = delete;
   ReplicationState(ReplicationState &&) = delete;
   ReplicationState &operator=(ReplicationState const &) = delete;
   ReplicationState &operator=(ReplicationState &&) = delete;
 
-  void SetRole(ReplicationRole role) { return replication_role_.store(role); }
-  auto GetRole() const -> ReplicationRole { return replication_role_.load(); }
-  bool IsMain() const { return replication_role_ == ReplicationRole::MAIN; }
-  bool IsReplica() const { return replication_role_ == ReplicationRole::REPLICA; }
-
-  auto GetEpoch() const -> const ReplicationEpoch & { return epoch_; }
-  auto GetEpoch() -> ReplicationEpoch & { return epoch_; }
-
   enum class FetchReplicationError : uint8_t {
     NOTHING_FETCHED,
     PARSE_ERROR,
   };
-  using ReplicationDataReplica = ReplicationServerConfig;
-  using ReplicationDataMain = std::vector<ReplicationClientConfig>;
-  using ReplicationData = std::variant<ReplicationDataMain, ReplicationDataReplica>;
-  using FetchReplicationResult = utils::BasicResult<FetchReplicationError, ReplicationData>;
-  auto FetchReplicationData() -> FetchReplicationResult;
+
+  using ReplicationData_t = std::variant<RoleMainData, RoleReplicaData>;
+  using FetchReplicationResult_t = utils::BasicResult<FetchReplicationError, ReplicationData_t>;
+  auto FetchReplicationData() -> FetchReplicationResult_t;
+
+  auto GetRole() const -> ReplicationRole {
+    return std::holds_alternative<RoleReplicaData>(replication_data_) ? ReplicationRole::REPLICA
+                                                                      : ReplicationRole::MAIN;
+  }
+  bool IsMain() const { return GetRole() == ReplicationRole::MAIN; }
+  bool IsReplica() const { return GetRole() == ReplicationRole::REPLICA; }
 
   bool ShouldPersist() const { return nullptr != durability_; }
-  bool TryPersistRoleMain();
+  bool TryPersistRoleMain(std::string new_epoch);
   bool TryPersistRoleReplica(const ReplicationServerConfig &config);
-  bool TryPersistUnregisterReplica(std::string_view &name);
+  bool TryPersistUnregisterReplica(std::string_view name);
   bool TryPersistRegisteredReplica(const ReplicationClientConfig &config);
 
+  // TODO: locked access
+  auto ReplicationData() -> ReplicationData_t & { return replication_data_; }
+  auto ReplicationData() const -> ReplicationData_t const & { return replication_data_; }
+  auto RegisterReplica(const ReplicationClientConfig &config) -> RegisterReplicaError;
+
+  bool SetReplicationRoleMain();
+
+  bool SetReplicationRoleReplica(const ReplicationServerConfig &config);
+
  private:
-  ReplicationEpoch epoch_;
-  std::atomic<ReplicationRole> replication_role_{ReplicationRole::MAIN};
+  bool HandleVersionMigration(durability::ReplicationRoleEntry &data) const;
+
   std::unique_ptr<kvstore::KVStore> durability_;
+  ReplicationData_t replication_data_;
   std::atomic<RolePersisted> role_persisted = RolePersisted::UNKNOWN_OR_NO;
 };
 

src/replication/include/replication/status.hpp

@@ -15,25 +15,56 @@
 #include <cstdint>
 #include <optional>
 #include <string>
+#include <variant>
 
 #include "json/json.hpp"
 
 #include "replication/config.hpp"
+#include "replication/epoch.hpp"
 #include "replication/role.hpp"
 
-namespace memgraph::replication {
-struct ReplicationStatus {
-  std::string name;
-  std::string ip_address;
-  uint16_t port;
-  ReplicationMode sync_mode;
-  std::chrono::seconds replica_check_frequency;
-  std::optional<ReplicationClientConfig::SSL> ssl;
-  std::optional<ReplicationRole> role;
+namespace memgraph::replication::durability {
 
-  friend bool operator==(const ReplicationStatus &, const ReplicationStatus &) = default;
+// Keys
+constexpr auto *kReplicationRoleName{"__replication_role"};
+constexpr auto *kReplicationReplicaPrefix{"__replication_replica:"};  // introduced in V2
+
+enum class DurabilityVersion : uint8_t {
+  V1,  // no distinct key for replicas
+  V2,  // this version, epoch, replica prefix introduced
 };
 
-nlohmann::json ReplicationStatusToJSON(ReplicationStatus &&status);
-std::optional<ReplicationStatus> JSONToReplicationStatus(nlohmann::json &&data);
-}  // namespace memgraph::replication
+// fragment of key: "__replication_role"
+struct MainRole {
+  ReplicationEpoch epoch{};
+  friend bool operator==(MainRole const &, MainRole const &) = default;
+};
+
+// fragment of key: "__replication_role"
+struct ReplicaRole {
+  ReplicationServerConfig config;
+  friend bool operator==(ReplicaRole const &, ReplicaRole const &) = default;
+};
+
+// from key: "__replication_role"
+struct ReplicationRoleEntry {
+  DurabilityVersion version =
+      DurabilityVersion::V2;  // if not latest then migration required for kReplicationReplicaPrefix
+  std::variant<MainRole, ReplicaRole> role;
+
+  friend bool operator==(ReplicationRoleEntry const &, ReplicationRoleEntry const &) = default;
+};
+
+// from key: "__replication_replica:"
+struct ReplicationReplicaEntry {
+  ReplicationClientConfig config;
+  friend bool operator==(ReplicationReplicaEntry const &, ReplicationReplicaEntry const &) = default;
+};
+
+void to_json(nlohmann::json &j, const ReplicationRoleEntry &p);
+void from_json(const nlohmann::json &j, ReplicationRoleEntry &p);
+
+void to_json(nlohmann::json &j, const ReplicationReplicaEntry &p);
+void from_json(const nlohmann::json &j, ReplicationReplicaEntry &p);
+
+}  // namespace memgraph::replication::durability

src/replication/replication_server.cpp

@@ -0,0 +1,96 @@
+// Copyright 2023 Memgraph Ltd.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+// License, and you may not use this file except in compliance with the Business Source License.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+#include "replication/replication_server.hpp"
+#include "rpc/messages.hpp"
+#include "slk/serialization.hpp"
+#include "slk/streams.hpp"
+
+namespace memgraph::slk {
+
+// Serialize code for FrequentHeartbeatRes
+void Save(const memgraph::replication::FrequentHeartbeatRes &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.success, builder);
+}
+void Load(memgraph::replication::FrequentHeartbeatRes *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->success, reader);
+}
+
+// Serialize code for FrequentHeartbeatReq
+void Save(const memgraph::replication::FrequentHeartbeatReq &self, memgraph::slk::Builder *builder) {}
+void Load(memgraph::replication::FrequentHeartbeatReq *self, memgraph::slk::Reader *reader) {}
+
+}  // namespace memgraph::slk
+
+namespace memgraph::replication {
+namespace {
+
+auto CreateServerContext(const memgraph::replication::ReplicationServerConfig &config) -> communication::ServerContext {
+  return (config.ssl) ? communication::ServerContext{config.ssl->key_file, config.ssl->cert_file, config.ssl->ca_file,
+                                                     config.ssl->verify_peer}
+                      : communication::ServerContext{};
+}
+
+void FrequentHeartbeatHandler(slk::Reader *req_reader, slk::Builder *res_builder) {
+  FrequentHeartbeatReq req;
+  memgraph::slk::Load(&req, req_reader);
+  FrequentHeartbeatRes res{true};
+  memgraph::slk::Save(res, res_builder);
+}
+
+// NOTE: The replication server must have a single thread for processing
+// because there is no need for more processing threads - each replica can
+// have only a single main server. Also, the single-threaded guarantee
+// simplifies the rest of the implementation.
+constexpr auto kReplicationServerThreads = 1;
+}  // namespace
+
+constexpr utils::TypeInfo FrequentHeartbeatReq::kType{utils::TypeId::REP_FREQUENT_HEARTBEAT_REQ, "FrequentHeartbeatReq",
+                                                      nullptr};
+
+constexpr utils::TypeInfo FrequentHeartbeatRes::kType{utils::TypeId::REP_FREQUENT_HEARTBEAT_RES, "FrequentHeartbeatRes",
+                                                      nullptr};
+
+void FrequentHeartbeatReq::Save(const FrequentHeartbeatReq &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self, builder);
+}
+void FrequentHeartbeatReq::Load(FrequentHeartbeatReq *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(self, reader);
+}
+void FrequentHeartbeatRes::Save(const FrequentHeartbeatRes &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self, builder);
+}
+void FrequentHeartbeatRes::Load(FrequentHeartbeatRes *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(self, reader);
+}
+
+ReplicationServer::ReplicationServer(const memgraph::replication::ReplicationServerConfig &config)
+    : rpc_server_context_{CreateServerContext(config)},
+      rpc_server_{io::network::Endpoint{config.ip_address, config.port}, &rpc_server_context_,
+                  kReplicationServerThreads} {
+  rpc_server_.Register<FrequentHeartbeatRpc>([](auto *req_reader, auto *res_builder) {
+    spdlog::debug("Received FrequentHeartbeatRpc");
+    FrequentHeartbeatHandler(req_reader, res_builder);
+  });
+}
+
+ReplicationServer::~ReplicationServer() {
+  if (rpc_server_.IsRunning()) {
+    auto const &endpoint = rpc_server_.endpoint();
+    spdlog::trace("Closing replication server on {}:{}", endpoint.address, endpoint.port);
+    rpc_server_.Shutdown();
+  }
+  rpc_server_.AwaitShutdown();
+}
+
+bool ReplicationServer::Start() { return rpc_server_.Start(); }
+
+}  // namespace memgraph::replication

src/replication/state.cpp

@@ -11,134 +11,265 @@
 
 #include "replication/state.hpp"
 
-#include "replication/status.hpp"  //TODO: don't use status for durability
+#include "replication/replication_server.hpp"
+#include "replication/status.hpp"
 #include "utils/file.hpp"
+#include "utils/variant_helpers.hpp"
 
 constexpr auto kReplicationDirectory = std::string_view{"replication"};
 
 namespace memgraph::replication {
 
+auto BuildReplicaKey(std::string_view name) -> std::string {
+  auto key = std::string{durability::kReplicationReplicaPrefix};
+  key.append(name);
+  return key;
+}
+
 ReplicationState::ReplicationState(std::optional<std::filesystem::path> durability_dir) {
   if (!durability_dir) return;
   auto repl_dir = *std::move(durability_dir);
   repl_dir /= kReplicationDirectory;
   utils::EnsureDirOrDie(repl_dir);
   durability_ = std::make_unique<kvstore::KVStore>(std::move(repl_dir));
+
+  auto replicationData = FetchReplicationData();
+  if (replicationData.HasError()) {
+    switch (replicationData.GetError()) {
+      using enum ReplicationState::FetchReplicationError;
+      case NOTHING_FETCHED: {
+        spdlog::debug("Cannot find data needed for restore replication role in persisted metadata.");
+        replication_data_ = RoleMainData{};
+        return;
+      }
+      case PARSE_ERROR: {
+        LOG_FATAL("Cannot parse previously saved configuration of replication role.");
+        return;
+      }
+    }
+  }
+  replication_data_ = std::move(replicationData).GetValue();
 }
+
 bool ReplicationState::TryPersistRoleReplica(const ReplicationServerConfig &config) {
   if (!ShouldPersist()) return true;
-  // Only thing that matters here is the role saved as REPLICA and the listening port
-  auto data = ReplicationStatusToJSON(ReplicationStatus{.name = kReservedReplicationRoleName,
-                                                        .ip_address = config.ip_address,
-                                                        .port = config.port,
-                                                        .sync_mode = ReplicationMode::SYNC,
-                                                        .replica_check_frequency = std::chrono::seconds(0),
-                                                        .ssl = std::nullopt,
-                                                        .role = ReplicationRole::REPLICA});
 
-  if (durability_->Put(kReservedReplicationRoleName, data.dump())) {
-    role_persisted = RolePersisted::YES;
-    return true;
+  auto data = durability::ReplicationRoleEntry{.role = durability::ReplicaRole{
+                                                   .config = config,
+                                               }};
+
+  if (!durability_->Put(durability::kReplicationRoleName, nlohmann::json(data).dump())) {
+    spdlog::error("Error when saving REPLICA replication role in settings.");
+    return false;
   }
-  spdlog::error("Error when saving REPLICA replication role in settings.");
-  return false;
-}
-bool ReplicationState::TryPersistRoleMain() {
-  if (!ShouldPersist()) return true;
-  // Only thing that matters here is the role saved as MAIN
-  auto data = ReplicationStatusToJSON(ReplicationStatus{.name = kReservedReplicationRoleName,
-                                                        .ip_address = "",
-                                                        .port = 0,
-                                                        .sync_mode = ReplicationMode::SYNC,
-                                                        .replica_check_frequency = std::chrono::seconds(0),
-                                                        .ssl = std::nullopt,
-                                                        .role = ReplicationRole::MAIN});
+  role_persisted = RolePersisted::YES;
 
-  if (durability_->Put(kReservedReplicationRoleName, data.dump())) {
+  // Cleanup remove registered replicas (assume successful delete)
+  // NOTE: we could do the alternative which would be on REPLICA -> MAIN we recover these registered replicas
+  auto b = durability_->begin(durability::kReplicationReplicaPrefix);
+  auto e = durability_->end(durability::kReplicationReplicaPrefix);
+  for (; b != e; ++b) {
+    durability_->Delete(b->first);
+  }
+
+  return true;
+}
+
+bool ReplicationState::TryPersistRoleMain(std::string new_epoch) {
+  if (!ShouldPersist()) return true;
+
+  auto data =
+      durability::ReplicationRoleEntry{.role = durability::MainRole{.epoch = ReplicationEpoch{std::move(new_epoch)}}};
+
+  if (durability_->Put(durability::kReplicationRoleName, nlohmann::json(data).dump())) {
     role_persisted = RolePersisted::YES;
     return true;
   }
   spdlog::error("Error when saving MAIN replication role in settings.");
   return false;
 }
-bool ReplicationState::TryPersistUnregisterReplica(std::string_view &name) {
+
+bool ReplicationState::TryPersistUnregisterReplica(std::string_view name) {
   if (!ShouldPersist()) return true;
-  if (durability_->Delete(name)) return true;
+
+  auto key = BuildReplicaKey(name);
+
+  if (durability_->Delete(key)) return true;
   spdlog::error("Error when removing replica {} from settings.", name);
   return false;
 }
-auto ReplicationState::FetchReplicationData() -> FetchReplicationResult {
+
+// TODO: FetchEpochData (agnostic of FetchReplicationData, but should be done before)
+
+auto ReplicationState::FetchReplicationData() -> FetchReplicationResult_t {
   if (!ShouldPersist()) return FetchReplicationError::NOTHING_FETCHED;
-  const auto replication_data = durability_->Get(kReservedReplicationRoleName);
+  const auto replication_data = durability_->Get(durability::kReplicationRoleName);
   if (!replication_data.has_value()) {
     return FetchReplicationError::NOTHING_FETCHED;
   }
 
-  const auto maybe_replication_status = JSONToReplicationStatus(nlohmann::json::parse(*replication_data));
-  if (!maybe_replication_status.has_value()) {
+  auto json = nlohmann::json::parse(*replication_data, nullptr, false);
+  if (json.is_discarded()) {
     return FetchReplicationError::PARSE_ERROR;
   }
+  try {
+    durability::ReplicationRoleEntry data = json.get<durability::ReplicationRoleEntry>();
 
-  // To get here this must be the case
-  role_persisted = memgraph::replication::RolePersisted::YES;
-
-  const auto replication_status = *maybe_replication_status;
-  auto role = replication_status.role.value_or(ReplicationRole::MAIN);
-  switch (role) {
-    case ReplicationRole::REPLICA: {
-      return {ReplicationServerConfig{
-          .ip_address = kDefaultReplicationServerIp,
-          .port = replication_status.port,
-      }};
+    if (!HandleVersionMigration(data)) {
+      return FetchReplicationError::PARSE_ERROR;
     }
-    case ReplicationRole::MAIN: {
-      auto res = ReplicationState::ReplicationDataMain{};
-      res.reserve(durability_->Size() - 1);
-      for (const auto &[replica_name, replica_data] : *durability_) {
-        if (replica_name == kReservedReplicationRoleName) {
+
+    // To get here this must be the case
+    role_persisted = memgraph::replication::RolePersisted::YES;
+
+    return std::visit(
+        utils::Overloaded{
+            [&](durability::MainRole &&r) -> FetchReplicationResult_t {
+              auto res = RoleMainData{
+                  .epoch_ = std::move(r.epoch),
+              };
+              auto b = durability_->begin(durability::kReplicationReplicaPrefix);
+              auto e = durability_->end(durability::kReplicationReplicaPrefix);
+              res.registered_replicas_.reserve(durability_->Size(durability::kReplicationReplicaPrefix));
+              for (; b != e; ++b) {
+                auto const &[replica_name, replica_data] = *b;
+                auto json = nlohmann::json::parse(replica_data, nullptr, false);
+                if (json.is_discarded()) return FetchReplicationError::PARSE_ERROR;
+                try {
+                  durability::ReplicationReplicaEntry data = json.get<durability::ReplicationReplicaEntry>();
+                  auto key_name = std::string_view{replica_name}.substr(strlen(durability::kReplicationReplicaPrefix));
+                  if (key_name != data.config.name) {
+                    return FetchReplicationError::PARSE_ERROR;
+                  }
+                  res.registered_replicas_.emplace_back(std::move(data.config));
+                } catch (...) {
+                  return FetchReplicationError::PARSE_ERROR;
+                }
+              }
+              return {std::move(res)};
+            },
+            [&](durability::ReplicaRole &&r) -> FetchReplicationResult_t {
+              return {RoleReplicaData{r.config, std::make_unique<ReplicationServer>(r.config)}};
+            },
+        },
+        std::move(data.role));
+  } catch (...) {
+    return FetchReplicationError::PARSE_ERROR;
+  }
+}
+
+bool ReplicationState::HandleVersionMigration(durability::ReplicationRoleEntry &data) const {
+  switch (data.version) {
+    case durability::DurabilityVersion::V1: {
+      // For each replica config, change key to use the prefix
+      std::map<std::string, std::string> to_put;
+      std::vector<std::string> to_delete;
+      for (auto [old_key, old_data] : *durability_) {
+        // skip reserved keys
+        if (old_key == durability::kReplicationRoleName) {
           continue;
         }
 
-        const auto maybe_replica_status = JSONToReplicationStatus(nlohmann::json::parse(replica_data));
-        if (!maybe_replica_status.has_value()) {
-          return FetchReplicationError::PARSE_ERROR;
-        }
+        // Turn old data to new data
+        auto old_json = nlohmann::json::parse(old_data, nullptr, false);
+        if (old_json.is_discarded()) return false;  // Can not read old_data as json
+        try {
+          durability::ReplicationReplicaEntry new_data = old_json.get<durability::ReplicationReplicaEntry>();
 
-        auto replica_status = *maybe_replica_status;
-        if (replica_status.name != replica_name) {
-          return FetchReplicationError::PARSE_ERROR;
+          // Migrate to using new key
+          to_put.emplace(BuildReplicaKey(old_key), nlohmann::json(new_data).dump());
+        } catch (...) {
+          return false;  // Can not parse as ReplicationReplicaEntry
         }
-        res.emplace_back(ReplicationClientConfig{
-            .name = replica_status.name,
-            .mode = replica_status.sync_mode,
-            .ip_address = replica_status.ip_address,
-            .port = replica_status.port,
-            .replica_check_frequency = replica_status.replica_check_frequency,
-            .ssl = replica_status.ssl,
-        });
+        to_delete.push_back(std::move(old_key));
       }
-      return {std::move(res)};
+      // Set version
+      data.version = durability::DurabilityVersion::V2;
+      // Re-serialise (to include version + epoch)
+      to_put.emplace(durability::kReplicationRoleName, nlohmann::json(data).dump());
+      if (!durability_->PutAndDeleteMultiple(to_put, to_delete)) return false;  // some reason couldn't persist
+      [[fallthrough]];
+    }
+    case durability::DurabilityVersion::V2: {
+      // do nothing - add code if V3 ever happens
+      break;
     }
   }
+  return true;
 }
+
 bool ReplicationState::TryPersistRegisteredReplica(const ReplicationClientConfig &config) {
   if (!ShouldPersist()) return true;
 
   // If any replicas are persisted then Role must be persisted
   if (role_persisted != RolePersisted::YES) {
     DMG_ASSERT(IsMain(), "MAIN is expected");
-    if (!TryPersistRoleMain()) return false;
+    auto epoch_str = std::string(std::get<RoleMainData>(replication_data_).epoch_.id());
+    if (!TryPersistRoleMain(std::move(epoch_str))) return false;
   }
 
-  auto data = ReplicationStatusToJSON(ReplicationStatus{.name = config.name,
-                                                        .ip_address = config.ip_address,
-                                                        .port = config.port,
-                                                        .sync_mode = config.mode,
-                                                        .replica_check_frequency = config.replica_check_frequency,
-                                                        .ssl = config.ssl,
-                                                        .role = ReplicationRole::REPLICA});
-  if (durability_->Put(config.name, data.dump())) return true;
+  auto data = durability::ReplicationReplicaEntry{.config = config};
+
+  auto key = BuildReplicaKey(config.name);
+  if (durability_->Put(key, nlohmann::json(data).dump())) return true;
   spdlog::error("Error when saving replica {} in settings.", config.name);
   return false;
 }
+
+bool ReplicationState::SetReplicationRoleMain() {
+  auto new_epoch = utils::GenerateUUID();
+  if (!TryPersistRoleMain(new_epoch)) {
+    return false;
+  }
+  replication_data_ = RoleMainData{.epoch_ = ReplicationEpoch{new_epoch}};
+  return true;
+}
+
+bool ReplicationState::SetReplicationRoleReplica(const ReplicationServerConfig &config) {
+  if (!TryPersistRoleReplica(config)) {
+    return false;
+  }
+  replication_data_ = RoleReplicaData{config, std::make_unique<ReplicationServer>(config)};
+  return true;
+}
+
+auto ReplicationState::RegisterReplica(const ReplicationClientConfig &config) -> RegisterReplicaError {
+  auto const replica_handler = [](RoleReplicaData const &) -> RegisterReplicaError {
+    return RegisterReplicaError::NOT_MAIN;
+  };
+  auto const main_handler = [this, &config](RoleMainData &mainData) -> RegisterReplicaError {
+    // name check
+    auto name_check = [&config](auto const &replicas) {
+      auto name_matches = [&name = config.name](ReplicationClientConfig const &registered_config) {
+        return registered_config.name == name;
+      };
+      return std::any_of(replicas.begin(), replicas.end(), name_matches);
+    };
+    if (name_check(mainData.registered_replicas_)) {
+      return RegisterReplicaError::NAME_EXISTS;
+    }
+
+    // endpoint check
+    auto endpoint_check = [&](auto const &replicas) {
+      auto endpoint_matches = [&config](ReplicationClientConfig const &registered_config) {
+        return registered_config.ip_address == config.ip_address && registered_config.port == config.port;
+      };
+      return std::any_of(replicas.begin(), replicas.end(), endpoint_matches);
+    };
+    if (endpoint_check(mainData.registered_replicas_)) {
+      return RegisterReplicaError::END_POINT_EXISTS;
+    }
+
+    // Durability
+    if (!TryPersistRegisteredReplica(config)) {
+      return RegisterReplicaError::COULD_NOT_BE_PERSISTED;
+    }
+
+    // set
+    mainData.registered_replicas_.emplace_back(config);
+    return RegisterReplicaError::SUCCESS;
+  };
+
+  return std::visit(utils::Overloaded{main_handler, replica_handler}, replication_data_);
+}
 }  // namespace memgraph::replication

src/replication/status.cpp

@@ -12,6 +12,9 @@
 
 #include "fmt/format.h"
 #include "utils/logging.hpp"
+#include "utils/variant_helpers.hpp"
+
+namespace memgraph::replication::durability {
 
 constexpr auto *kReplicaName = "replica_name";
 constexpr auto *kIpAddress = "replica_ip_address";
@@ -21,71 +24,87 @@ constexpr auto *kCheckFrequency = "replica_check_frequency";
 constexpr auto *kSSLKeyFile = "replica_ssl_key_file";
 constexpr auto *kSSLCertFile = "replica_ssl_cert_file";
 constexpr auto *kReplicationRole = "replication_role";
+constexpr auto *kEpoch = "epoch";
+constexpr auto *kVersion = "durability_version";
 
-namespace memgraph::replication {
-
-nlohmann::json ReplicationStatusToJSON(ReplicationStatus &&status) {
-  auto data = nlohmann::json::object();
-
-  data[kReplicaName] = std::move(status.name);
-  data[kIpAddress] = std::move(status.ip_address);
-  data[kPort] = status.port;
-  data[kSyncMode] = status.sync_mode;
-
-  data[kCheckFrequency] = status.replica_check_frequency.count();
-
-  if (status.ssl.has_value()) {
-    data[kSSLKeyFile] = std::move(status.ssl->key_file);
-    data[kSSLCertFile] = std::move(status.ssl->cert_file);
-  } else {
-    data[kSSLKeyFile] = nullptr;
-    data[kSSLCertFile] = nullptr;
-  }
-
-  if (status.role.has_value()) {
-    data[kReplicationRole] = *status.role;
-  }
-
-  return data;
-}
-std::optional<ReplicationStatus> JSONToReplicationStatus(nlohmann::json &&data) {
-  ReplicationStatus replica_status;
-
-  const auto get_failed_message = [](const std::string_view message, const std::string_view nested_message) {
-    return fmt::format("Failed to deserialize replica's configuration: {} : {}", message, nested_message);
+void to_json(nlohmann::json &j, const ReplicationRoleEntry &p) {
+  auto processMAIN = [&](MainRole const &main) {
+    j = nlohmann::json{{kVersion, p.version}, {kReplicationRole, ReplicationRole::MAIN}, {kEpoch, main.epoch.id()}};
   };
-
-  try {
-    data.at(kReplicaName).get_to(replica_status.name);
-    data.at(kIpAddress).get_to(replica_status.ip_address);
-    data.at(kPort).get_to(replica_status.port);
-    data.at(kSyncMode).get_to(replica_status.sync_mode);
-
-    replica_status.replica_check_frequency = std::chrono::seconds(data.at(kCheckFrequency));
-
-    const auto &key_file = data.at(kSSLKeyFile);
-    const auto &cert_file = data.at(kSSLCertFile);
-
-    MG_ASSERT(key_file.is_null() == cert_file.is_null());
-
-    if (!key_file.is_null()) {
-      replica_status.ssl = ReplicationClientConfig::SSL{};
-      data.at(kSSLKeyFile).get_to(replica_status.ssl->key_file);
-      data.at(kSSLCertFile).get_to(replica_status.ssl->cert_file);
-    }
-
-    if (data.find(kReplicationRole) != data.end()) {
-      replica_status.role = ReplicationRole::MAIN;
-      data.at(kReplicationRole).get_to(replica_status.role.value());
-    }
-  } catch (const nlohmann::json::type_error &exception) {
-    spdlog::error(get_failed_message("Invalid type conversion", exception.what()));
-    return std::nullopt;
-  } catch (const nlohmann::json::out_of_range &exception) {
-    spdlog::error(get_failed_message("Non existing field", exception.what()));
-    return std::nullopt;
-  }
-
-  return replica_status;
+  auto processREPLICA = [&](ReplicaRole const &replica) {
+    j = nlohmann::json{
+        {kVersion, p.version},
+        {kReplicationRole, ReplicationRole::REPLICA},
+        {kIpAddress, replica.config.ip_address},
+        {kPort, replica.config.port}
+        // TODO: SSL
+    };
+  };
+  std::visit(utils::Overloaded{processMAIN, processREPLICA}, p.role);
 }
-}  // namespace memgraph::replication
+
+void from_json(const nlohmann::json &j, ReplicationRoleEntry &p) {
+  // This value did not exist in V1, hence default DurabilityVersion::V1
+  DurabilityVersion version = j.value(kVersion, DurabilityVersion::V1);
+  // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
+  ReplicationRole role;
+  j.at(kReplicationRole).get_to(role);
+  switch (role) {
+    case ReplicationRole::MAIN: {
+      auto json_epoch = j.value(kEpoch, std::string{});
+      auto epoch = ReplicationEpoch{};
+      if (!json_epoch.empty()) epoch.SetEpoch(json_epoch);
+      p = ReplicationRoleEntry{.version = version, .role = MainRole{.epoch = std::move(epoch)}};
+      break;
+    }
+    case ReplicationRole::REPLICA: {
+      std::string ip_address;
+      // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
+      uint16_t port;
+      j.at(kIpAddress).get_to(ip_address);
+      j.at(kPort).get_to(port);
+      auto config = ReplicationServerConfig{.ip_address = std::move(ip_address), .port = port};
+      p = ReplicationRoleEntry{.version = version, .role = ReplicaRole{.config = std::move(config)}};
+      break;
+    }
+  }
+}
+
+void to_json(nlohmann::json &j, const ReplicationReplicaEntry &p) {
+  auto common = nlohmann::json{{kReplicaName, p.config.name},
+                               {kIpAddress, p.config.ip_address},
+                               {kPort, p.config.port},
+                               {kSyncMode, p.config.mode},
+                               {kCheckFrequency, p.config.replica_check_frequency.count()}};
+
+  if (p.config.ssl.has_value()) {
+    common[kSSLKeyFile] = p.config.ssl->key_file;
+    common[kSSLCertFile] = p.config.ssl->cert_file;
+  } else {
+    common[kSSLKeyFile] = nullptr;
+    common[kSSLCertFile] = nullptr;
+  }
+  j = std::move(common);
+}
+void from_json(const nlohmann::json &j, ReplicationReplicaEntry &p) {
+  const auto &key_file = j.at(kSSLKeyFile);
+  const auto &cert_file = j.at(kSSLCertFile);
+
+  MG_ASSERT(key_file.is_null() == cert_file.is_null());
+
+  auto seconds = j.at(kCheckFrequency).get<std::chrono::seconds::rep>();
+  auto config = ReplicationClientConfig{
+      .name = j.at(kReplicaName).get<std::string>(),
+      .mode = j.at(kSyncMode).get<ReplicationMode>(),
+      .ip_address = j.at(kIpAddress).get<std::string>(),
+      .port = j.at(kPort).get<uint16_t>(),
+      .replica_check_frequency = std::chrono::seconds{seconds},
+  };
+  if (!key_file.is_null()) {
+    config.ssl = ReplicationClientConfig::SSL{};
+    key_file.get_to(config.ssl->key_file);
+    cert_file.get_to(config.ssl->cert_file);
+  }
+  p = ReplicationReplicaEntry{.config = std::move(config)};
+}
+}  // namespace memgraph::replication::durability

src/rpc/CMakeLists.txt

@@ -7,5 +7,6 @@ find_package(fmt REQUIRED)
 find_package(gflags REQUIRED)
 
 add_library(mg-rpc STATIC ${rpc_src_files})
+add_library(mg::rpc ALIAS mg-rpc)
 target_link_libraries(mg-rpc Threads::Threads mg-communication mg-utils mg-io fmt::fmt gflags)
 target_link_libraries(mg-rpc mg-slk)

src/rpc/client.hpp

@@ -19,6 +19,7 @@
 #include "io/network/endpoint.hpp"
 #include "rpc/exceptions.hpp"
 #include "rpc/messages.hpp"
+#include "rpc/version.hpp"
 #include "slk/serialization.hpp"
 #include "slk/streams.hpp"
 #include "utils/logging.hpp"
@@ -43,7 +44,7 @@ class Client {
         : self_(self),
           guard_(std::move(guard)),
           req_builder_([self](const uint8_t *data, size_t size, bool have_more) {
-            if (!self->client_->Write(data, size, have_more)) throw RpcFailedException(self->endpoint_);
+            if (!self->client_->Write(data, size, have_more)) throw GenericRpcFailedException();
           }),
           res_load_(res_load) {}
 
@@ -69,11 +70,11 @@ class Client {
       while (true) {
         auto ret = slk::CheckStreamComplete(self_->client_->GetData(), self_->client_->GetDataSize());
         if (ret.status == slk::StreamStatus::INVALID) {
-          throw RpcFailedException(self_->endpoint_);
+          throw GenericRpcFailedException();
         } else if (ret.status == slk::StreamStatus::PARTIAL) {
           if (!self_->client_->Read(ret.stream_size - self_->client_->GetDataSize(),
                                     /* exactly_len = */ false)) {
-            throw RpcFailedException(self_->endpoint_);
+            throw GenericRpcFailedException();
           }
         } else {
           response_data_size = ret.stream_size;
@@ -88,11 +89,22 @@ class Client {
       utils::TypeId res_id{utils::TypeId::UNKNOWN};
       slk::Load(&res_id, &res_reader);
 
+      // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
+      rpc::Version version;
+      slk::Load(&version, &res_reader);
+
+      if (version != rpc::current_version) {
+        // V1 we introduced versioning with, absolutely no backwards compatibility,
+        // because it's impossible to provide backwards compatibility with pre versioning.
+        // Future versions this may require mechanism for graceful version handling.
+        throw VersionMismatchRpcFailedException();
+      }
+
       // Check the response ID.
       if (res_id != res_type.id && res_id != utils::TypeId::UNKNOWN) {
         spdlog::error("Message response was of unexpected type");
         self_->client_ = std::nullopt;
-        throw RpcFailedException(self_->endpoint_);
+        throw GenericRpcFailedException();
       }
 
       SPDLOG_TRACE("[RpcClient] received {}", res_type.name);
@@ -153,7 +165,7 @@ class Client {
       if (!client_->Connect(endpoint_)) {
         SPDLOG_ERROR("Couldn't connect to remote address {}", endpoint_);
         client_ = std::nullopt;
-        throw RpcFailedException(endpoint_);
+        throw GenericRpcFailedException();
       }
     }
 
@@ -162,6 +174,8 @@ class Client {
 
     // Build and send the request.
     slk::Save(req_type.id, handler.GetBuilder());
+    slk::Save(rpc::current_version, handler.GetBuilder());
+
     TRequestResponse::Request::Save(request, handler.GetBuilder());
 
     // Return the handler to the user.

src/rpc/exceptions.hpp

@@ -19,19 +19,30 @@ namespace memgraph::rpc {
 /// `utils::BasicException` is used for transient errors that should be reported
 /// to the user and `utils::StacktraceException` is used for fatal errors.
 /// This exception always requires explicit handling.
-class RpcFailedException final : public utils::BasicException {
+class RpcFailedException : public utils::BasicException {
  public:
-  RpcFailedException(const io::network::Endpoint &endpoint)
-      : utils::BasicException::BasicException(
-            "Couldn't communicate with the cluster! Please contact your "
-            "database administrator."),
-        endpoint_(endpoint) {}
-
-  /// Returns the endpoint associated with the error.
-  const io::network::Endpoint &endpoint() const { return endpoint_; }
-  SPECIALIZE_GET_EXCEPTION_NAME(RpcFailedException)
-
- private:
-  io::network::Endpoint endpoint_;
+  RpcFailedException(std::string_view msg) : utils::BasicException(msg) {}
+  SPECIALIZE_GET_EXCEPTION_NAME(RpcFailedException);
 };
+
+class VersionMismatchRpcFailedException : public RpcFailedException {
+ public:
+  VersionMismatchRpcFailedException()
+      : RpcFailedException(
+            "Couldn't communicate with the cluster! There was a version mismatch. "
+            "Please contact your database administrator.") {}
+
+  SPECIALIZE_GET_EXCEPTION_NAME(VersionMismatchRpcFailedException);
+};
+
+class GenericRpcFailedException : public RpcFailedException {
+ public:
+  GenericRpcFailedException()
+      : RpcFailedException(
+            "Couldn't communicate with the cluster! Please contact your "
+            "database administrator.") {}
+
+  SPECIALIZE_GET_EXCEPTION_NAME(GenericRpcFailedException);
+};
+
 }  // namespace memgraph::rpc

src/rpc/protocol.cpp

@@ -13,6 +13,7 @@
 
 #include "rpc/messages.hpp"
 #include "rpc/server.hpp"
+#include "rpc/version.hpp"
 #include "slk/serialization.hpp"
 #include "slk/streams.hpp"
 #include "utils/on_scope_exit.hpp"
@@ -44,6 +45,16 @@ void Session::Execute() {
   // Load the request ID.
   utils::TypeId req_id{utils::TypeId::UNKNOWN};
   slk::Load(&req_id, &req_reader);
+  // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
+  rpc::Version version;
+  slk::Load(&version, &req_reader);
+
+  if (version != rpc::current_version) {
+    // V1 we introduced versioning with, absolutely no backwards compatibility,
+    // because it's impossible to provide backwards compatibility with pre versioning.
+    // Future versions this may require mechanism for graceful version handling.
+    throw SessionException("Session trying to execute a RPC call of an incorrect version!");
+  }
 
   // Access to `callbacks_` and `extended_callbacks_` is done here without
   // acquiring the `mutex_` because we don't allow RPC registration after the
@@ -62,10 +73,12 @@ void Session::Execute() {
     }
     SPDLOG_TRACE("[RpcServer] received {}", extended_it->second.req_type.name);
     slk::Save(extended_it->second.res_type.id, &res_builder);
+    slk::Save(rpc::current_version, &res_builder);
     extended_it->second.callback(endpoint_, &req_reader, &res_builder);
   } else {
     SPDLOG_TRACE("[RpcServer] received {}", it->second.req_type.name);
     slk::Save(it->second.res_type.id, &res_builder);
+    slk::Save(rpc::current_version, &res_builder);
     it->second.callback(&req_reader, &res_builder);
   }
 

src/rpc/version.hpp

@@ -0,0 +1,27 @@
+// Copyright 2023 Memgraph Ltd.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+// License, and you may not use this file except in compliance with the Business Source License.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+#pragma once
+
+#include <cstdint>
+
+namespace memgraph::rpc {
+
+using Version = uint64_t;
+
+// versioning of RPC was/will be introduced in 2.13
+// We start the versioning with a strange number, to radically reduce the
+// probability of accidental match/conformance with pre 2.13 versions
+constexpr auto v1 = Version{2023'10'30'0'2'13};
+
+constexpr auto current_version = v1;
+
+}  // namespace memgraph::rpc

src/slk/CMakeLists.txt

@@ -4,5 +4,6 @@ set(slk_src_files
 find_package(gflags REQUIRED)
 
 add_library(mg-slk STATIC ${slk_src_files})
+add_library(mg::slk ALIAS mg-slk)
 target_link_libraries(mg-slk gflags)
 target_link_libraries(mg-slk mg-utils)

src/storage/v2/CMakeLists.txt

@@ -35,24 +35,10 @@ add_library(mg-storage-v2 STATIC
         disk/unique_constraints.cpp
         storage_mode.cpp
         replication/replication_client.cpp
-        replication/replication_server.cpp
         replication/serialization.cpp
         replication/slk.cpp
         replication/rpc.cpp
         replication/replication_storage_state.cpp
-        replication/replication_handler.cpp
-        inmemory/replication/replication_server.cpp
         inmemory/replication/replication_client.cpp
 )
 target_link_libraries(mg-storage-v2 mg::replication Threads::Threads mg-utils gflags absl::flat_hash_map mg-rpc mg-slk mg-events)
-
-# Until we get LTO there is an advantage to do some unity builds
-set_target_properties(mg-storage-v2
-        PROPERTIES
-        UNITY_BUILD ON
-        UNITY_BUILD_MODE GROUP
-)
-set_source_files_properties(
-        vertex_info_cache.cpp vertex_accessor.cpp
-        PROPERTIES UNITY_GROUP "ensure inline of vertex_info_cache"
-)

src/storage/v2/config.hpp

@@ -17,6 +17,7 @@
 #include "storage/v2/isolation_level.hpp"
 #include "storage/v2/storage_mode.hpp"
 #include "utils/exceptions.hpp"
+#include "utils/logging.hpp"
 
 namespace memgraph::storage {
 
@@ -34,10 +35,12 @@ struct Config {
 
     Type type{Type::PERIODIC};
     std::chrono::milliseconds interval{std::chrono::milliseconds(1000)};
+    friend bool operator==(const Gc &lrh, const Gc &rhs) = default;
   } gc;
 
   struct Items {
     bool properties_on_edges{true};
+    friend bool operator==(const Items &lrh, const Items &rhs) = default;
   } items;
 
   struct Durability {
@@ -62,10 +65,12 @@ struct Config {
     uint64_t recovery_thread_count{8};
 
     bool allow_parallel_index_creation{false};
+    friend bool operator==(const Durability &lrh, const Durability &rhs) = default;
   } durability;
 
   struct Transaction {
     IsolationLevel isolation_level{IsolationLevel::SNAPSHOT_ISOLATION};
+    friend bool operator==(const Transaction &lrh, const Transaction &rhs) = default;
   } transaction;
 
   struct DiskConfig {
@@ -77,13 +82,26 @@ struct Config {
     std::filesystem::path id_name_mapper_directory{"storage/rocksdb_id_name_mapper"};
     std::filesystem::path durability_directory{"storage/rocksdb_durability"};
     std::filesystem::path wal_directory{"storage/rocksdb_wal"};
+    friend bool operator==(const DiskConfig &lrh, const DiskConfig &rhs) = default;
   } disk;
 
   std::string name;
   bool force_on_disk{false};
   StorageMode storage_mode{StorageMode::IN_MEMORY_TRANSACTIONAL};
+
+  friend bool operator==(const Config &lrh, const Config &rhs) = default;
 };
 
+inline auto ReplicationStateRootPath(memgraph::storage::Config const &config) -> std::optional<std::filesystem::path> {
+  if (!config.durability.restore_replication_state_on_startup) {
+    spdlog::warn(
+        "Replication configuration will NOT be stored. When the server restarts, replication state will be "
+        "forgotten.");
+    return std::nullopt;
+  }
+  return {config.durability.storage_directory};
+}
+
 static inline void UpdatePaths(Config &config, const std::filesystem::path &storage_dir) {
   auto contained = [](const auto &path, const auto &base) -> std::optional<std::filesystem::path> {
     auto rel = std::filesystem::relative(path, base);

src/storage/v2/disk/label_index.hpp

@@ -19,6 +19,7 @@
 #include "storage/v2/indices/label_index.hpp"
 #include "storage/v2/vertex.hpp"
 #include "utils/rocksdb_serialization.hpp"
+#include "utils/synchronized.hpp"
 
 namespace memgraph::storage {
 class DiskLabelIndex : public storage::LabelIndex {

src/storage/v2/disk/label_property_index.hpp

@@ -13,6 +13,7 @@
 
 #include "storage/v2/disk/rocksdb_storage.hpp"
 #include "storage/v2/indices/label_property_index.hpp"
+#include "utils/synchronized.hpp"
 
 namespace memgraph::storage {
 

src/storage/v2/disk/storage.cpp

@@ -801,7 +801,7 @@ StorageInfo DiskStorage::GetBaseInfo(bool /* unused */) {
     // NOLINTNEXTLINE(bugprone-narrowing-conversions, cppcoreguidelines-narrowing-conversions)
     info.average_degree = 2.0 * static_cast<double>(info.edge_count) / info.vertex_count;
   }
-  info.memory_usage = utils::GetMemoryUsage();
+  info.memory_res = utils::GetMemoryRES();
   info.disk_usage = GetDiskSpaceUsage();
   return info;
 }
@@ -1551,7 +1551,7 @@ DiskStorage::CheckExistingVerticesBeforeCreatingUniqueConstraint(LabelId label,
 
 // NOLINTNEXTLINE(google-default-arguments)
 utils::BasicResult<StorageManipulationError, void> DiskStorage::DiskAccessor::Commit(
-    const std::optional<uint64_t> desired_commit_timestamp) {
+    const std::optional<uint64_t> desired_commit_timestamp, bool /*is_main*/) {
   MG_ASSERT(is_transaction_active_, "The transaction is already terminated!");
   MG_ASSERT(!transaction_.must_abort, "The transaction can't be committed!");
 
@@ -1957,7 +1957,7 @@ UniqueConstraints::DeletionStatus DiskStorage::DiskAccessor::DropUniqueConstrain
   return UniqueConstraints::DeletionStatus::SUCCESS;
 }
 
-Transaction DiskStorage::CreateTransaction(IsolationLevel isolation_level, StorageMode storage_mode) {
+Transaction DiskStorage::CreateTransaction(IsolationLevel isolation_level, StorageMode storage_mode, bool /*is_main*/) {
   /// We acquire the transaction engine lock here because we access (and
   /// modify) the transaction engine variables (`transaction_id` and
   /// `timestamp`) below.
@@ -1982,7 +1982,8 @@ uint64_t DiskStorage::CommitTimestamp(const std::optional<uint64_t> desired_comm
   return *desired_commit_timestamp;
 }
 
-std::unique_ptr<Storage::Accessor> DiskStorage::Access(std::optional<IsolationLevel> override_isolation_level) {
+std::unique_ptr<Storage::Accessor> DiskStorage::Access(std::optional<IsolationLevel> override_isolation_level,
+                                                       bool /*is_main*/) {
   auto isolation_level = override_isolation_level.value_or(isolation_level_);
   if (isolation_level != IsolationLevel::SNAPSHOT_ISOLATION) {
     throw utils::NotYetImplemented("Disk storage supports only SNAPSHOT isolation level.");
@@ -1990,7 +1991,8 @@ std::unique_ptr<Storage::Accessor> DiskStorage::Access(std::optional<IsolationLe
   return std::unique_ptr<DiskAccessor>(
       new DiskAccessor{Storage::Accessor::shared_access, this, isolation_level, storage_mode_});
 }
-std::unique_ptr<Storage::Accessor> DiskStorage::UniqueAccess(std::optional<IsolationLevel> override_isolation_level) {
+std::unique_ptr<Storage::Accessor> DiskStorage::UniqueAccess(std::optional<IsolationLevel> override_isolation_level,
+                                                             bool /*is_main*/) {
   auto isolation_level = override_isolation_level.value_or(isolation_level_);
   if (isolation_level != IsolationLevel::SNAPSHOT_ISOLATION) {
     throw utils::NotYetImplemented("Disk storage supports only SNAPSHOT isolation level.");

src/storage/v2/disk/storage.hpp

@@ -142,8 +142,8 @@ class DiskStorage final : public Storage {
     ConstraintsInfo ListAllConstraints() const override;
 
     // NOLINTNEXTLINE(google-default-arguments)
-    utils::BasicResult<StorageManipulationError, void> Commit(
-        std::optional<uint64_t> desired_commit_timestamp = {}) override;
+    utils::BasicResult<StorageManipulationError, void> Commit(std::optional<uint64_t> desired_commit_timestamp = {},
+                                                              bool is_main = true) override;
 
     void UpdateObjectsCountOnAbort();
 
@@ -172,9 +172,13 @@ class DiskStorage final : public Storage {
                                                            const std::set<PropertyId> &properties) override;
   };
 
-  std::unique_ptr<Storage::Accessor> Access(std::optional<IsolationLevel> override_isolation_level) override;
+  using Storage::Access;
+  std::unique_ptr<Storage::Accessor> Access(std::optional<IsolationLevel> override_isolation_level,
+                                            bool is_main) override;
 
-  std::unique_ptr<Storage::Accessor> UniqueAccess(std::optional<IsolationLevel> override_isolation_level) override;
+  using Storage::UniqueAccess;
+  std::unique_ptr<Storage::Accessor> UniqueAccess(std::optional<IsolationLevel> override_isolation_level,
+                                                  bool is_main) override;
 
   /// Flushing methods
   [[nodiscard]] utils::BasicResult<StorageManipulationError, void> FlushIndexCache(Transaction *transaction);
@@ -277,7 +281,8 @@ class DiskStorage final : public Storage {
 
   RocksDBStorage *GetRocksDBStorage() const { return kvstore_.get(); }
 
-  Transaction CreateTransaction(IsolationLevel isolation_level, StorageMode storage_mode) override;
+  using Storage::CreateTransaction;
+  Transaction CreateTransaction(IsolationLevel isolation_level, StorageMode storage_mode, bool is_main) override;
 
   void SetEdgeImportMode(EdgeImportMode edge_import_status);
 
@@ -304,22 +309,16 @@ class DiskStorage final : public Storage {
 
   void FreeMemory(std::unique_lock<utils::ResourceLock> /*lock*/) override {}
 
-  void PrepareForNewEpoch(std::string /*prev_epoch*/) override {
-    throw utils::BasicException("Disk storage mode does not support replication.");
-  }
+  void PrepareForNewEpoch() override { throw utils::BasicException("Disk storage mode does not support replication."); }
 
   uint64_t CommitTimestamp(std::optional<uint64_t> desired_commit_timestamp = {});
 
-  auto CreateReplicationClient(const memgraph::replication::ReplicationClientConfig & /*config*/)
+  auto CreateReplicationClient(const memgraph::replication::ReplicationClientConfig & /*config*/,
+                               const memgraph::replication::ReplicationEpoch * /*current_epoch*/)
       -> std::unique_ptr<ReplicationClient> override {
     throw utils::BasicException("Disk storage mode does not support replication.");
   }
 
-  auto CreateReplicationServer(const memgraph::replication::ReplicationServerConfig & /*config*/)
-      -> std::unique_ptr<ReplicationServer> override {
-    throw utils::BasicException("Disk storage mode does not support replication.");
-  }
-
   std::unique_ptr<RocksDBStorage> kvstore_;
   DurableMetadata durable_metadata_;
   EdgeImportMode edge_import_status_{EdgeImportMode::INACTIVE};

src/storage/v2/durability/durability.cpp

@@ -211,12 +211,10 @@ void RecoverIndicesAndConstraints(const RecoveredIndicesAndConstraints &indices_
 
 std::optional<RecoveryInfo> RecoverData(const std::filesystem::path &snapshot_directory,
                                         const std::filesystem::path &wal_directory, std::string *uuid,
-                                        memgraph::replication::ReplicationEpoch &epoch,
-                                        std::deque<std::pair<std::string, uint64_t>> *epoch_history,
-                                        utils::SkipList<Vertex> *vertices, utils::SkipList<Edge> *edges,
-                                        std::atomic<uint64_t> *edge_count, NameIdMapper *name_id_mapper,
-                                        Indices *indices, Constraints *constraints, const Config &config,
-                                        uint64_t *wal_seq_num) {
+                                        ReplicationStorageState &repl_storage_state, utils::SkipList<Vertex> *vertices,
+                                        utils::SkipList<Edge> *edges, std::atomic<uint64_t> *edge_count,
+                                        NameIdMapper *name_id_mapper, Indices *indices, Constraints *constraints,
+                                        const Config &config, uint64_t *wal_seq_num) {
   utils::MemoryTracker::OutOfMemoryExceptionEnabler oom_exception;
   spdlog::info("Recovering persisted data using snapshot ({}) and WAL directory ({}).", snapshot_directory,
                wal_directory);
@@ -226,6 +224,7 @@ std::optional<RecoveryInfo> RecoverData(const std::filesystem::path &snapshot_di
     return std::nullopt;
   }
 
+  auto *const epoch_history = &repl_storage_state.history;
   utils::Timer timer;
 
   auto snapshot_files = GetSnapshotFiles(snapshot_directory);
@@ -264,7 +263,7 @@ std::optional<RecoveryInfo> RecoverData(const std::filesystem::path &snapshot_di
     recovery_info = recovered_snapshot->recovery_info;
     indices_constraints = std::move(recovered_snapshot->indices_constraints);
     snapshot_timestamp = recovered_snapshot->snapshot_info.start_timestamp;
-    epoch.SetEpoch(std::move(recovered_snapshot->snapshot_info.epoch_id));
+    repl_storage_state.epoch_.SetEpoch(std::move(recovered_snapshot->snapshot_info.epoch_id));
 
     if (!utils::DirExists(wal_directory)) {
       const auto par_exec_info = config.durability.allow_parallel_index_creation
@@ -309,7 +308,7 @@ std::optional<RecoveryInfo> RecoverData(const std::filesystem::path &snapshot_di
     // UUID used for durability is the UUID of the last WAL file.
     // Same for the epoch id.
     *uuid = std::move(wal_files.back().uuid);
-    epoch.SetEpoch(std::move(wal_files.back().epoch_id));
+    repl_storage_state.epoch_.SetEpoch(std::move(wal_files.back().epoch_id));
   }
 
   auto maybe_wal_files = GetWalFiles(wal_directory, *uuid);
@@ -365,7 +364,7 @@ std::optional<RecoveryInfo> RecoverData(const std::filesystem::path &snapshot_di
       }
       previous_seq_num = wal_file.seq_num;
 
-      if (wal_file.epoch_id != epoch.id()) {
+      if (wal_file.epoch_id != repl_storage_state.epoch_.id()) {
         // This way we skip WALs finalized only because of role change.
         // We can also set the last timestamp to 0 if last loaded timestamp
         // is nullopt as this can only happen if the WAL file with seq = 0
@@ -373,7 +372,7 @@ std::optional<RecoveryInfo> RecoverData(const std::filesystem::path &snapshot_di
         if (last_loaded_timestamp) {
           epoch_history->emplace_back(wal_file.epoch_id, *last_loaded_timestamp);
         }
-        epoch.SetEpoch(std::move(wal_file.epoch_id));
+        repl_storage_state.epoch_.SetEpoch(std::move(wal_file.epoch_id));
       }
       try {
         auto info = LoadWal(wal_file.path, &indices_constraints, last_loaded_timestamp, vertices, edges, name_id_mapper,
@@ -398,7 +397,12 @@ std::optional<RecoveryInfo> RecoverData(const std::filesystem::path &snapshot_di
     spdlog::info("All necessary WAL files are loaded successfully.");
   }
 
-  RecoverIndicesAndConstraints(indices_constraints, indices, constraints, vertices);
+  const auto par_exec_info =
+      config.durability.allow_parallel_index_creation && !recovery_info.vertex_batches.empty()
+          ? std::make_optional(std::make_pair(recovery_info.vertex_batches, config.durability.recovery_thread_count))
+          : std::nullopt;
+
+  RecoverIndicesAndConstraints(indices_constraints, indices, constraints, vertices, par_exec_info);
 
   memgraph::metrics::Measure(memgraph::metrics::SnapshotRecoveryLatency_us,
                              std::chrono::duration_cast<std::chrono::microseconds>(timer.Elapsed()).count());

src/storage/v2/durability/durability.hpp

@@ -19,6 +19,7 @@
 #include <variant>
 
 #include "replication/epoch.hpp"
+#include "replication/state.hpp"
 #include "storage/v2/config.hpp"
 #include "storage/v2/constraints/constraints.hpp"
 #include "storage/v2/durability/metadata.hpp"
@@ -26,6 +27,7 @@
 #include "storage/v2/edge.hpp"
 #include "storage/v2/indices/indices.hpp"
 #include "storage/v2/name_id_mapper.hpp"
+#include "storage/v2/replication/replication_storage_state.hpp"
 #include "storage/v2/vertex.hpp"
 #include "utils/skip_list.hpp"
 
@@ -110,11 +112,9 @@ void RecoverIndicesAndConstraints(
 /// @throw std::bad_alloc
 std::optional<RecoveryInfo> RecoverData(const std::filesystem::path &snapshot_directory,
                                         const std::filesystem::path &wal_directory, std::string *uuid,
-                                        memgraph::replication::ReplicationEpoch &epoch,
-                                        std::deque<std::pair<std::string, uint64_t>> *epoch_history,
-                                        utils::SkipList<Vertex> *vertices, utils::SkipList<Edge> *edges,
-                                        std::atomic<uint64_t> *edge_count, NameIdMapper *name_id_mapper,
-                                        Indices *indices, Constraints *constraints, const Config &config,
-                                        uint64_t *wal_seq_num);
+                                        ReplicationStorageState &repl_storage_state, utils::SkipList<Vertex> *vertices,
+                                        utils::SkipList<Edge> *edges, std::atomic<uint64_t> *edge_count,
+                                        NameIdMapper *name_id_mapper, Indices *indices, Constraints *constraints,
+                                        const Config &config, uint64_t *wal_seq_num);
 
 }  // namespace memgraph::storage::durability

src/storage/v2/edge_accessor.hpp

@@ -110,6 +110,9 @@ class EdgeAccessor final {
 
 }  // namespace memgraph::storage
 
+static_assert(std::is_trivially_copyable<memgraph::storage::EdgeAccessor>::value,
+              "storage::EdgeAccessor must be trivially copyable!");
+
 namespace std {
 template <>
 struct hash<memgraph::storage::EdgeAccessor> {

src/storage/v2/inmemory/label_index.hpp

@@ -16,6 +16,7 @@
 #include "storage/v2/indices/label_index_stats.hpp"
 #include "storage/v2/vertex.hpp"
 #include "utils/rw_lock.hpp"
+#include "utils/synchronized.hpp"
 
 namespace memgraph::storage {
 

src/storage/v2/inmemory/label_property_index.hpp

@@ -15,6 +15,7 @@
 #include "storage/v2/indices/label_property_index.hpp"
 #include "storage/v2/indices/label_property_index_stats.hpp"
 #include "utils/rw_lock.hpp"
+#include "utils/synchronized.hpp"
 
 namespace memgraph::storage {
 

src/storage/v2/inmemory/replication/replication_client.cpp

@@ -44,7 +44,7 @@ class CurrentWalHandler {
 
 ////// CurrentWalHandler //////
 CurrentWalHandler::CurrentWalHandler(ReplicationClient *self)
-    : self_(self), stream_(self_->rpc_client_.Stream<replication::CurrentWalRpc>()) {}
+    : self_(self), stream_(self_->rpc_client_.Stream<replication::CurrentWalRpc>(self->GetStorageId())) {}
 
 void CurrentWalHandler::AppendFilename(const std::string &filename) {
   replication::Encoder encoder(stream_.GetBuilder());
@@ -70,9 +70,10 @@ replication::CurrentWalRes CurrentWalHandler::Finalize() { return stream_.AwaitR
 
 ////// ReplicationClient Helpers //////
 
-replication::WalFilesRes TransferWalFiles(rpc::Client &client, const std::vector<std::filesystem::path> &wal_files) {
+replication::WalFilesRes TransferWalFiles(std::string db_name, rpc::Client &client,
+                                          const std::vector<std::filesystem::path> &wal_files) {
   MG_ASSERT(!wal_files.empty(), "Wal files list is empty!");
-  auto stream = client.Stream<replication::WalFilesRpc>(wal_files.size());
+  auto stream = client.Stream<replication::WalFilesRpc>(std::move(db_name), wal_files.size());
   replication::Encoder encoder(stream.GetBuilder());
   for (const auto &wal : wal_files) {
     spdlog::debug("Sending wal file: {}", wal);
@@ -81,8 +82,8 @@ replication::WalFilesRes TransferWalFiles(rpc::Client &client, const std::vector
   return stream.AwaitResponse();
 }
 
-replication::SnapshotRes TransferSnapshot(rpc::Client &client, const std::filesystem::path &path) {
-  auto stream = client.Stream<replication::SnapshotRpc>();
+replication::SnapshotRes TransferSnapshot(std::string db_name, rpc::Client &client, const std::filesystem::path &path) {
+  auto stream = client.Stream<replication::SnapshotRpc>(std::move(db_name));
   replication::Encoder encoder(stream.GetBuilder());
   encoder.WriteFile(path);
   return stream.AwaitResponse();
@@ -115,19 +116,19 @@ void InMemoryReplicationClient::RecoverReplica(uint64_t replica_commit) {
 
     const auto steps = GetRecoverySteps(replica_commit, &file_locker);
     int i = 0;
-    for (const auto &recovery_step : steps) {
+    for (const InMemoryReplicationClient::RecoveryStep &recovery_step : steps) {
       spdlog::trace("Recovering in step: {}", i++);
       try {
         std::visit(
             [&, this]<typename T>(T &&arg) {
               using StepType = std::remove_cvref_t<T>;
-              if constexpr (std::is_same_v<StepType, RecoverySnapshot>) {
+              if constexpr (std::is_same_v<StepType, RecoverySnapshot>) {  // TODO: split into 3 overloads
                 spdlog::debug("Sending the latest snapshot file: {}", arg);
-                auto response = TransferSnapshot(rpc_client_, arg);
+                auto response = TransferSnapshot(storage->id(), rpc_client_, arg);
                 replica_commit = response.current_commit_timestamp;
               } else if constexpr (std::is_same_v<StepType, RecoveryWals>) {
                 spdlog::debug("Sending the latest wal files");
-                auto response = TransferWalFiles(rpc_client_, arg);
+                auto response = TransferWalFiles(storage->id(), rpc_client_, arg);
                 replica_commit = response.current_commit_timestamp;
                 spdlog::debug("Wal files successfully transferred.");
               } else if constexpr (std::is_same_v<StepType, RecoveryCurrentWal>) {

src/storage/v2/inmemory/replication/replication_server.hpp

@@ -1,52 +0,0 @@
-// Copyright 2023 Memgraph Ltd.
-//
-// Use of this software is governed by the Business Source License
-// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
-// License, and you may not use this file except in compliance with the Business Source License.
-//
-// As of the Change Date specified in that file, in accordance with
-// the Business Source License, use of this software will be governed
-// by the Apache License, Version 2.0, included in the file
-// licenses/APL.txt.
-
-#pragma once
-
-#include "replication/epoch.hpp"
-#include "storage/v2/replication/replication_server.hpp"
-#include "storage/v2/replication/serialization.hpp"
-
-namespace memgraph::storage {
-
-class InMemoryStorage;
-
-class InMemoryReplicationServer : public ReplicationServer {
- public:
-  explicit InMemoryReplicationServer(InMemoryStorage *storage,
-                                     const memgraph::replication::ReplicationServerConfig &config,
-                                     memgraph::replication::ReplicationEpoch *repl_epoch);
-
- private:
-  // RPC handlers
-  void HeartbeatHandler(slk::Reader *req_reader, slk::Builder *res_builder);
-
-  void AppendDeltasHandler(slk::Reader *req_reader, slk::Builder *res_builder);
-
-  void SnapshotHandler(slk::Reader *req_reader, slk::Builder *res_builder);
-
-  void WalFilesHandler(slk::Reader *req_reader, slk::Builder *res_builder);
-
-  void CurrentWalHandler(slk::Reader *req_reader, slk::Builder *res_builder);
-
-  void TimestampHandler(slk::Reader *req_reader, slk::Builder *res_builder);
-
-  static void LoadWal(InMemoryStorage *storage, memgraph::replication::ReplicationEpoch &epoch,
-                      replication::Decoder *decoder);
-
-  static uint64_t ReadAndApplyDelta(InMemoryStorage *storage, durability::BaseDecoder *decoder, uint64_t version);
-
-  InMemoryStorage *storage_;
-
-  memgraph::replication::ReplicationEpoch *repl_epoch_;
-};
-
-}  // namespace memgraph::storage

src/storage/v2/inmemory/storage.cpp

@@ -16,11 +16,11 @@
 #include "storage/v2/metadata_delta.hpp"
 
 /// REPLICATION ///
+#include "dbms/inmemory/replication_handlers.hpp"
 #include "storage/v2/inmemory/replication/replication_client.hpp"
-#include "storage/v2/inmemory/replication/replication_server.hpp"
 #include "storage/v2/inmemory/unique_constraints.hpp"
-#include "storage/v2/replication/replication_handler.hpp"
 #include "utils/resource_lock.hpp"
+#include "utils/stat.hpp"
 
 namespace memgraph::storage {
 
@@ -60,12 +60,10 @@ InMemoryStorage::InMemoryStorage(Config config, StorageMode storage_mode)
               "process!",
               config_.durability.storage_directory);
   }
-  auto &repl_state = repl_state_;
   if (config_.durability.recover_on_startup) {
-    auto &epoch = repl_state.GetEpoch();
-    auto info = durability::RecoverData(snapshot_directory_, wal_directory_, &uuid_, epoch,
-                                        &repl_storage_state_.history, &vertices_, &edges_, &edge_count_,
-                                        name_id_mapper_.get(), &indices_, &constraints_, config_, &wal_seq_num_);
+    auto info =
+        durability::RecoverData(snapshot_directory_, wal_directory_, &uuid_, repl_storage_state_, &vertices_, &edges_,
+                                &edge_count_, name_id_mapper_.get(), &indices_, &constraints_, config_, &wal_seq_num_);
     if (info) {
       vertex_id_ = info->next_vertex_id;
       edge_id_ = info->next_edge_id;
@@ -103,51 +101,14 @@ InMemoryStorage::InMemoryStorage(Config config, StorageMode storage_mode)
           "those files into a .backup directory inside the storage directory.");
     }
   }
-  if (config_.durability.snapshot_wal_mode != Config::Durability::SnapshotWalMode::DISABLED) {
-    snapshot_runner_.Run("Snapshot", config_.durability.snapshot_interval, [this] {
-      auto const &repl_state = repl_state_;
-      if (auto maybe_error = this->CreateSnapshot(repl_state, {true}); maybe_error.HasError()) {
-        switch (maybe_error.GetError()) {
-          case CreateSnapshotError::DisabledForReplica:
-            spdlog::warn(
-                utils::MessageWithLink("Snapshots are disabled for replicas.", "https://memgr.ph/replication"));
-            break;
-          case CreateSnapshotError::DisabledForAnalyticsPeriodicCommit:
-            spdlog::warn(utils::MessageWithLink("Periodic snapshots are disabled for analytical mode.",
-                                                "https://memgr.ph/durability"));
-            break;
-          case storage::InMemoryStorage::CreateSnapshotError::ReachedMaxNumTries:
-            spdlog::warn("Failed to create snapshot. Reached max number of tries. Please contact support");
-            break;
-        }
-      }
-    });
-  }
   if (config_.gc.type == Config::Gc::Type::PERIODIC) {
     gc_runner_.Run("Storage GC", config_.gc.interval, [this] { this->CollectGarbage<false>(); });
   }
-
   if (timestamp_ == kTimestampInitialId) {
     commit_log_.emplace();
   } else {
     commit_log_.emplace(timestamp_);
   }
-
-  if (config_.durability.restore_replication_state_on_startup) {
-    spdlog::info("Replication configuration will be stored and will be automatically restored in case of a crash.");
-    ReplicationHandler{repl_state, *this}.RestoreReplication();
-  } else {
-    spdlog::warn(
-        "Replication configuration will NOT be stored. When the server restarts, replication state will be "
-        "forgotten.");
-  }
-
-  if (config_.durability.snapshot_wal_mode == Config::Durability::SnapshotWalMode::DISABLED && repl_state.IsMain()) {
-    spdlog::warn(
-        "The instance has the MAIN replication role, but durability logs and snapshots are disabled. Please consider "
-        "enabling durability by using --storage-snapshot-interval-sec and --storage-wal-enabled flags because "
-        "without write-ahead logs this instance is not replicating any data.");
-  }
 }
 
 InMemoryStorage::InMemoryStorage(Config config) : InMemoryStorage(config, StorageMode::IN_MEMORY_TRANSACTIONAL) {}
@@ -167,29 +128,15 @@ InMemoryStorage::~InMemoryStorage() {
   if (config_.durability.snapshot_wal_mode != Config::Durability::SnapshotWalMode::DISABLED) {
     snapshot_runner_.Stop();
   }
-  if (config_.durability.snapshot_on_exit) {
-    auto const &repl_state = repl_state_;
-    if (auto maybe_error = this->CreateSnapshot(repl_state, {false}); maybe_error.HasError()) {
-      switch (maybe_error.GetError()) {
-        case CreateSnapshotError::DisabledForReplica:
-          spdlog::warn(utils::MessageWithLink("Snapshots are disabled for replicas.", "https://memgr.ph/replication"));
-          break;
-        case CreateSnapshotError::DisabledForAnalyticsPeriodicCommit:
-          spdlog::warn(utils::MessageWithLink("Periodic snapshots are disabled for analytical mode.",
-                                              "https://memgr.ph/replication"));
-          break;
-        case storage::InMemoryStorage::CreateSnapshotError::ReachedMaxNumTries:
-          spdlog::warn("Failed to create snapshot. Reached max number of tries. Please contact support");
-          break;
-      }
-    }
+  if (config_.durability.snapshot_on_exit && this->create_snapshot_handler) {
+    create_snapshot_handler();
   }
   committed_transactions_.WithLock([](auto &transactions) { transactions.clear(); });
 }
 
 InMemoryStorage::InMemoryAccessor::InMemoryAccessor(auto tag, InMemoryStorage *storage, IsolationLevel isolation_level,
-                                                    StorageMode storage_mode)
-    : Accessor(tag, storage, isolation_level, storage_mode), config_(storage->config_.items) {}
+                                                    StorageMode storage_mode, bool is_main)
+    : Accessor(tag, storage, isolation_level, storage_mode, is_main), config_(storage->config_.items) {}
 InMemoryStorage::InMemoryAccessor::InMemoryAccessor(InMemoryAccessor &&other) noexcept
     : Accessor(std::move(other)), config_(other.config_) {}
 
@@ -711,7 +658,7 @@ Result<EdgeAccessor> InMemoryStorage::InMemoryAccessor::EdgeChangeType(EdgeAcces
 
 // NOLINTNEXTLINE(google-default-arguments)
 utils::BasicResult<StorageManipulationError, void> InMemoryStorage::InMemoryAccessor::Commit(
-    const std::optional<uint64_t> desired_commit_timestamp) {
+    const std::optional<uint64_t> desired_commit_timestamp, bool is_main) {
   MG_ASSERT(is_transaction_active_, "The transaction is already terminated!");
   MG_ASSERT(!transaction_.must_abort, "The transaction can't be committed!");
 
@@ -719,7 +666,7 @@ utils::BasicResult<StorageManipulationError, void> InMemoryStorage::InMemoryAcce
 
   auto *mem_storage = static_cast<InMemoryStorage *>(storage_);
 
-  auto const &replState = mem_storage->repl_state_;
+  // TODO: duplicated transaction finalisation in md_deltas and deltas processing cases
   if (!transaction_.md_deltas.empty()) {
     // This is usually done by the MVCC, but it does not handle the metadata deltas
     transaction_.EnsureCommitTimestampExists();
@@ -739,14 +686,14 @@ utils::BasicResult<StorageManipulationError, void> InMemoryStorage::InMemoryAcce
     // modifications before they are written to disk.
     // Replica can log only the write transaction received from Main
     // so the Wal files are consistent
-    if (replState.IsMain() || desired_commit_timestamp.has_value()) {
+    if (is_main || desired_commit_timestamp.has_value()) {
       could_replicate_all_sync_replicas =
           mem_storage->AppendToWalDataDefinition(transaction_, *commit_timestamp_);  // protected by engine_guard
       // TODO: release lock, and update all deltas to have a local copy of the commit timestamp
       transaction_.commit_timestamp->store(*commit_timestamp_, std::memory_order_release);  // protected by engine_guard
       // Replica can only update the last commit timestamp with
       // the commits received from main.
-      if (replState.IsMain() || desired_commit_timestamp.has_value()) {
+      if (is_main || desired_commit_timestamp.has_value()) {
         // Update the last commit timestamp
         mem_storage->repl_storage_state_.last_commit_timestamp_.store(*commit_timestamp_);  // protected by engine_guard
       }
@@ -820,7 +767,7 @@ utils::BasicResult<StorageManipulationError, void> InMemoryStorage::InMemoryAcce
         // modifications before they are written to disk.
         // Replica can log only the write transaction received from Main
         // so the Wal files are consistent
-        if (replState.IsMain() || desired_commit_timestamp.has_value()) {
+        if (is_main || desired_commit_timestamp.has_value()) {
           could_replicate_all_sync_replicas =
               mem_storage->AppendToWalDataManipulation(transaction_, *commit_timestamp_);  // protected by engine_guard
         }
@@ -831,7 +778,7 @@ utils::BasicResult<StorageManipulationError, void> InMemoryStorage::InMemoryAcce
                                              std::memory_order_release);  // protected by engine_guard
         // Replica can only update the last commit timestamp with
         // the commits received from main.
-        if (replState.IsMain() || desired_commit_timestamp.has_value()) {
+        if (is_main || desired_commit_timestamp.has_value()) {
           // Update the last commit timestamp
           mem_storage->repl_storage_state_.last_commit_timestamp_.store(
               *commit_timestamp_);  // protected by engine_guard
@@ -1195,7 +1142,7 @@ VerticesIterable InMemoryStorage::InMemoryAccessor::Vertices(
       mem_label_property_index->Vertices(label, property, lower_bound, upper_bound, view, storage_, &transaction_));
 }
 
-Transaction InMemoryStorage::CreateTransaction(IsolationLevel isolation_level, StorageMode storage_mode) {
+Transaction InMemoryStorage::CreateTransaction(IsolationLevel isolation_level, StorageMode storage_mode, bool is_main) {
   // We acquire the transaction engine lock here because we access (and
   // modify) the transaction engine variables (`transaction_id` and
   // `timestamp`) below.
@@ -1210,16 +1157,34 @@ Transaction InMemoryStorage::CreateTransaction(IsolationLevel isolation_level, S
     // of any query on replica to the last commited transaction
     // which is timestamp_ as only commit of transaction with writes
     // can change the value of it.
-    auto const &replState = repl_state_;
-    if (replState.IsReplica()) {
-      start_timestamp = timestamp_;
-    } else {
+    if (is_main) {
       start_timestamp = timestamp_++;
+    } else {
+      start_timestamp = timestamp_;
     }
   }
   return {transaction_id, start_timestamp, isolation_level, storage_mode, false};
 }
 
+void InMemoryStorage::SetStorageMode(StorageMode new_storage_mode) {
+  std::unique_lock main_guard{main_lock_};
+  MG_ASSERT(
+      (storage_mode_ == StorageMode::IN_MEMORY_ANALYTICAL || storage_mode_ == StorageMode::IN_MEMORY_TRANSACTIONAL) &&
+      (new_storage_mode == StorageMode::IN_MEMORY_ANALYTICAL ||
+       new_storage_mode == StorageMode::IN_MEMORY_TRANSACTIONAL));
+  if (storage_mode_ != new_storage_mode) {
+    if (new_storage_mode == StorageMode::IN_MEMORY_ANALYTICAL) {
+      snapshot_runner_.Stop();
+    } else if (config_.durability.snapshot_wal_mode != Config::Durability::SnapshotWalMode::DISABLED) {
+      snapshot_runner_.Run("Snapshot", config_.durability.snapshot_interval,
+                           [this]() { this->create_snapshot_handler(); });
+    }
+
+    storage_mode_ = new_storage_mode;
+    FreeMemory(std::move(main_guard));
+  }
+}
+
 template <bool force>
 void InMemoryStorage::CollectGarbage(std::unique_lock<utils::ResourceLock> main_guard) {
   // NOTE: You do not need to consider cleanup of deleted object that occurred in
@@ -1557,7 +1522,7 @@ StorageInfo InMemoryStorage::GetBaseInfo(bool force_directory) {
     // NOLINTNEXTLINE(bugprone-narrowing-conversions, cppcoreguidelines-narrowing-conversions)
     info.average_degree = 2.0 * static_cast<double>(info.edge_count) / info.vertex_count;
   }
-  info.memory_usage = utils::GetMemoryUsage();
+  info.memory_res = utils::GetMemoryRES();
   // Special case for the default database
   auto update_path = [&](const std::filesystem::path &dir) {
     if (!force_directory && std::filesystem::is_directory(dir) && dir.has_filename()) {
@@ -1629,8 +1594,7 @@ void InMemoryStorage::FinalizeWalFile() {
 }
 
 bool InMemoryStorage::AppendToWalDataManipulation(const Transaction &transaction, uint64_t final_commit_timestamp) {
-  auto &replState = repl_state_;
-  if (!InitializeWalFile(replState.GetEpoch())) {
+  if (!InitializeWalFile(repl_storage_state_.epoch_)) {
     return true;
   }
   // Traverse deltas and append them to the WAL file.
@@ -1800,8 +1764,7 @@ bool InMemoryStorage::AppendToWalDataManipulation(const Transaction &transaction
 }
 
 bool InMemoryStorage::AppendToWalDataDefinition(const Transaction &transaction, uint64_t final_commit_timestamp) {
-  auto &replState = repl_state_;
-  if (!InitializeWalFile(replState.GetEpoch())) {
+  if (!InitializeWalFile(repl_storage_state_.epoch_)) {
     return true;
   }
 
@@ -1910,13 +1873,8 @@ void InMemoryStorage::AppendToWalDataDefinition(durability::StorageMetadataOpera
   return AppendToWalDataDefinition(operation, label, {}, {}, final_commit_timestamp);
 }
 
-utils::BasicResult<InMemoryStorage::CreateSnapshotError> InMemoryStorage::CreateSnapshot(
-    memgraph::replication::ReplicationState const &replicationState, std::optional<bool> is_periodic) {
-  if (replicationState.IsReplica()) {
-    return CreateSnapshotError::DisabledForReplica;
-  }
-
-  auto const &epoch = replicationState.GetEpoch();
+utils::BasicResult<InMemoryStorage::CreateSnapshotError> InMemoryStorage::CreateSnapshot() {
+  auto const &epoch = repl_storage_state_.epoch_;
   auto snapshot_creator = [this, &epoch]() {
     utils::Timer timer;
     auto transaction = CreateTransaction(IsolationLevel::SNAPSHOT_ISOLATION, storage_mode_);
@@ -1943,9 +1901,6 @@ utils::BasicResult<InMemoryStorage::CreateSnapshotError> InMemoryStorage::Create
     } else {
       std::unique_lock main_guard{main_lock_};
       if (storage_mode_ == memgraph::storage::StorageMode::IN_MEMORY_ANALYTICAL) {
-        if (is_periodic && *is_periodic) {
-          return CreateSnapshotError::DisabledForAnalyticsPeriodicCommit;
-        }
         snapshot_creator();
         return {};
       }
@@ -1976,13 +1931,13 @@ uint64_t InMemoryStorage::CommitTimestamp(const std::optional<uint64_t> desired_
   return *desired_commit_timestamp;
 }
 
-void InMemoryStorage::PrepareForNewEpoch(std::string prev_epoch) {
+void InMemoryStorage::PrepareForNewEpoch() {
   std::unique_lock engine_guard{engine_lock_};
   if (wal_file_) {
     wal_file_->FinalizeWal();
     wal_file_.reset();
   }
-  repl_storage_state_.AddEpochToHistory(std::move(prev_epoch));
+  repl_storage_state_.TrackLatestHistory();
 }
 
 utils::FileRetainer::FileLockerAccessor::ret_type InMemoryStorage::IsPathLocked() {
@@ -2010,26 +1965,45 @@ utils::FileRetainer::FileLockerAccessor::ret_type InMemoryStorage::UnlockPath()
   return true;
 }
 
-auto InMemoryStorage::CreateReplicationClient(const memgraph::replication::ReplicationClientConfig &config)
+auto InMemoryStorage::CreateReplicationClient(const memgraph::replication::ReplicationClientConfig &config,
+                                              const memgraph::replication::ReplicationEpoch *current_epoch)
     -> std::unique_ptr<ReplicationClient> {
-  auto &replState = this->repl_state_;
-  return std::make_unique<InMemoryReplicationClient>(this, config, &replState.GetEpoch());
+  return std::make_unique<InMemoryReplicationClient>(this, config, current_epoch);
 }
 
-std::unique_ptr<ReplicationServer> InMemoryStorage::CreateReplicationServer(
-    const memgraph::replication::ReplicationServerConfig &config) {
-  auto &replState = this->repl_state_;
-  return std::make_unique<InMemoryReplicationServer>(this, config, &replState.GetEpoch());
+std::unique_ptr<Storage::Accessor> InMemoryStorage::Access(std::optional<IsolationLevel> override_isolation_level,
+                                                           bool is_main) {
+  return std::unique_ptr<InMemoryAccessor>(new InMemoryAccessor{Storage::Accessor::shared_access, this,
+                                                                override_isolation_level.value_or(isolation_level_),
+                                                                storage_mode_, is_main});
+}
+std::unique_ptr<Storage::Accessor> InMemoryStorage::UniqueAccess(std::optional<IsolationLevel> override_isolation_level,
+                                                                 bool is_main) {
+  return std::unique_ptr<InMemoryAccessor>(new InMemoryAccessor{Storage::Accessor::unique_access, this,
+                                                                override_isolation_level.value_or(isolation_level_),
+                                                                storage_mode_, is_main});
 }
 
-std::unique_ptr<Storage::Accessor> InMemoryStorage::Access(std::optional<IsolationLevel> override_isolation_level) {
-  return std::unique_ptr<InMemoryAccessor>(new InMemoryAccessor{
-      Storage::Accessor::shared_access, this, override_isolation_level.value_or(isolation_level_), storage_mode_});
-}
-std::unique_ptr<Storage::Accessor> InMemoryStorage::UniqueAccess(
-    std::optional<IsolationLevel> override_isolation_level) {
-  return std::unique_ptr<InMemoryAccessor>(new InMemoryAccessor{
-      Storage::Accessor::unique_access, this, override_isolation_level.value_or(isolation_level_), storage_mode_});
+void InMemoryStorage::CreateSnapshotHandler(
+    std::function<utils::BasicResult<InMemoryStorage::CreateSnapshotError>()> cb) {
+  create_snapshot_handler = [cb]() {
+    if (auto maybe_error = cb(); maybe_error.HasError()) {
+      switch (maybe_error.GetError()) {
+        case CreateSnapshotError::DisabledForReplica:
+          spdlog::warn(utils::MessageWithLink("Snapshots are disabled for replicas.", "https://memgr.ph/replication"));
+          break;
+        case CreateSnapshotError::ReachedMaxNumTries:
+          spdlog::warn("Failed to create snapshot. Reached max number of tries. Please contact support");
+          break;
+      }
+    }
+  };
+
+  // Run the snapshot thread (if enabled)
+  if (config_.durability.snapshot_wal_mode != Config::Durability::SnapshotWalMode::DISABLED) {
+    snapshot_runner_.Run("Snapshot", config_.durability.snapshot_interval,
+                         [this]() { this->create_snapshot_handler(); });
+  }
 }
 IndicesInfo InMemoryStorage::InMemoryAccessor::ListAllIndices() const {
   auto *in_memory = static_cast<InMemoryStorage *>(storage_);

src/storage/v2/inmemory/storage.hpp

@@ -31,6 +31,10 @@
 #include "utils/resource_lock.hpp"
 #include "utils/synchronized.hpp"
 
+namespace memgraph::dbms {
+class InMemoryReplicationHandlers;
+}
+
 namespace memgraph::storage {
 
 // The storage is based on this paper:
@@ -39,15 +43,11 @@ namespace memgraph::storage {
 // only implement snapshot isolation for transactions.
 
 class InMemoryStorage final : public Storage {
-  friend class InMemoryReplicationServer;
+  friend class memgraph::dbms::InMemoryReplicationHandlers;
   friend class InMemoryReplicationClient;
 
  public:
-  enum class CreateSnapshotError : uint8_t {
-    DisabledForReplica,
-    DisabledForAnalyticsPeriodicCommit,
-    ReachedMaxNumTries
-  };
+  enum class CreateSnapshotError : uint8_t { DisabledForReplica, ReachedMaxNumTries };
 
   /// @throw std::system_error
   /// @throw std::bad_alloc
@@ -66,7 +66,7 @@ class InMemoryStorage final : public Storage {
     friend class InMemoryStorage;
 
     explicit InMemoryAccessor(auto tag, InMemoryStorage *storage, IsolationLevel isolation_level,
-                              StorageMode storage_mode);
+                              StorageMode storage_mode, bool is_main = true);
 
    public:
     InMemoryAccessor(const InMemoryAccessor &) = delete;
@@ -204,8 +204,8 @@ class InMemoryStorage final : public Storage {
     /// case the transaction is automatically aborted.
     /// @throw std::bad_alloc
     // NOLINTNEXTLINE(google-default-arguments)
-    utils::BasicResult<StorageManipulationError, void> Commit(
-        std::optional<uint64_t> desired_commit_timestamp = {}) override;
+    utils::BasicResult<StorageManipulationError, void> Commit(std::optional<uint64_t> desired_commit_timestamp = {},
+                                                              bool is_main = true) override;
 
     /// @throw std::bad_alloc
     void Abort() override;
@@ -311,9 +311,13 @@ class InMemoryStorage final : public Storage {
     Transaction &GetTransaction() { return transaction_; }
   };
 
-  std::unique_ptr<Storage::Accessor> Access(std::optional<IsolationLevel> override_isolation_level) override;
+  using Storage::Access;
+  std::unique_ptr<Storage::Accessor> Access(std::optional<IsolationLevel> override_isolation_level,
+                                            bool is_main) override;
 
-  std::unique_ptr<Storage::Accessor> UniqueAccess(std::optional<IsolationLevel> override_isolation_level) override;
+  using Storage::UniqueAccess;
+  std::unique_ptr<Storage::Accessor> UniqueAccess(std::optional<IsolationLevel> override_isolation_level,
+                                                  bool is_main) override;
 
   void FreeMemory(std::unique_lock<utils::ResourceLock> main_guard) override;
 
@@ -321,16 +325,18 @@ class InMemoryStorage final : public Storage {
   utils::FileRetainer::FileLockerAccessor::ret_type LockPath();
   utils::FileRetainer::FileLockerAccessor::ret_type UnlockPath();
 
-  utils::BasicResult<InMemoryStorage::CreateSnapshotError> CreateSnapshot(
-      memgraph::replication::ReplicationState const &replicationState, std::optional<bool> is_periodic);
+  utils::BasicResult<InMemoryStorage::CreateSnapshotError> CreateSnapshot();
 
-  Transaction CreateTransaction(IsolationLevel isolation_level, StorageMode storage_mode) override;
+  void CreateSnapshotHandler(std::function<utils::BasicResult<InMemoryStorage::CreateSnapshotError>()> cb);
 
-  auto CreateReplicationClient(const memgraph::replication::ReplicationClientConfig &config)
+  using Storage::CreateTransaction;
+  Transaction CreateTransaction(IsolationLevel isolation_level, StorageMode storage_mode, bool is_main) override;
+
+  auto CreateReplicationClient(const memgraph::replication::ReplicationClientConfig &config,
+                               const memgraph::replication::ReplicationEpoch *current_epoch)
       -> std::unique_ptr<ReplicationClient> override;
 
-  auto CreateReplicationServer(const memgraph::replication::ReplicationServerConfig &config)
-      -> std::unique_ptr<ReplicationServer> override;
+  void SetStorageMode(StorageMode storage_mode);
 
  private:
   /// The force parameter determines the behaviour of the garbage collector.
@@ -377,7 +383,7 @@ class InMemoryStorage final : public Storage {
 
   uint64_t CommitTimestamp(std::optional<uint64_t> desired_commit_timestamp = {});
 
-  void PrepareForNewEpoch(std::string prev_epoch) override;
+  void PrepareForNewEpoch() override;
 
   // Main object storage
   utils::SkipList<storage::Vertex> vertices_;
@@ -444,6 +450,9 @@ class InMemoryStorage final : public Storage {
   // Flags to inform CollectGarbage that it needs to do the more expensive full scans
   std::atomic<bool> gc_full_scan_vertices_delete_ = false;
   std::atomic<bool> gc_full_scan_edges_delete_ = false;
+
+  // Moved the create snapshot to a user defined handler so we can remove the global replication state from the storage
+  std::function<void()> create_snapshot_handler{};
 };
 
 }  // namespace memgraph::storage

src/storage/v2/replication/replication_client.cpp

@@ -49,8 +49,8 @@ uint64_t ReplicationClient::LastCommitTimestamp() const {
 void ReplicationClient::InitializeClient() {
   uint64_t current_commit_timestamp{kTimestampInitialId};
 
-  auto stream{rpc_client_.Stream<replication::HeartbeatRpc>(storage_->repl_storage_state_.last_commit_timestamp_,
-                                                            std::string{repl_epoch_->id()})};
+  auto stream{rpc_client_.Stream<replication::HeartbeatRpc>(
+      storage_->id(), storage_->repl_storage_state_.last_commit_timestamp_, std::string{repl_epoch_->id()})};
 
   const auto replica = stream.AwaitResponse();
   std::optional<uint64_t> branching_point;
@@ -98,7 +98,7 @@ TimestampInfo ReplicationClient::GetTimestampInfo() {
   info.current_number_of_timestamp_behind_master = 0;
 
   try {
-    auto stream{rpc_client_.Stream<replication::TimestampRpc>()};
+    auto stream{rpc_client_.Stream<replication::TimestampRpc>(storage_->id())};
     const auto response = stream.AwaitResponse();
     const auto is_success = response.success;
     if (!is_success) {
@@ -135,8 +135,15 @@ void ReplicationClient::TryInitializeClientAsync() {
 void ReplicationClient::TryInitializeClientSync() {
   try {
     InitializeClient();
+  } catch (const rpc::VersionMismatchRpcFailedException &) {
+    std::unique_lock client_guard{client_lock_};
+    replica_state_.store(replication::ReplicaState::INVALID);
+    spdlog::error(
+        utils::MessageWithLink("Failed to connect to replica {} at the endpoint {}. Because the replica "
+                               "deployed is not a compatible version.",
+                               name_, rpc_client_.Endpoint(), "https://memgr.ph/replication"));
   } catch (const rpc::RpcFailedException &) {
-    std::unique_lock client_guarde{client_lock_};
+    std::unique_lock client_guard{client_lock_};
     replica_state_.store(replication::ReplicaState::INVALID);
     spdlog::error(utils::MessageWithLink("Failed to connect to replica {} at the endpoint {}.", name_,
                                          rpc_client_.Endpoint(), "https://memgr.ph/replication"));
@@ -222,7 +229,7 @@ bool ReplicationClient::FinalizeTransactionReplication() {
 void ReplicationClient::FrequentCheck() {
   const auto is_success = std::invoke([this]() {
     try {
-      auto stream{rpc_client_.Stream<replication::FrequentHeartbeatRpc>()};
+      auto stream{rpc_client_.Stream<memgraph::replication::FrequentHeartbeatRpc>()};
       const auto response = stream.AwaitResponse();
       return response.success;
     } catch (const rpc::RpcFailedException &) {
@@ -280,7 +287,8 @@ void ReplicationClient::IfStreamingTransaction(const std::function<void(ReplicaS
 ReplicaStream::ReplicaStream(ReplicationClient *self, const uint64_t previous_commit_timestamp,
                              const uint64_t current_seq_num)
     : self_(self),
-      stream_(self_->rpc_client_.Stream<replication::AppendDeltasRpc>(previous_commit_timestamp, current_seq_num)) {
+      stream_(self_->rpc_client_.Stream<replication::AppendDeltasRpc>(self->GetStorageId(), previous_commit_timestamp,
+                                                                      current_seq_num)) {
   replication::Encoder encoder{stream_.GetBuilder()};
 
   encoder.WriteString(self->repl_epoch_->id());
@@ -312,4 +320,5 @@ void ReplicaStream::AppendOperation(durability::StorageMetadataOperation operati
 
 replication::AppendDeltasRes ReplicaStream::Finalize() { return stream_.AwaitResponse(); }
 
+auto ReplicationClient::GetStorageId() const -> std::string { return storage_->id(); }
 }  // namespace memgraph::storage

src/storage/v2/replication/replication_client.hpp

@@ -29,6 +29,7 @@
 #include <optional>
 #include <set>
 #include <string>
+#include <variant>
 
 namespace memgraph::storage {
 
@@ -86,6 +87,8 @@ class ReplicationClient {
   auto State() const -> replication::ReplicaState { return replica_state_.load(); }
   auto GetTimestampInfo() -> TimestampInfo;
 
+  auto GetStorageId() const -> std::string;
+
   void Start();
   void StartTransactionReplication(const uint64_t current_wal_seq_num);
   // Replication clients can be removed at any point

src/storage/v2/replication/replication_handler.cpp

@@ -1,209 +0,0 @@
-// Copyright 2023 Memgraph Ltd.
-//
-// Use of this software is governed by the Business Source License
-// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
-// License, and you may not use this file except in compliance with the Business Source License.
-//
-// As of the Change Date specified in that file, in accordance with
-// the Business Source License, use of this software will be governed
-// by the Apache License, Version 2.0, included in the file
-// licenses/APL.txt.
-
-#include "storage/v2/replication/replication_handler.hpp"
-
-#include "replication/state.hpp"
-#include "storage/v2/storage.hpp"
-
-namespace memgraph::storage {
-
-namespace {
-
-std::string RegisterReplicaErrorToString(RegisterReplicaError error) {
-  switch (error) {
-    using enum RegisterReplicaError;
-    case NAME_EXISTS:
-      return "NAME_EXISTS";
-    case END_POINT_EXISTS:
-      return "END_POINT_EXISTS";
-    case CONNECTION_FAILED:
-      return "CONNECTION_FAILED";
-    case COULD_NOT_BE_PERSISTED:
-      return "COULD_NOT_BE_PERSISTED";
-  }
-}
-}  // namespace
-
-bool ReplicationHandler::SetReplicationRoleMain() {
-  // We don't want to generate new epoch_id and do the
-  // cleanup if we're already a MAIN
-  // TODO: under lock
-  if (repl_state_.IsMain()) {
-    return false;
-  }
-
-  // STEP 1) bring down all REPLICA servers
-  auto current_epoch = std::string(repl_state_.GetEpoch().id());
-  {  // TODO: foreach storage
-    // ensure replica server brought down
-    storage_.repl_storage_state_.replication_server_.reset(nullptr);
-    // Remember old epoch + storage timestamp association
-    storage_.PrepareForNewEpoch(current_epoch);
-  }
-
-  // STEP 2) Change to MAIN
-  repl_state_.GetEpoch().NewEpoch();
-  if (!repl_state_.TryPersistRoleMain()) {
-    // TODO: On failure restore old epoch? restore replication servers?
-    return false;
-  }
-  repl_state_.SetRole(memgraph::replication::ReplicationRole::MAIN);
-  return true;
-}
-memgraph::utils::BasicResult<RegisterReplicaError> ReplicationHandler::RegisterReplica(
-    const RegistrationMode registration_mode, const memgraph::replication::ReplicationClientConfig &config) {
-  MG_ASSERT(repl_state_.IsMain(), "Only main instance can register a replica!");
-
-  auto name_check = [&config](auto &clients) {
-    auto name_matches = [&name = config.name](const auto &client) { return client->Name() == name; };
-    return std::any_of(clients.begin(), clients.end(), name_matches);
-  };
-
-  io::network::Endpoint desired_endpoint;
-  if (io::network::Endpoint::GetIpFamily(config.ip_address) == io::network::Endpoint::IpFamily::NONE) {
-    desired_endpoint = io::network::Endpoint{io::network::Endpoint::needs_resolving, config.ip_address, config.port};
-  } else {
-    desired_endpoint = io::network::Endpoint{config.ip_address, config.port};
-  }
-  auto endpoint_check = [&](auto &clients) {
-    auto endpoint_matches = [&](const auto &client) { return client->Endpoint() == desired_endpoint; };
-    return std::any_of(clients.begin(), clients.end(), endpoint_matches);
-  };
-
-  auto task = [&](auto &clients) -> utils::BasicResult<RegisterReplicaError> {
-    if (name_check(clients)) {
-      return RegisterReplicaError::NAME_EXISTS;
-    }
-
-    if (endpoint_check(clients)) {
-      return RegisterReplicaError::END_POINT_EXISTS;
-    }
-
-    using enum RegistrationMode;
-    if (registration_mode != RESTORE && !repl_state_.TryPersistRegisteredReplica(config)) {
-      return RegisterReplicaError::COULD_NOT_BE_PERSISTED;
-    }
-
-    auto client = storage_.CreateReplicationClient(config);
-    client->Start();
-
-    if (client->State() == replication::ReplicaState::INVALID) {
-      if (registration_mode != RESTORE) {
-        return RegisterReplicaError::CONNECTION_FAILED;
-      }
-
-      spdlog::warn("Connection failed when registering replica {}. Replica will still be registered.", client->Name());
-    }
-
-    clients.push_back(std::move(client));
-    return {};
-  };
-
-  return storage_.repl_storage_state_.replication_clients_.WithLock(task);
-}
-bool ReplicationHandler::SetReplicationRoleReplica(const memgraph::replication::ReplicationServerConfig &config) {
-  // We don't want to restart the server if we're already a REPLICA
-  if (repl_state_.IsReplica()) {
-    return false;
-  }
-
-  std::unique_ptr<ReplicationServer> replication_server = storage_.CreateReplicationServer(config);
-  bool res = replication_server->Start();
-  if (!res) {
-    spdlog::error("Unable to start the replication server.");
-    return false;
-  }
-  storage_.repl_storage_state_.replication_server_ = std::move(replication_server);
-
-  if (!repl_state_.TryPersistRoleReplica(config)) {
-    return false;
-  }
-
-  repl_state_.SetRole(memgraph::replication::ReplicationRole::REPLICA);
-  return true;
-}
-auto ReplicationHandler::UnregisterReplica(std::string_view name) -> UnregisterReplicaResult {
-  if (repl_state_.IsReplica()) {
-    return UnregisterReplicaResult::NOT_MAIN;
-  }
-
-  if (!repl_state_.TryPersistUnregisterReplica(name)) {
-    return UnregisterReplicaResult::COULD_NOT_BE_PERSISTED;
-  }
-
-  auto const n_unregistered = storage_.repl_storage_state_.replication_clients_.WithLock([&](auto &clients) {
-    return std::erase_if(clients, [&](const auto &client) { return client->Name() == name; });
-  });
-  return (n_unregistered != 0) ? UnregisterReplicaResult::SUCCESS : UnregisterReplicaResult::CAN_NOT_UNREGISTER;
-}
-void ReplicationHandler::RestoreReplication() {
-  if (!repl_state_.ShouldPersist()) {
-    return;
-  }
-
-  spdlog::info("Restoring replication role.");
-
-  using memgraph::replication::ReplicationState;
-
-  auto replicationData = repl_state_.FetchReplicationData();
-  if (replicationData.HasError()) {
-    switch (replicationData.GetError()) {
-      using enum ReplicationState::FetchReplicationError;
-      case NOTHING_FETCHED: {
-        spdlog::debug("Cannot find data needed for restore replication role in persisted metadata.");
-        return;
-      }
-      case PARSE_ERROR: {
-        LOG_FATAL("Cannot parse previously saved configuration of replication role.");
-        return;
-      }
-    }
-  }
-
-  /// MAIN
-  auto const recover_main = [this](ReplicationState::ReplicationDataMain const &configs) {
-    storage_.repl_storage_state_.replication_server_.reset();
-    repl_state_.SetRole(memgraph::replication::ReplicationRole::MAIN);
-    for (const auto &config : configs) {
-      spdlog::info("Replica {} restored for {}.", config.name, storage_.id());
-      auto ret = RegisterReplica(RegistrationMode::RESTORE, config);
-      if (ret.HasError()) {
-        MG_ASSERT(RegisterReplicaError::CONNECTION_FAILED != ret.GetError());
-        LOG_FATAL("Failure when restoring replica {}: {}.", config.name, RegisterReplicaErrorToString(ret.GetError()));
-      }
-      spdlog::info("Replica {} restored for {}.", config.name, storage_.id());
-    }
-    spdlog::info("Replication role restored to MAIN.");
-  };
-
-  /// REPLICA
-  auto const recover_replica = [this](ReplicationState::ReplicationDataReplica const &config) {
-    auto replication_server = storage_.CreateReplicationServer(config);
-    if (!replication_server->Start()) {
-      LOG_FATAL("Unable to start the replication server.");
-    }
-    storage_.repl_storage_state_.replication_server_ = std::move(replication_server);
-    repl_state_.SetRole(memgraph::replication::ReplicationRole::REPLICA);
-    spdlog::info("Replication role restored to REPLICA.");
-  };
-
-  std::visit(
-      utils::Overloaded{
-          recover_main,
-          recover_replica,
-      },
-      *replicationData);
-}
-auto ReplicationHandler::GetRole() const -> memgraph::replication::ReplicationRole { return repl_state_.GetRole(); }
-bool ReplicationHandler::IsMain() const { return repl_state_.IsMain(); }
-bool ReplicationHandler::IsReplica() const { return repl_state_.IsReplica(); }
-}  // namespace memgraph::storage

src/storage/v2/replication/replication_server.cpp

@@ -1,61 +0,0 @@
-// Copyright 2023 Memgraph Ltd.
-//
-// Use of this software is governed by the Business Source License
-// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
-// License, and you may not use this file except in compliance with the Business Source License.
-//
-// As of the Change Date specified in that file, in accordance with
-// the Business Source License, use of this software will be governed
-// by the Apache License, Version 2.0, included in the file
-// licenses/APL.txt.
-
-#include "replication_server.hpp"
-#include "io/network/endpoint.hpp"
-#include "replication/config.hpp"
-#include "rpc.hpp"
-
-namespace memgraph::storage {
-namespace {
-
-auto CreateServerContext(const memgraph::replication::ReplicationServerConfig &config) -> communication::ServerContext {
-  return (config.ssl) ? communication::ServerContext{config.ssl->key_file, config.ssl->cert_file, config.ssl->ca_file,
-                                                     config.ssl->verify_peer}
-                      : communication::ServerContext{};
-}
-
-// NOTE: The replication server must have a single thread for processing
-// because there is no need for more processing threads - each replica can
-// have only a single main server. Also, the single-threaded guarantee
-// simplifies the rest of the implementation.
-constexpr auto kReplictionServerThreads = 1;
-}  // namespace
-
-ReplicationServer::ReplicationServer(const memgraph::replication::ReplicationServerConfig &config)
-    : rpc_server_context_{CreateServerContext(config)},
-      rpc_server_{io::network::Endpoint{config.ip_address, config.port}, &rpc_server_context_,
-                  kReplictionServerThreads} {
-  rpc_server_.Register<replication::FrequentHeartbeatRpc>([](auto *req_reader, auto *res_builder) {
-    spdlog::debug("Received FrequentHeartbeatRpc");
-    FrequentHeartbeatHandler(req_reader, res_builder);
-  });
-}
-
-ReplicationServer::~ReplicationServer() {
-  if (rpc_server_.IsRunning()) {
-    auto const &endpoint = rpc_server_.endpoint();
-    spdlog::trace("Closing replication server on {}:{}", endpoint.address, endpoint.port);
-    rpc_server_.Shutdown();
-  }
-  rpc_server_.AwaitShutdown();
-}
-
-bool ReplicationServer::Start() { return rpc_server_.Start(); }
-
-void ReplicationServer::FrequentHeartbeatHandler(slk::Reader *req_reader, slk::Builder *res_builder) {
-  replication::FrequentHeartbeatReq req;
-  slk::Load(&req, req_reader);
-  replication::FrequentHeartbeatRes res{true};
-  slk::Save(res, res_builder);
-}
-
-}  // namespace memgraph::storage

src/storage/v2/replication/replication_storage_state.cpp

@@ -11,8 +11,8 @@
 
 #include "storage/v2/replication/replication_storage_state.hpp"
 
+#include "replication/replication_server.hpp"
 #include "storage/v2/replication/replication_client.hpp"
-#include "storage/v2/replication/replication_server.hpp"
 
 namespace memgraph::storage {
 
@@ -91,17 +91,16 @@ std::vector<ReplicaInfo> ReplicationStorageState::ReplicasInfo() const {
 }
 
 void ReplicationStorageState::Reset() {
-  replication_server_.reset();
   replication_clients_.WithLock([](auto &clients) { clients.clear(); });
 }
 
-void ReplicationStorageState::AddEpochToHistory(std::string prev_epoch) {
+void ReplicationStorageState::TrackLatestHistory() {
   constexpr uint16_t kEpochHistoryRetention = 1000;
   // Generate new epoch id and save the last one to the history.
   if (history.size() == kEpochHistoryRetention) {
     history.pop_front();
   }
-  history.emplace_back(std::move(prev_epoch), last_commit_timestamp_);
+  history.emplace_back(epoch_.id(), last_commit_timestamp_);
 }
 
 void ReplicationStorageState::AddEpochToHistoryForce(std::string prev_epoch) {

src/storage/v2/replication/replication_storage_state.hpp

@@ -27,11 +27,12 @@
 #include "storage/v2/replication/global.hpp"
 #include "storage/v2/replication/rpc.hpp"
 #include "storage/v2/replication/serialization.hpp"
+#include "utils/synchronized.hpp"
 
 namespace memgraph::storage {
 
 class Storage;
-class ReplicationServer;
+
 class ReplicationClient;
 
 struct ReplicationStorageState {
@@ -49,7 +50,7 @@ struct ReplicationStorageState {
   auto ReplicasInfo() const -> std::vector<ReplicaInfo>;
 
   // History
-  void AddEpochToHistory(std::string prev_epoch);
+  void TrackLatestHistory();
   void AddEpochToHistoryForce(std::string prev_epoch);
 
   void Reset();
@@ -76,10 +77,9 @@ struct ReplicationStorageState {
   using ReplicationClientPtr = std::unique_ptr<ReplicationClient>;
   using ReplicationClientList = utils::Synchronized<std::vector<ReplicationClientPtr>, utils::RWSpinLock>;
 
-  // NOTE: Server is not in MAIN it is in REPLICA
-  std::unique_ptr<ReplicationServer> replication_server_{nullptr};
-
   ReplicationClientList replication_clients_;
+
+  memgraph::replication::ReplicationEpoch epoch_;
 };
 
 }  // namespace memgraph::storage

src/storage/v2/replication/rpc.cpp

@@ -34,18 +34,6 @@ void HeartbeatRes::Save(const HeartbeatRes &self, memgraph::slk::Builder *builde
   memgraph::slk::Save(self, builder);
 }
 void HeartbeatRes::Load(HeartbeatRes *self, memgraph::slk::Reader *reader) { memgraph::slk::Load(self, reader); }
-void FrequentHeartbeatReq::Save(const FrequentHeartbeatReq &self, memgraph::slk::Builder *builder) {
-  memgraph::slk::Save(self, builder);
-}
-void FrequentHeartbeatReq::Load(FrequentHeartbeatReq *self, memgraph::slk::Reader *reader) {
-  memgraph::slk::Load(self, reader);
-}
-void FrequentHeartbeatRes::Save(const FrequentHeartbeatRes &self, memgraph::slk::Builder *builder) {
-  memgraph::slk::Save(self, builder);
-}
-void FrequentHeartbeatRes::Load(FrequentHeartbeatRes *self, memgraph::slk::Reader *reader) {
-  memgraph::slk::Load(self, reader);
-}
 void SnapshotReq::Save(const SnapshotReq &self, memgraph::slk::Builder *builder) { memgraph::slk::Save(self, builder); }
 void SnapshotReq::Load(SnapshotReq *self, memgraph::slk::Reader *reader) { memgraph::slk::Load(self, reader); }
 void SnapshotRes::Save(const SnapshotRes &self, memgraph::slk::Builder *builder) { memgraph::slk::Save(self, builder); }
@@ -86,12 +74,6 @@ constexpr utils::TypeInfo storage::replication::HeartbeatReq::kType{utils::TypeI
 constexpr utils::TypeInfo storage::replication::HeartbeatRes::kType{utils::TypeId::REP_HEARTBEAT_RES, "HeartbeatRes",
                                                                     nullptr};
 
-constexpr utils::TypeInfo storage::replication::FrequentHeartbeatReq::kType{utils::TypeId::REP_FREQUENT_HEARTBEAT_REQ,
-                                                                            "FrequentHeartbeatReq", nullptr};
-
-constexpr utils::TypeInfo storage::replication::FrequentHeartbeatRes::kType{utils::TypeId::REP_FREQUENT_HEARTBEAT_RES,
-                                                                            "FrequentHeartbeatRes", nullptr};
-
 constexpr utils::TypeInfo storage::replication::SnapshotReq::kType{utils::TypeId::REP_SNAPSHOT_REQ, "SnapshotReq",
                                                                    nullptr};
 
@@ -121,47 +103,61 @@ namespace slk {
 // Serialize code for TimestampRes
 
 void Save(const memgraph::storage::replication::TimestampRes &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
   memgraph::slk::Save(self.success, builder);
   memgraph::slk::Save(self.current_commit_timestamp, builder);
 }
 
 void Load(memgraph::storage::replication::TimestampRes *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
   memgraph::slk::Load(&self->success, reader);
   memgraph::slk::Load(&self->current_commit_timestamp, reader);
 }
 
 // Serialize code for TimestampReq
 
-void Save(const memgraph::storage::replication::TimestampReq &self, memgraph::slk::Builder *builder) {}
+void Save(const memgraph::storage::replication::TimestampReq &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
+}
 
-void Load(memgraph::storage::replication::TimestampReq *self, memgraph::slk::Reader *reader) {}
+void Load(memgraph::storage::replication::TimestampReq *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
+}
 
 // Serialize code for CurrentWalRes
 
 void Save(const memgraph::storage::replication::CurrentWalRes &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
   memgraph::slk::Save(self.success, builder);
   memgraph::slk::Save(self.current_commit_timestamp, builder);
 }
 
 void Load(memgraph::storage::replication::CurrentWalRes *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
   memgraph::slk::Load(&self->success, reader);
   memgraph::slk::Load(&self->current_commit_timestamp, reader);
 }
 
 // Serialize code for CurrentWalReq
 
-void Save(const memgraph::storage::replication::CurrentWalReq &self, memgraph::slk::Builder *builder) {}
+void Save(const memgraph::storage::replication::CurrentWalReq &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
+}
 
-void Load(memgraph::storage::replication::CurrentWalReq *self, memgraph::slk::Reader *reader) {}
+void Load(memgraph::storage::replication::CurrentWalReq *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
+}
 
 // Serialize code for WalFilesRes
 
 void Save(const memgraph::storage::replication::WalFilesRes &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
   memgraph::slk::Save(self.success, builder);
   memgraph::slk::Save(self.current_commit_timestamp, builder);
 }
 
 void Load(memgraph::storage::replication::WalFilesRes *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
   memgraph::slk::Load(&self->success, reader);
   memgraph::slk::Load(&self->current_commit_timestamp, reader);
 }
@@ -169,56 +165,50 @@ void Load(memgraph::storage::replication::WalFilesRes *self, memgraph::slk::Read
 // Serialize code for WalFilesReq
 
 void Save(const memgraph::storage::replication::WalFilesReq &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
   memgraph::slk::Save(self.file_number, builder);
 }
 
 void Load(memgraph::storage::replication::WalFilesReq *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
   memgraph::slk::Load(&self->file_number, reader);
 }
 
 // Serialize code for SnapshotRes
 
 void Save(const memgraph::storage::replication::SnapshotRes &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
   memgraph::slk::Save(self.success, builder);
   memgraph::slk::Save(self.current_commit_timestamp, builder);
 }
 
 void Load(memgraph::storage::replication::SnapshotRes *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
   memgraph::slk::Load(&self->success, reader);
   memgraph::slk::Load(&self->current_commit_timestamp, reader);
 }
 
 // Serialize code for SnapshotReq
 
-void Save(const memgraph::storage::replication::SnapshotReq &self, memgraph::slk::Builder *builder) {}
-
-void Load(memgraph::storage::replication::SnapshotReq *self, memgraph::slk::Reader *reader) {}
-
-// Serialize code for FrequentHeartbeatRes
-
-void Save(const memgraph::storage::replication::FrequentHeartbeatRes &self, memgraph::slk::Builder *builder) {
-  memgraph::slk::Save(self.success, builder);
+void Save(const memgraph::storage::replication::SnapshotReq &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
 }
 
-void Load(memgraph::storage::replication::FrequentHeartbeatRes *self, memgraph::slk::Reader *reader) {
-  memgraph::slk::Load(&self->success, reader);
+void Load(memgraph::storage::replication::SnapshotReq *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
 }
 
-// Serialize code for FrequentHeartbeatReq
-
-void Save(const memgraph::storage::replication::FrequentHeartbeatReq &self, memgraph::slk::Builder *builder) {}
-
-void Load(memgraph::storage::replication::FrequentHeartbeatReq *self, memgraph::slk::Reader *reader) {}
-
 // Serialize code for HeartbeatRes
 
 void Save(const memgraph::storage::replication::HeartbeatRes &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
   memgraph::slk::Save(self.success, builder);
   memgraph::slk::Save(self.current_commit_timestamp, builder);
   memgraph::slk::Save(self.epoch_id, builder);
 }
 
 void Load(memgraph::storage::replication::HeartbeatRes *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
   memgraph::slk::Load(&self->success, reader);
   memgraph::slk::Load(&self->current_commit_timestamp, reader);
   memgraph::slk::Load(&self->epoch_id, reader);
@@ -227,11 +217,13 @@ void Load(memgraph::storage::replication::HeartbeatRes *self, memgraph::slk::Rea
 // Serialize code for HeartbeatReq
 
 void Save(const memgraph::storage::replication::HeartbeatReq &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
   memgraph::slk::Save(self.main_commit_timestamp, builder);
   memgraph::slk::Save(self.epoch_id, builder);
 }
 
 void Load(memgraph::storage::replication::HeartbeatReq *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
   memgraph::slk::Load(&self->main_commit_timestamp, reader);
   memgraph::slk::Load(&self->epoch_id, reader);
 }
@@ -239,11 +231,13 @@ void Load(memgraph::storage::replication::HeartbeatReq *self, memgraph::slk::Rea
 // Serialize code for AppendDeltasRes
 
 void Save(const memgraph::storage::replication::AppendDeltasRes &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
   memgraph::slk::Save(self.success, builder);
   memgraph::slk::Save(self.current_commit_timestamp, builder);
 }
 
 void Load(memgraph::storage::replication::AppendDeltasRes *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
   memgraph::slk::Load(&self->success, reader);
   memgraph::slk::Load(&self->current_commit_timestamp, reader);
 }
@@ -251,11 +245,13 @@ void Load(memgraph::storage::replication::AppendDeltasRes *self, memgraph::slk::
 // Serialize code for AppendDeltasReq
 
 void Save(const memgraph::storage::replication::AppendDeltasReq &self, memgraph::slk::Builder *builder) {
+  memgraph::slk::Save(self.db_name, builder);
   memgraph::slk::Save(self.previous_commit_timestamp, builder);
   memgraph::slk::Save(self.seq_num, builder);
 }
 
 void Load(memgraph::storage::replication::AppendDeltasReq *self, memgraph::slk::Reader *reader) {
+  memgraph::slk::Load(&self->db_name, reader);
   memgraph::slk::Load(&self->previous_commit_timestamp, reader);
   memgraph::slk::Load(&self->seq_num, reader);
 }

src/storage/v2/replication/rpc.hpp

@@ -32,9 +32,10 @@ struct AppendDeltasReq {
   static void Load(AppendDeltasReq *self, memgraph::slk::Reader *reader);
   static void Save(const AppendDeltasReq &self, memgraph::slk::Builder *builder);
   AppendDeltasReq() {}
-  AppendDeltasReq(uint64_t previous_commit_timestamp, uint64_t seq_num)
-      : previous_commit_timestamp(previous_commit_timestamp), seq_num(seq_num) {}
+  AppendDeltasReq(std::string name, uint64_t previous_commit_timestamp, uint64_t seq_num)
+      : db_name(std::move(name)), previous_commit_timestamp(previous_commit_timestamp), seq_num(seq_num) {}
 
+  std::string db_name;
   uint64_t previous_commit_timestamp;
   uint64_t seq_num;
 };
@@ -46,9 +47,10 @@ struct AppendDeltasRes {
   static void Load(AppendDeltasRes *self, memgraph::slk::Reader *reader);
   static void Save(const AppendDeltasRes &self, memgraph::slk::Builder *builder);
   AppendDeltasRes() {}
-  AppendDeltasRes(bool success, uint64_t current_commit_timestamp)
-      : success(success), current_commit_timestamp(current_commit_timestamp) {}
+  AppendDeltasRes(std::string name, bool success, uint64_t current_commit_timestamp)
+      : db_name(std::move(name)), success(success), current_commit_timestamp(current_commit_timestamp) {}
 
+  std::string db_name;
   bool success;
   uint64_t current_commit_timestamp;
 };
@@ -62,9 +64,10 @@ struct HeartbeatReq {
   static void Load(HeartbeatReq *self, memgraph::slk::Reader *reader);
   static void Save(const HeartbeatReq &self, memgraph::slk::Builder *builder);
   HeartbeatReq() {}
-  HeartbeatReq(uint64_t main_commit_timestamp, std::string epoch_id)
-      : main_commit_timestamp(main_commit_timestamp), epoch_id(std::move(epoch_id)) {}
+  HeartbeatReq(std::string name, uint64_t main_commit_timestamp, std::string epoch_id)
+      : db_name(std::move(name)), main_commit_timestamp(main_commit_timestamp), epoch_id(std::move(epoch_id)) {}
 
+  std::string db_name;
   uint64_t main_commit_timestamp;
   std::string epoch_id;
 };
@@ -76,9 +79,13 @@ struct HeartbeatRes {
   static void Load(HeartbeatRes *self, memgraph::slk::Reader *reader);
   static void Save(const HeartbeatRes &self, memgraph::slk::Builder *builder);
   HeartbeatRes() {}
-  HeartbeatRes(bool success, uint64_t current_commit_timestamp, std::string epoch_id)
-      : success(success), current_commit_timestamp(current_commit_timestamp), epoch_id(epoch_id) {}
+  HeartbeatRes(std::string name, bool success, uint64_t current_commit_timestamp, std::string epoch_id)
+      : db_name(std::move(name)),
+        success(success),
+        current_commit_timestamp(current_commit_timestamp),
+        epoch_id(epoch_id) {}
 
+  std::string db_name;
   bool success;
   uint64_t current_commit_timestamp;
   std::string epoch_id;
@@ -86,29 +93,6 @@ struct HeartbeatRes {
 
 using HeartbeatRpc = rpc::RequestResponse<HeartbeatReq, HeartbeatRes>;
 
-struct FrequentHeartbeatReq {
-  static const utils::TypeInfo kType;
-  static const utils::TypeInfo &GetTypeInfo() { return kType; }
-
-  static void Load(FrequentHeartbeatReq *self, memgraph::slk::Reader *reader);
-  static void Save(const FrequentHeartbeatReq &self, memgraph::slk::Builder *builder);
-  FrequentHeartbeatReq() {}
-};
-
-struct FrequentHeartbeatRes {
-  static const utils::TypeInfo kType;
-  static const utils::TypeInfo &GetTypeInfo() { return kType; }
-
-  static void Load(FrequentHeartbeatRes *self, memgraph::slk::Reader *reader);
-  static void Save(const FrequentHeartbeatRes &self, memgraph::slk::Builder *builder);
-  FrequentHeartbeatRes() {}
-  explicit FrequentHeartbeatRes(bool success) : success(success) {}
-
-  bool success;
-};
-
-using FrequentHeartbeatRpc = rpc::RequestResponse<FrequentHeartbeatReq, FrequentHeartbeatRes>;
-
 struct SnapshotReq {
   static const utils::TypeInfo kType;
   static const utils::TypeInfo &GetTypeInfo() { return kType; }
@@ -116,6 +100,9 @@ struct SnapshotReq {
   static void Load(SnapshotReq *self, memgraph::slk::Reader *reader);
   static void Save(const SnapshotReq &self, memgraph::slk::Builder *builder);
   SnapshotReq() {}
+  explicit SnapshotReq(std::string name) : db_name(std::move(name)) {}
+
+  std::string db_name;
 };
 
 struct SnapshotRes {
@@ -125,9 +112,10 @@ struct SnapshotRes {
   static void Load(SnapshotRes *self, memgraph::slk::Reader *reader);
   static void Save(const SnapshotRes &self, memgraph::slk::Builder *builder);
   SnapshotRes() {}
-  SnapshotRes(bool success, uint64_t current_commit_timestamp)
-      : success(success), current_commit_timestamp(current_commit_timestamp) {}
+  SnapshotRes(std::string name, bool success, uint64_t current_commit_timestamp)
+      : db_name(std::move(name)), success(success), current_commit_timestamp(current_commit_timestamp) {}
 
+  std::string db_name;
   bool success;
   uint64_t current_commit_timestamp;
 };
@@ -141,8 +129,9 @@ struct WalFilesReq {
   static void Load(WalFilesReq *self, memgraph::slk::Reader *reader);
   static void Save(const WalFilesReq &self, memgraph::slk::Builder *builder);
   WalFilesReq() {}
-  explicit WalFilesReq(uint64_t file_number) : file_number(file_number) {}
+  explicit WalFilesReq(std::string name, uint64_t file_number) : db_name(std::move(name)), file_number(file_number) {}
 
+  std::string db_name;
   uint64_t file_number;
 };
 
@@ -153,9 +142,10 @@ struct WalFilesRes {
   static void Load(WalFilesRes *self, memgraph::slk::Reader *reader);
   static void Save(const WalFilesRes &self, memgraph::slk::Builder *builder);
   WalFilesRes() {}
-  WalFilesRes(bool success, uint64_t current_commit_timestamp)
-      : success(success), current_commit_timestamp(current_commit_timestamp) {}
+  WalFilesRes(std::string name, bool success, uint64_t current_commit_timestamp)
+      : db_name(std::move(name)), success(success), current_commit_timestamp(current_commit_timestamp) {}
 
+  std::string db_name;
   bool success;
   uint64_t current_commit_timestamp;
 };
@@ -169,6 +159,9 @@ struct CurrentWalReq {
   static void Load(CurrentWalReq *self, memgraph::slk::Reader *reader);
   static void Save(const CurrentWalReq &self, memgraph::slk::Builder *builder);
   CurrentWalReq() {}
+  explicit CurrentWalReq(std::string name) : db_name(std::move(name)) {}
+
+  std::string db_name;
 };
 
 struct CurrentWalRes {
@@ -178,9 +171,10 @@ struct CurrentWalRes {
   static void Load(CurrentWalRes *self, memgraph::slk::Reader *reader);
   static void Save(const CurrentWalRes &self, memgraph::slk::Builder *builder);
   CurrentWalRes() {}
-  CurrentWalRes(bool success, uint64_t current_commit_timestamp)
-      : success(success), current_commit_timestamp(current_commit_timestamp) {}
+  CurrentWalRes(std::string name, bool success, uint64_t current_commit_timestamp)
+      : db_name(std::move(name)), success(success), current_commit_timestamp(current_commit_timestamp) {}
 
+  std::string db_name;
   bool success;
   uint64_t current_commit_timestamp;
 };
@@ -194,6 +188,9 @@ struct TimestampReq {
   static void Load(TimestampReq *self, memgraph::slk::Reader *reader);
   static void Save(const TimestampReq &self, memgraph::slk::Builder *builder);
   TimestampReq() {}
+  explicit TimestampReq(std::string name) : db_name(std::move(name)) {}
+
+  std::string db_name;
 };
 
 struct TimestampRes {
@@ -203,9 +200,10 @@ struct TimestampRes {
   static void Load(TimestampRes *self, memgraph::slk::Reader *reader);
   static void Save(const TimestampRes &self, memgraph::slk::Builder *builder);
   TimestampRes() {}
-  TimestampRes(bool success, uint64_t current_commit_timestamp)
-      : success(success), current_commit_timestamp(current_commit_timestamp) {}
+  TimestampRes(std::string name, bool success, uint64_t current_commit_timestamp)
+      : db_name(std::move(name)), success(success), current_commit_timestamp(current_commit_timestamp) {}
 
+  std::string db_name;
   bool success;
   uint64_t current_commit_timestamp;
 };
@@ -251,14 +249,6 @@ void Save(const memgraph::storage::replication::SnapshotReq &self, memgraph::slk
 
 void Load(memgraph::storage::replication::SnapshotReq *self, memgraph::slk::Reader *reader);
 
-void Save(const memgraph::storage::replication::FrequentHeartbeatRes &self, memgraph::slk::Builder *builder);
-
-void Load(memgraph::storage::replication::FrequentHeartbeatRes *self, memgraph::slk::Reader *reader);
-
-void Save(const memgraph::storage::replication::FrequentHeartbeatReq &self, memgraph::slk::Builder *builder);
-
-void Load(memgraph::storage::replication::FrequentHeartbeatReq *self, memgraph::slk::Reader *reader);
-
 void Save(const memgraph::storage::replication::HeartbeatRes &self, memgraph::slk::Builder *builder);
 
 void Load(memgraph::storage::replication::HeartbeatRes *self, memgraph::slk::Reader *reader);

src/storage/v2/storage.cpp

@@ -31,14 +31,8 @@ namespace memgraph::storage {
 
 class InMemoryStorage;
 
-auto ReplicationStateHelper(Config const &config) -> std::optional<std::filesystem::path> {
-  if (!config.durability.restore_replication_state_on_startup) return std::nullopt;
-  return {config.durability.storage_directory};
-}
-
 Storage::Storage(Config config, StorageMode storage_mode)
-    : repl_state_(ReplicationStateHelper(config)),
-      name_id_mapper_(std::invoke([config, storage_mode]() -> std::unique_ptr<NameIdMapper> {
+    : name_id_mapper_(std::invoke([config, storage_mode]() -> std::unique_ptr<NameIdMapper> {
         if (storage_mode == StorageMode::ON_DISK_TRANSACTIONAL) {
           return std::make_unique<DiskNameIdMapper>(config.disk.name_id_mapper_directory,
                                                     config.disk.id_name_mapper_directory);
@@ -55,26 +49,26 @@ Storage::Storage(Config config, StorageMode storage_mode)
 }
 
 Storage::Accessor::Accessor(SharedAccess /* tag */, Storage *storage, IsolationLevel isolation_level,
-                            StorageMode storage_mode)
+                            StorageMode storage_mode, bool is_main)
     : storage_(storage),
       // The lock must be acquired before creating the transaction object to
       // prevent freshly created transactions from dangling in an active state
       // during exclusive operations.
       storage_guard_(storage_->main_lock_),
       unique_guard_(storage_->main_lock_, std::defer_lock),
-      transaction_(storage->CreateTransaction(isolation_level, storage_mode)),
+      transaction_(storage->CreateTransaction(isolation_level, storage_mode, is_main)),
       is_transaction_active_(true),
       creation_storage_mode_(storage_mode) {}
 
 Storage::Accessor::Accessor(UniqueAccess /* tag */, Storage *storage, IsolationLevel isolation_level,
-                            StorageMode storage_mode)
+                            StorageMode storage_mode, bool is_main)
     : storage_(storage),
       // The lock must be acquired before creating the transaction object to
       // prevent freshly created transactions from dangling in an active state
       // during exclusive operations.
       storage_guard_(storage_->main_lock_, std::defer_lock),
       unique_guard_(storage_->main_lock_),
-      transaction_(storage->CreateTransaction(isolation_level, storage_mode)),
+      transaction_(storage->CreateTransaction(isolation_level, storage_mode, is_main)),
       is_transaction_active_(true),
       creation_storage_mode_(storage_mode) {}
 
@@ -91,18 +85,6 @@ Storage::Accessor::Accessor(Accessor &&other) noexcept
   other.commit_timestamp_.reset();
 }
 
-/// Main lock is taken by the caller.
-void Storage::SetStorageMode(StorageMode storage_mode) {
-  std::unique_lock main_guard{main_lock_};
-  MG_ASSERT(
-      (storage_mode_ == StorageMode::IN_MEMORY_ANALYTICAL || storage_mode_ == StorageMode::IN_MEMORY_TRANSACTIONAL) &&
-      (storage_mode == StorageMode::IN_MEMORY_ANALYTICAL || storage_mode == StorageMode::IN_MEMORY_TRANSACTIONAL));
-  if (storage_mode_ != storage_mode) {
-    storage_mode_ = storage_mode;
-    FreeMemory(std::move(main_guard));
-  }
-}
-
 StorageMode Storage::GetStorageMode() const { return storage_mode_; }
 
 IsolationLevel Storage::GetIsolationLevel() const noexcept { return isolation_level_; }

src/storage/v2/storage.hpp

@@ -20,6 +20,7 @@
 #include "kvstore/kvstore.hpp"
 #include "query/exceptions.hpp"
 #include "replication/config.hpp"
+#include "replication/replication_server.hpp"
 #include "storage/v2/all_vertices_iterable.hpp"
 #include "storage/v2/commit_log.hpp"
 #include "storage/v2/config.hpp"
@@ -30,7 +31,6 @@
 #include "storage/v2/mvcc.hpp"
 #include "storage/v2/replication/enums.hpp"
 #include "storage/v2/replication/replication_client.hpp"
-#include "storage/v2/replication/replication_server.hpp"
 #include "storage/v2/replication/replication_storage_state.hpp"
 #include "storage/v2/storage_error.hpp"
 #include "storage/v2/storage_mode.hpp"
@@ -69,7 +69,7 @@ struct StorageInfo {
   uint64_t vertex_count;
   uint64_t edge_count;
   double average_degree;
-  uint64_t memory_usage;
+  uint64_t memory_res;
   uint64_t disk_usage;
   uint64_t label_indices;
   uint64_t label_property_indices;
@@ -86,7 +86,7 @@ static inline nlohmann::json ToJson(const StorageInfo &info) {
 
   res["edges"] = info.edge_count;
   res["vertices"] = info.vertex_count;
-  res["memory"] = info.memory_usage;
+  res["memory"] = info.memory_res;
   res["disk"] = info.disk_usage;
   res["label_indices"] = info.label_indices;
   res["label_prop_indices"] = info.label_property_indices;
@@ -130,8 +130,10 @@ class Storage {
     static constexpr struct UniqueAccess {
     } unique_access;
 
-    Accessor(SharedAccess /* tag */, Storage *storage, IsolationLevel isolation_level, StorageMode storage_mode);
-    Accessor(UniqueAccess /* tag */, Storage *storage, IsolationLevel isolation_level, StorageMode storage_mode);
+    Accessor(SharedAccess /* tag */, Storage *storage, IsolationLevel isolation_level, StorageMode storage_mode,
+             bool is_main = true);
+    Accessor(UniqueAccess /* tag */, Storage *storage, IsolationLevel isolation_level, StorageMode storage_mode,
+             bool is_main = true);
     Accessor(const Accessor &) = delete;
     Accessor &operator=(const Accessor &) = delete;
     Accessor &operator=(Accessor &&other) = delete;
@@ -211,7 +213,7 @@ class Storage {
 
     // NOLINTNEXTLINE(google-default-arguments)
     virtual utils::BasicResult<StorageManipulationError, void> Commit(
-        std::optional<uint64_t> desired_commit_timestamp = {}) = 0;
+        std::optional<uint64_t> desired_commit_timestamp = {}, bool is_main = true) = 0;
 
     virtual void Abort() = 0;
 
@@ -299,19 +301,26 @@ class Storage {
     return EdgeTypeId::FromUint(name_id_mapper_->NameToId(name));
   }
 
-  void SetStorageMode(StorageMode storage_mode);
-
   StorageMode GetStorageMode() const;
 
   virtual void FreeMemory(std::unique_lock<utils::ResourceLock> main_guard) = 0;
 
   void FreeMemory() { FreeMemory({}); }
 
-  virtual std::unique_ptr<Accessor> Access(std::optional<IsolationLevel> override_isolation_level) = 0;
-  std::unique_ptr<Accessor> Access() { return Access(std::optional<IsolationLevel>{}); }
+  virtual std::unique_ptr<Accessor> Access(std::optional<IsolationLevel> override_isolation_level, bool is_main) = 0;
+  std::unique_ptr<Accessor> Access(bool is_main = true) { return Access(std::optional<IsolationLevel>{}, is_main); }
+  std::unique_ptr<Accessor> Access(std::optional<IsolationLevel> override_isolation_level) {
+    return Access(std::move(override_isolation_level), true);
+  }
 
-  virtual std::unique_ptr<Accessor> UniqueAccess(std::optional<IsolationLevel> override_isolation_level) = 0;
-  std::unique_ptr<Accessor> UniqueAccess() { return UniqueAccess(std::optional<IsolationLevel>{}); }
+  virtual std::unique_ptr<Accessor> UniqueAccess(std::optional<IsolationLevel> override_isolation_level,
+                                                 bool is_main) = 0;
+  std::unique_ptr<Accessor> UniqueAccess(bool is_main = true) {
+    return UniqueAccess(std::optional<IsolationLevel>{}, is_main);
+  }
+  std::unique_ptr<Accessor> UniqueAccess(std::optional<IsolationLevel> override_isolation_level) {
+    return UniqueAccess(std::move(override_isolation_level), true);
+  }
 
   enum class SetIsolationLevelError : uint8_t { DisabledForAnalyticalMode };
 
@@ -338,23 +347,24 @@ class Storage {
     return GetInfo(force_dir);
   }
 
-  virtual Transaction CreateTransaction(IsolationLevel isolation_level, StorageMode storage_mode) = 0;
+  Transaction CreateTransaction(IsolationLevel isolation_level, StorageMode storage_mode) {
+    return CreateTransaction(isolation_level, storage_mode, true);
+  }
 
-  virtual void PrepareForNewEpoch(std::string prev_epoch) = 0;
+  virtual Transaction CreateTransaction(IsolationLevel isolation_level, StorageMode storage_mode, bool is_main) = 0;
 
-  virtual auto CreateReplicationClient(const memgraph::replication::ReplicationClientConfig &config)
+  virtual void PrepareForNewEpoch() = 0;
+
+  virtual auto CreateReplicationClient(const memgraph::replication::ReplicationClientConfig &config,
+                                       const memgraph::replication::ReplicationEpoch *current_epoch)
       -> std::unique_ptr<ReplicationClient> = 0;
 
-  virtual auto CreateReplicationServer(const memgraph::replication::ReplicationServerConfig &config)
-      -> std::unique_ptr<ReplicationServer> = 0;
-
   auto ReplicasInfo() const { return repl_storage_state_.ReplicasInfo(); }
   auto GetReplicaState(std::string_view name) const -> std::optional<replication::ReplicaState> {
     return repl_storage_state_.GetReplicaState(name);
   }
 
   // TODO: make non-public
-  memgraph::replication::ReplicationState repl_state_;
   ReplicationStorageState repl_storage_state_;
 
   // Main storage lock.

src/storage/v2/vertex_accessor.hpp

@@ -127,6 +127,9 @@ class VertexAccessor final {
   bool for_deleted_{false};
 };
 
+static_assert(std::is_trivially_copyable<memgraph::storage::VertexAccessor>::value,
+              "storage::VertexAccessor must be trivially copyable!");
+
 struct EdgesVertexAccessorResult {
   std::vector<EdgeAccessor> edges;
   int64_t expanded_count;

src/telemetry/collectors.cpp

@@ -78,8 +78,12 @@ const nlohmann::json GetResourceUsage(std::filesystem::path root_directory) {
   }
   auto cpu_total = GetCpuUsage(pid);
   cpu["usage"] = cpu_total.second;
-
-  return {{"cpu", cpu}, {"memory", utils::GetMemoryUsage()}, {"disk", utils::GetDirDiskUsage(root_directory)}};
+  const auto vm_max_map_count = utils::GetVmMaxMapCount();
+  return {{"cpu", cpu},
+          {"memory", utils::GetMemoryRES()},
+          {"disk", utils::GetDirDiskUsage(root_directory)},
+          {"vm_max_map_count",
+           vm_max_map_count.has_value() ? *vm_max_map_count : memgraph::utils::VM_MAX_MAP_COUNT_DEFAULT}};
 }
 
 }  // namespace memgraph::telemetry

src/telemetry/telemetry.cpp

@@ -15,6 +15,7 @@
 
 #include <fmt/format.h>
 
+#include "communication/bolt/metrics.hpp"
 #include "requests/requests.hpp"
 #include "telemetry/collectors.hpp"
 #include "utils/event_counter.hpp"
@@ -154,7 +155,6 @@ void Telemetry::AddDatabaseCollector(dbms::DbmsHandler &dbms_handler) {
 #else
 #endif
 
-#ifdef MG_ENTERPRISE
 void Telemetry::AddStorageCollector(
     dbms::DbmsHandler &dbms_handler,
     memgraph::utils::Synchronized<memgraph::auth::Auth, memgraph::utils::WritePrioritizedRWLock> &auth) {
@@ -164,33 +164,6 @@ void Telemetry::AddStorageCollector(
     return ToJson(stats);
   });
 }
-#else
-void Telemetry::AddStorageCollector(
-    memgraph::utils::Gatekeeper<memgraph::dbms::Database> &db_gatekeeper,
-    memgraph::utils::Synchronized<memgraph::auth::Auth, memgraph::utils::WritePrioritizedRWLock> &auth) {
-  AddCollector("storage", [&db_gatekeeper, &auth]() -> nlohmann::json {
-    memgraph::dbms::Statistics stats;
-    auto db_acc_opt = db_gatekeeper.access();
-    MG_ASSERT(db_acc_opt, "Failed to get access to the default database");
-    auto &db_acc = *db_acc_opt;
-    const auto &info = db_acc->GetInfo();
-    const auto &storage_info = info.storage_info;
-    stats.num_vertex = storage_info.vertex_count;
-    stats.num_edges = storage_info.edge_count;
-    stats.triggers = info.triggers;
-    stats.streams = info.streams;
-    stats.num_databases = 1;
-    stats.indices += storage_info.label_indices + storage_info.label_property_indices;
-    stats.constraints += storage_info.existence_constraints + storage_info.unique_constraints;
-    ++stats.storage_modes[(int)storage_info.storage_mode];
-    ++stats.isolation_levels[(int)storage_info.isolation_level];
-    stats.snapshot_enabled = storage_info.durability_snapshot_enabled;
-    stats.wal_enabled = storage_info.durability_wal_enabled;
-    stats.users = auth->AllUsers().size();
-    return ToJson(stats);
-  });
-}
-#endif
 
 void Telemetry::AddExceptionCollector() {
   AddCollector("exception", []() -> nlohmann::json { return memgraph::metrics::global_counters_map.ToJson(); });

src/telemetry/telemetry.hpp

@@ -43,15 +43,9 @@ class Telemetry final {
   void AddCollector(const std::string &name, const std::function<const nlohmann::json(void)> &func);
 
   // Specialized collectors
-#ifdef MG_ENTERPRISE
   void AddStorageCollector(
       dbms::DbmsHandler &dbms_handler,
       memgraph::utils::Synchronized<memgraph::auth::Auth, memgraph::utils::WritePrioritizedRWLock> &auth);
-#else
-  void AddStorageCollector(
-      memgraph::utils::Gatekeeper<memgraph::dbms::Database> &db_gatekeeper,
-      memgraph::utils::Synchronized<memgraph::auth::Auth, memgraph::utils::WritePrioritizedRWLock> &auth);
-#endif
 
 #ifdef MG_ENTERPRISE
   void AddDatabaseCollector(dbms::DbmsHandler &dbms_handler);

src/utils/CMakeLists.txt

@@ -14,7 +14,8 @@ set(utils_src_files
     tsc.cpp
     system_info.cpp
     uuid.cpp
-    build_info.cpp)
+    build_info.cpp
+    query_memory_tracker.cpp)
 
 find_package(Boost REQUIRED)
 find_package(fmt REQUIRED)

src/utils/file.cpp

@@ -519,7 +519,7 @@ void OutputFile::FlushBufferInternal() {
   auto *buffer = buffer_;
   auto buffer_position = buffer_position_.load();
   while (buffer_position > 0) {
-    auto written = write(fd_, buffer, buffer_position_);
+    auto written = write(fd_, buffer, buffer_position);
     if (written == -1 && errno == EINTR) {
       continue;
     }

src/utils/gatekeeper.hpp

@@ -81,7 +81,7 @@ EvalResult(run_t, Func &&, T &) -> EvalResult<std::invoke_result_t<Func, T &>>;
 template <typename T>
 struct Gatekeeper {
   template <typename... Args>
-  explicit Gatekeeper(Args &&...args) : value_{std::forward<Args>(args)...} {}
+  explicit Gatekeeper(Args &&...args) : value_{std::in_place, std::forward<Args>(args)...} {}
 
   Gatekeeper(Gatekeeper const &) = delete;
   Gatekeeper(Gatekeeper &&) noexcept = delete;

src/utils/logging.hpp

@@ -47,8 +47,12 @@ inline void AssertFailed(const char *file_name, int line_num, const char *expr,
 #define GET_MESSAGE(...) \
   BOOST_PP_IF(BOOST_PP_EQUAL(BOOST_PP_VARIADIC_SIZE(__VA_ARGS__), 0), "", fmt::format(__VA_ARGS__))
 
-#define MG_ASSERT(expr, ...) \
-  [[likely]] !!(expr) ? (void)0 : ::memgraph::logging::AssertFailed(__FILE__, __LINE__, #expr, GET_MESSAGE(__VA_ARGS__))
+#define MG_ASSERT(expr, ...)                                                                \
+  if (expr) [[likely]] {                                                                    \
+    (void)0;                                                                                \
+  } else {                                                                                  \
+    ::memgraph::logging::AssertFailed(__FILE__, __LINE__, #expr, GET_MESSAGE(__VA_ARGS__)); \
+  }
 
 #ifndef NDEBUG
 #define DMG_ASSERT(expr, ...) MG_ASSERT(expr, __VA_ARGS__)

src/utils/lru_cache.hpp

@@ -0,0 +1,68 @@
+// Copyright 2023 Memgraph Ltd.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+// License, and you may not use this file except in compliance with the Business Source License.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+#pragma once
+
+#include <list>
+#include <optional>
+#include <unordered_map>
+#include <utility>
+
+namespace memgraph::utils {
+
+/// A simple LRU cache implementation.
+/// It is not thread-safe.
+
+template <class TKey, class TVal>
+class LRUCache {
+ public:
+  LRUCache(int cache_size_) : cache_size(cache_size_){};
+
+  void put(const TKey &key, const TVal &val) {
+    auto it = item_map.find(key);
+    if (it != item_map.end()) {
+      item_list.erase(it->second);
+      item_map.erase(it);
+    }
+    item_list.push_front(std::make_pair(key, val));
+    item_map.insert(std::make_pair(key, item_list.begin()));
+    try_clean();
+  };
+  std::optional<TVal> get(const TKey &key) {
+    if (!exists(key)) {
+      return std::nullopt;
+    }
+    auto it = item_map.find(key);
+    item_list.splice(item_list.begin(), item_list, it->second);
+    return it->second->second;
+  }
+  void reset() {
+    item_list.clear();
+    item_map.clear();
+  };
+  std::size_t size() { return item_map.size(); };
+
+ private:
+  void try_clean() {
+    while (item_map.size() > cache_size) {
+      auto last_it_elem_it = item_list.end();
+      last_it_elem_it--;
+      item_map.erase(last_it_elem_it->first);
+      item_list.pop_back();
+    }
+  };
+  bool exists(const TKey &key) { return (item_map.count(key) > 0); };
+
+  std::list<std::pair<TKey, TVal>> item_list;
+  std::unordered_map<TKey, decltype(item_list.begin())> item_map;
+  std::size_t cache_size;
+};
+}  // namespace memgraph::utils

src/utils/memory.hpp

@@ -539,9 +539,9 @@ class SynchronizedPoolResource final : public MemoryResource {
   bool DoIsEqual(const MemoryResource &other) const noexcept override { return this == &other; }
 };
 
-class LimitedMemoryResource final : public utils::MemoryResource {
+class MemoryTrackingResource final : public utils::MemoryResource {
  public:
-  explicit LimitedMemoryResource(utils::MemoryResource *memory, size_t max_allocated_bytes)
+  explicit MemoryTrackingResource(utils::MemoryResource *memory, size_t max_allocated_bytes)
       : memory_(memory), max_allocated_bytes_(max_allocated_bytes) {}
 
   size_t GetAllocatedBytes() const noexcept { return max_allocated_bytes_ - available_bytes_; }
@@ -552,13 +552,11 @@ class LimitedMemoryResource final : public utils::MemoryResource {
   size_t available_bytes_{max_allocated_bytes_};
 
   void *DoAllocate(size_t bytes, size_t alignment) override {
-    if (bytes > available_bytes_) throw utils::BadAlloc("Memory allocation limit exceeded!");
     available_bytes_ -= bytes;
     return memory_->Allocate(bytes, alignment);
   }
 
   void DoDeallocate(void *p, size_t bytes, size_t alignment) override {
-    MG_ASSERT(available_bytes_ + bytes > available_bytes_, "Failed deallocation");
     available_bytes_ += bytes;
     return memory_->Deallocate(p, bytes, alignment);
   }

src/utils/memory_tracker.cpp

@@ -111,7 +111,7 @@ void MemoryTracker::Alloc(const int64_t size) {
 
   const auto current_hard_limit = hard_limit_.load(std::memory_order_relaxed);
 
-  if (UNLIKELY(current_hard_limit && will_be > current_hard_limit && MemoryTrackerCanThrow())) {
+  if (current_hard_limit && will_be > current_hard_limit && MemoryTrackerCanThrow()) [[unlikely]] {
     MemoryTracker::OutOfMemoryExceptionBlocker exception_blocker;
 
     amount_.fetch_sub(size, std::memory_order_relaxed);
@@ -121,7 +121,6 @@ void MemoryTracker::Alloc(const int64_t size) {
                     "use to {}, while the maximum allowed size for allocation is set to {}.",
                     GetReadableSize(size), GetReadableSize(will_be), GetReadableSize(current_hard_limit)));
   }
-
   UpdatePeak(will_be);
 }
 

src/utils/memory_tracker.hpp

@@ -25,17 +25,6 @@ class OutOfMemoryException : public utils::BasicException {
 };
 
 class MemoryTracker final {
- private:
-  std::atomic<int64_t> amount_{0};
-  std::atomic<int64_t> peak_{0};
-  std::atomic<int64_t> hard_limit_{0};
-  // Maximum possible value of a hard limit. If it's set to 0, no upper bound on the hard limit is set.
-  int64_t maximum_hard_limit_{0};
-
-  void UpdatePeak(int64_t will_be);
-
-  static void LogMemoryUsage(int64_t current);
-
  public:
   void LogPeakMemoryUsage() const;
 
@@ -73,6 +62,13 @@ class MemoryTracker final {
 
   void ResetTrackings();
 
+  bool IsProcedureTracked();
+
+  void SetProcTrackingLimit(size_t limit);
+
+  void StartProcTracking();
+  void StopProcTracking();
+
   // By creating an object of this class, every allocation in its scope that goes over
   // the set hard limit produces an OutOfMemoryException.
   class OutOfMemoryExceptionEnabler final {
@@ -109,6 +105,17 @@ class MemoryTracker final {
    private:
     static thread_local uint64_t counter_;
   };
+
+ private:
+  std::atomic<int64_t> amount_{0};
+  std::atomic<int64_t> peak_{0};
+  std::atomic<int64_t> hard_limit_{0};
+  // Maximum possible value of a hard limit. If it's set to 0, no upper bound on the hard limit is set.
+  int64_t maximum_hard_limit_{0};
+
+  void UpdatePeak(int64_t will_be);
+
+  static void LogMemoryUsage(int64_t current);
 };
 
 // Global memory tracker which tracks every allocation in the application.

src/utils/query_memory_tracker.cpp

@@ -0,0 +1,78 @@
+// Copyright 2023 Memgraph Ltd.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+// License, and you may not use this file except in compliance with the Business Source License.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+#include "utils/query_memory_tracker.hpp"
+#include <atomic>
+#include <optional>
+#include "memory/query_memory_control.hpp"
+#include "utils/memory_tracker.hpp"
+
+namespace memgraph::utils {
+
+void QueryMemoryTracker::TrackAlloc(size_t size) {
+  if (query_tracker_.has_value()) [[likely]] {
+    query_tracker_->Alloc(static_cast<int64_t>(size));
+  }
+
+  auto *proc_tracker = GetActiveProc();
+
+  if (proc_tracker == nullptr) {
+    return;
+  }
+
+  proc_tracker->Alloc(static_cast<int64_t>(size));
+}
+void QueryMemoryTracker::TrackFree(size_t size) {
+  if (query_tracker_.has_value()) [[likely]] {
+    query_tracker_->Free(static_cast<int64_t>(size));
+  }
+
+  auto *proc_tracker = GetActiveProc();
+
+  if (proc_tracker == nullptr) {
+    return;
+  }
+
+  proc_tracker->Free(static_cast<int64_t>(size));
+}
+
+void QueryMemoryTracker::SetQueryLimit(size_t size) {
+  if (size == memgraph::memory::UNLIMITED_MEMORY) {
+    return;
+  }
+  InitializeQueryTracker();
+  query_tracker_->SetMaximumHardLimit(static_cast<int64_t>(size));
+  query_tracker_->SetHardLimit(static_cast<int64_t>(size));
+}
+
+memgraph::utils::MemoryTracker *QueryMemoryTracker::GetActiveProc() {
+  if (active_proc_id == NO_PROCEDURE) [[likely]] {
+    return nullptr;
+  }
+  return &proc_memory_trackers_[active_proc_id];
+}
+
+void QueryMemoryTracker::SetActiveProc(int64_t new_active_proc) { active_proc_id = new_active_proc; }
+
+void QueryMemoryTracker::StopProcTracking() { active_proc_id = QueryMemoryTracker::NO_PROCEDURE; }
+
+void QueryMemoryTracker::TryCreateProcTracker(int64_t procedure_id, size_t limit) {
+  if (proc_memory_trackers_.contains(procedure_id)) {
+    return;
+  }
+  auto [it, inserted] = proc_memory_trackers_.emplace(procedure_id, utils::MemoryTracker{});
+  it->second.SetMaximumHardLimit(static_cast<int64_t>(limit));
+  it->second.SetHardLimit(static_cast<int64_t>(limit));
+}
+
+void QueryMemoryTracker::InitializeQueryTracker() { query_tracker_.emplace(MemoryTracker{}); }
+
+}  // namespace memgraph::utils