Added proper truncation support for RocksDB indexes.

arangod/RocksDBEngine/RocksDBCollection.cpp

@@ -644,52 +644,7 @@ void RocksDBCollection::truncate(transaction::Methods* trx,
   // don't do anything beyond deleting their contents
   for (std::shared_ptr<Index> const& index : _indexes) {
     RocksDBIndex* rindex = static_cast<RocksDBIndex*>(index.get());
-
+    rindex->truncate(trx);
-    RocksDBKeyBounds indexBounds = RocksDBKeyBounds::Empty();
-    switch (rindex->type()) {
-      case RocksDBIndex::TRI_IDX_TYPE_PRIMARY_INDEX:
-        indexBounds = RocksDBKeyBounds::PrimaryIndex(rindex->objectId());
-        break;
-      case RocksDBIndex::TRI_IDX_TYPE_EDGE_INDEX:
-        indexBounds = RocksDBKeyBounds::EdgeIndex(rindex->objectId());
-        break;
-      case RocksDBIndex::TRI_IDX_TYPE_HASH_INDEX:
-      case RocksDBIndex::TRI_IDX_TYPE_SKIPLIST_INDEX:
-      case RocksDBIndex::TRI_IDX_TYPE_PERSISTENT_INDEX:
-        if (rindex->unique()) {
-          indexBounds = RocksDBKeyBounds::UniqueIndex(rindex->objectId());
-        } else {
-          indexBounds = RocksDBKeyBounds::IndexEntries(rindex->objectId());
-        }
-        break;
-      case RocksDBIndex::TRI_IDX_TYPE_FULLTEXT_INDEX:
-        indexBounds = RocksDBKeyBounds::FulltextIndex(rindex->objectId());
-        break;
-      case RocksDBIndex::TRI_IDX_TYPE_GEO1_INDEX:
-      case RocksDBIndex::TRI_IDX_TYPE_GEO2_INDEX:
-        indexBounds = RocksDBKeyBounds::GeoIndex(rindex->objectId());
-        break;
-      case RocksDBIndex::TRI_IDX_TYPE_UNKNOWN:
-        THROW_ARANGO_EXCEPTION(TRI_ERROR_NOT_IMPLEMENTED);
-    }
-
-    rocksdb::ReadOptions options = state->readOptions();
-    options.iterate_upper_bound = &(indexBounds.end());
-    iter.reset(rtrx->GetIterator(options));
-
-    iter->Seek(indexBounds.start());
-    rindex->disableCache();  // TODO: proper blacklisting of keys?
-    TRI_DEFER(rindex->createCache());
-
-    while (iter->Valid()) {
-      rocksdb::Status s = rtrx->Delete(iter->key());
-      if (!s.ok()) {
-        auto converted = convertStatus(s);
-        THROW_ARANGO_EXCEPTION(converted);
-      }
-
-      iter->Next();
-    }
   }
 }
 

arangod/RocksDBEngine/RocksDBIndex.cpp

@@ -27,13 +27,16 @@
 #include "Cache/TransactionalCache.h"
 #include "Cache/Common.h"
 #include "Cache/Manager.h"
+#include "RocksDBEngine/RocksDBCommon.h"
 #include "RocksDBEngine/RocksDBComparator.h"
 #include "RocksDBEngine/RocksDBEngine.h"
+#include "RocksDBEngine/RocksDBTransactionState.h"
 #include "StorageEngine/EngineSelectorFeature.h"
 #include "VocBase/LogicalCollection.h"
 #include "VocBase/ticks.h"
 
 using namespace arangodb;
