src/java/org/apache/cassandra/io/tries/IncrementalDeepTrieWriterPageAware.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.cassandra.io.tries;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import javax.annotation.concurrent.NotThreadSafe;

import org.apache.cassandra.io.util.DataOutputPlus;

/**
 * This class is a variant of {@link IncrementalTrieWriterPageAware} which is able to build even very deep
 * tries. While the parent class uses recursion for clarity, it may end up with stack overflow for tries with
 * very long keys. This implementation can switch processing from stack to heap at a certain depth (provided
 * as a constructor param).
 * <p>
 * This class intentionally repeats code present in the parent class, both in the in-stack and on-heap versions
 * of each of the three implemented recursive operations. Removing this repetition can cause higher stack usage
 * and thus stack overflow failures.
 */
@NotThreadSafe
public class IncrementalDeepTrieWriterPageAware<VALUE> extends IncrementalTrieWriterPageAware<VALUE>
{
    private final int maxRecursionDepth;

    public IncrementalDeepTrieWriterPageAware(TrieSerializer<VALUE, ? super DataOutputPlus> trieSerializer, DataOutputPlus dest, int maxRecursionDepth)
    {
        super(trieSerializer, dest);
        this.maxRecursionDepth = maxRecursionDepth;
    }

    public IncrementalDeepTrieWriterPageAware(TrieSerializer<VALUE, ? super DataOutputPlus> trieSerializer, DataOutputPlus dest)
    {
        this(trieSerializer, dest, 64);
    }

    /**
     * Simple framework for executing recursion using on-heap linked trace to avoid stack overruns.
     */
    static abstract class Recursion<NODE>
    {
        final Recursion<NODE> parent;
        final NODE node;
        final Iterator<NODE> childIterator;

        Recursion(NODE node, Iterator<NODE> childIterator, Recursion<NODE> parent)
        {
            this.parent = parent;
            this.node = node;
            this.childIterator = childIterator;
        }

        /**
         * Make a child Recursion object for the given node and initialize it as necessary to continue processing
         * with it.
         * <p>
         * May return null if the recursion does not need to continue inside the child branch.
         */
        abstract Recursion<NODE> makeChild(NODE child);

        /**
         * Complete the processing this Recursion object.
         * <p>
         * Note: this method is not called for the nodes for which makeChild() returns null.
         */
        abstract void complete() throws IOException;

        /**
         * Complete processing of the given child (possibly retrieve data to apply to any accumulation performed
         * in this Recursion object).
         * <p>
         * This is called when processing a child completes, including when recursion inside the child branch
         * is skipped by makeChild() returning null.
         */
        void completeChild(NODE child)
        {}

        /**
         * Recursive process, in depth-first order, the branch rooted at this recursion node.
         * <p>
         * Returns this.
         */
        Recursion<NODE> process() throws IOException
        {
            Recursion<NODE> curr = this;

            while (true)
            {
                if (curr.childIterator.hasNext())
                {
                    NODE child = curr.childIterator.next();
                    Recursion<NODE> childRec = curr.makeChild(child);
                    if (childRec != null)
                        curr = childRec;
                    else
                        curr.completeChild(child);
                }
                else
                {
                    curr.complete();
                    Recursion<NODE> currParent = curr.parent;
                    if (currParent == null)
                        return curr;
                    currParent.completeChild(curr.node);
                    curr = currParent;
                }
            }
        }
    }

    @Override
    protected int recalcTotalSize(Node<VALUE> node, long nodePosition) throws IOException
    {
        return recalcTotalSizeRecursiveOnStack(node, nodePosition, 0);
    }

    @SuppressWarnings("DuplicatedCode") // intentionally duplicates IncrementalTrieWriterPageAware
    private int recalcTotalSizeRecursiveOnStack(Node<VALUE> node, long nodePosition, int depth) throws IOException
    {
        if (node.hasOutOfPageInBranch)
        {
            int sz = 0;
            for (Node<VALUE> child : node.children)
            {
                if (depth < maxRecursionDepth)
                    sz += recalcTotalSizeRecursiveOnStack(child, nodePosition + sz, depth + 1);
                else
                    sz += recalcTotalSizeRecursiveOnHeap(child, nodePosition + sz);
            }
            node.branchSize = sz;
        }

        // The sizing below will use the branch size calculated above. Since that can change on out-of-page in branch,
        // we need to recalculate the size if either flag is set.
        if (node.hasOutOfPageChildren || node.hasOutOfPageInBranch)
            node.nodeSize = serializer.sizeofNode(node, nodePosition + node.branchSize);

        return node.branchSize + node.nodeSize;
    }

