Check-in initial RsmEngine class

Co-authored-by: János Benjamin Antal <antaljanosbenjamin@users.noreply.github.com>

src/CMakeLists.txt

@@ -18,6 +18,7 @@ add_subdirectory(query/v2)
 add_subdirectory(slk)
 add_subdirectory(rpc)
 add_subdirectory(auth)
+add_subdirectory(coordinator)
 
 if (MG_ENTERPRISE)
   add_subdirectory(audit)

src/coordinator/CMakeLists.txt

@@ -0,0 +1,10 @@
+set(coordinator_src_files
+    coordinator.hpp
+    shard_map.hpp
+    hybrid_logical_clock.hpp)
+
+find_package(fmt REQUIRED)
+find_package(Threads REQUIRED)
+
+add_library(mg-coordinator STATIC ${coordinator_src_files})
+target_link_libraries(mg-coordinator stdc++fs Threads::Threads fmt::fmt mg-utils)

src/coordinator/coordinator.hpp

@@ -0,0 +1,243 @@
+// Copyright 2022 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 <unordered_set>
+
+#include "coordinator/hybrid_logical_clock.hpp"
+#include "coordinator/shard_map.hpp"
+#include "io/simulator/simulator.hpp"
+#include "io/time.hpp"
+#include "io/transport.hpp"
+
+namespace memgraph::coordinator {
+
+using Address = memgraph::io::Address;
+using SimT = memgraph::io::simulator::SimulatorTransport;
+
+struct HlcRequest {
+  Hlc last_shard_map_version;
+};
+
+struct HlcResponse {
+  Hlc new_hlc;
+  std::optional<ShardMap> fresher_shard_map;
+};
+
+struct GetShardMapRequest {
+  // No state
+};
+
+struct GetShardMapResponse {
+  ShardMap shard_map;
+};
+
+struct AllocateHlcBatchRequest {
+  Hlc low;
+  Hlc high;
+};
+
+struct AllocateHlcBatchResponse {
+  bool success;
+  Hlc low;
+  Hlc high;
+};
+
+struct AllocateEdgeIdBatchRequest {
+  size_t batch_size;
+};
+
+struct AllocateEdgeIdBatchResponse {
+  uint64_t low;
+  uint64_t high;
+};
+
+struct SplitShardRequest {
+  Hlc previous_shard_map_version;
+  Label label;
+  CompoundKey split_key;
+};
+
+struct SplitShardResponse {
+  bool success;
+};
+
+struct RegisterStorageEngineRequest {
+  Address address;
+};
+
+struct RegisterStorageEngineResponse {
+  bool success;
+};
+
+struct DeregisterStorageEngineRequest {
+  Address address;
+};
+
+struct DeregisterStorageEngineResponse {
+  bool success;
+};
+
+struct InitializeLabelRequest {
+  std::string label_name;
+  Hlc last_shard_map_version;
+};
+
+struct InitializeLabelResponse {
+  bool success;
+  std::optional<ShardMap> fresher_shard_map;
+};
+
+using WriteRequests =
+    std::variant<AllocateHlcBatchRequest, AllocateEdgeIdBatchRequest, SplitShardRequest, RegisterStorageEngineRequest,
+                 DeregisterStorageEngineRequest, InitializeLabelRequest>;
+using WriteResponses =
+    std::variant<AllocateHlcBatchResponse, AllocateEdgeIdBatchResponse, SplitShardResponse,
+                 RegisterStorageEngineResponse, DeregisterStorageEngineResponse, InitializeLabelResponse>;
+
+using ReadRequests = std::variant<HlcRequest, GetShardMapRequest>;
+using ReadResponses = std::variant<HlcResponse, GetShardMapResponse>;
+
+class Coordinator {
+  ShardMap shard_map_;
+  /// The highest reserved timestamp / highest allocated timestamp
+  /// is a way for minimizing communication involved in query engines
+  /// reserving Hlc's for their transaction processing.
+  /// Periodically, the coordinator will allocate a batch of timestamps
+  /// and this will need to go over consensus. From that point forward,
+  /// each timestamp in that batch can be given out to "readers" who issue
+  /// HlcRequest without blocking on consensus first. But if
+  /// highest_allocated_timestamp_ approaches highest_reserved_timestamp_,
+  /// it is time to allocate another batch, so that we can keep guaranteeing
+  /// forward progress.
+  /// Any time a coordinator becomes a new leader, it will need to issue
+  /// a new AllocateHlcBatchRequest to create a pool of IDs to allocate.
+  uint64_t highest_allocated_timestamp_;
+  uint64_t highest_reserved_timestamp_;
+
+  /// Query engines need to periodically request batches of unique edge IDs.
+  uint64_t highest_allocated_edge_id_;
+
+  /// Increment our
+  ReadResponses HandleRead(HlcRequest &&hlc_request) {
+    HlcResponse res{};
+
+    auto hlc_shard_map = shard_map_.GetHlc();
+
+    MG_ASSERT(!(hlc_request.last_shard_map_version.logical_id > hlc_shard_map.logical_id));
+
+    res.new_hlc = shard_map_.IncrementShardMapVersion();
+
+    // res.fresher_shard_map = hlc_request.last_shard_map_version.logical_id < hlc_shard_map.logical_id
+    //                             ? std::make_optional(shard_map_)
+    //                             : std::nullopt;
+
+    // Allways return fresher shard_map for now.
+    res.fresher_shard_map = std::make_optional(shard_map_);
+
+    return res;
+  }
+
+  ReadResponses HandleRead(GetShardMapRequest &&get_shard_map_request) {
+    GetShardMapResponse res;
+    res.shard_map = shard_map_;
+    return res;
+  }
+
+  WriteResponses ApplyWrite(AllocateHlcBatchRequest &&ahr) {
+    AllocateHlcBatchResponse res{};
+
+    return res;
+  }
+
+  WriteResponses ApplyWrite(AllocateEdgeIdBatchRequest &&ahr) {
+    AllocateEdgeIdBatchResponse res{};
+
+    uint64_t low = highest_allocated_edge_id_;
+
+    highest_allocated_edge_id_ += ahr.batch_size;
+
+    uint64_t high = highest_allocated_edge_id_;
+
+    res.low = low;
+    res.high = high;
+
+    return res;
+  }
+
+  /// This splits the shard immediately beneath the provided
+  /// split key, keeping the assigned peers identical for now,
+  /// but letting them be gradually migrated over time.
+  WriteResponses ApplyWrite(SplitShardRequest &&split_shard_request) {
+    SplitShardResponse res{};
+
+    if (split_shard_request.previous_shard_map_version != shard_map_.shard_map_version) {
+      res.success = false;
+    } else {
+      res.success = shard_map_.SplitShard(split_shard_request.previous_shard_map_version, split_shard_request.label,
+                                          split_shard_request.split_key);
+    }
+
+    return res;
+  }
+
+  /// This adds the provided storage engine to the standby storage engine pool,
+  /// which can be used to rebalance storage over time.
+  WriteResponses ApplyWrite(RegisterStorageEngineRequest &&register_storage_engine_request) {
+    RegisterStorageEngineResponse res{};
+    // TODO
+
+    return res;
+  }
+
+  /// This begins the process of draining the provided storage engine from all raft
+  /// clusters that it might be participating in.
+  WriteResponses ApplyWrite(DeregisterStorageEngineRequest &&register_storage_engine_request) {
+    DeregisterStorageEngineResponse res{};
+    // TODO
+    // const Address &address = register_storage_engine_request.address;
+    // storage_engine_pool_.erase(address);
+    // res.success = true;
+
+    return res;
+  }
+
+  WriteResponses ApplyWrite(InitializeLabelRequest &&initialize_label_request) {
+    InitializeLabelResponse res{};
+
+    bool success = shard_map_.InitializeNewLabel(initialize_label_request.label_name,
+                                                 initialize_label_request.last_shard_map_version);
+
+    if (success) {
+      res.fresher_shard_map = shard_map_;
+      res.success = false;
+    } else {
+      res.fresher_shard_map = std::nullopt;
+      res.success = true;
+    }
+
+    return res;
+  }
+
+ public:
+  explicit Coordinator(ShardMap sm) : shard_map_{(sm)} {}
+
+  ReadResponses Read(ReadRequests requests) {
+    return std::visit([&](auto &&requests) { return HandleRead(std::move(requests)); }, std::move(requests));
+  }
+
+  WriteResponses Apply(WriteRequests requests) {
+    return std::visit([&](auto &&requests) { return ApplyWrite(std::move(requests)); }, std::move(requests));
+  }
+};
+
+}  // namespace memgraph::coordinator

src/coordinator/hybrid_logical_clock.hpp

@@ -0,0 +1,30 @@
+// Copyright 2022 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 "io/time.hpp"
+
+namespace memgraph::coordinator {
+
+using Time = memgraph::io::Time;
+
+/// Hybrid-logical clock
+struct Hlc {
+  uint64_t logical_id;
+  Time coordinator_wall_clock;
+
+  bool operator==(const Hlc &other) const {
+    return (logical_id == other.logical_id) && (coordinator_wall_clock == other.coordinator_wall_clock);
+  }
+};
+
+}  // namespace memgraph::coordinator

src/coordinator/shard_map.hpp

