Track document sizes in bucket db.

Extend document meta store save/load to handle document sizes.

3 files changed, 105 insertions, 33 deletions

diff --git a/searchcore/src/tests/proton/bucketdb/bucketdb/bucketdb_test.cpp b/searchcore/src/tests/proton/bucketdb/bucketdb/bucketdb_test.cpp
index 6895e469319..8845cb633e5 100644
--- a/searchcore/src/tests/proton/bucketdb/bucketdb/bucketdb_test.cpp
+++ b/searchcore/src/tests/proton/bucketdb/bucketdb/bucketdb_test.cpp
@@ -23,6 +23,7 @@ const BucketId BUCKET_1(MIN_NUM_BITS, GID_1.convertToBucketId().getRawId());
 const Timestamp TIME_1(1u);
 const Timestamp TIME_2(2u);
 const Timestamp TIME_3(3u);
+constexpr uint32_t DOCSIZE_1(4096u);
 
 typedef BucketInfo::ReadyState RS;
 typedef SubDbType SDT;
@@ -42,6 +43,17 @@ assertDocCount(uint32_t ready,
 }
 
 void
+assertDocSizes(size_t ready,
+               size_t notReady,
+               size_t removed,
+               const BucketState &state)
+{
+    EXPECT_EQUAL(ready, state.getReadyDocSizes());
+    EXPECT_EQUAL(notReady, state.getNotReadyDocSizes());
+    EXPECT_EQUAL(removed, state.getRemovedDocSizes());
+}
+
+void
 assertReady(bool expReady,
             const BucketInfo &info)
 {
@@ -54,37 +66,55 @@ struct Fixture
     Fixture()
         : _db()
     {}
-    const BucketState &add(const Timestamp &timestamp,
-                           SubDbType subDbType) {
-        return _db.add(GID_1, BUCKET_1, timestamp, subDbType);
+    const BucketState &add(const Timestamp &timestamp, uint32_t docSize, SubDbType subDbType) {
+        return _db.add(GID_1, BUCKET_1, timestamp, docSize, subDbType);
+    }
+    const BucketState &add(const Timestamp &timestamp, SubDbType subDbType) {
+        return add(timestamp, DOCSIZE_1, subDbType);
     }
-    BucketState remove(const Timestamp &timestamp,
-                SubDbType subDbType) {
-        _db.remove(GID_1, BUCKET_1, timestamp, subDbType);
+    BucketState remove(const Timestamp &timestamp, uint32_t docSize, SubDbType subDbType) {
+        _db.remove(GID_1, BUCKET_1, timestamp, docSize, subDbType);
         return get();
     }
+    BucketState remove(const Timestamp &timestamp, SubDbType subDbType) {
+        return remove(timestamp, DOCSIZE_1, subDbType);
+    }
     BucketState get() const {
         return _db.get(BUCKET_1);
     }
-    BucketChecksum getChecksum(const Timestamp &timestamp,
-                               SubDbType subDbType) {
+    BucketChecksum getChecksum(const Timestamp &timestamp, uint32_t docSize, SubDbType subDbType) {
         BucketDB db;
-        BucketChecksum retval = db.add(GID_1, BUCKET_1, timestamp, subDbType).getChecksum();
+        BucketChecksum retval = db.add(GID_1, BUCKET_1, timestamp, docSize, subDbType).getChecksum();
         // Must ensure empty bucket db before destruction.
-        db.remove(GID_1, BUCKET_1, timestamp, subDbType);
+        db.remove(GID_1, BUCKET_1, timestamp, docSize, subDbType);
         return retval;
     }
+    BucketChecksum getChecksum(const Timestamp &timestamp, SubDbType subDbType) {
+        return getChecksum(timestamp, DOCSIZE_1, subDbType);
+    }
 };
 
 TEST_F("require that bucket db tracks doc counts per sub db type", Fixture)
 {
-    assertDocCount(0, 0, 0, f.get());
-    assertDocCount(1, 0, 0, f.add(TIME_1, SDT::READY));
-    assertDocCount(1, 1, 0, f.add(TIME_2, SDT::NOTREADY));
-    assertDocCount(1, 1, 1, f.add(TIME_3, SDT::REMOVED));
-    assertDocCount(0, 1, 1, f.remove(TIME_1, SDT::READY));
-    assertDocCount(0, 0, 1, f.remove(TIME_2, SDT::NOTREADY));
-    assertDocCount(0, 0, 0, f.remove(TIME_3, SDT::REMOVED));
+    TEST_DO(assertDocCount(0, 0, 0, f.get()));
+    TEST_DO(assertDocCount(1, 0, 0, f.add(TIME_1, SDT::READY)));
+    TEST_DO(assertDocCount(1, 1, 0, f.add(TIME_2, SDT::NOTREADY)));
+    TEST_DO(assertDocCount(1, 1, 1, f.add(TIME_3, SDT::REMOVED)));
+    TEST_DO(assertDocCount(0, 1, 1, f.remove(TIME_1, SDT::READY)));
+    TEST_DO(assertDocCount(0, 0, 1, f.remove(TIME_2, SDT::NOTREADY)));
+    TEST_DO(assertDocCount(0, 0, 0, f.remove(TIME_3, SDT::REMOVED)));
+}
+
+TEST_F("require that bucket db tracks doc sizes per sub db type", Fixture)
+{
+    constexpr size_t S = DOCSIZE_1;
+    TEST_DO(assertDocSizes(0, 0, 0, f.get()));
+    TEST_DO(assertDocSizes(S, 0, 0, f.add(TIME_1, DOCSIZE_1, SDT::READY)));
+    TEST_DO(assertDocSizes(S, S, 0, f.add(TIME_2, DOCSIZE_1, SDT::NOTREADY)));
+    TEST_DO(assertDocSizes(S, S, S, f.add(TIME_3, DOCSIZE_1, SDT::REMOVED)));
+    TEST_DO(assertDocSizes(0, S, S, f.remove(TIME_1, DOCSIZE_1, SDT::READY)));
+    TEST_DO(assertDocSizes(0, 0, S, f.remove(TIME_2, DOCSIZE_1, SDT::NOTREADY)));
+    TEST_DO(assertDocSizes(0, 0, 0, f.remove(TIME_3, DOCSIZE_1, SDT::REMOVED)));
 }
 
 TEST_F("require that bucket checksum is a combination of sub db types", Fixture)
@@ -120,13 +150,13 @@ TEST_F("require that bucket can be cached", Fixture)
     f._db.cacheBucket(BUCKET_1);
     EXPECT_TRUE(f._db.isCachedBucket(BUCKET_1));
 
-    assertDocCount(1, 0, 0, f._db.cachedGet(BUCKET_1));
+    TEST_DO(assertDocCount(1, 0, 0, f._db.cachedGet(BUCKET_1)));
     f.add(TIME_2, SDT::NOTREADY);
-    assertDocCount(1, 0, 0, f._db.cachedGet(BUCKET_1));
+    TEST_DO(assertDocCount(1, 0, 0, f._db.cachedGet(BUCKET_1)));
 
     f._db.uncacheBucket();
     EXPECT_FALSE(f._db.isCachedBucket(BUCKET_1));
-    assertDocCount(1, 1, 0, f._db.cachedGet(BUCKET_1));
+    TEST_DO(assertDocCount(1, 1, 0, f._db.cachedGet(BUCKET_1)));
 
     // Must ensure empty bucket db before destruction.
     f.remove(TIME_1, SDT::READY);
@@ -136,7 +166,7 @@ TEST_F("require that bucket can be cached", Fixture)
 TEST("require that bucket db can be explored")
 {
     BucketDBOwner db;
-    db.takeGuard()->add(GID_1, BUCKET_1, TIME_1, SDT::READY);
+    db.takeGuard()->add(GID_1, BUCKET_1, TIME_1, DOCSIZE_1, SDT::READY);
     {
         BucketDBExplorer explorer(db.takeGuard());
         Slime expectSlime;
@@ -146,10 +176,13 @@ TEST("require that bucket db can be explored")
             "  buckets: ["
             "    {"
             "      id: '0x2000000000000031',"
-            "      checksum: '0x93939394',"
+            "      checksum: '0x9393a394',"
             "      readyCount: 1,"
             "      notReadyCount: 0,"
             "      removedCount: 0,"
+            "      readyDocSizes: 4096,"
+            "      notReadyDocSizes: 0,"
+            "      removedDocSizes: 0,"
             "      active: false"
             "    }"
             "  ]"
@@ -163,7 +196,7 @@ TEST("require that bucket db can be explored")
     }
 
     // Must ensure empty bucket db before destruction.
-    db.takeGuard()->remove(GID_1, BUCKET_1, TIME_1, SDT::READY);
+    db.takeGuard()->remove(GID_1, BUCKET_1, TIME_1, DOCSIZE_1, SDT::READY);
 }
 
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/searchcore/src/tests/proton/documentdb/buckethandler/buckethandler_test.cpp b/searchcore/src/tests/proton/documentdb/buckethandler/buckethandler_test.cpp
index ba7c17f1c8d..03690ead45a 100644
--- a/searchcore/src/tests/proton/documentdb/buckethandler/buckethandler_test.cpp
+++ b/searchcore/src/tests/proton/documentdb/buckethandler/buckethandler_test.cpp
@@ -22,6 +22,7 @@ const PartitionId PART_ID(0);
 const GlobalId GID_1("111111111111");
 const BucketId BUCKET_1(8, GID_1.convertToBucketId().getRawId());
 const Timestamp TIME_1(1u);
+const uint32_t DOCSIZE_1(4096u);
 
 struct MySubDb
 {
@@ -182,9 +183,9 @@ TEST_F("require that handleGetBucketInfo() can get cached bucket", Fixture)
 {
     {
         BucketDBOwner::Guard db = f._bucketDB->takeGuard();
-        db->add(GID_1, BUCKET_1, TIME_1, SubDbType::READY);
+        db->add(GID_1, BUCKET_1, TIME_1, DOCSIZE_1, SubDbType::READY);
         db->cacheBucket(BUCKET_1);
-        db->add(GID_1, BUCKET_1, TIME_1, SubDbType::NOTREADY);
+        db->add(GID_1, BUCKET_1, TIME_1, DOCSIZE_1, SubDbType::NOTREADY);
     }
     f.handleGetBucketInfo(BUCKET_1);
     EXPECT_TRUE(expectEqual(1, 1, f._bucketInfo.getInfo()));
@@ -196,8 +197,8 @@ TEST_F("require that handleGetBucketInfo() can get cached bucket", Fixture)
     {
         // Must ensure empty bucket db before destruction.
         BucketDBOwner::Guard db = f._bucketDB->takeGuard();
-        db->remove(GID_1, BUCKET_1, TIME_1, SubDbType::READY);
-        db->remove(GID_1, BUCKET_1, TIME_1, SubDbType::NOTREADY);
+        db->remove(GID_1, BUCKET_1, TIME_1, DOCSIZE_1, SubDbType::READY);
+        db->remove(GID_1, BUCKET_1, TIME_1, DOCSIZE_1, SubDbType::NOTREADY);
     }
 }
 
diff --git a/searchcore/src/tests/proton/documentmetastore/documentmetastore_test.cpp b/searchcore/src/tests/proton/documentmetastore/documentmetastore_test.cpp
index 1958d79b1f9..e99b09c8b81 100644
--- a/searchcore/src/tests/proton/documentmetastore/documentmetastore_test.cpp
+++ b/searchcore/src/tests/proton/documentmetastore/documentmetastore_test.cpp
@@ -263,11 +263,10 @@ uint32_t docSize4 = 1;
 uint32_t docSize5 = 1;
 
 uint32_t
-addGid(DocumentMetaStore &dms, const GlobalId &gid, const BucketId &bid, Timestamp timestamp)
+addGid(DocumentMetaStore &dms, const GlobalId &gid, const BucketId &bid, Timestamp timestamp, uint32_t docSize = 1)
 {
     Result inspect = dms.inspect(gid);
     PutRes putRes;
-    uint32_t docSize = 1;
     EXPECT_TRUE((putRes = dms.put(gid, bid, timestamp, docSize, inspect.getLid())).ok());
     return putRes.getLid();
 }
@@ -812,10 +811,11 @@ TEST("requireThatBasicBucketInfoWorks")
     uint32_t maxcnt = 0u;
     BucketDBOwner::Guard bucketDB = dms.getBucketDB().takeGuard();
     for (Map::const_iterator i = m.begin(), ie = m.end(); i != ie; ++i) {
+        uint32_t docSize = 1;
         if (i->first.first == prevBucket) {
             cksum = BucketChecksum(cksum +
                                    BucketState::calcChecksum(i->first.second,
-                                           i->second));
+                                                             i->second, docSize));
             ++cnt;
         } else {
             BucketInfo bi = bucketDB->get(prevBucket);
@@ -823,7 +823,7 @@ TEST("requireThatBasicBucketInfoWorks")
             EXPECT_EQUAL(cksum, bi.getChecksum());
             prevBucket = i->first.first;
             cksum = BucketState::calcChecksum(i->first.second,
-                    i->second);
+                                              i->second, docSize);
             maxcnt = std::max(maxcnt, cnt);
             cnt = 1u;
         }
@@ -1906,16 +1906,24 @@ TEST("requireThatShrinkViaFlushTargetWorks")
 namespace {
 
 void
-addLid(DocumentMetaStore &dms, uint32_t lid)
+addLid(DocumentMetaStore &dms, uint32_t lid, uint32_t docSize = 1)
 {
     GlobalId gid = createGid(lid);
     BucketId bucketId(gid.convertToBucketId());
     bucketId.setUsedBits(numBucketBits);
-    uint32_t addedLid = addGid(dms, gid, bucketId, Timestamp(lid + timestampBias));
+    uint32_t addedLid = addGid(dms, gid, bucketId, Timestamp(lid + timestampBias), docSize);
     EXPECT_EQUAL(lid, addedLid);
 }
 
 void
+assertSize(DocumentMetaStore &dms, uint32_t lid, uint32_t expSize)
+{
+    EXPECT_TRUE(dms.validLid(lid));
+    const auto &metaData = dms.getRawMetaData(lid);
+    EXPECT_EQUAL(expSize, metaData.getSize());
+}
+
+void
 removeLid(DocumentMetaStore &dms, uint32_t lid)
 {
     dms.remove(lid);
@@ -1965,6 +1973,36 @@ TEST("requireThatSecondShrinkWorksAfterCompactAndInactiveInsert")
     TEST_DO(assertShrink(dms, 2, 1));
 }
 
+TEST("require that document sizes are saved")
+{
+    DocumentMetaStore dms1(createBucketDB());
+    dms1.constructFreeList();
+    TEST_DO(addLid(dms1, 1, 100));
+    TEST_DO(addLid(dms1, 2, 10000));
+    TEST_DO(addLid(dms1, 3, 100000000));
+
+    TuneFileAttributes tuneFileAttributes;
+    DummyFileHeaderContext fileHeaderContext;
+    AttributeFileSaveTarget saveTarget(tuneFileAttributes, fileHeaderContext);
+    EXPECT_TRUE(dms1.saveAs("documentmetastore3", saveTarget));
+    dms1.setTrackDocumentSizes(false);
+    EXPECT_TRUE(dms1.saveAs("documentmetastore4", saveTarget));
+
+    DocumentMetaStore dms3(createBucketDB(), "documentmetastore3");
+    EXPECT_TRUE(dms3.load());
+    dms3.constructFreeList();
+    TEST_DO(assertSize(dms3, 1, 100));
+    TEST_DO(assertSize(dms3, 2, 10000));
+    TEST_DO(assertSize(dms3, 3, (1u << 24) - 1));
+
+    DocumentMetaStore dms4(createBucketDB(), "documentmetastore4");
+    EXPECT_TRUE(dms4.load());
+    dms4.constructFreeList();
+    TEST_DO(assertSize(dms4, 1, 1));
+    TEST_DO(assertSize(dms4, 2, 1));
+    TEST_DO(assertSize(dms4, 3, 1));
+}
+
 }
 
 TEST_MAIN()