kostas test case

src/integrations/kafka/consumer.cpp

@@ -35,9 +35,9 @@ constexpr std::chrono::milliseconds kMinimumInterval{1};
 constexpr int64_t kMinimumSize{1};
 
 namespace {
-utils::BasicResult<std::string, std::vector<Message>> GetBatch(RdKafka::KafkaConsumer &consumer,
-                                                               const ConsumerInfo &info,
-                                                               std::atomic<bool> &is_running) {
+utils::BasicResult<std::string, std::pair<int64_t, std::vector<Message>>> GetBatch(RdKafka::KafkaConsumer &consumer,
+                                                                                   const ConsumerInfo &info,
+                                                                                   std::atomic<bool> &is_running) {
   std::vector<Message> batch{};
 
   int64_t batch_size = info.batch_size.value_or(kDefaultBatchSize);
@@ -47,6 +47,7 @@ utils::BasicResult<std::string, std::vector<Message>> GetBatch(RdKafka::KafkaCon
   auto start = std::chrono::steady_clock::now();
 
   bool run_batch = true;
+  int64_t offset = 0;
   for (int64_t i = 0; remaining_timeout_in_ms > 0 && i < batch_size && is_running.load(); ++i) {
     std::unique_ptr<RdKafka::Message> msg(consumer.consume(remaining_timeout_in_ms));
     switch (msg->err()) {
@@ -55,6 +56,7 @@ utils::BasicResult<std::string, std::vector<Message>> GetBatch(RdKafka::KafkaCon
         break;
 
       case RdKafka::ERR_NO_ERROR:
+        offset = msg->offset();
         batch.emplace_back(std::move(msg));
         break;
       case RdKafka::ERR__MAX_POLL_EXCEEDED:
@@ -78,7 +80,7 @@ utils::BasicResult<std::string, std::vector<Message>> GetBatch(RdKafka::KafkaCon
     start = now;
   }
 
-  return {std::move(batch)};
+  return std::make_pair(offset, std::move(batch));
 }
 }  // namespace
 
@@ -131,6 +133,10 @@ Consumer::Consumer(const std::string &bootstrap_servers, ConsumerInfo info, Cons
     throw ConsumerFailedToInitializeException(info_.consumer_name, error);
   }
 
+  if (conf->set("rebalance_cb", &cb_, error) != RdKafka::Conf::CONF_OK) {
+    throw ConsumerFailedToInitializeException(info_.consumer_name, error);
+  }
+
   if (conf->set("enable.partition.eof", "false", error) != RdKafka::Conf::CONF_OK) {
     throw ConsumerFailedToInitializeException(info_.consumer_name, error);
   }
@@ -269,13 +275,13 @@ void Consumer::Check(std::optional<std::chrono::milliseconds> timeout, std::opti
 
     const auto &batch = maybe_batch.GetValue();
 
-    if (batch.empty()) {
+    if (batch.second.empty()) {
       continue;
     }
     ++i;
 
     try {
-      check_consumer_function(batch);
+      check_consumer_function(batch.second);
     } catch (const std::exception &e) {
       spdlog::warn("Kafka consumer {} check failed with error {}", info_.consumer_name, e.what());
       throw ConsumerCheckFailedException(info_.consumer_name, e.what());
@@ -333,13 +339,25 @@ void Consumer::StartConsuming() {
       }
       const auto &batch = maybe_batch.GetValue();
 
-      if (batch.empty()) continue;
+      if (batch.second.empty()) continue;
 
       spdlog::info("Kafka consumer {} is processing a batch", info_.consumer_name);
 
       try {
-        consumer_function_(batch);
-        if (const auto err = consumer_->commitSync(); err != RdKafka::ERR_NO_ERROR) {
+        consumer_function_(batch.second);
+        std::vector<RdKafka::TopicPartition *> partitions;
+        if (const auto err = consumer_->assignment(partitions); err != RdKafka::ERR_NO_ERROR) {
+          spdlog::warn("Saving the commited offset of consumer {} failed: {}", info_.consumer_name,
+                       RdKafka::err2str(err));
+          throw ConsumerCheckFailedException(
+              info_.consumer_name, fmt::format("Couldn't save commited offsets: '{}'", RdKafka::err2str(err)));
+        }
+
+        for (auto *partition : partitions) {
+          partition->set_offset(batch.first + 1);
+        }
+
+        if (const auto err = consumer_->commitSync(partitions); err != RdKafka::ERR_NO_ERROR) {
           spdlog::warn("Committing offset of consumer {} failed: {}", info_.consumer_name, RdKafka::err2str(err));
           break;
         }
@@ -358,4 +376,21 @@ void Consumer::StopConsuming() {
   if (thread_.joinable()) thread_.join();
 }
 
+std::string Consumer::SetConsumerOffsets(const std::string_view stream_name, int64_t offset) {
+  std::vector<RdKafka::TopicPartition *> partitions;
+  auto maybe_error = consumer_->assignment(partitions);
+  if (maybe_error != RdKafka::ErrorCode::ERR_NO_ERROR) {
+    return fmt::format("Can't access assigned topic partitions to the consumer: {}", maybe_error);
+  }
+
+  std::vector<std::unique_ptr<RdKafka::TopicPartition>> owners(partitions.begin(), partitions.end());
+  if (offset == -1) {
+    offset = RD_KAFKA_OFFSET_BEGINNING;
+  }
+
+  cb_.set_offset(offset);
+  consumer_->subscribe(info_.topics);
+
+  return "";
+}
 }  // namespace integrations::kafka

src/integrations/kafka/consumer.hpp

@@ -137,6 +137,9 @@ class Consumer final : public RdKafka::EventCb {
   /// Returns true if the consumer is actively consuming messages.
   bool IsRunning() const;
 
+  ///
+  std::string SetConsumerOffsets(const std::string_view stream_name, int64_t offset);
+
   const ConsumerInfo &Info() const;
 
  private:
@@ -153,5 +156,35 @@ class Consumer final : public RdKafka::EventCb {
   std::optional<int64_t> limit_batches_{std::nullopt};
   std::unique_ptr<RdKafka::KafkaConsumer, std::function<void(RdKafka::KafkaConsumer *)>> consumer_;
   std::thread thread_;
+  class ExampleRebalanceCb : public RdKafka::RebalanceCb {
+   public:
+    void rebalance_cb(RdKafka::KafkaConsumer *consumer, RdKafka::ErrorCode err,
+                      std::vector<RdKafka::TopicPartition *> &partitions) {
+      if (offset_) {
+        for (auto partition : partitions) {
+          partition->set_offset(*offset_);
+        }
+        offset_.reset();
+      }
+      consumer->assign(partitions);
+      consumer->commitSync(partitions);
+    }
+    void set_offset(int64_t offset) { offset_ = offset; }
+
+   private:
+    std::optional<int64_t> offset_ = -1;
+  };
+
+  ExampleRebalanceCb cb_;
+
+  class ExampleCommitCb : public RdKafka::OffsetCommitCb {
+    void offset_commit_cb(RdKafka::ErrorCode err, std::vector<RdKafka::TopicPartition *> &offsets) override {
+      for (auto *partition : offsets) {
+        spdlog::critical("Trying to commit {}", partition->offset());
+      }
+    }
+  };
+
+  ExampleCommitCb commit_cb_;
 };
 }  // namespace integrations::kafka

src/query/procedure/mg_procedure_impl.hpp

@@ -676,6 +676,17 @@ struct mgp_proc {
         results(memory),
         is_write_procedure(is_write_procedure) {}
 
+  /// @throw std::bad_alloc
+  /// @throw std::length_error
+  mgp_proc(const std::string_view name, std::function<void(mgp_list *, mgp_graph *, mgp_result *, mgp_memory *)> cb,
+           utils::MemoryResource *memory, bool is_write_procedure)
+      : name(name, memory),
+        cb(cb),
+        args(memory),
+        opt_args(memory),
+        results(memory),
+        is_write_procedure(is_write_procedure) {}
+
   /// @throw std::bad_alloc
   /// @throw std::length_error
   mgp_proc(const mgp_proc &other, utils::MemoryResource *memory)

src/query/procedure/module.cpp

@@ -326,14 +326,16 @@ bool SharedLibraryModule::Load(const std::filesystem::path &file_path) {
   handle_ = dlopen(file_path.c_str(), RTLD_NOW | RTLD_LOCAL);
   if (!handle_) {
     // NOLINTNEXTLINE(concurrency-mt-unsafe)
-    spdlog::error(utils::MessageWithLink("Unable to load module {}; {}.", file_path, dlerror(), "https://memgr.ph/modules"));
+    spdlog::error(
+        utils::MessageWithLink("Unable to load module {}; {}.", file_path, dlerror(), "https://memgr.ph/modules"));
     return false;
   }
   // Get required mgp_init_module
   init_fn_ = reinterpret_cast<int (*)(mgp_module *, mgp_memory *)>(dlsym(handle_, "mgp_init_module"));
   char *dl_errored = dlerror();
   if (!init_fn_ || dl_errored) {
-    spdlog::error(utils::MessageWithLink("Unable to load module {}; {}.", file_path, dl_errored, "https://memgr.ph/modules"));
+    spdlog::error(
+        utils::MessageWithLink("Unable to load module {}; {}.", file_path, dl_errored, "https://memgr.ph/modules"));
     dlclose(handle_);
     handle_ = nullptr;
     return false;
@@ -385,7 +387,8 @@ bool SharedLibraryModule::Close() {
   }
   if (dlclose(handle_) != 0) {
     // NOLINTNEXTLINE(concurrency-mt-unsafe)
-    spdlog::error(utils::MessageWithLink("Failed to close module {}; {}.", file_path_, dlerror(), "https://memgr.ph/modules"));
+    spdlog::error(
+        utils::MessageWithLink("Failed to close module {}; {}.", file_path_, dlerror(), "https://memgr.ph/modules"));
     return false;
   }
   spdlog::info("Closed module {}", file_path_);
@@ -444,7 +447,8 @@ bool PythonModule::Load(const std::filesystem::path &file_path) {
   auto gil = py::EnsureGIL();
   auto maybe_exc = py::AppendToSysPath(file_path.parent_path().c_str());
   if (maybe_exc) {
-    spdlog::error(utils::MessageWithLink("Unable to load module {}; {}.", file_path, *maybe_exc, "https://memgr.ph/modules"));
+    spdlog::error(
+        utils::MessageWithLink("Unable to load module {}; {}.", file_path, *maybe_exc, "https://memgr.ph/modules"));
     return false;
   }
   bool succ = true;
@@ -469,7 +473,8 @@ bool PythonModule::Load(const std::filesystem::path &file_path) {
     return true;
   }
   auto exc_info = py::FetchError().value();
-  spdlog::error(utils::MessageWithLink("Unable to load module {}; {}.", file_path, exc_info, "https://memgr.ph/modules"));
+  spdlog::error(
+      utils::MessageWithLink("Unable to load module {}; {}.", file_path, exc_info, "https://memgr.ph/modules"));
   return false;
 }
 
@@ -534,7 +539,8 @@ bool ModuleRegistry::RegisterModule(const std::string_view &name, std::unique_pt
   MG_ASSERT(!name.empty(), "Module name cannot be empty");
   MG_ASSERT(module, "Tried to register an invalid module");
   if (modules_.find(name) != modules_.end()) {
-    spdlog::error(utils::MessageWithLink("Unable to overwrite an already loaded module {}.", name, "https://memgr.ph/modules"));
+    spdlog::error(
+        utils::MessageWithLink("Unable to overwrite an already loaded module {}.", name, "https://memgr.ph/modules"));
     return false;
   }
   modules_.emplace(name, std::move(module));
@@ -564,7 +570,8 @@ void ModuleRegistry::SetModulesDirectory(std::vector<std::filesystem::path> modu
 
 bool ModuleRegistry::LoadModuleIfFound(const std::filesystem::path &modules_dir, const std::string_view name) {
   if (!utils::DirExists(modules_dir)) {
-    spdlog::error(utils::MessageWithLink("Module directory {} doesn't exist.", modules_dir, "https://memgr.ph/modules"));
+    spdlog::error(
+        utils::MessageWithLink("Module directory {} doesn't exist.", modules_dir, "https://memgr.ph/modules"));
     return false;
   }
   for (const auto &entry : std::filesystem::directory_iterator(modules_dir)) {
@@ -601,7 +608,8 @@ bool ModuleRegistry::LoadOrReloadModuleFromName(const std::string_view name) {
 void ModuleRegistry::LoadModulesFromDirectory(const std::filesystem::path &modules_dir) {
   if (modules_dir.empty()) return;
   if (!utils::DirExists(modules_dir)) {
-    spdlog::error(utils::MessageWithLink("Module directory {} doesn't exist.", modules_dir, "https://memgr.ph/modules"));
+    spdlog::error(
+        utils::MessageWithLink("Module directory {} doesn't exist.", modules_dir, "https://memgr.ph/modules"));
     return;
   }
   for (const auto &entry : std::filesystem::directory_iterator(modules_dir)) {
@@ -638,6 +646,19 @@ void ModuleRegistry::UnloadAllModules() {
 
 utils::MemoryResource &ModuleRegistry::GetSharedMemoryResource() noexcept { return *shared_; }
 
+bool ModuleRegistry::RegisterProcedure(const std::string_view name,
+                                       std::function<void(mgp_list *, mgp_graph *, mgp_result *, mgp_memory *)> f) {
+  std::unique_lock<utils::RWLock> guard(lock_);
+  auto module = modules_.find("mg");
+  if (module != modules_.end()) {
+    auto *builtin_module = dynamic_cast<BuiltinModule *>(module->second.get());
+    mgp_proc proc(name, std::move(f), utils::NewDeleteResource(), false);
+    builtin_module->AddProcedure(name, std::move(proc));
+    return true;
+  }
+  return false;
+}
+
 namespace {
 
 /// This function returns a pair of either

src/query/procedure/module.hpp

@@ -117,6 +117,9 @@ class ModuleRegistry final {
   /// Returns the shared memory allocator used by modules
   utils::MemoryResource &GetSharedMemoryResource() noexcept;
 
+  bool RegisterProcedure(const std::string_view name,
+                         std::function<void(mgp_list *, mgp_graph *, mgp_result *, mgp_memory *)> f);
+
  private:
   std::vector<std::filesystem::path> modules_dirs_;
 };

src/query/streams.cpp

@@ -20,6 +20,7 @@
 #include "query/db_accessor.hpp"
 #include "query/discard_value_stream.hpp"
 #include "query/interpreter.hpp"
+#include "query/procedure//mg_procedure_helpers.hpp"
 #include "query/procedure/mg_procedure_impl.hpp"
 #include "query/procedure/module.hpp"
 #include "query/typed_value.hpp"
@@ -185,7 +186,28 @@ Streams::Streams(InterpreterContext *interpreter_context, std::string bootstrap_
                  std::filesystem::path directory)
     : interpreter_context_(interpreter_context),
       bootstrap_servers_(std::move(bootstrap_servers)),
-      storage_(std::move(directory)) {}
+      storage_(std::move(directory)) {
+  auto set_stream_offset_procedure = [ictx = interpreter_context](mgp_list *args, mgp_graph * /*graph*/,
+                                                                  mgp_result *result, mgp_memory * /*memory*/) {
+    MG_ASSERT(procedure::Call<size_t>(mgp_list_size, args) == 2U, "Should have been type checked already");
+    auto *arg = procedure::Call<mgp_value *>(mgp_list_at, args, 0);
+    MG_ASSERT(procedure::CallBool(mgp_value_is_string, arg), "Should have been type checked already");
+    bool succ = false;
+    const char *arg_as_string{nullptr};
+    if (const auto err = mgp_value_get_string(arg, &arg_as_string); err != MGP_ERROR_NO_ERROR) {
+      succ = false;
+    }
+    auto *offset_value = procedure::Call<mgp_value *>(mgp_list_at, args, 1);
+    int64_t offset{0};
+    auto res = mgp_value_get_int(offset_value, &offset);
+    std::string error = ictx->streams.SetStreamOffset(arg_as_string, offset);
+    if (!error.empty()) {
+      MG_ASSERT(mgp_result_set_error_msg(result, "something") == MGP_ERROR_NO_ERROR);
+    }
+  };
+
+  procedure::gModuleRegistry.RegisterProcedure("set_stream_offset", set_stream_offset_procedure);
+}
 
 void Streams::RestoreStreams() {
   spdlog::info("Loading streams...");
@@ -434,4 +456,14 @@ void Streams::Persist(StreamStatus &&status) {
 }
 
 std::string_view Streams::BootstrapServers() const { return bootstrap_servers_; }
+
+std::string Streams::SetStreamOffset(const std::string_view stream_name, int64_t offset) {
+  auto lock_ptr = streams_.Lock();
+  if (auto it = lock_ptr->find(std::string(stream_name)); it != lock_ptr->end()) {
+    auto consumer_lock_ptr = it->second.consumer->Lock();
+    return consumer_lock_ptr->SetConsumerOffsets(stream_name, offset);
+  }
+  return fmt::format("Stream: {} not found", stream_name);
+}
+
 }  // namespace query

src/query/streams.hpp

@@ -143,6 +143,8 @@ class Streams final {
   /// Return the configuration value passed to memgraph.
   std::string_view BootstrapServers() const;
 
+  std::string SetStreamOffset(const std::string_view stream_name, int64_t offset);
+
  private:
   using StreamsMap = std::unordered_map<std::string, StreamData>;
   using SynchronizedStreamsMap = utils::Synchronized<StreamsMap, utils::WritePrioritizedRWLock>;

tests/unit/integrations_kafka_consumer.cpp

@@ -1,3 +1,14 @@
+// Copyright 2021 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 <optional>
 #include <string>
@@ -45,7 +56,7 @@ struct ConsumerTest : public ::testing::Test {
     };
   };
 
-  std::unique_ptr<Consumer> CreateConsumer(ConsumerInfo &&info, ConsumerFunction consumer_function) {
+  std::unique_ptr<Consumer> CreateConsumer(ConsumerInfo &&info, ConsumerFunction consumer_function, long &sent) {
     EXPECT_EQ(1, info.topics.size());
     EXPECT_EQ(info.topics.at(0), kTopicName);
     auto last_received_message = std::make_shared<std::atomic<int>>(0);
@@ -89,6 +100,7 @@ struct ConsumerTest : public ::testing::Test {
 
     consumer->Stop();
     std::this_thread::sleep_for(std::chrono::seconds(4));
+    sent = sent_messages;
     return consumer;
   }
 
@@ -104,7 +116,7 @@ struct ConsumerTest : public ::testing::Test {
 };
 
 const std::string ConsumerTest::kTopicName{"FirstTopic"};
-
+/*
 TEST_F(ConsumerTest, BatchInterval) {
   // There might be ~300ms delay in message delivery with librdkafka mock, thus the batch interval cannot be too small.
   constexpr auto kBatchInterval = std::chrono::milliseconds{500};
@@ -483,3 +495,59 @@ TEST_F(ConsumerTest, ConsumerStatus) {
   consumer.StopIfRunning();
   check_info(consumer.Info());
 }
+*/
+TEST_F(ConsumerTest, SetOffset) {
+  constexpr auto kBatchInterval = std::chrono::milliseconds{1000};
+  constexpr auto kBatchSize = 3;
+  auto info = CreateDefaultConsumerInfo();
+  info.batch_interval = kBatchInterval;
+  info.batch_size = kBatchSize;
+  constexpr std::string_view kMessage = "BatchSizeTestMessage";
+  std::vector<std::string> messages_received;
+  std::vector<std::string> expected_messages_received;
+  auto consumer_function = [&](const std::vector<Message> &messages) mutable {
+    for (const auto &message : messages) {
+      messages_received.push_back(std::string(message.Payload().data(), message.Payload().size()));
+      spdlog::info("Message Received {}", messages_received.back());
+    }
+  };
+
+  constexpr auto kLastBatchMessageCount = 1;
+  constexpr auto kMessageCount = 3 * kBatchSize + kLastBatchMessageCount;
+  for (auto sent_messages = 0; sent_messages < kMessageCount; ++sent_messages) {
+    auto message = fmt::format("{}, {}", sent_messages, kMessage);
+    cluster.SeedTopic(kTopicName, std::string_view(message));
+    expected_messages_received.push_back(std::move(message));
+  }
+  /*
+  cluster.SeedTopic(kTopicName, std::string_view{"final message"});
+  std::this_thread::sleep_for(kBatchInterval * 2);
+  consumer->Stop();
+  consumer->SetConsumerOffsets("Test stream", -1);
+
+  consumer->Start();
+  std::this_thread::sleep_for(kBatchInterval * 2);
+  consumer->Stop();
+
+  cluster.SeedTopic(kTopicName, std::string_view{"after final message"});
+  consumer->Start();
+  std::this_thread::sleep_for(kBatchInterval * 2);
+  ASSERT_TRUE((messages_received.size() + 3) == expected_messages_received.size());
+  ASSERT_TRUE(std::equal(messages_received.begin() + 4, messages_received.end(), expected_messages_received.begin()));
+
+*/
+  long msgs = 0;
+  auto consumer = CreateConsumer(std::move(info), std::move(consumer_function), msgs);
+  std::this_thread::sleep_for(kBatchInterval * 8);
+  auto err = consumer->SetConsumerOffsets("Test stream", msgs);
+  std::this_thread::sleep_for(kBatchInterval * 2);
+  messages_received.clear();
+  consumer->Start();
+  std::this_thread::sleep_for(kBatchInterval * 2);
+  cluster.SeedTopic(kTopicName, std::string_view{"final message"});
+  std::this_thread::sleep_for(kBatchInterval * 2);
+  consumer->Stop();
+  for (auto message : messages_received) {
+    std::cout << message << std::endl;
+  }
+}