@@ -0,0 +1,135 @@
+// Copyright 2022 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 <map>
+#include <vector>
+
+#include "coordinator/hybrid_logical_clock.hpp"
+#include "io/address.hpp"
+#include "storage/v3/property_value.hpp"
+
+namespace memgraph::coordinator {
+
+enum class Status : uint8_t {
+  CONSENSUS_PARTICIPANT,
+  INITIALIZING,
+  // TODO(tyler) this will possibly have more states,
+  // depending on the reconfiguration protocol that we
+  // implement.
+};
+
+struct AddressAndStatus {
+  memgraph::io::Address address;
+  Status status;
+};
+
+using memgraph::io::Address;
+
+using CompoundKey = std::vector<memgraph::storage::v3::PropertyValue>;
+using Shard = std::vector<AddressAndStatus>;
+using Shards = std::map<CompoundKey, Shard>;
+
+// use string for intermachine communication and NameIdMapper within the machine
+using Label = std::string;
+
+struct ShardMap {
+  Hlc shard_map_version;
+  std::map<Label, Shards> shards;
+
+  // TODO(gabor) later we will want to update the wallclock time with
+  // the given Io<impl>'s time as well
+  Hlc IncrementShardMapVersion() noexcept {
+    ++shard_map_version.logical_id;
+    return shard_map_version;
+  }
+
+  Hlc GetHlc() const noexcept { return shard_map_version; }
+
+  bool SplitShard(Hlc previous_shard_map_version, Label label, CompoundKey key) {
+    if (CompareShardMapVersions(previous_shard_map_version, shard_map_version)) {
+      MG_ASSERT(shards.contains(label));
+      auto &shards_in_map = shards[label];
+      MG_ASSERT(!shards_in_map.contains(key));
+
+      // Finding the Shard that the new CompoundKey should map to.
+      Shard shard_to_map_to;
+      CompoundKey prev_key = ((*shards_in_map.begin()).first);
+
+      for (auto iter = std::next(shards_in_map.begin()); iter != shards_in_map.end(); ++iter) {
+        const auto &current_key = (*iter).first;
+        if (key > prev_key && key < current_key) {
+          shard_to_map_to = shards_in_map[prev_key];
+        }
+
+        prev_key = (*iter).first;
+      }
+
+      // Apply the split
+      shards_in_map[key] = shard_to_map_to;
+
+      return true;
+    }
+
+    return false;
+  }
+
+  bool InitializeNewLabel(std::string label_name, Hlc last_shard_map_version) {
+    if (shard_map_version != last_shard_map_version) {
+      return false;
+    }
+
+    if (shards.contains(label_name)) {
+      return false;
+    }
+
+    shards.emplace(label_name, Shards{});
+
+    IncrementShardMapVersion();
+
+    return true;
+  }
+
+  void AddServer(Address server_address) {
+    // Find a random place for the server to plug in
+  }
+
+  std::map<Label, Shards> &GetShards() noexcept { return shards; }
+
+  Shards GetShardsForRange(Label label, CompoundKey start, CompoundKey end);
+
+  Shard GetShardForKey(Label label, CompoundKey key) {
+    auto shard_for_label = shards.at(label);
+
+    auto max = (--shard_for_label.end())->first;
+
+    if (key > max) {
+      return shard_for_label[max];
+    }
+
+    for (auto it = shard_for_label.lower_bound(key);; --it) {
+      MG_ASSERT(it->first <= key);
+      return it->second;
+    }
+
+    MG_ASSERT(false, "failed to find shard with a key that is less than or equal to the provided key");
+  }
+
+ private:
+  // TODO(gabor) later we will want to update the wallclock time with
+  // the given Io<impl>'s time as well. This function should just be
+  // replaced with operator== since it is already overloaded for Hlc
+  // objects.
+  bool CompareShardMapVersions(Hlc one, Hlc two) { return one.logical_id == two.logical_id; }
+};
+
+}  // namespace memgraph::coordinator

src/io/rsm/coordinator_rsm.hpp

@@ -0,0 +1,25 @@
+// Copyright 2022 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 "coordinator/coordinator.hpp"
+#include "io/rsm/raft.hpp"
+
+namespace memgraph::io::rsm {
+
+// TODO(tyler) don't
+using namespace memgraph::coordinator;
+
+template <typename IoImpl>
+using CoordinatorRsm = io::rsm::Raft<IoImpl, Coordinator, WriteRequests, WriteResponses, ReadRequests, ReadResponses>;
+
+}  // namespace memgraph::io::rsm

src/io/rsm/raft.hpp

