Walk trees for compaction without allocating per page

database/src/main/java/io/hstore/db/property/PropertySlots.java

@@ -54,7 +54,13 @@ public boolean holdsRefs() {
 
                 @Override
                 public IndexRoot mapRefs(IndexRoot value, RefMapper mapper) {
-                    return new IndexRoot(mapper.map(value.root(), VALUE_POSTINGS.schema()));
+                    Ref mapped = mapper.map(value.root(), VALUE_POSTINGS.schema());
+                    return mapped == value.root() ? value : new IndexRoot(mapped);
+                }
+
+                @Override
+                public void forEachRef(IndexRoot value, RefVisitor visitor) {
+                    visitor.visit(value.root(), VALUE_POSTINGS.schema());
                 }
             },
             EntryMeasure.keyed(root -> root.root().summary().fingerprint()));

engine/src/main/java/io/hstore/engine/catalog/AtomCodec.java

@@ -80,10 +80,20 @@ public boolean holdsRefs() {
     public AtomRecord mapRefs(AtomRecord record, RefMapper mapper) {
         return switch (record) {
             case NodeRecord node -> node;
-            case EdgeRecord edge -> new EdgeRecord(edge.type(), edge.kind(),
-                    mapper.map(edge.members(), TopologySchemas.members(edge.kind())),
-                    mapper.map(edge.order(), TopologySchemas.ORDER_INDEX),
-                    edge.version(), edge.dataRef(), edge.tenant(), edge.isolation());
+            case EdgeRecord edge -> {
+                Ref members = mapper.map(edge.members(), TopologySchemas.members(edge.kind()));
+                Ref order = mapper.map(edge.order(), TopologySchemas.ORDER_INDEX);
+                yield members == edge.members() && order == edge.order() ? edge
+                        : new EdgeRecord(edge.type(), edge.kind(), members, order, edge.version(), edge.dataRef(), edge.tenant(), edge.isolation());
+            }
         };
     }
+
+    @Override
+    public void forEachRef(AtomRecord record, RefVisitor visitor) {
+        if (record instanceof EdgeRecord edge) {
+            visitor.visit(edge.members(), TopologySchemas.members(edge.kind()));
+            visitor.visit(edge.order(), TopologySchemas.ORDER_INDEX);
+        }
+    }
 }

engine/src/main/java/io/hstore/engine/index/PostingsCodec.java

@@ -89,6 +89,17 @@ public boolean holdsRefs() {
 
     @Override
     public Postings<P> mapRefs(Postings<P> value, RefMapper mapper) {
-        return value instanceof Postings.Promoted<P>(Ref root) ? new Postings.Promoted<>(mapper.map(root, postings)) : value;
+        if (!(value instanceof Postings.Promoted<P>(Ref root))) {
+            return value;
+        }
+        Ref mapped = mapper.map(root, postings);
+        return mapped == root ? value : new Postings.Promoted<>(mapped);
+    }
+
+    @Override
+    public void forEachRef(Postings<P> value, RefVisitor visitor) {
+        if (value instanceof Postings.Promoted<P>(Ref root)) {
+            visitor.visit(root, postings);
+        }
     }
 }

engine/src/main/java/io/hstore/engine/maintenance/Compactor.java

@@ -7,17 +7,16 @@
 import io.hstore.engine.page.SegmentInfo;
 import io.hstore.engine.page.SegmentState;
 import io.hstore.engine.tree.TreeWalker;
