MINIFICPP-2845 - Monotonic provenance event identifiers, iterations support, refactor

core-framework/common/src/utils/Id.cpp

@@ -43,6 +43,18 @@ bool Identifier::isNil() const {
   return *this == Identifier{};
 }
 
+Identifier& Identifier::operator++() {
+  for (int byte_idx = 15; byte_idx >= 0 && ++data_[byte_idx] == 0; --byte_idx) {}
+  return *this;
+}
+
+Identifier Identifier::operator++(int) {
+  auto result = *this;
+  ++*this;
+  return result;
+}
+
+
 bool Identifier::operator!=(const Identifier& other) const {
   return !(*this == other);
 }

core-framework/include/core/Repository.h

@@ -106,12 +106,6 @@ class RepositoryImpl : public core::CoreComponentImpl, public core::RepositoryMe
     return false;
   }
 
-  bool getElements(std::vector<std::shared_ptr<core::SerializableComponent>>& /*store*/, size_t& /*max_size*/) override {
-    return true;
-  }
-
-  bool storeElement(const std::shared_ptr<core::SerializableComponent>& element) override;
-
   void loadComponent(const std::shared_ptr<core::ContentRepository>& /*content_repo*/) override {
   }
 

core-framework/src/core/Repository.cpp

@@ -27,20 +27,4 @@ bool RepositoryImpl::Delete(std::vector<std::shared_ptr<core::SerializableCompon
   return found;
 }
 
-bool RepositoryImpl::storeElement(const std::shared_ptr<core::SerializableComponent>& element) {
-  if (!element) {
-    return false;
-  }
-
-  org::apache::nifi::minifi::io::BufferStream stream;
-
-  element->serialize(stream);
-
-  if (!Put(element->getUUIDStr(), reinterpret_cast<const uint8_t*>(stream.getBuffer().data()), stream.size())) {
-    logger_->log_error("NiFi Provenance Store event {} size {} fail", element->getUUIDStr(), stream.size());
-    return false;
-  }
-  return true;
-}
-
 }  // namespace org::apache::nifi::minifi::core

extensions/grafana-loki/PushGrafanaLokiGrpc.cpp