@@ -0,0 +1,845 @@
+// Copyright 2022 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.
+
+// TODO(tyler) buffer out-of-order Append buffers on the Followers to reassemble more quickly
+// TODO(tyler) handle granular batch sizes based on simple flow control
+
+#pragma once
+
+#include <deque>
+#include <iostream>
+#include <map>
+#include <set>
+#include <thread>
+#include <unordered_map>
+#include <vector>
+
+#include "io/simulator/simulator.hpp"
+#include "io/transport.hpp"
+
+namespace memgraph::io::rsm {
+
+using memgraph::io::Address;
+using memgraph::io::Duration;
+using memgraph::io::Io;
+using memgraph::io::ResponseFuture;
+using memgraph::io::ResponseResult;
+using memgraph::io::Time;
+
+using Term = uint64_t;
+using LogIndex = uint64_t;
+using RequestId = uint64_t;
+
+template <typename WriteOperation>
+struct WriteRequest {
+  WriteOperation operation;
+};
+
+/// WriteResponse is returned to a client after
+/// their WriteRequest was entered in to the raft
+/// log and it reached consensus.
+///
+/// WriteReturn is the result of applying the WriteRequest to
+/// ReplicatedState, and if the ReplicatedState::write
+/// method is deterministic, all replicas will
+/// have the same ReplicatedState after applying
+/// the submitted WriteRequest.
+template <typename WriteReturn>
+struct WriteResponse {
+  bool success;
+  WriteReturn write_return;
+  std::optional<Address> retry_leader;
+};
+
+template <typename ReadOperation>
+struct ReadRequest {
+  ReadOperation operation;
+};
+
+template <typename ReadReturn>
+struct ReadResponse {
+  bool success;
+  ReadReturn read_return;
+  std::optional<Address> retry_leader;
+};
+
+/// AppendRequest is a raft-level message that the Leader
+/// periodically broadcasts to all Follower peers. This
+/// serves three main roles:
+/// 1. acts as a heartbeat from the Leader to the Follower
+/// 2. replicates new data that the Leader has received to the Follower
+/// 3. informs Follower peers when the commit index has increased,
+///    signalling that it is now safe to apply log items to the
+///    replicated state machine
+template <typename WriteRequest>
+struct AppendRequest {
+  Term term = 0;
+  LogIndex last_log_index;
+  Term last_log_term;
+  std::vector<std::pair<Term, WriteRequest>> entries;
+  LogIndex leader_commit;
+};
+
+struct AppendResponse {
+  bool success;
+  Term term;
+  Term last_log_term;
+  // a small optimization over the raft paper, tells
+  // the leader the offset that we are interested in
+  // to send log offsets from for us. This will only
+  // be useful at the beginning of a leader's term.
+  LogIndex last_log_index;
+};
+
+struct VoteRequest {
+  Term term = 0;
+  LogIndex last_log_index;
+  Term last_log_term;
+};
+
+struct VoteResponse {
+  Term term = 0;
+  LogIndex committed_log_size;
+  bool vote_granted = false;
+};
+
+template <typename WriteRequest>
+struct CommonState {
+  Term term = 0;
+  std::vector<std::pair<Term, WriteRequest>> log;
+  LogIndex committed_log_size = 0;
+  LogIndex applied_size = 0;
+};
+
+struct FollowerTracker {
+  LogIndex next_index = 0;
+  LogIndex confirmed_contiguous_index = 0;
+};
+
+struct PendingClientRequest {
+  RequestId request_id;
+  Address address;
+  Time received_at;
+};
+
+struct Leader {
+  std::map<Address, FollowerTracker> followers;
+  std::unordered_map<LogIndex, PendingClientRequest> pending_client_requests;
+  Time last_broadcast = Time::min();
+
+  std::string ToString() { return "\tLeader   \t"; }
+};
+
+struct Candidate {
+  std::map<Address, LogIndex> successful_votes;
+  Time election_began = Time::min();
+  std::set<Address> outstanding_votes;
+
+  std::string ToString() { return "\tCandidate\t"; }
+};
+
+struct Follower {
+  Time last_received_append_entries_timestamp;
+  Address leader_address;
+
+  std::string ToString() { return "\tFollower \t"; }
+};
+
+using Role = std::variant<Candidate, Leader, Follower>;
+
+/*
+
+all ReplicatedState classes should have an Apply method
+that returns our WriteResponseValue after consensus, and
+a Read method that returns our ReadResponseValue without
+requiring consensus.
+
+ReadResponse Read(ReadOperation);
+WriteResponseValue ReplicatedState::Apply(WriteRequest);
+
+For example:
+If the state is uint64_t, and WriteRequest is `struct PlusOne {};`,
+and WriteResponseValue is also uint64_t (the new value), then
+each call to state.Apply(PlusOne{}) will return the new value
+after incrementing it. 0, 1, 2, 3... and this will be sent back
+to the client that requested the mutation.
+
+In practice, these mutations will usually be predicated on some
+previous value, so that they are idempotent, functioning similarly
+to a CAS operation.
+*/
+template <typename WriteOperation, typename ReadOperation, typename ReplicatedState, typename WriteResponseValue,
+          typename ReadResponseValue>
+concept Rsm = requires(ReplicatedState state, WriteOperation w, ReadOperation r) {
+  { state.Read(r) } -> std::same_as<ReadResponseValue>;
+  { state.Apply(w) } -> std::same_as<WriteResponseValue>;
+};
+
+/// Parameter           Purpose
+/// --------------------------
+/// IoImpl              the concrete Io provider - SimulatorTransport, ThriftTransport, etc...
+/// ReplicatedState     the high-level data structure that is managed by the raft-backed replicated state machine
+/// WriteOperation      the individual operation type that is applied to the ReplicatedState in identical order
+///                     across each replica
+/// WriteResponseValue  the return value of calling ReplicatedState::Apply(WriteOperation), which is executed in
+///                     identical order across all replicas after an WriteOperation reaches consensus.
+/// ReadOperation       the type of operations that do not require consensus before executing directly
+///                     on a const ReplicatedState &
+/// ReadResponseValue   the return value of calling ReplicatedState::Read(ReadOperation), which is executed directly
+///                     without going through consensus first
+template <typename IoImpl, typename ReplicatedState, typename WriteOperation, typename WriteResponseValue,
+          typename ReadOperation, typename ReadResponseValue>
+requires Rsm<WriteOperation, ReadOperation, ReplicatedState, WriteResponseValue, ReadResponseValue>
+class Raft {
+  CommonState<WriteOperation> state_;
+  Role role_ = Candidate{};
+  Io<IoImpl> io_;
+  std::vector<Address> peers_;
+  ReplicatedState replicated_state_;
+
+ public:
+  Raft(Io<IoImpl> &&io, std::vector<Address> peers, ReplicatedState &&replicated_state)
+      : io_(std::move(io)), peers_(peers), replicated_state_(std::move(replicated_state)) {}
+
+  void Run() {
+    Time last_cron = io_.Now();
+
+    while (!io_.ShouldShutDown()) {
+      const auto now = io_.Now();
+      const Duration random_cron_interval = RandomTimeout(1000, 2000);
+      if (now - last_cron > random_cron_interval) {
+        Cron();
+        last_cron = now;
+      }
+
+      Duration receive_timeout = RandomTimeout(10000, 50000);
+
+      auto request_result =
+          io_.template ReceiveWithTimeout<ReadRequest<ReadOperation>, AppendRequest<WriteOperation>, AppendResponse,
+                                          WriteRequest<WriteOperation>, VoteRequest, VoteResponse>(receive_timeout);
+      if (request_result.HasError()) {
+        continue;
+      }
+
+      auto request = std::move(request_result.GetValue());
+
+      Handle(std::move(request.message), request.request_id, request.from_address);
+    }
+  }
+
+ private:
+  // Raft paper - 5.3
+  // When the entry has been safely replicated, the leader applies the
+  // entry to its state machine and returns the result of that
+  // execution to the client.
+  //
+  // "Safely replicated" is defined as being known to be present
+  // on at least a majority of all peers (inclusive of the Leader).
+  void BumpCommitIndexAndReplyToClients(Leader &leader) {
+    auto indices = std::vector<LogIndex>{};
+
+    // We include our own log size in the calculation of the log
+    // index that is present on at least a majority of all peers.
+    indices.push_back(state_.log.size());
+
+    for (const auto &[addr, f] : leader.followers) {
+      indices.push_back(f.confirmed_contiguous_index);
+      Log("at port ", addr.last_known_port, " has confirmed contiguous index of: ", f.confirmed_contiguous_index);
+    }
+
+    // reverse sort from highest to lowest (using std::ranges::greater)
+    std::ranges::sort(indices, std::ranges::greater());
+
+    // This is a particularly correctness-critical calculation because it
+    // determines which index we will consider to be the committed index.
+    //
+    // If the following indexes are recorded for clusters of different sizes,
+    // these are the expected indexes that are considered to have reached
+    // consensus:
+    // state           | expected value | (indices.size() / 2)
+    // [1]              1                 (1 / 2) => 0
+    // [2, 1]           1                 (2 / 2) => 1
+    // [3, 2, 1]        2                 (3 / 2) => 1
+    // [4, 3, 2, 1]     2                 (4 / 2) => 2
+    // [5, 4, 3, 2, 1]  3                 (5 / 2) => 2
+    size_t index_present_on_majority = indices.size() / 2;
+    LogIndex new_committed_log_size = indices[index_present_on_majority];
+
+    // We never go backwards in history.
+    MG_ASSERT(state_.committed_log_size <= new_committed_log_size);
+
+    state_.committed_log_size = new_committed_log_size;
+
+    // For each index between the old index and the new one (inclusive),
+    // Apply that log's WriteOperation to our replicated_state_,
+    // and use the specific return value of the ReplicatedState::Apply
+    // method (WriteResponseValue) to respond to the requester.
+    for (; state_.applied_size < state_.committed_log_size; state_.applied_size++) {
+      const LogIndex apply_index = state_.applied_size;
+      const auto &write_request = state_.log[apply_index].second;
+      WriteResponseValue write_return = replicated_state_.Apply(write_request);
+
+      if (leader.pending_client_requests.contains(apply_index)) {
+        PendingClientRequest client_request = std::move(leader.pending_client_requests.at(apply_index));
+
+        WriteResponse<WriteResponseValue> resp;
+        resp.success = true;
+        resp.write_return = std::move(write_return);
+
+        io_.Send(client_request.address, client_request.request_id, std::move(resp));
+        leader.pending_client_requests.erase(apply_index);
+      }
+    }
+
+    Log("committed_log_size is now ", state_.committed_log_size);
+  }
+
+  // Raft paper - 5.1
+  // AppendEntries RPCs are initiated by leaders to replicate log entries and to provide a form of heartbeat
+  void BroadcastAppendEntries(std::map<Address, FollowerTracker> &followers) {
+    for (auto &[address, follower] : followers) {
+      const LogIndex index = follower.confirmed_contiguous_index;
+
+      std::vector<std::pair<Term, WriteOperation>> entries;
+
+      if (state_.log.size() > index) {
+        entries.insert(entries.begin(), state_.log.begin() + index, state_.log.end());
+      }
+
+      const Term previous_term_from_index = PreviousTermFromIndex(index);
+
+      Log("sending ", entries.size(), " entries to Follower ", address.last_known_port,
+          " which are above its known index of ", index);
+
+      AppendRequest<WriteOperation> ar{
+          .term = state_.term,
+          .last_log_index = index,
+          .last_log_term = previous_term_from_index,
+          .entries = entries,
+          .leader_commit = state_.committed_log_size,
+      };
+
+      // request_id not necessary to set because it's not a Future-backed Request.
+      static constexpr RequestId request_id = 0;
+
+      io_.Send(address, request_id, ar);
+    }
+  }
+
+  // Raft paper - 5.2
+  // Raft uses randomized election timeouts to ensure that split votes are rare and that they are resolved quickly
+  Duration RandomTimeout(Duration min, Duration max) {
+    std::uniform_int_distribution time_distrib(min.count(), max.count());
+
+    auto rand_micros = io_.Rand(time_distrib);
+
+    return std::chrono::microseconds{rand_micros};
+  }
+
+  Duration RandomTimeout(int min_micros, int max_micros) {
+    std::uniform_int_distribution time_distrib(min_micros, max_micros);
+
+    int rand_micros = io_.Rand(time_distrib);
+
+    return std::chrono::microseconds{rand_micros};
+  }
+
+  Term PreviousTermFromIndex(LogIndex index) const {
+    if (index == 0 || state_.log.size() + 1 <= index) {
+      return 0;
+    }
+
+    const auto &[term, data] = state_.log.at(index - 1);
+    return term;
+  }
+
+  LogIndex CommittedLogIndex() { return state_.committed_log_size; }
+
+  Term CommittedLogTerm() {
+    MG_ASSERT(state_.log.size() >= state_.committed_log_size);
+    if (state_.log.empty() || state_.committed_log_size == 0) {
+      return 0;
+    }
+
+    auto &[term, data] = state_.log.at(state_.committed_log_size - 1);
+    return term;
+  }
+
+  LogIndex LastLogIndex() { return state_.log.size(); }
+
+  Term LastLogTerm() const {
+    if (state_.log.empty()) {
+      return 0;
+    }
+
+    const auto &[term, data] = state_.log.back();
+    return term;
+  }
+
+  template <typename... Ts>
+  void Log(Ts &&...args) {
+    const Time now = io_.Now();
+    auto micros = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()).count();
+
+    const Term term = state_.term;
+
+    std::ostringstream out;
+
+    out << '\t' << (int)micros << "\t" << term << "\t" << io_.GetAddress().last_known_port;
+
+    std::string role_string = std::visit([&](auto &&role) { return role.ToString(); }, role_);
+
+    out << role_string;
+
+    (out << ... << args);
+
+    spdlog::debug(out.str());
+  }
+
+  /////////////////////////////////////////////////////////////
+  /// Raft-related Cron methods
+  ///
+  /// Cron + std::visit is how events are dispatched
+  /// to certain code based on Raft role.
+  ///
+  /// Cron(role) takes as the first argument a reference to its
+  /// role, and as the second argument, the message that has
+  /// been received.
+  /////////////////////////////////////////////////////////////
+
+  /// Periodic protocol maintenance.
+  void Cron() {
+    // dispatch periodic logic based on our role to a specific Cron method.
+    std::optional<Role> new_role = std::visit([&](auto &&role) { return Cron(role); }, role_);
+
+    if (new_role) {
+      role_ = std::move(new_role).value();
+    }
+  }
+
+  // Raft paper - 5.2
+  // Candidates keep sending Vote to peers until:
+  // 1. receiving Append with a higher term (become Follower)
+  // 2. receiving Vote with a higher term (become a Follower)
+  // 3. receiving a quorum of responses to our last batch of Vote (become a Leader)
+  std::optional<Role> Cron(Candidate &candidate) {
+    const auto now = io_.Now();
+    const Duration election_timeout = RandomTimeout(100000, 200000);
+    const auto election_timeout_us = std::chrono::duration_cast<std::chrono::milliseconds>(election_timeout).count();
+
+    if (now - candidate.election_began > election_timeout) {
+      state_.term++;
+      Log("becoming Candidate for term ", state_.term, " after leader timeout of ", election_timeout_us,
+          " elapsed since last election attempt");
+
+      const VoteRequest request{
+          .term = state_.term,
+          .last_log_index = LastLogIndex(),
+          .last_log_term = LastLogTerm(),
+      };
+
+      auto outstanding_votes = std::set<Address>();
+
+      for (const auto &peer : peers_) {
+        // request_id not necessary to set because it's not a Future-backed Request.
+        static constexpr auto request_id = 0;
+        io_.template Send<VoteRequest>(peer, request_id, request);
+        outstanding_votes.insert(peer);
+      }
+
+      return Candidate{
+          .successful_votes = std::map<Address, LogIndex>(),
+          .election_began = now,
+          .outstanding_votes = outstanding_votes,
+      };
+    }
+    return std::nullopt;
+  }
+
+  // Raft paper - 5.2
+  // Followers become candidates if we haven't heard from the leader
+  // after a randomized timeout.
+  std::optional<Role> Cron(Follower &follower) {
+    const auto now = io_.Now();
+    const auto time_since_last_append_entries = now - follower.last_received_append_entries_timestamp;
+    Duration election_timeout = RandomTimeout(100000, 200000);
+
+    // randomized follower timeout with a range of 100-150ms.
+    if (time_since_last_append_entries > election_timeout) {
+      // become a Candidate if we haven't heard from the Leader after this timeout
+      return Candidate{};
+    }
+
+    return std::nullopt;
+  }
+
+  // Leaders (re)send AppendRequest to followers.
+  std::optional<Role> Cron(Leader &leader) {
+    const Time now = io_.Now();
+    const Duration broadcast_timeout = RandomTimeout(40000, 60000);
+
+    if (now - leader.last_broadcast > broadcast_timeout) {
+      BroadcastAppendEntries(leader.followers);
+      leader.last_broadcast = now;
+    }
+
+    return std::nullopt;
+  }
+
+  /////////////////////////////////////////////////////////////
+  /// Raft-related Handle methods
+  ///
+  /// Handle + std::visit is how events are dispatched
+  /// to certain code based on Raft role.
+  ///
+  /// Handle(role, message, ...)
+  /// takes as the first argument a reference
+  /// to its role, and as the second argument, the
+  /// message that has been received.
+  /////////////////////////////////////////////////////////////
+
+  void Handle(std::variant<ReadRequest<ReadOperation>, AppendRequest<WriteOperation>, AppendResponse,
+                           WriteRequest<WriteOperation>, VoteRequest, VoteResponse> &&message_variant,
+              RequestId request_id, Address from_address) {
+    // dispatch the message to a handler based on our role,
+    // which can be specified in the Handle first argument,
+    // or it can be `auto` if it's a handler for several roles
+    // or messages.
+    std::optional<Role> new_role =
+        std::visit([&](auto &&msg, auto &&role) { return Handle(role, std::move(msg), request_id, from_address); },
+                   std::move(message_variant), role_);
+
+    // TODO(tyler) (M3) maybe replace std::visit with get_if for explicit prioritized matching, [[likely]] etc...
+    if (new_role) {
+      role_ = std::move(new_role).value();
+    }
+  }
+
+  // all roles can receive Vote and possibly become a follower
+  template <typename AllRoles>
+  std::optional<Role> Handle(AllRoles &, VoteRequest &&req, RequestId request_id, Address from_address) {
+    Log("received Vote from ", from_address.last_known_port, " with term ", req.term);
+    const bool last_log_term_dominates = req.last_log_term >= LastLogTerm();
+    const bool term_dominates = req.term > state_.term;
+    const bool last_log_index_dominates = req.last_log_index >= LastLogIndex();
+    const bool new_leader = last_log_term_dominates && term_dominates && last_log_index_dominates;
+
+    if (new_leader) {
+      MG_ASSERT(req.term > state_.term);
+      MG_ASSERT(std::max(req.term, state_.term) == req.term);
+    }
+
+    const VoteResponse res{
+        .term = std::max(req.term, state_.term),
+        .committed_log_size = state_.committed_log_size,
+        .vote_granted = new_leader,
+    };
+
+    io_.Send(from_address, request_id, res);
+
+    if (new_leader) {
+      // become a follower
+      state_.term = req.term;
+      return Follower{
+          .last_received_append_entries_timestamp = io_.Now(),
+          .leader_address = from_address,
+      };
+    } else if (term_dominates) {
+      Log("received a vote from an inferior candidate. Becoming Candidate");
+      state_.term = std::max(state_.term, req.term) + 1;
+      return Candidate{};
+    }
+
+    return std::nullopt;
+  }
+
+  std::optional<Role> Handle(Candidate &candidate, VoteResponse &&res, RequestId, Address from_address) {
+    Log("received VoteResponse");
+
+    if (!res.vote_granted || res.term != state_.term) {
+      Log("received unsuccessful VoteResponse from term ", res.term, " when our candidacy term is ", state_.term);
+      // we received a delayed VoteResponse from the past, which has to do with an election that is
+      // no longer valid. We can simply drop this.
+      return std::nullopt;
+    }
+
+    MG_ASSERT(candidate.outstanding_votes.contains(from_address),
+              "Received unexpected VoteResponse from server not present in Candidate's outstanding_votes!");
+    candidate.outstanding_votes.erase(from_address);
+
+    MG_ASSERT(!candidate.successful_votes.contains(from_address),
+              "Received unexpected VoteResponse from server already in Candidate's successful_votes!");
+    candidate.successful_votes.insert({from_address, res.committed_log_size});
+
+    if (candidate.successful_votes.size() >= candidate.outstanding_votes.size()) {
+      std::map<Address, FollowerTracker> followers{};
+
+      for (const auto &[address, committed_log_size] : candidate.successful_votes) {
+        FollowerTracker follower{
+            .next_index = committed_log_size,
+            .confirmed_contiguous_index = committed_log_size,
+        };
+        followers.insert({address, follower});
+      }
+      for (const auto &address : candidate.outstanding_votes) {
+        FollowerTracker follower{
+            .next_index = state_.log.size(),
+            .confirmed_contiguous_index = 0,
+        };
+        followers.insert({address, follower});
+      }
+
+      Log("becoming Leader at term ", state_.term);
+
+      BroadcastAppendEntries(followers);
+
+      return Leader{
+          .followers = std::move(followers),
+          .pending_client_requests = std::unordered_map<LogIndex, PendingClientRequest>(),
+      };
+    }
+
+    return std::nullopt;
+  }
+
+  template <typename AllRoles>
+  std::optional<Role> Handle(AllRoles &, VoteResponse &&, RequestId, Address) {
+    Log("non-Candidate received VoteResponse");
+    return std::nullopt;
+  }
+
+  template <typename AllRoles>
+  std::optional<Role> Handle(AllRoles &role, AppendRequest<WriteOperation> &&req, RequestId request_id,
+                             Address from_address) {
+    AppendResponse res{
+        .success = false,
+        .term = state_.term,
+        .last_log_term = CommittedLogTerm(),
+        .last_log_index = CommittedLogIndex(),
+    };
+
+    if constexpr (std::is_same<AllRoles, Leader>()) {
+      MG_ASSERT(req.term != state_.term, "Multiple leaders are acting under the term ", req.term);
+    }
+
+    const bool is_candidate = std::is_same<AllRoles, Candidate>();
+    const bool is_failed_competitor = is_candidate && req.term == state_.term;
+    const Time now = io_.Now();
+
+    // Raft paper - 5.2
+    // While waiting for votes, a candidate may receive an
+    // AppendEntries RPC from another server claiming to be leader. If
+    // the leader’s term (included in its RPC) is at least as large as
+    // the candidate’s current term, then the candidate recognizes the
+    // leader as legitimate and returns to follower state.
+    if (req.term > state_.term || is_failed_competitor) {
+      // become follower of this leader, reply with our log status
+      state_.term = req.term;
+
+      io_.Send(from_address, request_id, res);
+
+      Log("becoming Follower of Leader ", from_address.last_known_port, " at term ", req.term);
+      return Follower{
+          .last_received_append_entries_timestamp = now,
+          .leader_address = from_address,
+      };
+    } else if (req.term < state_.term) {
+      // nack this request from an old leader
+      io_.Send(from_address, request_id, res);
+
+      return std::nullopt;
+    }
+
+    // at this point, we're dealing with our own leader
+    if constexpr (std::is_same<AllRoles, Follower>()) {
+      // small specialization for when we're already a Follower
+      MG_ASSERT(role.leader_address == from_address, "Multiple Leaders are acting under the same term number!");
+      role.last_received_append_entries_timestamp = now;
+    } else {
+      Log("Somehow entered Follower-specific logic as a non-Follower");
+      MG_ASSERT(false, "Somehow entered Follower-specific logic as a non-Follower");
+    }
+
+    res.last_log_term = LastLogTerm();
+    res.last_log_index = LastLogIndex();
+
+    Log("returning last_log_index of ", res.last_log_index);
+
+    // Handle steady-state conditions.
+    if (req.last_log_index != LastLogIndex()) {
+      Log("req.last_log_index is above our last applied log index");
+    } else if (req.last_log_term != LastLogTerm()) {
+      Log("req.last_log_term differs from our leader term at that slot, expected: ", LastLogTerm(), " but got ",
+          req.last_log_term);
+    } else {
+      // happy path - Apply log
+      Log("applying batch of entries to log of size ", req.entries.size());
+
+      MG_ASSERT(req.last_log_index >= state_.committed_log_size,
+                "Applied history from Leader which goes back in time from our commit_index");
+
+      // possibly chop-off stuff that was replaced by
+      // things with different terms (we got data that
+      // hasn't reached consensus yet, which is normal)
+      state_.log.resize(req.last_log_index);
+
+      state_.log.insert(state_.log.end(), req.entries.begin(), req.entries.end());
+
+      MG_ASSERT(req.leader_commit >= state_.committed_log_size);
+      state_.committed_log_size = std::min(req.leader_commit, LastLogIndex());
+
+      for (; state_.applied_size < state_.committed_log_size; state_.applied_size++) {
+        const auto &write_request = state_.log[state_.applied_size].second;
+        replicated_state_.Apply(write_request);
+      }
+
+      res.success = true;
+    }
+
+    io_.Send(from_address, request_id, res);
+
+    return std::nullopt;
+  }
+
+  std::optional<Role> Handle(Leader &leader, AppendResponse &&res, RequestId, Address from_address) {
+    if (res.term != state_.term) {
+    } else if (!leader.followers.contains(from_address)) {
+      Log("received AppendResponse from unknown Follower");
+      MG_ASSERT(false, "received AppendResponse from unknown Follower");
+    } else {
+      if (res.success) {
+        Log("got successful AppendResponse from ", from_address.last_known_port, " with last_log_index of ",
+            res.last_log_index);
+      } else {
+        Log("got unsuccessful AppendResponse from ", from_address.last_known_port, " with last_log_index of ",
+            res.last_log_index);
+      }
+      FollowerTracker &follower = leader.followers.at(from_address);
+      follower.next_index = std::max(follower.next_index, res.last_log_index);
+      follower.confirmed_contiguous_index = std::max(follower.confirmed_contiguous_index, res.last_log_index);
+
+      BumpCommitIndexAndReplyToClients(leader);
+    }
+    return std::nullopt;
+  }
+
+  template <typename AllRoles>
+  std::optional<Role> Handle(AllRoles &, AppendResponse &&, RequestId, Address) {
+    // we used to be the leader, and are getting old delayed responses
+    return std::nullopt;
+  }
+
+  /////////////////////////////////////////////////////////////
+  /// RSM-related handle methods
+  /////////////////////////////////////////////////////////////
+
+  // Leaders are able to immediately respond to the requester (with a ReadResponseValue) applied to the ReplicatedState
+  std::optional<Role> Handle(Leader &, ReadRequest<ReadOperation> &&req, RequestId request_id, Address from_address) {
+    Log("handling ReadOperation");
+    ReadOperation read_operation = req.operation;
+
+    ReadResponseValue read_return = replicated_state_.Read(read_operation);
+
+    ReadResponse<ReadResponseValue> resp{
+        .success = true,
+        .read_return = std::move(read_return),
+        .retry_leader = std::nullopt,
+    };
+
+    io_.Send(from_address, request_id, resp);
+
+    return std::nullopt;
+  }
+
+  // Candidates should respond with a failure, similar to the Candidate + WriteRequest failure below
+  std::optional<Role> Handle(Candidate &, ReadRequest<ReadOperation> &&, RequestId request_id, Address from_address) {
+    Log("received ReadOperation - not redirecting because no Leader is known");
+    auto res = ReadResponse<ReadResponseValue>{};
+
+    res.success = false;
+
+    Cron();
+
+    io_.Send(from_address, request_id, res);
+
+    return std::nullopt;
+  }
+
+  // Followers should respond with a redirection, similar to the Follower + WriteRequest response below
+  std::optional<Role> Handle(Follower &follower, ReadRequest<ReadOperation> &&, RequestId request_id,
+                             Address from_address) {
+    auto res = ReadResponse<ReadResponseValue>{};
+
+    res.success = false;
+    Log("redirecting client to known Leader with port ", follower.leader_address.last_known_port);
+    res.retry_leader = follower.leader_address;
+
+    io_.Send(from_address, request_id, res);
+
+    return std::nullopt;
+  }
+
+  // Raft paper - 8
+  // When a client first starts up, it connects to a randomly chosen
+  // server. If the client’s first choice is not the leader, that
+  // server will reject the client’s request and supply information
+  // about the most recent leader it has heard from.
+  std::optional<Role> Handle(Follower &follower, WriteRequest<WriteOperation> &&, RequestId request_id,
+                             Address from_address) {
+    auto res = WriteResponse<WriteResponseValue>{};
+
+    res.success = false;
+    Log("redirecting client to known Leader with port ", follower.leader_address.last_known_port);
+    res.retry_leader = follower.leader_address;
+
+    io_.Send(from_address, request_id, res);
+
+    return std::nullopt;
+  }
+
+  std::optional<Role> Handle(Candidate &, WriteRequest<WriteOperation> &&, RequestId request_id, Address from_address) {
+    Log("received WriteRequest - not redirecting because no Leader is known");
+    auto res = WriteResponse<WriteResponseValue>{};
+
+    res.success = false;
+
+    Cron();
+
+    io_.Send(from_address, request_id, res);
+
+    return std::nullopt;
+  }
+
+  // only leaders actually handle replication requests from clients
+  std::optional<Role> Handle(Leader &leader, WriteRequest<WriteOperation> &&req, RequestId request_id,
+                             Address from_address) {
+    Log("handling WriteRequest");
+
+    // we are the leader. add item to log and send Append to peers
+    state_.log.emplace_back(std::pair(state_.term, std::move(req.operation)));
+
+    LogIndex log_index = state_.log.size() - 1;
+
+    PendingClientRequest pcr{
+        .request_id = request_id,
+        .address = from_address,
+        .received_at = io_.Now(),
+    };
+
+    leader.pending_client_requests.emplace(log_index, pcr);
+
+    BroadcastAppendEntries(leader.followers);
+
+    return std::nullopt;
+  }
+};
+
+};  // namespace memgraph::io::rsm

