initial implementation of count()

arangod/RocksDBEngine/RocksDBCollection.cpp

@@ -409,37 +409,37 @@ int RocksDBCollection::insert(arangodb::transaction::Methods* trx,
   }
 
   transaction::BuilderLeaser builder(trx);
-  VPackSlice newSlice;
+  int res = newObjectForInsert(trx, slice, fromSlice, toSlice, isEdgeCollection,
-  int res = TRI_ERROR_NO_ERROR;
-  if (options.recoveryData == nullptr) {
-    res = newObjectForInsert(trx, slice, fromSlice, toSlice, isEdgeCollection,
                                *builder.get(), options.isRestore);
   if (res != TRI_ERROR_NO_ERROR) {
     return res;
   }
-    newSlice = builder->slice();
+  VPackSlice newSlice = builder->slice();
-  } else {
-    TRI_ASSERT(slice.isObject());
-    // we can get away with the fast hash function here, as key values are
-    // restricted to strings
-    newSlice = slice;
-  }
 
   TRI_voc_rid_t revisionId =
       transaction::helpers::extractRevFromDocument(newSlice);
 
-  res = insertDocument(trx, revisionId, newSlice, options.waitForSync);
+  RocksDBSavePoint guard(rocksTransaction(trx));
-  arangodb::Result result(res);
 
-  if (result.ok()) {
+  res = insertDocument(trx, revisionId, newSlice, options.waitForSync);
-    result = lookupRevisionVPack(revisionId, trx, mdr);
+
-  }
+  if (res == TRI_ERROR_NO_ERROR) {
+    Result result = lookupRevisionVPack(revisionId, trx, mdr);
+
+    if (!result.ok()) {
       return result.errorNumber();
     }
 
+    static_cast<RocksDBTransactionState*>(trx->state())->addOperation(_logicalCollection->cid(), TRI_VOC_DOCUMENT_OPERATION_INSERT, newSlice.byteSize());
+    guard.commit();
+  }
+
+  return res;
+}
+
 int RocksDBCollection::update(arangodb::transaction::Methods* trx,
                               arangodb::velocypack::Slice const newSlice,
-                              arangodb::ManagedDocumentResult& result,
+                              arangodb::ManagedDocumentResult& mdr,
                               OperationOptions& options,
                               TRI_voc_tick_t& resultMarkerTick, bool /*lock*/,
                               TRI_voc_rid_t& prevRev,
@@ -475,9 +475,11 @@ int RocksDBCollection::update(arangodb::transaction::Methods* trx,
   }
 
   if (newSlice.length() <= 1) {
-    // no need to do anything
+    // shortcut. no need to do anything
-    result = std::move(previous);
+    previous.clone(mdr);
+
     if (_logicalCollection->waitForSync()) {
+      trx->state()->waitForSync(true);
       options.waitForSync = true;
     }
     return TRI_ERROR_NO_ERROR;
@@ -499,13 +501,22 @@ int RocksDBCollection::update(arangodb::transaction::Methods* trx,
     }
   }
 
+  RocksDBSavePoint guard(rocksTransaction(trx));
+
   VPackSlice const newDoc(builder->slice());
 
   res = updateDocument(trx, oldRevisionId, oldDoc, revisionId, newDoc,
                        options.waitForSync);
 
   if (res == TRI_ERROR_NO_ERROR) {
-    lookupRevisionVPack(revisionId, trx, result);
+    Result result = lookupRevisionVPack(revisionId, trx, mdr);
+
+    if (!result.ok()) {
+      return result.errorNumber();
+    }
+    
+    static_cast<RocksDBTransactionState*>(trx->state())->addOperation(_logicalCollection->cid(), TRI_VOC_DOCUMENT_OPERATION_UPDATE, newDoc.byteSize());
+    guard.commit();
   }
 
   return res;
@@ -520,7 +531,6 @@ int RocksDBCollection::replace(
     arangodb::velocypack::Slice const toSlice) {
   resultMarkerTick = 0;
 
-  arangodb::Result result;
   bool const isEdgeCollection =
       (_logicalCollection->type() == TRI_COL_TYPE_EDGE);
 
@@ -571,15 +581,25 @@ int RocksDBCollection::replace(
     }
   }
 
+  RocksDBSavePoint guard(rocksTransaction(trx));
+
   res = updateDocument(trx, oldRevisionId, oldDoc, revisionId,
                        VPackSlice(builder->slice()), options.waitForSync);
-  result = Result(res);
+
-  if (result.ok()) {
+  if (res == TRI_ERROR_NO_ERROR) {
-    result = lookupRevisionVPack(revisionId, trx, mdr);
+    Result result = lookupRevisionVPack(revisionId, trx, mdr);
-  }
+
+    if (!result.ok()) {
       return result.errorNumber();
     }
 
+    static_cast<RocksDBTransactionState*>(trx->state())->addOperation(_logicalCollection->cid(), TRI_VOC_DOCUMENT_OPERATION_REPLACE, VPackSlice(builder->slice()).byteSize());
+    guard.commit();
+  }
+
+  return res;
+}
+
 int RocksDBCollection::remove(arangodb::transaction::Methods* trx,
                               arangodb::velocypack::Slice const slice,
                               arangodb::ManagedDocumentResult& previous,
@@ -626,8 +646,15 @@ int RocksDBCollection::remove(arangodb::transaction::Methods* trx,
     }
   }
 