+import io.hstore.engine.tree.VisitedPages;
 import io.hstore.engine.txn.TransactionManager;
 
 import java.util.Comparator;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.OptionalLong;
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.function.LongPredicate;
-import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 public final class Compactor {
@@ -36,14 +35,17 @@ public Compactor(TransactionManager transactions, PageStore pages, SlotRegistry
     }
 
     public Map<Integer, Long> liveness() {
-        Map<Long, Integer> visited = new HashMap<>();
+        VisitedPages visited = new VisitedPages(pages.segments().stream().mapToLong(SegmentInfo::pages).sum());
         TreeWalker walker = new TreeWalker(transactions.source());
         Stream.concat(transactions.history().stream(), Stream.of(transactions.current()))
                 .flatMap(generation -> generation.branches().values().stream())
                 .forEach(branch -> branch.roots().roots().forEach((slot, ref) ->
                         walker.visit(ref, slots.slot(slot).schema(), visited)));
-        return visited.entrySet().stream().collect(Collectors.groupingBy(entry -> PageId.segmentOf(entry.getKey()), TreeMap::new,
-                Collectors.summingLong(entry -> (long) entry.getValue() * PageId.UNIT_BYTES)));
+        Map<Integer, long[]> bySegment = new TreeMap<>();
+        visited.forEach((pageId, units) -> bySegment.computeIfAbsent(PageId.segmentOf(pageId), _ -> new long[1])[0] += (long) units * PageId.UNIT_BYTES);
+        Map<Integer, Long> live = new TreeMap<>();
+        bySegment.forEach((segment, bytes) -> live.put(segment, bytes[0]));
+        return live;
     }
 
     public Report compact(double liveThreshold) {

engine/src/main/java/io/hstore/engine/tree/TreeWalker.java

@@ -17,29 +17,50 @@ public TreeWalker(NodeSource source) {
         this.source = source;
     }
 
+    @FunctionalInterface
+    private interface FirstVisit {
+        boolean first(long pageId, int units);
+    }
+
     public void visit(Ref root, TreeSchema<?> schema, Map<Long, Integer> visited) {
-        visit(root, schema, UNKNOWN_HEIGHT, visited);
+        new Walk((pageId, units) -> visited.putIfAbsent(pageId, units) == null).walk(root, schema, UNKNOWN_HEIGHT);
     }
 
-    private void visit(Ref ref, TreeSchema<?> schema, int height, Map<Long, Integer> visited) {
-        if (!(ref instanceof Ref.Stored(long pageId, int units, Summary _)) || visited.putIfAbsent(pageId, units) != null) {
-            return;
+    public void visit(Ref root, TreeSchema<?> schema, VisitedPages visited) {
+        new Walk(visited::add).walk(root, schema, UNKNOWN_HEIGHT);
+    }
+
+    private final class Walk implements ValueCodec.RefVisitor {
+        private final FirstVisit visited;
+
+        Walk(FirstVisit visited) {
+            this.visited = visited;
         }
-        if (height == 0 && !schema.codec().holdsRefs()) {
-            return;
+
+        @Override
+        public void visit(Ref ref, TreeSchema<?> schema) {
+            walk(ref, schema, UNKNOWN_HEIGHT);
         }
-        switch (source.load(pageId, schema)) {
-            case Leaf leaf -> {
-                for (int i = 0; i < leaf.size(); i++) {
-                    mapNested(schema, leaf.value(i), (nested, nestedSchema) -> {
-                        visit(nested, nestedSchema, UNKNOWN_HEIGHT, visited);
-                        return nested;
-                    });
-                }
+
+        void walk(Ref ref, TreeSchema<?> schema, int height) {
+            if (!(ref instanceof Ref.Stored(long pageId, int units, Summary _)) || !visited.first(pageId, units)) {
+                return;
             }
-            case Branch branch -> {
-                for (int i = 0; i < branch.size(); i++) {
-                    visit(branch.child(i), schema, branch.height() - 1, visited);
+            if (height == 0 && !schema.codec().holdsRefs()) {
+                return;
+            }
+            switch (source.load(pageId, schema)) {
+                case Leaf leaf -> {
+                    if (schema.codec().holdsRefs()) {
+                        for (int i = 0; i < leaf.size(); i++) {
+                            forEachNested(schema, leaf.value(i), this);
+                        }
+                    }
+                }
+                case Branch branch -> {
+                    for (int i = 0; i < branch.size(); i++) {
+                        walk(branch.child(i), schema, branch.height() - 1);
+                    }
                 }
             }
         }
@@ -94,7 +115,11 @@ private void pairMoves(Ref before, Ref after, TreeSchema<?> schema, int height,
         if (height == 0 && !schema.codec().holdsRefs()) {
             return;
         }
-        switch (source.load(oldId, schema)) {
+        Node original = source.load(oldId, schema);
+        if (original instanceof Leaf && !schema.codec().holdsRefs()) {
+            return;
+        }
+        switch (original) {
             case Branch branch -> {
                 Branch copy = (Branch) source.load(newId, schema);
                 for (int i = 0; i < branch.size(); i++) {
@@ -107,15 +132,11 @@ private void pairMoves(Ref before, Ref after, TreeSchema<?> schema, int height,
                     List<Ref> oldRefs = new ArrayList<>();
                     List<Ref> newRefs = new ArrayList<>();
                     List<TreeSchema<?>> schemas = new ArrayList<>();
-                    mapNested(schema, leaf.value(i), (nested, nestedSchema) -> {
+                    forEachNested(schema, leaf.value(i), (nested, nestedSchema) -> {
                         oldRefs.add(nested);
                         schemas.add(nestedSchema);
-                        return nested;
-                    });
-                    mapNested(schema, copy.value(i), (nested, _) -> {
-                        newRefs.add(nested);
-                        return nested;
                     });
+                    forEachNested(schema, copy.value(i), (nested, _) -> newRefs.add(nested));
                     for (int k = 0; k < oldRefs.size(); k++) {
                         pairMoves(oldRefs.get(k), newRefs.get(k), schemas.get(k), UNKNOWN_HEIGHT, moved);
                     }
@@ -175,6 +196,12 @@ private Ref rebuild(Ref ref, long pageId, TreeSchema<?> schema, boolean force, W
         };
     }
 
+    private static <V> void forEachNested(TreeSchema<V> schema, Object value, ValueCodec.RefVisitor visitor) {
+        if (schema.codec().holdsRefs()) {
+            schema.codec().forEachRef(schema.cast(value), visitor);
+        }
+    }
+
     private static <V> Object mapNested(TreeSchema<V> schema, Object value, ValueCodec.RefMapper mapper) {
         if (!schema.codec().holdsRefs()) {
             return value;

engine/src/main/java/io/hstore/engine/tree/ValueCodec.java

@@ -27,6 +27,13 @@ default V mapRefs(V value, RefMapper mapper) {
         return value;
     }
 
+    default void forEachRef(V value, RefVisitor visitor) {
+        mapRefs(value, (ref, schema) -> {
+            visitor.visit(ref, schema);
+            return ref;
+        });
+    }
+
     static <V> ValueCodec<V> rows(ToIntFunction<V> size, BiConsumer<ByteCursor, V> writer, Function<ByteCursor, V> reader) {
         return new ValueCodec<>() {
             @Override
@@ -52,4 +59,8 @@ public void decode(ByteCursor in, long[] keys, Object[] into) {
     interface RefMapper {
         Ref map(Ref ref, TreeSchema<?> schema);
     }
+
+    interface RefVisitor {
+        void visit(Ref ref, TreeSchema<?> schema);
+    }
 }

engine/src/main/java/io/hstore/engine/tree/VisitedPages.java

@@ -0,0 +1,89 @@
+package io.hstore.engine.tree;
+
+public final class VisitedPages {
+
+    @FunctionalInterface
+    public interface Consumer {
+        void accept(long pageId, int units);
+    }
+
+    private static final long EMPTY = 0L;
+
+    private long[] keys;
+    private int[] units;
+    private int size;
+
+    public VisitedPages() {
+        this(1 << 10);
+    }
+
+    public VisitedPages(long expected) {
+        int capacity = Integer.highestOneBit((int) Math.min(1 << 30, Math.max(16, expected * 10 / 6))) << 1;
+        keys = new long[capacity];
+        units = new int[capacity];
+    }
+
+    public boolean add(long pageId, int pageUnits) {
+        if (pageId == EMPTY) {
+            throw new IllegalArgumentException("page id 0 is not a page");
+        }
+        if ((size + 1) * 10L > keys.length * 6L) {
+            grow();
+        }
+        int slot = slot(pageId, keys.length);
+        while (keys[slot] != EMPTY) {
+            if (keys[slot] == pageId) {
+                return false;
+            }
+            slot = (slot + 1) & (keys.length - 1);
+        }
+        keys[slot] = pageId;
+        units[slot] = pageUnits;
+        size++;
+        return true;
+    }
+
+    public boolean contains(long pageId) {
+        int slot = slot(pageId, keys.length);
+        while (keys[slot] != EMPTY) {
+            if (keys[slot] == pageId) {
+                return true;
+            }
+            slot = (slot + 1) & (keys.length - 1);
+        }
+        return false;
+    }
+
+    public int size() {
+        return size;
+    }
+
+    public void forEach(Consumer consumer) {
+        for (int i = 0; i < keys.length; i++) {
+            if (keys[i] != EMPTY) {
+                consumer.accept(keys[i], units[i]);
+            }
+        }
+    }
+
+    private void grow() {
+        long[] oldKeys = keys;
+        int[] oldUnits = units;
+        keys = new long[oldKeys.length * 2];
+        units = new int[oldKeys.length * 2];
+        for (int i = 0; i < oldKeys.length; i++) {
+            if (oldKeys[i] != EMPTY) {
+                int slot = slot(oldKeys[i], keys.length);
+                while (keys[slot] != EMPTY) {
+                    slot = (slot + 1) & (keys.length - 1);
+                }
+                keys[slot] = oldKeys[i];
+                units[slot] = oldUnits[i];
+            }
+        }
+    }
+
+    private static int slot(long pageId, int capacity) {
+        return (int) (Hashing.mix(pageId) >>> 32) & (capacity - 1);
+    }
+}

engine/src/main/java/io/hstore/engine/txn/TransactionManager.java

@@ -24,6 +24,7 @@
 import io.hstore.engine.tree.Tree;
 import io.hstore.engine.tree.TreeSchema;
 import io.hstore.engine.tree.TreeWalker;
+import io.hstore.engine.tree.VisitedPages;
 import io.hstore.engine.tree.WriteScope;
 import io.hstore.engine.wal.WalRecord;
 import io.hstore.engine.wal.WriteAheadLog;
@@ -322,7 +323,7 @@ public long relocate(LongPredicate moving) {
             commitLock.unlock();
         }
         TreeWalker walker = new TreeWalker(storage.source());
-        Set<Long> reached = new HashSet<>();
+        VisitedPages reached = new VisitedPages();
         Map<Long, Ref> moved = new HashMap<>();
         List<WalRecord> records = new ArrayList<>();
         List<PageWrite> writes = new ArrayList<>();
@@ -331,7 +332,7 @@ public long relocate(LongPredicate moving) {
         List<Runnable> pairings = new ArrayList<>();
         snapshot.branches().values().stream().filter(Branch::isActive).forEach(branch -> branch.roots().roots().forEach((slot, ref) -> {
             TreeSchema<?> schema = storage.slots().slot(slot).schema();
-            Ref copy = materializer.materialize(walker.relocate(ref, schema, moving, copies, reached::add,
+            Ref copy = materializer.materialize(walker.relocate(ref, schema, moving, copies, pageId -> reached.add(pageId, 0),
                     leaf -> copyPage(leaf, snapshot.id() + 1, txnId, records, writes)), schema);
             pairings.add(() -> walker.pairMoves(ref, copy, schema, moved));
         }));

engine/src/test/java/io/hstore/engine/CompactionTest.java

@@ -31,6 +31,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 import java.util.TreeSet;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -39,6 +40,7 @@
 import java.util.random.RandomGenerator;
 import java.util.random.RandomGeneratorFactory;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -114,6 +116,22 @@ private static Map<Long, TreeSet<Long>> copy(Map<Long, TreeSet<Long>> oracle) {
         return oracle.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> new TreeSet<>(entry.getValue())));
     }
 
+    @Test
+    void livenessCountsEveryPageOfEveryRetainedRootOnce() {
+        try (StorageEngine engine = StorageEngine.open(directory, options())) {
+            Map<Long, TreeSet<Long>> oracle = load(engine);
+            churn(engine, Branch.MAIN, oracle, RandomGeneratorFactory.of("L64X128MixRandom").create(6), 30);
+            Map<Long, Integer> pages = new HashMap<>();
+            TreeWalker walker = new TreeWalker(engine.transactions().source());
+            Stream.concat(engine.transactions().history().stream(), Stream.of(engine.transactions().current()))
+                    .flatMap(generation -> generation.branches().values().stream())
+                    .forEach(branch -> branch.roots().roots().forEach((slot, ref) -> walker.visit(ref, engine.slots().slot(slot).schema(), pages)));
+            Map<Integer, Long> expected = new TreeMap<>();
+            pages.forEach((pageId, units) -> expected.merge(PageId.segmentOf(pageId), (long) units * PageId.UNIT_BYTES, Long::sum));
+            assertEquals(expected, engine.liveness());
+        }
+    }
+
     @Test
     void commitsGoThroughWhileCompactionCopiesPages() throws Exception {
         try (StorageEngine engine = StorageEngine.open(directory, options())) {