    private int recalcTotalSizeRecursiveOnHeap(Node<VALUE> node, long nodePosition) throws IOException
    {
        if (node.hasOutOfPageInBranch)
            new RecalcTotalSizeRecursion(node, null, nodePosition).process();

        if (node.hasOutOfPageChildren || node.hasOutOfPageInBranch)
            node.nodeSize = serializer.sizeofNode(node, nodePosition + node.branchSize);

        return node.branchSize + node.nodeSize;
    }

    class RecalcTotalSizeRecursion extends Recursion<Node<VALUE>>
    {
        final long nodePosition;
        int sz;

        RecalcTotalSizeRecursion(Node<VALUE> node, Recursion<Node<VALUE>> parent, long nodePosition)
        {
            super(node, node.children.iterator(), parent);
            sz = 0;
            this.nodePosition = nodePosition;
        }

        @Override
        Recursion<Node<VALUE>> makeChild(Node<VALUE> child)
        {
            if (child.hasOutOfPageInBranch)
                return new RecalcTotalSizeRecursion(child, this, nodePosition + sz);
            else
                return null;
        }

        @Override
        void complete()
        {
            node.branchSize = sz;
        }

        @SuppressWarnings("DuplicatedCode") // intentionally duplicates IncrementalTrieWriterPageAware and onStack code
        @Override
        void completeChild(Node<VALUE> child)
        {
            // This will be called for nodes that were recursively processed as well as the ones that weren't.

            // The sizing below will use the branch size calculated above. Since that can change on out-of-page in branch,
            // we need to recalculate the size if either flag is set.
            if (child.hasOutOfPageChildren || child.hasOutOfPageInBranch)
            {
                long childPosition = this.nodePosition + sz;
                child.nodeSize = serializer.sizeofNode(child, childPosition + child.branchSize);
            }

            sz += child.branchSize + child.nodeSize;
        }
    }

    @Override
    protected long write(Node<VALUE> node) throws IOException
    {
        return writeRecursiveOnStack(node, 0);
    }

    @SuppressWarnings("DuplicatedCode") // intentionally duplicates IncrementalTrieWriterPageAware
    private long writeRecursiveOnStack(Node<VALUE> node, int depth) throws IOException
    {
        long nodePosition = dest.position();
        for (Node<VALUE> child : node.children)
            if (child.filePos == -1)
            {
                if (depth < maxRecursionDepth)
                    child.filePos = writeRecursiveOnStack(child, depth + 1);
                else
                    child.filePos = writeRecursiveOnHeap(child);
            }

        nodePosition += node.branchSize;
        assert dest.position() == nodePosition
        : "Expected node position to be " + nodePosition + " but got " + dest.position() + " after writing children.\n" + dumpNode(node, dest.position());

        serializer.write(dest, node, nodePosition);

        assert dest.position() == nodePosition + node.nodeSize
               || dest.paddedPosition() == dest.position() // For PartitionIndexTest.testPointerGrowth where position may jump on page boundaries.
        : "Expected node position to be " + (nodePosition + node.nodeSize) + " but got " + dest.position() + " after writing node, nodeSize " + node.nodeSize + ".\n" + dumpNode(node, nodePosition);
        return nodePosition;
    }

    private long writeRecursiveOnHeap(Node<VALUE> node) throws IOException
    {
        return new WriteRecursion(node, null).process().node.filePos;
    }

    class WriteRecursion extends Recursion<Node<VALUE>>
    {
        long nodePosition;

        WriteRecursion(Node<VALUE> node, Recursion<Node<VALUE>> parent)
        {
            super(node, node.children.iterator(), parent);
            nodePosition = dest.position();
        }

        @Override
        Recursion<Node<VALUE>> makeChild(Node<VALUE> child)
        {
            if (child.filePos == -1)
                return new WriteRecursion(child, this);
            else
                return null;
        }

        @SuppressWarnings("DuplicatedCode") // intentionally duplicates IncrementalTrieWriterPageAware and onStack code
        @Override
        void complete() throws IOException
        {
            nodePosition = nodePosition + node.branchSize;
            assert dest.position() == nodePosition
                    : "Expected node position to be " + nodePosition + " but got " + dest.position() + " after writing children.\n" + dumpNode(node, dest.position());

            serializer.write(dest, node, nodePosition);

            assert dest.position() == nodePosition + node.nodeSize
                   || dest.paddedPosition() == dest.position() // For PartitionIndexTest.testPointerGrowth where position may jump on page boundaries.
                    : "Expected node position to be " + (nodePosition + node.nodeSize) + " but got " + dest.position() + " after writing node, nodeSize " + node.nodeSize + ".\n" + dumpNode(node, nodePosition);

            node.filePos = nodePosition;
        }
    }