+  RocksDBSavePoint guard(rocksTransaction(trx));
+
   res = removeDocument(trx, oldRevisionId, oldDoc, options.waitForSync);
 
+  if (res == TRI_ERROR_NO_ERROR) {
+    static_cast<RocksDBTransactionState*>(trx->state())->addOperation(_logicalCollection->cid(), TRI_VOC_DOCUMENT_OPERATION_REMOVE, oldDoc.byteSize());
+    guard.commit();
+  }
+
   return res;
 }
 
@@ -752,7 +779,6 @@ int RocksDBCollection::insertDocument(arangodb::transaction::Methods* trx,
   RocksDBValue value(RocksDBValue::Document(doc));
 
   rocksdb::Transaction* rtrx = rocksTransaction(trx);
-  RocksDBSavePoint guard(rtrx);
 
   LOG_TOPIC(ERR, Logger::FIXME)
       << "INSERT DOCUMENT. COLLECTION '" << _logicalCollection->name()
@@ -788,7 +814,7 @@ int RocksDBCollection::insertDocument(arangodb::transaction::Methods* trx,
     }
   }
 
-  if (result != TRI_ERROR_NO_ERROR) {
+  if (result == TRI_ERROR_NO_ERROR) {
     if (_logicalCollection->waitForSync()) {
       waitForSync = true;
     }
@@ -813,8 +839,6 @@ int RocksDBCollection::removeDocument(arangodb::transaction::Methods* trx,
 
   rocksdb::Transaction* rtrx = rocksTransaction(trx);
 
-  RocksDBSavePoint guard(rtrx);
-
   rtrx->Delete(key.string());
 
   auto indexes = _indexes;
@@ -833,7 +857,7 @@ int RocksDBCollection::removeDocument(arangodb::transaction::Methods* trx,
     }
   }
 
-  if (result != TRI_ERROR_NO_ERROR) {
+  if (result == TRI_ERROR_NO_ERROR) {
     if (_logicalCollection->waitForSync()) {
       waitForSync = true;
     }
@@ -850,7 +874,6 @@ int RocksDBCollection::removeDocument(arangodb::transaction::Methods* trx,
 /// the key must be a string slice, no revision check is performed
 int RocksDBCollection::lookupDocument(transaction::Methods* trx, VPackSlice key,
                                       ManagedDocumentResult& mdr) {
-  arangodb::Result result;
   if (!key.isString()) {
     return TRI_ERROR_ARANGO_DOCUMENT_KEY_BAD;
   }
@@ -859,7 +882,7 @@ int RocksDBCollection::lookupDocument(transaction::Methods* trx, VPackSlice key,
   TRI_voc_rid_t revisionId = token.revisionId();
 
   if (revisionId > 0) {
-    result = lookupRevisionVPack(revisionId, trx, mdr);
+    Result result = lookupRevisionVPack(revisionId, trx, mdr);
     return result.errorNumber();
   }
 
@@ -876,9 +899,6 @@ int RocksDBCollection::updateDocument(transaction::Methods* trx,
   TRI_ASSERT(trx->state()->isRunning());
   TRI_ASSERT(!ServerState::instance()->isCoordinator());
 
-  rocksdb::Transaction* rtrx = rocksTransaction(trx);
-  RocksDBSavePoint guard(rtrx);
-
   LOG_TOPIC(ERR, Logger::FIXME)
       << "UPDATE DOCUMENT. COLLECTION '" << _logicalCollection->name()
       << "', OBJECTID: " << _objectId << ", OLDREVISIONID: " << oldRevisionId
@@ -935,3 +955,11 @@ arangodb::Result RocksDBCollection::lookupRevisionVPack(
   }
   return result;
 }
+
+void RocksDBCollection::adjustNumberDocuments(int64_t adjustment) {
+  if (adjustment < 0) {
+    _numberDocuments -= static_cast<uint64_t>(-adjustment);
+  } else if (adjustment > 0) {
+    _numberDocuments += static_cast<uint64_t>(adjustment);
+  }
+}

arangod/RocksDBEngine/RocksDBCollection.h

@@ -170,6 +170,7 @@ class RocksDBCollection final : public PhysicalCollection {
   void deferDropCollection(
       std::function<bool(LogicalCollection*)> callback) override;
 
+  void adjustNumberDocuments(int64_t adjustment);
   uint64_t objectId() const { return _objectId; }
 
   Result lookupDocumentToken(transaction::Methods* trx, arangodb::StringRef key,

arangod/RocksDBEngine/RocksDBTransactionCollection.cpp

@@ -39,7 +39,10 @@ RocksDBTransactionCollection::RocksDBTransactionCollection(TransactionState* trx
                                                            AccessMode::Type accessType) 
     : TransactionCollection(trx, cid),
       _accessType(accessType), 
-      _numOperations(0) {} 
+      _operationSize(0),
+      _numInserts(0),
+      _numUpdates(0),
+      _numRemoves(0) {}
 
 RocksDBTransactionCollection::~RocksDBTransactionCollection() {}
 
@@ -86,7 +89,7 @@ bool RocksDBTransactionCollection::isLocked() const {
 
 /// @brief whether or not any write operations for the collection happened
 bool RocksDBTransactionCollection::hasOperations() const {
-  return (_numOperations > 0);
+  return (_numInserts > 0 || _numRemoves > 0 || _numUpdates > 0);
 }
   
 void RocksDBTransactionCollection::freeOperations(transaction::Methods* /*activeTrx*/, bool /*mustRollback*/) {}
@@ -153,3 +156,24 @@ void RocksDBTransactionCollection::release() {
     _collection = nullptr;
   }
 }
+  
+/// @brief add an operation for a transaction collection
+void RocksDBTransactionCollection::addOperation(TRI_voc_document_operation_e operationType,
+                                                uint64_t operationSize) {
+  switch (operationType) {
+    case TRI_VOC_DOCUMENT_OPERATION_UNKNOWN:
+      break;
+    case TRI_VOC_DOCUMENT_OPERATION_INSERT:
+      ++_numInserts;
+      break;
+    case TRI_VOC_DOCUMENT_OPERATION_UPDATE:
+    case TRI_VOC_DOCUMENT_OPERATION_REPLACE:
+      ++_numUpdates;
+      break;
+    case TRI_VOC_DOCUMENT_OPERATION_REMOVE:
+      ++_numRemoves;
+      break;
+  }
+
+  _operationSize += operationSize;
+}

arangod/RocksDBEngine/RocksDBTransactionCollection.h

@@ -69,9 +69,19 @@ class RocksDBTransactionCollection final : public TransactionCollection {
   void unuse(int nestingLevel) override;
   void release() override;
   
+  uint64_t numInserts() const { return _numInserts; }
+  uint64_t numUpdates() const { return _numUpdates; }
+  uint64_t numRemoves() const { return _numRemoves; }
+  
+  /// @brief add an operation for a transaction collection
+  void addOperation(TRI_voc_document_operation_e operationType, uint64_t operationSize);
+
  private:
   AccessMode::Type _accessType;  // access type (read|write)
-  uint64_t _numOperations;
+  uint64_t _operationSize;
+  uint64_t _numInserts;
+  uint64_t _numUpdates;
+  uint64_t _numRemoves;
 };
 
 }

arangod/RocksDBEngine/RocksDBTransactionState.cpp

@@ -32,6 +32,7 @@
 #include "RocksDBEngine/RocksDBCollection.h"
 #include "RocksDBEngine/RocksDBCommon.h"
 #include "RocksDBEngine/RocksDBEngine.h"
+#include "RocksDBEngine/RocksDBTransactionCollection.h"
 #include "StorageEngine/EngineSelectorFeature.h"
 #include "StorageEngine/StorageEngine.h"
 #include "StorageEngine/TransactionCollection.h"
@@ -52,7 +53,9 @@ using namespace arangodb;
 struct RocksDBTransactionData final : public TransactionData {};
 
 RocksDBSavePoint::RocksDBSavePoint(rocksdb::Transaction* trx)
-    : _trx(trx), _committed(false) {}
+    : _trx(trx), _committed(false) {
+  _trx->SetSavePoint();
+}
 
 RocksDBSavePoint::~RocksDBSavePoint() {
   if (!_committed) {
@@ -160,6 +163,13 @@ Result RocksDBTransactionState::commitTransaction(
         abortTransaction(activeTrx);
         return result;
       }
+      
+      for (auto& trxCollection : _collections) {
+        RocksDBTransactionCollection* collection = static_cast<RocksDBTransactionCollection*>(trxCollection);
+        int64_t adjustment = collection->numInserts() - collection->numRemoves();
+        static_cast<RocksDBCollection*>(trxCollection->collection()->getPhysical())->adjustNumberDocuments(adjustment);
+      }
+  
       _rocksTransaction.reset();
       if (_cacheTx != nullptr) {
         CacheManagerFeature::MANAGER->endTransaction(_cacheTx);
@@ -216,8 +226,17 @@ Result RocksDBTransactionState::abortTransaction(
 }
 
 /// @brief add an operation for a transaction collection
-void RocksDBTransactionState::addOperation(TRI_voc_document_operation_e operationType,
+void RocksDBTransactionState::addOperation(TRI_voc_cid_t cid, 
+                                           TRI_voc_document_operation_e operationType,
                                            uint64_t operationSize) {
+  auto collection = static_cast<RocksDBTransactionCollection*>(findCollection(cid));
+  
+  if (collection == nullptr) {
+    THROW_ARANGO_EXCEPTION_MESSAGE(TRI_ERROR_INTERNAL, "collection not found in transaction state");
+  }
+
+  collection->addOperation(operationType, operationSize);
+
   switch (operationType) {
     case TRI_VOC_DOCUMENT_OPERATION_UNKNOWN:
       break;

arangod/RocksDBEngine/RocksDBTransactionState.h

@@ -95,7 +95,7 @@ class RocksDBTransactionState final : public TransactionState {
   }
 
   /// @brief add an operation for a transaction collection
-  void addOperation(TRI_voc_document_operation_e operationType, uint64_t operationSize);
+  void addOperation(TRI_voc_cid_t collectionId, TRI_voc_document_operation_e operationType, uint64_t operationSize);
 
   rocksdb::Transaction* rocksTransaction() {
     TRI_ASSERT(_rocksTransaction != nullptr);

arangod/VocBase/ManagedDocumentResult.h

@@ -65,6 +65,19 @@ class ManagedDocumentResult {
 
   ManagedDocumentResult(ManagedDocumentResult&& other) = delete;
 
+  void clone(ManagedDocumentResult& cloned) const {
+    cloned.reset();
+    if (_useString) {
+      cloned._useString = true;
+      cloned._string = _string;
+      cloned._lastRevisionId = _lastRevisionId;
+    } else if (_managed) {
+      cloned.setManaged(_vpack, _lastRevisionId);
+    } else {
+      cloned.setUnmanaged(_vpack, _lastRevisionId);
+    }
+  }
+
   inline uint8_t const* vpack() const {
     TRI_ASSERT(_vpack != nullptr);
     return _vpack;