@@ -116,7 +116,7 @@ std::expected<void, std::string> PushGrafanaLokiGrpc::submitRequest(const std::v
 
   for (const auto& flow_file : batched_flow_files) {
     logproto::EntryAdapter *entry = stream->add_entries();
-    auto timestamp_str = std::to_string(flow_file->getlineageStartDate().time_since_epoch() / std::chrono::nanoseconds(1));
+    auto timestamp_str = std::to_string(flow_file->getLineageStartDate().time_since_epoch() / std::chrono::nanoseconds(1));
     auto timestamp_nanos = std::stoll(timestamp_str);
     *entry->mutable_timestamp() = google::protobuf::util::TimeUtil::NanosecondsToTimestamp(timestamp_nanos);
 

extensions/grafana-loki/PushGrafanaLokiREST.cpp

@@ -133,7 +133,7 @@ std::string PushGrafanaLokiREST::createLokiJson(const std::vector<std::shared_pt
   for (const auto& flow_file : batched_flow_files) {
     rapidjson::Value log_line(rapidjson::kArrayType);
 
-    auto timestamp_str = std::to_string(flow_file->getlineageStartDate().time_since_epoch() / std::chrono::nanoseconds(1));
+    auto timestamp_str = std::to_string(flow_file->getLineageStartDate().time_since_epoch() / std::chrono::nanoseconds(1));
     rapidjson::Value timestamp;
     timestamp.SetString(timestamp_str.c_str(), gsl::narrow<rapidjson::SizeType>(timestamp_str.length()), allocator);
     rapidjson::Value log_line_value;

extensions/rocksdb-repos/RocksDbProvenanceRepository.cpp

@@ -0,0 +1,207 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "RocksDbProvenanceRepository.h"
+
+#include <string>
+
+#include "core/Resource.h"
+
+namespace org::apache::nifi::minifi::provenance {
+
+namespace {
+class EventCursor : public ProvenanceRepository::Cursor {
+public:
+  explicit EventCursor(std::string event_id): event_id_(std::move(event_id)) {}
+  [[nodiscard]]
+  std::string toString() const override {
+    return event_id_;
+  }
+  ~EventCursor() override = default;
+
+  std::string event_id_;
+};
+}  // namespace
+
+static const std::string_view NEXT_EVENT_UUID_KEY = "next_event_uuid";
+
+bool RocksDbProvenanceRepository::initialize(const std::shared_ptr<org::apache::nifi::minifi::Configure> &config) {
+  if (!RocksDbRepository::initialize(config)) {
+    return false;
+  }
+  std::string value;
+  if (config->get(Configure::nifi_provenance_repository_directory_default, value) && !value.empty()) {
+    directory_ = value;
+  }
+  logger_->log_debug("MiNiFi Provenance Repository Directory {}", directory_);
+  if (config->get(Configure::nifi_provenance_repository_max_storage_size, value)) {
+    max_partition_bytes_ = gsl::narrow<int64_t>(parsing::parseDataSize(value) | utils::orThrow("expected parsable data size"));
+  }
+  logger_->log_debug("MiNiFi Provenance Max Partition Bytes {}", max_partition_bytes_);
+  if (config->get(Configure::nifi_provenance_repository_max_storage_time, value)) {
+    if (auto max_partition = utils::timeutils::StringToDuration<std::chrono::milliseconds>(value))
+      max_partition_millis_ = *max_partition;
+  }
+  logger_->log_debug("MiNiFi Provenance Max Storage Time: [{}]", max_partition_millis_);
+
+  verify_checksums_in_rocksdb_reads_ = (config->get(Configure::nifi_provenance_repository_rocksdb_read_verify_checksums) | utils::andThen(&utils::string::toBool)).value_or(false);
+  logger_->log_debug("{} checksum verification in RocksDbProvenanceRepository", verify_checksums_in_rocksdb_reads_ ? "Using" : "Not using");
+
+  auto db_options = [] (minifi::internal::Writable<rocksdb::DBOptions>& db_opts) {
+    minifi::internal::setCommonRocksDbOptions(db_opts);
+  };
+
+  // Rocksdb write buffers act as a log of database operation: grow till reaching the limit, serialized after
+  // This shouldn't go above 16MB and the configured total size of the db should cap it as well
+  auto cf_options = [this] (rocksdb::ColumnFamilyOptions& cf_opts) {
+    int64_t max_buffer_size = 16 << 20;
+    cf_opts.write_buffer_size = gsl::narrow<size_t>(std::min(max_buffer_size, max_partition_bytes_));
+    cf_opts.max_write_buffer_number = 4;
+    cf_opts.min_write_buffer_number_to_merge = 1;
+
+    cf_opts.compaction_style = rocksdb::CompactionStyle::kCompactionStyleFIFO;
+    cf_opts.compaction_options_fifo = rocksdb::CompactionOptionsFIFO(max_partition_bytes_, false);
+    if (max_partition_millis_ > std::chrono::milliseconds(0)) {
+      cf_opts.ttl = std::chrono::duration_cast<std::chrono::seconds>(max_partition_millis_).count();
+    }
+  };
+
+  db_ = minifi::internal::RocksDatabase::create(db_options, cf_options, directory_,
+    minifi::internal::getRocksDbOptionsToOverride(config, Configure::nifi_provenance_repository_rocksdb_options));
+  std::string internal_state_db_uri = [&] {
+    const std::string_view minifidb_scheme = "minifidb://";
+    if (directory_.starts_with(minifidb_scheme)) {
+      return directory_ + "-internal-state";
+    }
+    std::string uri = utils::string::join_pack(minifidb_scheme, directory_);
+    if (uri.ends_with("/") || uri.ends_with("\\")) {
+      uri.pop_back();
+    }
+    return uri + "/internal-state";
+  }();
+  internal_state_db_ = minifi::internal::RocksDatabase::create(db_options, {}, internal_state_db_uri, {});
+  if (auto open_state_db = internal_state_db_->open()) {
+    rocksdb::ReadOptions options;
+    options.verify_checksums = verify_checksums_in_rocksdb_reads_;
+    std::string next_event_uuid_str;
+    if (open_state_db->Get(options, NEXT_EVENT_UUID_KEY, &next_event_uuid_str).ok()) {
+      next_event_id_ = next_event_uuid_str;
+    } else {
+      logger_->log_error("Could not find '{}'", NEXT_EVENT_UUID_KEY);
+      next_event_id_ = utils::IdGenerator::getIdGenerator()->generate();
+    }
+    logger_->log_trace("Using next event uuid: {}", next_event_id_.to_string());
+  } else {
+    logger_->log_error("Could not open internal state column in provenance repository {}", internal_state_db_uri);
+    return false;
+  }
+  if (db_->open()) {
+    logger_->log_debug("MiNiFi Provenance Repository database open {} success", directory_);
+  } else {
+    logger_->log_error("MiNiFi Provenance Repository database open {} failed", directory_);
+    return false;
+  }
+
+  return true;
+}
+
+void RocksDbProvenanceRepository::destroy() {
+  db_.reset();
+}
+
+std::unique_ptr<ProvenanceRepository::Cursor> RocksDbProvenanceRepository::cursorFromString(std::optional<std::string> cursor_str) {
+  if (cursor_str.has_value()) {
+    return std::make_unique<EventCursor>(cursor_str.value());
+  }
+  return std::make_unique<EventCursor>("");
+}
+
+std::expected<std::vector<std::shared_ptr<provenance::ProvenanceEventRecord>>, std::string> RocksDbProvenanceRepository::getEvents(size_t max_size, Cursor* cursor) {
+  auto* event_cursor = dynamic_cast<EventCursor*>(cursor);
+  if (cursor && !event_cursor) {
+    return std::unexpected{"Invalid cursor"};
+  }
+  if (max_size == 0) {
+    return {};
+  }
+  auto opendb = db_->open();
+  if (!opendb) {
+    return std::unexpected{"Failed to open database"};
+  }
+  std::vector<std::shared_ptr<provenance::ProvenanceEventRecord>> records;
+  rocksdb::ReadOptions options;
+  options.verify_checksums = verify_checksums_in_rocksdb_reads_;
+  std::unique_ptr<rocksdb::Iterator> it(opendb->NewIterator(options));
+  std::string last_event_id;
+  if (event_cursor) {
+    last_event_id = event_cursor->event_id_;
+    it->Seek(event_cursor->event_id_);
+    if (it->Valid() && it->key() == event_cursor->event_id_) {
+      it->Next();
+    }
+  } else {
+    it->SeekToFirst();
+  }
+  for (; it->Valid(); it->Next()) {
+    last_event_id = it->key().ToString();
+    auto eventRead = ProvenanceEventRecord::create();
+    const auto slice = it->value();
+    io::BufferStream stream(std::as_bytes(std::span(slice.data(), slice.size())));
+    if (eventRead->deserialize(stream)) {
+      records.push_back(eventRead);
+      if (--max_size == 0) {
+        break;
+      }
+    }
+  }
+  if (event_cursor) {
+    event_cursor->event_id_ = last_event_id;
+  }
+  return records;
+}
+
+std::expected<void, std::string> RocksDbProvenanceRepository::appendEvents(const std::vector<std::shared_ptr<ProvenanceEventRecord>>& events) {
+  std::vector<std::pair<std::string, std::unique_ptr<io::BufferStream>>> data;
+  data.reserve(events.size());
+  std::lock_guard guard(next_event_id_mtx_);
+  for (auto& event : events) {
+    event->setUUID(next_event_id_++);
+  }
+  {
+    auto open_state_db = internal_state_db_->open();
+    if (!open_state_db) {
+      return std::unexpected{"Failed to open internal state column in provenance database"};
+    }
+    auto operation = [this, &open_state_db]() { return open_state_db->Put(rocksdb::WriteOptions(), NEXT_EVENT_UUID_KEY, next_event_id_.to_string().view()); };
+    if (!ExecuteWithRetry(operation)) {
+      return std::unexpected{"Failed to update next provenance event id"};
+    }
+  }
+  for (auto& event : events) {
+    data.emplace_back(event->getUUIDStr(), std::make_unique<io::BufferStream>());
+    event->serialize(*data.back().second);
+  }
+  if (MultiPut(data)) {
+    return {};
+  }
+
+  return std::unexpected{"Failed to append provenance events"};
+}
+
+REGISTER_RESOURCE_AS(RocksDbProvenanceRepository, InternalResource, ("RocksDbProvenanceRepository", "ProvenanceRepository", "provenancerepository"));
+
+}  // namespace org::apache::nifi::minifi::provenance