src/io/rsm/rsm_client.hpp

@@ -0,0 +1,137 @@
+// Copyright 2022 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 <iostream>
+#include <optional>
+#include <vector>
+
+#include "io/address.hpp"
+#include "io/rsm/raft.hpp"
+#include "utils/result.hpp"
+
+namespace memgraph::io::rsm {
+
+using memgraph::io::Address;
+using memgraph::io::Duration;
+using memgraph::io::ResponseEnvelope;
+using memgraph::io::ResponseFuture;
+using memgraph::io::ResponseResult;
+using memgraph::io::Time;
+using memgraph::io::TimedOut;
+using memgraph::io::rsm::ReadRequest;
+using memgraph::io::rsm::ReadResponse;
+using memgraph::io::rsm::WriteRequest;
+using memgraph::io::rsm::WriteResponse;
+using memgraph::utils::BasicResult;
+
+template <typename IoImpl, typename WriteRequestT, typename WriteResponseT, typename ReadRequestT,
+          typename ReadResponseT>
+class RsmClient {
+  using ServerPool = std::vector<Address>;
+
+  IoImpl io_;
+  Address leader_;
+
+  std::mt19937 cli_rng_{0};
+  ServerPool server_addrs_;
+
+  template <typename ResponseT>
+  void PossiblyRedirectLeader(const ResponseT &response) {
+    if (response.retry_leader) {
+      MG_ASSERT(!response.success, "retry_leader should never be set for successful responses");
+      leader_ = response.retry_leader.value();
+      spdlog::debug("client redirected to leader server {}", leader_.ToString());
+    } else if (!response.success) {
+      std::uniform_int_distribution<size_t> addr_distrib(0, (server_addrs_.size() - 1));
+      size_t addr_index = addr_distrib(cli_rng_);
+      leader_ = server_addrs_[addr_index];
+
+      spdlog::debug(
+          "client NOT redirected to leader server despite our success failing to be processed (it probably was sent to "
+          "a RSM Candidate) trying a random one at index {} with address {}",
+          addr_index, leader_.ToString());
+    }
+  }
+
+ public:
+  RsmClient(IoImpl io, Address leader, ServerPool server_addrs)
+      : io_{io}, leader_{leader}, server_addrs_{server_addrs} {}
+
+  RsmClient() = delete;
+
+  BasicResult<TimedOut, WriteResponseT> SendWriteRequest(WriteRequestT req) {
+    WriteRequest<WriteRequestT> client_req;
+    client_req.operation = req;
+
+    const Duration overall_timeout = io_.GetDefaultTimeout();
+    const Time before = io_.Now();
+
+    do {
+      spdlog::debug("client sending WriteRequest to Leader {}", leader_.ToString());
+      ResponseFuture<WriteResponse<WriteResponseT>> response_future =
+          io_.template Request<WriteRequest<WriteRequestT>, WriteResponse<WriteResponseT>>(leader_, client_req);
+      ResponseResult<WriteResponse<WriteResponseT>> response_result = std::move(response_future).Wait();
+
+      if (response_result.HasError()) {
+        spdlog::debug("client timed out while trying to communicate with leader server {}", leader_.ToString());
+        // continue;
+        return response_result.GetError();
+      }
+
+      ResponseEnvelope<WriteResponse<WriteResponseT>> &&response_envelope = std::move(response_result.GetValue());
+      WriteResponse<WriteResponseT> &&write_response = std::move(response_envelope.message);
+
+      if (write_response.success) {
+        return std::move(write_response.write_return);
+      }
+
+      PossiblyRedirectLeader(write_response);
+    } while (io_.Now() < before + overall_timeout);
+
+    return TimedOut{};
+  }
+
+  BasicResult<TimedOut, ReadResponseT> SendReadRequest(ReadRequestT req) {
+    ReadRequest<ReadRequestT> read_req;
+    read_req.operation = req;
+
+    const Duration overall_timeout = io_.GetDefaultTimeout();
+    const Time before = io_.Now();
+
+    do {
+      spdlog::debug("client sending ReadRequest to Leader {}", leader_.ToString());
+
+      ResponseFuture<ReadResponse<ReadResponseT>> get_response_future =
+          io_.template Request<ReadRequest<ReadRequestT>, ReadResponse<ReadResponseT>>(leader_, read_req);
+
+      // receive response
+      ResponseResult<ReadResponse<ReadResponseT>> get_response_result = std::move(get_response_future).Wait();
+
+      if (get_response_result.HasError()) {
+        spdlog::debug("client timed out while trying to communicate with leader server {}", leader_.ToString());
+        return get_response_result.GetError();
+      }
+
+      ResponseEnvelope<ReadResponse<ReadResponseT>> &&get_response_envelope = std::move(get_response_result.GetValue());
+      ReadResponse<ReadResponseT> &&read_get_response = std::move(get_response_envelope.message);
+
+      if (read_get_response.success) {
+        return std::move(read_get_response.read_return);
+      }
+
+      PossiblyRedirectLeader(read_get_response);
+    } while (io_.Now() < before + overall_timeout);
+
+    return TimedOut{};
+  }
+};
+
+}  // namespace memgraph::io::rsm

