From b3dd29599e94d18b0a2674391fedac051e166040 Mon Sep 17 00:00:00 2001
From: Sunjeet Singh <sunjeets@netflix.com>
Date: Wed, 11 Oct 2023 07:44:38 -0700
Subject: [PATCH] Cleanup of object read state

---
 .../object/HollowObjectTypeReadState.java     |  17 +-
 .../HollowObjectTypeReadStateShard.java       |   6 +-
 .../write/HollowObjectTypeWriteState.java     |   2 +
 ...ollowObjectTypeDataElementsJoinerTest.java | 150 +++++++++---------
 4 files changed, 88 insertions(+), 87 deletions(-)

diff --git a/hollow/src/main/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeReadState.java b/hollow/src/main/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeReadState.java
index c2786e76a0..b0316e0a71 100644
--- a/hollow/src/main/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeReadState.java
+++ b/hollow/src/main/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeReadState.java
@@ -45,8 +45,8 @@ public class HollowObjectTypeReadState extends HollowTypeReadState implements Ho
     private final HollowObjectSchema unfilteredSchema;
     private final HollowObjectSampler sampler;
 
-    public static class ShardsHolder {   // SNAP: TODO: remove
-        public final HollowObjectTypeReadStateShard shards[];   // SNAP: TODO: remove
+    static class ShardsHolder {
+        final HollowObjectTypeReadStateShard shards[];
         final int shardNumberMask;
 
         private ShardsHolder(HollowSchema schema, int numShards) {
@@ -71,7 +71,7 @@ private ShardsHolder(HollowSchema schema, HollowObjectTypeDataElements[] dataEle
         }
     }
 
-    volatile ShardsHolder shardsVolatile;    // SNAP: TODO: remove
+    volatile ShardsHolder shardsVolatile;
 
     private int maxOrdinal;
 
@@ -189,7 +189,7 @@ static int shardingFactor(int oldNumShards, int newNumShards) {
      *
      * @param newNumShards The desired number of shards
      */
-    public void reshard(int newNumShards) {   // TODO: package private
+    void reshard(int newNumShards) {
         int prevNumShards = shardsVolatile.shards.length;
         int shardingFactor = shardingFactor(prevNumShards, newNumShards);
         HollowObjectTypeDataElements[] newDataElements;
@@ -558,17 +558,16 @@ HollowObjectTypeDataElements[] currentDataElements() {
 
     @Override
     protected void applyToChecksum(HollowChecksum checksum, HollowSchema withSchema) {
-        final HollowObjectTypeReadStateShard[] shards = this.shardsVolatile.shards;
+        final ShardsHolder shardsHolder = this.shardsVolatile;
+        final HollowObjectTypeReadStateShard[] shards = shardsHolder.shards;
+        int shardNumberMask = shardsHolder.shardNumberMask;
         if(!(withSchema instanceof HollowObjectSchema))
             throw new IllegalArgumentException("HollowObjectTypeReadState can only calculate checksum with a HollowObjectSchema: " + getSchema().getName());
 
         BitSet populatedOrdinals = getPopulatedOrdinals();
 
         for(int i=0;i<shards.length;i++) {
-            if (shards[i].shardOrdinalShift != (31 - Integer.numberOfLeadingZeros(shards.length))) {
-                throw new UnsupportedOperationException("applyToChecksum called for virtual shard, unexpected");  // SNAP: TODO: remove this altogether, or support applyToChecksum for virtual shards
-            }
-            shards[i].applyToChecksum(checksum, withSchema, populatedOrdinals, i, shards.length);
+            shards[i].applyToChecksum(checksum, withSchema, populatedOrdinals, i, shardNumberMask);
         }
     }
 
diff --git a/hollow/src/main/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeReadStateShard.java b/hollow/src/main/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeReadStateShard.java
index 1b52eb7704..ce78774cd9 100644
--- a/hollow/src/main/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeReadStateShard.java
+++ b/hollow/src/main/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeReadStateShard.java
@@ -393,7 +393,7 @@ void setCurrentData(HollowObjectTypeDataElements data) {
         this.currentDataVolatile = data;
     }
 
-    protected void applyToChecksum(HollowChecksum checksum, HollowSchema withSchema, BitSet populatedOrdinals, int shardNumber, int numShards) {
+    protected void applyToChecksum(HollowChecksum checksum, HollowSchema withSchema, BitSet populatedOrdinals, int shardNumber, int shardNumberMask) {
         if(!(withSchema instanceof HollowObjectSchema))
             throw new IllegalArgumentException("HollowObjectTypeReadState can only calculate checksum with a HollowObjectSchema: " + schema.getName());
 
@@ -412,8 +412,8 @@ protected void applyToChecksum(HollowChecksum checksum, HollowSchema withSchema,
         HollowObjectTypeDataElements currentData = currentDataVolatile;
         int ordinal = populatedOrdinals.nextSetBit(0);
         while(ordinal != ORDINAL_NONE) {
-            if((ordinal & (numShards - 1)) == shardNumber) {
-                int shardOrdinal = ordinal / numShards;
+            if((ordinal & shardNumberMask) == shardNumber) {
+                int shardOrdinal = ordinal >> shardOrdinalShift;
                 checksum.applyInt(ordinal);
                 for(int i=0;i<fieldIndexes.length;i++) {
                     int fieldIdx = fieldIndexes[i];
diff --git a/hollow/src/main/java/com/netflix/hollow/core/write/HollowObjectTypeWriteState.java b/hollow/src/main/java/com/netflix/hollow/core/write/HollowObjectTypeWriteState.java
index 3372ec4c50..1d0d6de141 100644
--- a/hollow/src/main/java/com/netflix/hollow/core/write/HollowObjectTypeWriteState.java
+++ b/hollow/src/main/java/com/netflix/hollow/core/write/HollowObjectTypeWriteState.java
@@ -62,6 +62,8 @@ public HollowObjectSchema getSchema() {
     }
 
     private boolean isDynamicTypeShardingEnabled() {   // SNAP: TODO: configurable
+        // SNAP: TODO: enforce integrity check
+        // SNAP: TODO: publish metric
         return true;
     }
 
diff --git a/hollow/src/test/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeDataElementsJoinerTest.java b/hollow/src/test/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeDataElementsJoinerTest.java
index 335ada9112..2f0dfcea03 100644
--- a/hollow/src/test/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeDataElementsJoinerTest.java
+++ b/hollow/src/test/java/com/netflix/hollow/core/read/engine/object/HollowObjectTypeDataElementsJoinerTest.java
@@ -78,81 +78,81 @@ public void testJoinDifferentFieldWidths() throws IOException {
     }
 
 //    TODO: manually invoked, depends on producer side changes for supporting changing numShards in a delta chain
-//    @Test
-//    public void testLopsidedShards() {
-//        InMemoryBlobStore blobStore = new InMemoryBlobStore();
-//        HollowProducer p = HollowProducer.withPublisher(blobStore)
-//                .withBlobStager(new HollowInMemoryBlobStager())
-//                .build();
-//
-//        p.initializeDataModel(schema);
-//        int targetSize = 64;
-//        p.getWriteEngine().setTargetMaxTypeShardSize(targetSize);
-//        long v1 = oneRunCycle(p, new int[] {0, 1, 2, 3, 4, 5, 6, 7});
-//
-//        HollowConsumer c = HollowConsumer
-//                .withBlobRetriever(blobStore)
-//                .withDoubleSnapshotConfig(new HollowConsumer.DoubleSnapshotConfig() {
-//                    @Override
-//                    public boolean allowDoubleSnapshot() {
-//                        return false;
-//                    }
-//
-//                    @Override
-//                    public int maxDeltasBeforeDoubleSnapshot() {
-//                        return Integer.MAX_VALUE;
-//                    }
-//                })
-//                .withSkipTypeShardUpdateWithNoAdditions()
-//                .build();
-//        c.triggerRefreshTo(v1);
-//
-//        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards());
-//        assertEquals(true, c.getStateEngine().isSkipTypeShardUpdateWithNoAdditions());
-//
-//        long v2 = oneRunCycle(p, new int[] {0, 1, 2, 3, 5, 7});
-//        c.triggerRefreshTo(v2);
-//        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards());
-//
-//        long v3 = oneRunCycle(p, new int[] { 0, 1, 3, 5}); // drop to 1 ordinal per shard, skipTypeShardWithNoAdds will make it so that maxOrdinal is adjusted
-//        c.triggerRefreshTo(v3);
-//        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards());
-//
-//        long v4 = oneRunCycle(p, new int[] { 0, 1, 2, 3}); // now add another ordinal to one shard, maxOrdinals will be lopsided
-//        c.triggerRefreshTo(v4);
-//        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards());
-//
-//        readStateEngine = c.getStateEngine();
-//        assertDataUnchanged(3);
-//
-//        long v5 = oneRunCycle(p, new int[] {0, 1});
-//
-//        // assert lopsided shards before join
-//        assertEquals(2, ((HollowObjectTypeReadState) c.getStateEngine().getTypeState("TestObject")).shardsVolatile.shards[0].currentDataElements().maxOrdinal);
-//        assertEquals(3, ((HollowObjectTypeReadState) c.getStateEngine().getTypeState("TestObject")).shardsVolatile.shards[1].currentDataElements().maxOrdinal);
-//        c.triggerRefreshTo(v5);
-//        assertEquals(1, c.getStateEngine().getTypeState("TestObject").numShards()); // joined to 1 shard
-//        readStateEngine = c.getStateEngine();
-//        assertDataUnchanged(2);
-//
-//        long v6 = oneRunCycle(p, new int[] {0, 1, 2, 3, 4, 5 });
-//        c.triggerRefreshTo(v6);
-//        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards()); // split to 2 shards
-//
-//        long v7 = oneRunCycle(p, new int[] {8, 9});
-//        c.triggerRefreshTo(v7);
-//        assertEquals(4, c.getStateEngine().getTypeState("TestObject").numShards()); // still 2 shards
-//
-//        long v8 = oneRunCycle(p, new int[] {8});
-//        c.triggerRefreshTo(v8);
-//        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards()); // down to 1 shard
-//
-//        c.triggerRefreshTo(v1);
-//        assertEquals(v1, c.getCurrentVersionId());
-//
-//        c.triggerRefreshTo(v8);
-//        assertEquals(v8, c.getCurrentVersionId());
-//    }
+    @Test
+    public void testLopsidedShards() {
+        InMemoryBlobStore blobStore = new InMemoryBlobStore();
+        HollowProducer p = HollowProducer.withPublisher(blobStore)
+                .withBlobStager(new HollowInMemoryBlobStager())
+                .build();
+
+        p.initializeDataModel(schema);
+        int targetSize = 64;
+        p.getWriteEngine().setTargetMaxTypeShardSize(targetSize);
+        long v1 = oneRunCycle(p, new int[] {0, 1, 2, 3, 4, 5, 6, 7});
+
+        HollowConsumer c = HollowConsumer
+                .withBlobRetriever(blobStore)
+                .withDoubleSnapshotConfig(new HollowConsumer.DoubleSnapshotConfig() {
+                    @Override
+                    public boolean allowDoubleSnapshot() {
+                        return false;
+                    }
+
+                    @Override
+                    public int maxDeltasBeforeDoubleSnapshot() {
+                        return Integer.MAX_VALUE;
+                    }
+                })
+                .withSkipTypeShardUpdateWithNoAdditions()
+                .build();
+        c.triggerRefreshTo(v1);
+
+        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards());
+        assertEquals(true, c.getStateEngine().isSkipTypeShardUpdateWithNoAdditions());
+
+        long v2 = oneRunCycle(p, new int[] {0, 1, 2, 3, 5, 7});
+        c.triggerRefreshTo(v2);
+        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards());
+
+        long v3 = oneRunCycle(p, new int[] { 0, 1, 3, 5}); // drop to 1 ordinal per shard, skipTypeShardWithNoAdds will make it so that maxOrdinal is adjusted
+        c.triggerRefreshTo(v3);
+        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards());
+
+        long v4 = oneRunCycle(p, new int[] { 0, 1, 2, 3}); // now add another ordinal to one shard, maxOrdinals will be lopsided
+        c.triggerRefreshTo(v4);
+        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards());
+
+        readStateEngine = c.getStateEngine();
+        assertDataUnchanged(3);
+
+        long v5 = oneRunCycle(p, new int[] {0, 1});
+
+        // assert lopsided shards before join
+        assertEquals(2, ((HollowObjectTypeReadState) c.getStateEngine().getTypeState("TestObject")).shardsVolatile.shards[0].currentDataElements().maxOrdinal);
+        assertEquals(3, ((HollowObjectTypeReadState) c.getStateEngine().getTypeState("TestObject")).shardsVolatile.shards[1].currentDataElements().maxOrdinal);
+        c.triggerRefreshTo(v5);
+        assertEquals(1, c.getStateEngine().getTypeState("TestObject").numShards()); // joined to 1 shard
+        readStateEngine = c.getStateEngine();
+        assertDataUnchanged(2);
+
+        long v6 = oneRunCycle(p, new int[] {0, 1, 2, 3, 4, 5 });
+        c.triggerRefreshTo(v6);
+        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards()); // split to 2 shards
+
+        long v7 = oneRunCycle(p, new int[] {8, 9});
+        c.triggerRefreshTo(v7);
+        assertEquals(4, c.getStateEngine().getTypeState("TestObject").numShards()); // still 2 shards
+
+        long v8 = oneRunCycle(p, new int[] {8});
+        c.triggerRefreshTo(v8);
+        assertEquals(2, c.getStateEngine().getTypeState("TestObject").numShards()); // down to 1 shard
+
+        c.triggerRefreshTo(v1);
+        assertEquals(v1, c.getCurrentVersionId());
+
+        c.triggerRefreshTo(v8);
+        assertEquals(v8, c.getCurrentVersionId());
+    }
 
     private long oneRunCycle(HollowProducer p, int recordIds[]) {
         return p.runCycle(state -> {