    @Override
    protected long writePartial(Node<VALUE> node, DataOutputPlus dest, long baseOffset) throws IOException
    {
        return writePartialRecursiveOnStack(node, dest, baseOffset, 0);
    }

    @SuppressWarnings("DuplicatedCode") // intentionally duplicates IncrementalTrieWriterPageAware
    private long writePartialRecursiveOnStack(Node<VALUE> node, DataOutputPlus dest, long baseOffset, int depth) throws IOException
    {
        long startPosition = dest.position() + baseOffset;

        List<Node<VALUE>> childrenToClear = new ArrayList<>();
        for (Node<VALUE> child : node.children)
        {
            if (child.filePos == -1)
            {
                childrenToClear.add(child);
                if (depth < maxRecursionDepth)
                    child.filePos = writePartialRecursiveOnStack(child, dest, baseOffset, depth + 1);
                else
                    child.filePos = writePartialRecursiveOnHeap(child, dest, baseOffset);
            }
        }

        long nodePosition = dest.position() + baseOffset;

        if (node.hasOutOfPageInBranch)
        {
            // Update the branch size with the size of what we have just written. This may be used by the node's
            // maxPositionDelta, and it's a better approximation for later fitting calculations.
            node.branchSize = (int) (nodePosition - startPosition);
        }

        serializer.write(dest, node, nodePosition);

        if (node.hasOutOfPageChildren || node.hasOutOfPageInBranch)
        {
            // Update the node size with what we have just seen. It's a better approximation for later fitting
            // calculations.
            long endPosition = dest.position() + baseOffset;
            node.nodeSize = (int) (endPosition - nodePosition);
        }

        for (Node<VALUE> child : childrenToClear)
            child.filePos = -1;
        return nodePosition;
    }

    private long writePartialRecursiveOnHeap(Node<VALUE> node, DataOutputPlus dest, long baseOffset) throws IOException
    {
        new WritePartialRecursion(node, dest, baseOffset).process();
        long pos = node.filePos;
        node.filePos = -1;
        return pos;
    }

    class WritePartialRecursion extends Recursion<Node<VALUE>>
    {
        final DataOutputPlus dest;
        final long baseOffset;
        final long startPosition;
        final List<Node<VALUE>> childrenToClear;

        WritePartialRecursion(Node<VALUE> node, WritePartialRecursion parent)
        {
            super(node, node.children.iterator(), parent);
            this.dest = parent.dest;
            this.baseOffset = parent.baseOffset;
            this.startPosition = dest.position() + baseOffset;
            childrenToClear = new ArrayList<>();
        }

        WritePartialRecursion(Node<VALUE> node, DataOutputPlus dest, long baseOffset)
        {
            super(node, node.children.iterator(), null);
            this.dest = dest;
            this.baseOffset = baseOffset;
            this.startPosition = dest.position() + baseOffset;
            childrenToClear = new ArrayList<>();
        }

        @SuppressWarnings("DuplicatedCode") // intentionally duplicates IncrementalTrieWriterPageAware and onStack code
        @Override
        Recursion<Node<VALUE>> makeChild(Node<VALUE> child)
        {
            if (child.filePos == -1)
            {
                childrenToClear.add(child);
                return new WritePartialRecursion(child, this);
            }
            else
                return null;
        }

        @Override
        void complete() throws IOException
        {
            long nodePosition = dest.position() + baseOffset;

            if (node.hasOutOfPageInBranch)
            {
                // Update the branch size with the size of what we have just written. This may be used by the node's
                // maxPositionDelta, and it's a better approximation for later fitting calculations.
                node.branchSize = (int) (nodePosition - startPosition);
            }

            serializer.write(dest, node, nodePosition);

            if (node.hasOutOfPageChildren || node.hasOutOfPageInBranch)
            {
                // Update the node size with what we have just seen. It's a better approximation for later fitting
                // calculations.
                long endPosition = dest.position() + baseOffset;
                node.nodeSize = (int) (endPosition - nodePosition);
            }

            for (Node<VALUE> child : childrenToClear)
                child.filePos = -1;

            node.filePos = nodePosition;
        }
    }
}