src/io/rsm/shard_rsm.hpp

@@ -0,0 +1,165 @@
+// Copyright 2022 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
+
+/// The StorageRsm is a simple in-memory raft-backed kv store that can be used for simple testing
+/// and implementation of some query engine logic before storage engines are fully implemented.
+///
+/// To implement multiple read and write commands, change the StorageRead* and StorageWrite* requests
+/// and responses to a std::variant of the different options, and route them to specific handlers in
+/// the StorageRsm's Read and Apply methods. Remember that Read is called immediately when the Raft
+/// leader receives the request, and does not replicate anything over Raft. Apply is called only
+/// AFTER the StorageWriteRequest is replicated to a majority of Raft peers, and the result of calling
+/// StorageRsm::Apply(StorageWriteRequest) is returned to the client that submitted the request.
+
+#include <deque>
+#include <iostream>
+#include <map>
+#include <optional>
+#include <set>
+#include <thread>
+#include <vector>
+
+#include "coordinator/hybrid_logical_clock.hpp"
+#include "io/address.hpp"
+#include "io/rsm/raft.hpp"
+#include "io/simulator/simulator.hpp"
+#include "io/simulator/simulator_transport.hpp"
+#include "storage/v2/property_value.hpp"
+#include "utils/logging.hpp"
+
+using memgraph::coordinator::Hlc;
+using memgraph::io::Address;
+using memgraph::io::Io;
+using memgraph::io::ResponseEnvelope;
+using memgraph::io::ResponseFuture;
+using memgraph::io::ResponseResult;
+using memgraph::io::rsm::Raft;
+using memgraph::io::rsm::ReadRequest;
+using memgraph::io::rsm::ReadResponse;
+using memgraph::io::rsm::WriteRequest;
+using memgraph::io::rsm::WriteResponse;
+using memgraph::io::simulator::Simulator;
+using memgraph::io::simulator::SimulatorConfig;
+using memgraph::io::simulator::SimulatorStats;
+using memgraph::io::simulator::SimulatorTransport;
+using memgraph::storage::PropertyValue;
+
+namespace memgraph::io::rsm {
+
+using ShardRsmKey = std::vector<PropertyValue>;
+
+struct StorageWriteRequest {
+  ShardRsmKey key;
+  std::optional<int> value;
+};
+
+struct StorageWriteResponse {
+  bool shard_rsm_success;
+  std::optional<int> last_value;
+  // Only has a value if the given shard does not contain the requested key
+  std::optional<Hlc> latest_known_shard_map_version{std::nullopt};
+};
+
+struct StorageReadRequest {
+  ShardRsmKey key;
+};
+
+struct StorageReadResponse {
+  bool shard_rsm_success;
+  std::optional<int> value;
+  // Only has a value if the given shard does not contain the requested key
+  std::optional<Hlc> latest_known_shard_map_version{std::nullopt};
+};
+
+class StorageRsm {
+  std::map<ShardRsmKey, int> state_;
+  ShardRsmKey minimum_key_;
+  std::optional<ShardRsmKey> maximum_key_{std::nullopt};
+  Hlc shard_map_version_;
+
+  // The key is not located in this shard
+  bool IsKeyInRange(const ShardRsmKey &key) {
+    if (maximum_key_) [[likely]] {
+      return (key >= minimum_key_ && key <= maximum_key_);
+    }
+    return key >= minimum_key_;
+  }
+
+ public:
+  StorageReadResponse Read(StorageReadRequest request) {
+    StorageReadResponse ret;
+
+    if (!IsKeyInRange(request.key)) {
+      ret.latest_known_shard_map_version = shard_map_version_;
+      ret.shard_rsm_success = false;
+    } else if (state_.contains(request.key)) {
+      ret.value = state_[request.key];
+      ret.shard_rsm_success = true;
+    } else {
+      ret.shard_rsm_success = false;
+      ret.value = std::nullopt;
+    }
+    return ret;
+  }
+
+  StorageWriteResponse Apply(StorageWriteRequest request) {
+    StorageWriteResponse ret;
+
+    // Key is outside the prohibited range
+    if (!IsKeyInRange(request.key)) {
+      ret.latest_known_shard_map_version = shard_map_version_;
+      ret.shard_rsm_success = false;
+    }
+    // Key exist
+    else if (state_.contains(request.key)) {
+      auto &val = state_[request.key];
+
+      /*
+       *   Delete
+       */
+      if (!request.value) {
+        ret.shard_rsm_success = true;
+        ret.last_value = val;
+        state_.erase(state_.find(request.key));
+      }
+
+      /*
+       *   Update
+       */
+      // Does old_value match?
+      if (request.value == val) {
+        ret.last_value = val;
+        ret.shard_rsm_success = true;
+
+        val = request.value.value();
+
+      } else {
+        ret.last_value = val;
+        ret.shard_rsm_success = false;
+      }
+    }
+    /*
+     *   Create
+     */
+    else {
+      ret.last_value = std::nullopt;
+      ret.shard_rsm_success = true;
+
+      state_.emplace(request.key, std::move(request.value).value());
+    }
+
+    return ret;
+  }
+};
+
+}  // namespace memgraph::io::rsm