+using namespace arangodb::rocksutils;
 
 RocksDBIndex::RocksDBIndex(
     TRI_idx_iid_t id, LogicalCollection* collection,
@@ -153,6 +156,40 @@ int RocksDBIndex::drop() {
   return TRI_ERROR_NO_ERROR;
 }
 
+void RocksDBIndex::truncate(transaction::Methods* trx) {
+  RocksDBTransactionState* state = rocksutils::toRocksTransactionState(trx);
+  rocksdb::Transaction* rtrx = state->rocksTransaction();
+  RocksDBKeyBounds indexBounds = getBounds();
+
+  rocksdb::ReadOptions options = state->readOptions();
+  options.iterate_upper_bound = &(indexBounds.end());
+
+  std::unique_ptr<rocksdb::Iterator> iter(rtrx->GetIterator(options));
+  iter->Seek(indexBounds.start());
+
+  while (iter->Valid()) {
+    rocksdb::Status s = rtrx->Delete(iter->key());
+    if (!s.ok()) {
+      auto converted = convertStatus(s);
+      THROW_ARANGO_EXCEPTION(converted);
+    }
+
+    Result r = postprocessRemove(trx, iter->key(), iter->value());
+    if (!r.ok()) {
+      THROW_ARANGO_EXCEPTION(r);
+    }
+
+    iter->Next();
+  }
+}
+
+Result RocksDBIndex::postprocessRemove(transaction::Methods* trx,
+                                       rocksdb::Slice const& key,
+                                      rocksdb::Slice const& value) {
+  return {TRI_ERROR_NO_ERROR};
+}
+
+
 // blacklist given key from transactional cache
 void RocksDBIndex::blackListKey(char const* data, std::size_t len){
   if (useCache()) {
@@ -170,3 +207,27 @@ void RocksDBIndex::blackListKey(char const* data, std::size_t len){
     }
   }
 }
+
+RocksDBKeyBounds RocksDBIndex::getBounds() const {
+  switch (type()) {
+    case RocksDBIndex::TRI_IDX_TYPE_PRIMARY_INDEX:
+      return RocksDBKeyBounds::PrimaryIndex(objectId());
+    case RocksDBIndex::TRI_IDX_TYPE_EDGE_INDEX:
+      return RocksDBKeyBounds::EdgeIndex(objectId());
+    case RocksDBIndex::TRI_IDX_TYPE_HASH_INDEX:
+    case RocksDBIndex::TRI_IDX_TYPE_SKIPLIST_INDEX:
+    case RocksDBIndex::TRI_IDX_TYPE_PERSISTENT_INDEX:
+      if (unique()) {
+        return RocksDBKeyBounds::UniqueIndex(objectId());
+      }
+      return RocksDBKeyBounds::IndexEntries(objectId());
+    case RocksDBIndex::TRI_IDX_TYPE_FULLTEXT_INDEX:
+      return RocksDBKeyBounds::FulltextIndex(objectId());
+    case RocksDBIndex::TRI_IDX_TYPE_GEO1_INDEX:
+    case RocksDBIndex::TRI_IDX_TYPE_GEO2_INDEX:
+      return RocksDBKeyBounds::GeoIndex(objectId());
+    case RocksDBIndex::TRI_IDX_TYPE_UNKNOWN:
+    default:
+      THROW_ARANGO_EXCEPTION(TRI_ERROR_NOT_IMPLEMENTED);
+  }
+}

arangod/RocksDBEngine/RocksDBIndex.h

@@ -68,6 +68,8 @@ class RocksDBIndex : public Index {
 
   int unload() override;
 
+  virtual void truncate(transaction::Methods*);
+
   /// @brief provides a size hint for the index
   int sizeHint(transaction::Methods* /*trx*/, size_t /*size*/) override final {
     // nothing to do here
@@ -88,12 +90,19 @@ class RocksDBIndex : public Index {
   void disableCache();
 
  protected:
+  // Will be called during truncate to allow the index to update selectivity
+  // estimates, blacklist keys, etc.
+  virtual Result postprocessRemove(transaction::Methods* trx,
+                                   rocksdb::Slice const& key, rocksdb::Slice const& value);
+
   inline bool useCache() const { return (_useCache && _cachePresent); }
   void blackListKey(char const* data, std::size_t len);
   void blackListKey(StringRef& ref){
     blackListKey(ref.data(), ref.size());
   };
 
+  RocksDBKeyBounds getBounds() const;
+
  protected:
   uint64_t _objectId;
   RocksDBComparator* _cmp;

arangod/RocksDBEngine/RocksDBPrimaryIndex.cpp

@@ -560,6 +560,13 @@ void RocksDBPrimaryIndex::invokeOnAllElements(
   }
 }
 
+Result RocksDBPrimaryIndex::postprocessRemove(transaction::Methods* trx,
+                                              rocksdb::Slice const& key,
+                                              rocksdb::Slice const& value) {
+  blackListKey(key.data(), key.size());
+  return {TRI_ERROR_NO_ERROR};
+}
+
 /// @brief create the iterator, for a single attribute, IN operator
 IndexIterator* RocksDBPrimaryIndex::createInIterator(
     transaction::Methods* trx, ManagedDocumentResult* mmdr,

arangod/RocksDBEngine/RocksDBPrimaryIndex.h

@@ -211,6 +211,10 @@ class RocksDBPrimaryIndex final : public RocksDBIndex {
 
   int cleanup() override;
 
+ protected:
+  Result postprocessRemove(transaction::Methods* trx, rocksdb::Slice const& key,
+                           rocksdb::Slice const& value) override;
+
  private:
   /// @brief create the iterator, for a single attribute, IN operator
   IndexIterator* createInIterator(transaction::Methods*, ManagedDocumentResult*,

arangod/RocksDBEngine/RocksDBVPackIndex.cpp

@@ -1412,3 +1412,12 @@ int RocksDBVPackIndex::cleanup() {
   db->CompactRange(opts, &b, &e);
   return TRI_ERROR_NO_ERROR;
 }
+
+Result RocksDBVPackIndex::postprocessRemove(transaction::Methods* trx,
+                                            rocksdb::Slice const& key,
+                                            rocksdb::Slice const& value) {
+  if (!unique()) {
+    // TODO: update selectivity estimate
+  }
+  return {TRI_ERROR_NO_ERROR};
+}

arangod/RocksDBEngine/RocksDBVPackIndex.h

@@ -179,6 +179,10 @@ class RocksDBVPackIndex : public RocksDBIndex {
 
   int cleanup() override;
 
+protected:
+ Result postprocessRemove(transaction::Methods* trx, rocksdb::Slice const& key,
+                          rocksdb::Slice const& value) override;
+
  private:
   bool isDuplicateOperator(arangodb::aql::AstNode const*,
                            std::unordered_set<int> const&) const;