src/io/simulator/simulator.hpp

@@ -23,6 +23,7 @@ namespace memgraph::io::simulator {
 class Simulator {
   std::mt19937 rng_;
   std::shared_ptr<SimulatorHandle> simulator_handle_;
+  uint16_t auto_port_ = 0;
 
  public:
   explicit Simulator(SimulatorConfig config)
@@ -30,6 +31,11 @@ class Simulator {
 
   void ShutDown() { simulator_handle_->ShutDown(); }
 
+  Io<SimulatorTransport> RegisterNew() {
+    Address address = Address::TestAddress(auto_port_++);
+    return Register(address);
+  }
+
   Io<SimulatorTransport> Register(Address address) {
     std::uniform_int_distribution<uint64_t> seed_distrib;
     uint64_t seed = seed_distrib(rng_);

src/io/transport.hpp

@@ -72,6 +72,9 @@ class Io {
   /// without an explicit timeout set.
   void SetDefaultTimeout(Duration timeout) { default_timeout_ = timeout; }
 
+  /// Returns the current default timeout for this Io instance.
+  Duration GetDefaultTimeout() { return default_timeout_; }
+
   /// Issue a request with an explicit timeout in microseconds provided. This tends to be used by clients.
   template <Message Request, Message Response>
   ResponseFuture<Response> RequestWithTimeout(Address address, Request request, Duration timeout) {

src/storage/v3/rsm_engine.hpp

@@ -0,0 +1,41 @@
+// Copyright 2022 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 <boost/uuid/uuid.hpp>
+#include "io/rsm/coordinator_rsm.hpp"
+#include "io/rsm/rsm_client.hpp"
+
+namespace memgraph::storage::v3 {
+
+using boost::uuids::uuid;
+using memgraph::coordinator;
+using memgraph::io::rsm::RsmClient;
+
+/// The RsmEngine is responsible for:
+/// * reconciling the storage engine's local configuration with the Coordinator's
+///   intentions for how it should participate in multiple raft clusters
+/// * replying to heartbeat requests to the Coordinator
+/// * routing incoming messages from the Io interface to the appropriate
+///   storage RSM
+///
+/// Every storage engine has exactly one RsmEngine.
+template <typename IoImpl>
+class RsmEngine {
+  Io<IoImpl> io_;
+  std::map<uuid, StorageRsm> rsm_map_;
+
+ public:
+  RsmEngine(Io<IoImpl> io) : io_(io) {}
+};
+
+}  // namespace memgraph::storage::v3

tests/simulation/CMakeLists.txt

@@ -27,4 +27,8 @@ endfunction(add_simulation_test)
 
 add_simulation_test(basic_request.cpp address)
 
+add_simulation_test(raft.cpp address)
+
 add_simulation_test(trial_query_storage/query_storage_test.cpp address)
+
+add_simulation_test(sharded_map.cpp address)

tests/simulation/raft.cpp

@@ -0,0 +1,265 @@
+// Copyright 2022 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 <chrono>
+#include <deque>
+#include <iostream>
+#include <map>
+#include <optional>
+#include <set>
+#include <thread>
+#include <vector>
+
+#include "io/address.hpp"
+#include "io/rsm/raft.hpp"
+#include "io/rsm/rsm_client.hpp"
+#include "io/simulator/simulator.hpp"
+#include "io/simulator/simulator_transport.hpp"
+
+using memgraph::io::Address;
+using memgraph::io::Duration;
+using memgraph::io::Io;
+using memgraph::io::ResponseEnvelope;
+using memgraph::io::ResponseFuture;
+using memgraph::io::ResponseResult;
+using memgraph::io::Time;
+using memgraph::io::rsm::Raft;
+using memgraph::io::rsm::ReadRequest;
+using memgraph::io::rsm::ReadResponse;
+using memgraph::io::rsm::RsmClient;
+using memgraph::io::rsm::WriteRequest;
+using memgraph::io::rsm::WriteResponse;
+using memgraph::io::simulator::Simulator;
+using memgraph::io::simulator::SimulatorConfig;
+using memgraph::io::simulator::SimulatorStats;
+using memgraph::io::simulator::SimulatorTransport;
+
+struct CasRequest {
+  int key;
+  std::optional<int> old_value;
+  std::optional<int> new_value;
+};
+
+struct CasResponse {
+  bool cas_success;
+  std::optional<int> last_value;
+};
+
+struct GetRequest {
+  int key;
+};
+
+struct GetResponse {
+  std::optional<int> value;
+};
+
+class TestState {
+  std::map<int, int> state_;
+
+ public:
+  GetResponse Read(GetRequest request) {
+    GetResponse ret;
+    if (state_.contains(request.key)) {
+      ret.value = state_[request.key];
+    }
+    return ret;
+  }
+
+  CasResponse Apply(CasRequest request) {
+    CasResponse ret;
+
+    // Key exist
+    if (state_.contains(request.key)) {
+      auto &val = state_[request.key];
+
+      /*
+       *   Delete
+       */
+      if (!request.new_value) {
+        ret.last_value = val;
+        ret.cas_success = true;
+
+        state_.erase(state_.find(request.key));
+      }
+
+      /*
+       *   Update
+       */
+      // Does old_value match?
+      if (request.old_value == val) {
+        ret.last_value = val;
+        ret.cas_success = true;
+
+        val = request.new_value.value();
+      } else {
+        ret.last_value = val;
+        ret.cas_success = false;
+      }
+    }
+    /*
+     *   Create
+     */
+    else {
+      ret.last_value = std::nullopt;
+      ret.cas_success = true;
+
+      state_.emplace(request.key, std::move(request.new_value).value());
+    }
+
+    return ret;
+  }
+};
+
+template <typename IoImpl>
+void RunRaft(Raft<IoImpl, TestState, CasRequest, CasResponse, GetRequest, GetResponse> server) {
+  server.Run();
+}
+
+void RunSimulation() {
+  SimulatorConfig config{
+      .drop_percent = 5,
+      .perform_timeouts = true,
+      .scramble_messages = true,
+      .rng_seed = 0,
+      .start_time = Time::min() + std::chrono::microseconds{256 * 1024},
+      .abort_time = Time::min() + std::chrono::microseconds{8 * 1024 * 1024},
+  };
+
+  auto simulator = Simulator(config);
+
+  auto cli_addr = Address::TestAddress(1);
+  auto srv_addr_1 = Address::TestAddress(2);
+  auto srv_addr_2 = Address::TestAddress(3);
+  auto srv_addr_3 = Address::TestAddress(4);
+
+  Io<SimulatorTransport> cli_io = simulator.Register(cli_addr);
+  Io<SimulatorTransport> srv_io_1 = simulator.Register(srv_addr_1);
+  Io<SimulatorTransport> srv_io_2 = simulator.Register(srv_addr_2);
+  Io<SimulatorTransport> srv_io_3 = simulator.Register(srv_addr_3);
+
+  std::vector<Address> srv_1_peers = {srv_addr_2, srv_addr_3};
+  std::vector<Address> srv_2_peers = {srv_addr_1, srv_addr_3};
+  std::vector<Address> srv_3_peers = {srv_addr_1, srv_addr_2};
+
+  using RaftClass = Raft<SimulatorTransport, TestState, CasRequest, CasResponse, GetRequest, GetResponse>;
+  RaftClass srv_1{std::move(srv_io_1), srv_1_peers, TestState{}};
+  RaftClass srv_2{std::move(srv_io_2), srv_2_peers, TestState{}};
+  RaftClass srv_3{std::move(srv_io_3), srv_3_peers, TestState{}};
+
+  auto srv_thread_1 = std::jthread(RunRaft<SimulatorTransport>, std::move(srv_1));
+  simulator.IncrementServerCountAndWaitForQuiescentState(srv_addr_1);
+
+  auto srv_thread_2 = std::jthread(RunRaft<SimulatorTransport>, std::move(srv_2));
+  simulator.IncrementServerCountAndWaitForQuiescentState(srv_addr_2);
+
+  auto srv_thread_3 = std::jthread(RunRaft<SimulatorTransport>, std::move(srv_3));
+  simulator.IncrementServerCountAndWaitForQuiescentState(srv_addr_3);
+
+  spdlog::debug("beginning test after servers have become quiescent");
+
+  std::mt19937 cli_rng_{0};
+  std::vector<Address> server_addrs{srv_addr_1, srv_addr_2, srv_addr_3};
+  Address leader = server_addrs[0];
+
+  RsmClient<Io<SimulatorTransport>, CasRequest, CasResponse, GetRequest, GetResponse> client(cli_io, leader,
+                                                                                             server_addrs);
+
+  const int key = 0;
+  std::optional<int> last_known_value;
+
+  bool success = false;
+
+  for (int i = 0; !success; i++) {
+    /*
+     * Write Request
+     */
+    CasRequest cas_req;
+    cas_req.key = key;
+
+    cas_req.old_value = last_known_value;
+
+    cas_req.new_value = i;
+
+    auto write_cas_response_result = client.SendWriteRequest(cas_req);
+    if (write_cas_response_result.HasError()) {
+      // timed out
+      continue;
+    }
+    CasResponse cas_response = write_cas_response_result.GetValue();
+
+    bool cas_succeeded = cas_response.cas_success;
+
+    spdlog::debug("Client received CasResponse! success: {} last_known_value {}", cas_succeeded,
+                  (int)*last_known_value);
+
+    if (cas_succeeded) {
+      last_known_value = i;
+    } else {
+      last_known_value = cas_response.last_value;
+      continue;
+    }
+
+    /*
+     * Get Request
+     */
+    GetRequest get_req;
+    get_req.key = key;
+
+    auto read_get_response_result = client.SendReadRequest(get_req);
+    if (read_get_response_result.HasError()) {
+      // timed out
+      continue;
+    }
+    GetResponse get_response = read_get_response_result.GetValue();
+
+    MG_ASSERT(get_response.value == i);
+
+    spdlog::debug("client successfully cas'd a value and read it back! value: {}", i);
+
+    success = true;
+  }
+
+  MG_ASSERT(success);
+
+  simulator.ShutDown();
+
+  SimulatorStats stats = simulator.Stats();
+
+  spdlog::debug("total messages:     ", stats.total_messages);
+  spdlog::debug("dropped messages:   ", stats.dropped_messages);
+  spdlog::debug("timed out requests: ", stats.timed_out_requests);
+  spdlog::debug("total requests:     ", stats.total_requests);
+  spdlog::debug("total responses:    ", stats.total_responses);
+  spdlog::debug("simulator ticks:    ", stats.simulator_ticks);
+
+  spdlog::debug("========================== SUCCESS :) ==========================");
+
+  /*
+  this is implicit in jthread's dtor
+  srv_thread_1.join();
+  srv_thread_2.join();
+  srv_thread_3.join();
+  */
+}
+
+int main() {
+  int n_tests = 50;
+
+  for (int i = 0; i < n_tests; i++) {
+    spdlog::debug("========================== NEW SIMULATION {} ==========================", i);
+    spdlog::debug("\tTime\tTerm\tPort\tRole\t\tMessage\n");
+    RunSimulation();
+  }
+
+  spdlog::debug("passed {} tests!", n_tests);
+
+  return 0;
+}

tests/simulation/sharded_map.cpp

@@ -0,0 +1,340 @@
+// Copyright 2022 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 <chrono>
+#include <deque>
+#include <iostream>
+#include <map>
+#include <optional>
+#include <set>
+#include <thread>
+#include <vector>
+
+#include "io/address.hpp"
+#include "io/errors.hpp"
+#include "io/rsm/coordinator_rsm.hpp"
+#include "io/rsm/raft.hpp"
+#include "io/rsm/rsm_client.hpp"
+#include "io/rsm/shard_rsm.hpp"
+#include "io/simulator/simulator.hpp"
+#include "io/simulator/simulator_transport.hpp"
+#include "utils/result.hpp"
+
+using memgraph::coordinator::Address;
+using memgraph::coordinator::AddressAndStatus;
+using memgraph::coordinator::CompoundKey;
+using memgraph::coordinator::Coordinator;
+using memgraph::coordinator::HlcRequest;
+using memgraph::coordinator::HlcResponse;
+using memgraph::coordinator::Shard;
+using memgraph::coordinator::ShardMap;
+using memgraph::coordinator::Shards;
+using memgraph::coordinator::Status;
+using memgraph::io::Address;
+using memgraph::io::Io;
+using memgraph::io::ResponseEnvelope;
+using memgraph::io::ResponseFuture;
+using memgraph::io::Time;
+using memgraph::io::TimedOut;
+using memgraph::io::rsm::CoordinatorRsm;
+using memgraph::io::rsm::Raft;
+using memgraph::io::rsm::ReadRequest;
+using memgraph::io::rsm::ReadResponse;
+using memgraph::io::rsm::RsmClient;
+using memgraph::io::rsm::StorageReadRequest;
+using memgraph::io::rsm::StorageReadResponse;
+using memgraph::io::rsm::StorageRsm;
+using memgraph::io::rsm::StorageWriteRequest;
+using memgraph::io::rsm::StorageWriteResponse;
+using memgraph::io::rsm::WriteRequest;
+using memgraph::io::rsm::WriteResponse;
+using memgraph::io::simulator::Simulator;
+using memgraph::io::simulator::SimulatorConfig;
+using memgraph::io::simulator::SimulatorStats;
+using memgraph::io::simulator::SimulatorTransport;
+using memgraph::utils::BasicResult;
+
+using StorageClient = RsmClient<Io<SimulatorTransport>, StorageWriteRequest, StorageWriteResponse, StorageReadRequest,
+                                StorageReadResponse>;
+namespace {
+
+ShardMap CreateDummyShardmap(memgraph::coordinator::Address a_io_1, memgraph::coordinator::Address a_io_2,
+                             memgraph::coordinator::Address a_io_3, memgraph::coordinator::Address b_io_1,
+                             memgraph::coordinator::Address b_io_2, memgraph::coordinator::Address b_io_3) {
+  ShardMap sm1;
+  auto &shards = sm1.GetShards();
+
+  // 1
+  std::string label1 = std::string("label1");
+  auto key1 = memgraph::storage::v3::PropertyValue(3);
+  auto key2 = memgraph::storage::v3::PropertyValue(4);
+  CompoundKey cm1 = {key1, key2};
+  AddressAndStatus aas1_1{.address = a_io_1, .status = Status::CONSENSUS_PARTICIPANT};
+  AddressAndStatus aas1_2{.address = a_io_2, .status = Status::CONSENSUS_PARTICIPANT};
+  AddressAndStatus aas1_3{.address = a_io_3, .status = Status::CONSENSUS_PARTICIPANT};
+
+  Shard shard1 = {aas1_1, aas1_2, aas1_3};
+  Shards shards1;
+  shards1[cm1] = shard1;
+
+  // 2
+  std::string label2 = std::string("label2");
+  auto key3 = memgraph::storage::v3::PropertyValue(12);
+  auto key4 = memgraph::storage::v3::PropertyValue(13);
+  CompoundKey cm2 = {key3, key4};
+  AddressAndStatus aas2_1{.address = b_io_1, .status = Status::CONSENSUS_PARTICIPANT};
+  AddressAndStatus aas2_2{.address = b_io_2, .status = Status::CONSENSUS_PARTICIPANT};
+  AddressAndStatus aas2_3{.address = b_io_3, .status = Status::CONSENSUS_PARTICIPANT};
+
+  Shard shard2 = {aas2_1, aas2_2, aas2_3};
+  Shards shards2;
+  shards2[cm2] = shard2;
+
+  shards[label1] = shards1;
+  shards[label2] = shards2;
+
+  return sm1;
+}
+
+std::optional<StorageClient> DetermineShardLocation(Shard target_shard, const std::vector<Address> &a_addrs,
+                                                    StorageClient a_client, const std::vector<Address> &b_addrs,
+                                                    StorageClient b_client) {
+  for (const auto &addr : target_shard) {
+    if (addr.address == b_addrs[0]) {
+      return b_client;
+    }
+    if (addr.address == a_addrs[0]) {
+      return a_client;
+    }
+  }
+  return {};
+}
+
+}  // namespace
+
+using ConcreteCoordinatorRsm = CoordinatorRsm<SimulatorTransport>;
+using ConcreteStorageRsm = Raft<SimulatorTransport, StorageRsm, StorageWriteRequest, StorageWriteResponse,
+                                StorageReadRequest, StorageReadResponse>;
+
+template <typename IoImpl>
+void RunStorageRaft(
+    Raft<IoImpl, StorageRsm, StorageWriteRequest, StorageWriteResponse, StorageReadRequest, StorageReadResponse>
+        server) {
+  server.Run();
+}
+
+int main() {
+  SimulatorConfig config{
+      .drop_percent = 5,
+      .perform_timeouts = true,
+      .scramble_messages = true,
+      .rng_seed = 0,
+      .start_time = Time::min() + std::chrono::microseconds{256 * 1024},
+      .abort_time = Time::min() + std::chrono::microseconds{2 * 8 * 1024 * 1024},
+  };
+
+  auto simulator = Simulator(config);
+
+  Io<SimulatorTransport> cli_io = simulator.RegisterNew();
+
+  // Register
+  Io<SimulatorTransport> a_io_1 = simulator.RegisterNew();
+  Io<SimulatorTransport> a_io_2 = simulator.RegisterNew();
+  Io<SimulatorTransport> a_io_3 = simulator.RegisterNew();
+
+  Io<SimulatorTransport> b_io_1 = simulator.RegisterNew();
+  Io<SimulatorTransport> b_io_2 = simulator.RegisterNew();
+  Io<SimulatorTransport> b_io_3 = simulator.RegisterNew();
+
+  // Preconfigure coordinator with kv shard 'A' and 'B'
+  auto sm1 = CreateDummyShardmap(a_io_1.GetAddress(), a_io_2.GetAddress(), a_io_3.GetAddress(), b_io_1.GetAddress(),
+                                 b_io_2.GetAddress(), b_io_3.GetAddress());
+  auto sm2 = CreateDummyShardmap(a_io_1.GetAddress(), a_io_2.GetAddress(), a_io_3.GetAddress(), b_io_1.GetAddress(),
+                                 b_io_2.GetAddress(), b_io_3.GetAddress());
+  auto sm3 = CreateDummyShardmap(a_io_1.GetAddress(), a_io_2.GetAddress(), a_io_3.GetAddress(), b_io_1.GetAddress(),
+                                 b_io_2.GetAddress(), b_io_3.GetAddress());
+
+  // Spin up shard A
+  std::vector<Address> a_addrs = {a_io_1.GetAddress(), a_io_2.GetAddress(), a_io_3.GetAddress()};
+
+  std::vector<Address> a_1_peers = {a_addrs[1], a_addrs[2]};
+  std::vector<Address> a_2_peers = {a_addrs[0], a_addrs[2]};
+  std::vector<Address> a_3_peers = {a_addrs[0], a_addrs[1]};
+
+  ConcreteStorageRsm a_1{std::move(a_io_1), a_1_peers, StorageRsm{}};
+  ConcreteStorageRsm a_2{std::move(a_io_2), a_2_peers, StorageRsm{}};
+  ConcreteStorageRsm a_3{std::move(a_io_3), a_3_peers, StorageRsm{}};
+
+  auto a_thread_1 = std::jthread(RunStorageRaft<SimulatorTransport>, std::move(a_1));
+  simulator.IncrementServerCountAndWaitForQuiescentState(a_addrs[0]);
+
+  auto a_thread_2 = std::jthread(RunStorageRaft<SimulatorTransport>, std::move(a_2));
+  simulator.IncrementServerCountAndWaitForQuiescentState(a_addrs[1]);
+
+  auto a_thread_3 = std::jthread(RunStorageRaft<SimulatorTransport>, std::move(a_3));
+  simulator.IncrementServerCountAndWaitForQuiescentState(a_addrs[2]);
+
+  // Spin up shard B
+  std::vector<Address> b_addrs = {b_io_1.GetAddress(), b_io_2.GetAddress(), b_io_3.GetAddress()};
+
+  std::vector<Address> b_1_peers = {b_addrs[1], b_addrs[2]};
+  std::vector<Address> b_2_peers = {b_addrs[0], b_addrs[2]};
+  std::vector<Address> b_3_peers = {b_addrs[0], b_addrs[1]};
+
+  ConcreteStorageRsm b_1{std::move(b_io_1), b_1_peers, StorageRsm{}};
+  ConcreteStorageRsm b_2{std::move(b_io_2), b_2_peers, StorageRsm{}};
+  ConcreteStorageRsm b_3{std::move(b_io_3), b_3_peers, StorageRsm{}};
+
+  auto b_thread_1 = std::jthread(RunStorageRaft<SimulatorTransport>, std::move(b_1));
+  simulator.IncrementServerCountAndWaitForQuiescentState(b_addrs[0]);
+
+  auto b_thread_2 = std::jthread(RunStorageRaft<SimulatorTransport>, std::move(b_2));
+  simulator.IncrementServerCountAndWaitForQuiescentState(b_addrs[1]);
+
+  auto b_thread_3 = std::jthread(RunStorageRaft<SimulatorTransport>, std::move(b_3));
+  simulator.IncrementServerCountAndWaitForQuiescentState(b_addrs[2]);
+
+  // Spin up coordinators
+
+  Io<SimulatorTransport> c_io_1 = simulator.RegisterNew();
+  Io<SimulatorTransport> c_io_2 = simulator.RegisterNew();
+  Io<SimulatorTransport> c_io_3 = simulator.RegisterNew();
+
+  std::vector<Address> c_addrs = {c_io_1.GetAddress(), c_io_2.GetAddress(), c_io_3.GetAddress()};
+
+  std::vector<Address> c_1_peers = {c_addrs[1], c_addrs[2]};
+  std::vector<Address> c_2_peers = {c_addrs[0], c_addrs[2]};
+  std::vector<Address> c_3_peers = {c_addrs[0], c_addrs[1]};
+
+  ConcreteCoordinatorRsm c_1{std::move(c_io_1), c_1_peers, Coordinator{(sm1)}};
+  ConcreteCoordinatorRsm c_2{std::move(c_io_2), c_2_peers, Coordinator{(sm2)}};
+  ConcreteCoordinatorRsm c_3{std::move(c_io_3), c_3_peers, Coordinator{(sm3)}};
+
+  auto c_thread_1 = std::jthread([c_1]() mutable { c_1.Run(); });
+  simulator.IncrementServerCountAndWaitForQuiescentState(c_addrs[0]);
+
+  auto c_thread_2 = std::jthread([c_2]() mutable { c_2.Run(); });
+  simulator.IncrementServerCountAndWaitForQuiescentState(c_addrs[1]);
+
+  auto c_thread_3 = std::jthread([c_3]() mutable { c_3.Run(); });
+  simulator.IncrementServerCountAndWaitForQuiescentState(c_addrs[2]);
+
+  std::cout << "beginning test after servers have become quiescent" << std::endl;
+
+  // Have client contact coordinator RSM for a new transaction ID and
+  // also get the current shard map
+  using CoordinatorClient =
+      RsmClient<Io<SimulatorTransport>, memgraph::coordinator::WriteRequests, memgraph::coordinator::WriteResponses,
+                memgraph::coordinator::ReadRequests, memgraph::coordinator::ReadResponses>;
+  CoordinatorClient coordinator_client(cli_io, c_addrs[0], c_addrs);
+
+  StorageClient shard_a_client(cli_io, a_addrs[0], a_addrs);
+  StorageClient shard_b_client(cli_io, b_addrs[0], b_addrs);
+
+  memgraph::coordinator::HlcRequest req;
+
+  // Last ShardMap Version The query engine knows about.
+  ShardMap client_shard_map;
+  req.last_shard_map_version = client_shard_map.GetHlc();
+
+  while (true) {
+    // Create CompoundKey
+    const auto cm_key_1 = memgraph::storage::v3::PropertyValue(3);
+    const auto cm_key_2 = memgraph::storage::v3::PropertyValue(4);
+
+    const CompoundKey cm_k = {cm_key_1, cm_key_2};
+
+    // Look for Shard
+    BasicResult<TimedOut, memgraph::coordinator::ReadResponses> read_res = coordinator_client.SendReadRequest(req);
+
+    if (read_res.HasError()) {
+      // timeout
+      continue;
+    }
+
+    auto coordinator_read_response = read_res.GetValue();
+    HlcResponse hlc_response = std::get<HlcResponse>(coordinator_read_response);
+
+    // Transaction ID to be used later...
+    auto transaction_id = hlc_response.new_hlc;
+
+    if (hlc_response.fresher_shard_map) {
+      client_shard_map = hlc_response.fresher_shard_map.value();
+    }
+
+    // TODO(gabor) check somewhere in the call chain if the entries are actually valid
+    // for (auto &[key, val] : client_shard_map.GetShards()) {
+    //   std::cout << "key: " << key << std::endl;
+    // }
+
+    auto target_shard = client_shard_map.GetShardForKey(std::string("label1"), cm_k);
+
+    // Determine which shard to send the requests to
+    auto storage_client_opt = DetermineShardLocation(target_shard, a_addrs, shard_a_client, b_addrs, shard_b_client);
+    MG_ASSERT(storage_client_opt);
+
+    auto storage_client = storage_client_opt.value();
+
+    // Have client use shard map to decide which shard to communicate
+    // with in order to write a new value
+    // client_shard_map.
+    auto write_key_1 = memgraph::storage::PropertyValue(3);
+    auto write_key_2 = memgraph::storage::PropertyValue(4);
+
+    StorageWriteRequest storage_req;
+    storage_req.key = {write_key_1, write_key_2};
+    storage_req.value = 1000;
+
+    auto write_response_result = storage_client.SendWriteRequest(storage_req);
+    if (write_response_result.HasError()) {
+      // timed out
+      continue;
+    }
+    auto write_response = write_response_result.GetValue();
+
+    bool cas_succeeded = write_response.shard_rsm_success;
+
+    if (!cas_succeeded) {
+      continue;
+    }
+    // Have client use shard map to decide which shard to communicate
+    // with to read that same value back
+
+    StorageReadRequest storage_get_req;
+    storage_get_req.key = {write_key_1, write_key_2};
+
+    auto get_response_result = storage_client.SendReadRequest(storage_get_req);
+    if (get_response_result.HasError()) {
+      // timed out
+      continue;
+    }
+    auto get_response = get_response_result.GetValue();
+    auto val = get_response.value.value();
+
+    MG_ASSERT(val == 1000);
+    break;
+  }
+
+  simulator.ShutDown();
+
+  SimulatorStats stats = simulator.Stats();
+
+  std::cout << "total messages:     " << stats.total_messages << std::endl;
+  std::cout << "dropped messages:   " << stats.dropped_messages << std::endl;
+  std::cout << "timed out requests: " << stats.timed_out_requests << std::endl;
+  std::cout << "total requests:     " << stats.total_requests << std::endl;
+  std::cout << "total responses:    " << stats.total_responses << std::endl;
+  std::cout << "simulator ticks:    " << stats.simulator_ticks << std::endl;
+
+  std::cout << "========================== SUCCESS :) ==========================" << std::endl;
+
+  return 0;
+}

tests/unit/CMakeLists.txt

@@ -391,12 +391,6 @@ add_custom_target(test_lcp ALL DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/test_lcp)
 add_test(test_lcp ${CMAKE_CURRENT_BINARY_DIR}/test_lcp)
 add_dependencies(memgraph__unit test_lcp)
 
-# Test websocket
-find_package(Boost REQUIRED)
-
-add_unit_test(websocket.cpp)
-target_link_libraries(${test_prefix}websocket mg-communication Boost::headers)
-
 # Test future
 add_unit_test(future.cpp)
-target_link_libraries(${test_prefix}future mg-io)
+target_link_libraries(${test_prefix}future mg-io)