ActiveMQ example source code file (BTreeNode.java)

This example ActiveMQ source code file (BTreeNode.java) is included in the DevDaily.com "Java Source Code Warehouse" project. The intent of this project is to help you "Learn Java by Example" ^TM.

Java - ActiveMQ tags/keywords

btreeiterator, btreenode, btreenode, illegalargumentexception, io, ioexception, ioexception, key, key, object, suppresswarnings, t, transaction, util, value, value

The ActiveMQ BTreeNode.java source code

/**
 * 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.kahadb.index;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Arrays;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Map.Entry;

import org.apache.kahadb.index.BTreeIndex.Prefixer;
import org.apache.kahadb.page.Page;
import org.apache.kahadb.page.Transaction;
import org.apache.kahadb.util.VariableMarshaller;


/**
 * The BTreeNode class represents a node in the BTree object graph.  It is stored in 
 * one Page of a PageFile.
 */
public final class BTreeNode<Key,Value> {

    // The index that this node is part of.
    private final BTreeIndex<Key,Value> index;
    // The parent node or null if this is the root node of the BTree
    private BTreeNode<Key,Value> parent;
    // The page associated with this node
    private Page<BTreeNode page;
    
    // Order list of keys in the node
    private Key[] keys;
    // Values associated with the Keys. Null if this is a branch node.
    private Value[] values;
    // nodeId pointers to children BTreeNodes. Null if this is a leaf node.
    private long[] children;
    // The next leaf node after this one.  Used for fast iteration of the entries.
    private long next = -1;
    
    private final class KeyValueEntry implements Map.Entry<Key, Value> {
        private final Key key;
        private final Value value;

        public KeyValueEntry(Key key, Value value) {
            this.key = key;
            this.value = value;
        }

        public Key getKey() {
            return key;
        }

        public Value getValue() {
            return value;
        }

        public Value setValue(Value value) {
            throw new UnsupportedOperationException();
        }

    }

    private final class BTreeIterator implements Iterator<Map.Entry {
        
        private final Transaction tx;
        BTreeNode<Key,Value> current;
        int nextIndex;
        Map.Entry<Key,Value> nextEntry;

        private BTreeIterator(Transaction tx, BTreeNode<Key,Value> current, int nextIndex) {
            this.tx = tx;
            this.current = current;
            this.nextIndex=nextIndex;
        }

        synchronized private void findNextPage() {
            if( nextEntry!=null ) {
                return;
            }
            
            try {
                while( current!=null ) {
                    if( nextIndex >= current.keys.length ) {
                        // we need to roll to the next leaf..
                        if( current.next >= 0 ) {
                            current = index.loadNode(tx, current.next, null);
                            assert !current.isBranch() : "Should have linked to the next leaf node.";
                            nextIndex=0;
                        } else {
                            break;
                        }
                    }  else {
                        nextEntry = new KeyValueEntry(current.keys[nextIndex], current.values[nextIndex]);
                        nextIndex++;
                        break;
                    }
                    
                }
            } catch (IOException e) {
            }
        }

        public boolean hasNext() {
            findNextPage();
            return nextEntry !=null;
        }

        public Entry<Key, Value> next() {
            findNextPage(); 
            if( nextEntry !=null ) {
                Entry<Key, Value> lastEntry = nextEntry;
                nextEntry=null;
                return lastEntry;
            } else {
                throw new NoSuchElementException();
            }
        }

        public void remove() {
            throw new UnsupportedOperationException();
        }
    }

    /**
     * The Marshaller is used to store and load the data in the BTreeNode into a Page.
     *  
     * @param <Key>
     * @param <Value>
     */
    static public class Marshaller<Key,Value> extends VariableMarshaller> {
        private final BTreeIndex<Key,Value> index;
        
        public Marshaller(BTreeIndex<Key,Value> index) {
            this.index = index;
        }

        public void writePayload(BTreeNode<Key,Value> node, DataOutput os) throws IOException {
            // Write the keys
            short count = (short)node.keys.length; // cast may truncate value...
            if( count != node.keys.length ) {
                throw new IOException("Too many keys");
            }
            
            os.writeShort(count);
            for (int i = 0; i < node.keys.length; i++) {
                index.getKeyMarshaller().writePayload(node.keys[i], os);
            }
            
            if( node.isBranch() ) {
                // If this is a branch...
                os.writeBoolean(true);
                for (int i = 0; i < count+1; i++) {
                    os.writeLong(node.children[i]);
                }
                
            } else {
                // If this is a leaf
                os.writeBoolean(false);
                for (int i = 0; i < count; i++) {
                    index.getValueMarshaller().writePayload(node.values[i], os);
                }
                os.writeLong(node.next);
            }
        }

        @SuppressWarnings("unchecked")
        public BTreeNode<Key,Value> readPayload(DataInput is) throws IOException {
            BTreeNode<Key,Value>  node = new BTreeNode(index);
            int count = is.readShort();
            
            node.keys = (Key[])new Object[count];
            for (int i = 0; i < count; i++) {
                node.keys[i] = index.getKeyMarshaller().readPayload(is);
            }
            
            if( is.readBoolean() ) {
                node.children = new long[count+1];
                for (int i = 0; i < count+1; i++) {
                    node.children[i] = is.readLong();
                }
            } else {
                node.values = (Value[])new Object[count];
                for (int i = 0; i < count; i++) {
                    node.values[i] = index.getValueMarshaller().readPayload(is);
                }
                node.next = is.readLong();
            }
            return node;
        }
    }

    public BTreeNode(BTreeIndex<Key,Value> index) {
        this.index = index;
    }
    
    public void setEmpty() {
        setLeafData(createKeyArray(0), createValueArray(0));
    }
    

    /**
     * Internal (to the BTreeNode) method. Because this method is called only by
     * BTreeNode itself, no synchronization done inside of this method.
     * @throws IOException 
     */
    private BTreeNode<Key,Value> getChild(Transaction tx, int idx) throws IOException {
        if (isBranch() && idx >= 0 && idx < children.length) {
            BTreeNode<Key, Value> result = this.index.loadNode(tx, children[idx], this);
            return result;
        } else {
            return null;
        }
    }


    /**
     * Returns the right most leaf from the current btree graph.
     * @throws IOException
     */
    private BTreeNode<Key,Value> getRightLeaf(Transaction tx) throws IOException {
        BTreeNode<Key,Value> cur = this;
        while(cur.isBranch()) {
            cur = cur.getChild(tx, cur.keys.length);
        }
        return cur;
    }

    /**
     * Returns the left most leaf from the current btree graph.
     * @throws IOException
     */
    private BTreeNode<Key,Value> getLeftLeaf(Transaction tx) throws IOException {
        BTreeNode<Key,Value> cur = this;
        while(cur.isBranch()) {
            cur = cur.getChild(tx, 0);
        }
        return cur;
    }

    /**
     * Returns the left most leaf from the current btree graph.
     * @throws IOException
     */
    private BTreeNode<Key,Value> getLeftPeer(Transaction tx, BTreeNode x) throws IOException {
        BTreeNode<Key,Value> cur = x;
        while( cur.parent !=null ) {
            if( cur.parent.children[0] == cur.getPageId() ) {
                cur = cur.parent;
            } else {
                for( int i=0; i < cur.parent.children.length; i ++) {
                    if( cur.parent.children[i]==cur.getPageId() ) {
                        return  cur.parent.getChild(tx, i-1);
                    }
                }
                throw new AssertionError("page "+x+" was decendent of "+cur.getPageId());
            }
        }
        return null;
    }

    public Value remove(Transaction tx, Key key) throws IOException {

        if(isBranch()) {
            int idx = Arrays.binarySearch(keys, key);
            idx = idx < 0 ? -(idx + 1) : idx + 1;
            BTreeNode<Key, Value> child = getChild(tx, idx);
            if( child.getPageId() == index.getPageId() ) {
                throw new IOException("BTree corrupted: Cylce detected.");
            }
            Value rc = child.remove(tx, key);
            
            // child node is now empty.. remove it from the branch node.
            if( child.keys.length == 0 ) {
                
                // If the child node is a branch, promote
                if( child.isBranch() ) {
                    // This is cause branches are never really empty.. they just go down to 1 child..
                    children[idx] = child.children[0];
                } else {
                    
                    // The child was a leaf. Then we need to actually remove it from this branch node..
                    // and relink the previous leaf to skip to the next leaf.

                    BTreeNode<Key, Value> previousLeaf = null;
                    if( idx > 0 ) {
                        // easy if we this node hold the previous child.
                        previousLeaf = getChild(tx, idx-1).getRightLeaf(tx);
                    } else {
                        // less easy if we need to go to the parent to find the previous child.
                        BTreeNode<Key, Value> lp = getLeftPeer(tx, this);
                        if( lp!=null ) {
                            previousLeaf = lp.getRightLeaf(tx);
                        }
                        // lp will be null if there was no previous child.
                    }

                    if( previousLeaf !=null ) {
                        previousLeaf.next = child.next;
                        index.storeNode(tx, previousLeaf, true);
                    }

                    if( idx < children.length-1 ) {
                        // Delete it and key to the right.
                        setBranchData(arrayDelete(keys, idx), arrayDelete(children, idx));
                    } else {
                        // It was the last child.. Then delete it and key to the left
                        setBranchData(arrayDelete(keys, idx-1), arrayDelete(children, idx));
                    }
                    
                    // If we are the root node, and only have 1 child left.  Then 
                    // make the root be the leaf node.
                    if( children.length == 1 && parent==null ) {
                        child = getChild(tx, 0);
                        keys = child.keys;
                        children = child.children;
                        values = child.values;
                        // free up the page..
                        tx.free(child.getPage());
                    }
                    
                }
                index.storeNode(tx, this, true);
            }
            
            return rc;
        } else {
            int idx = Arrays.binarySearch(keys, key);
            if (idx < 0) {
                return null;
            } else {
                Value oldValue = values[idx];
                setLeafData(arrayDelete(keys, idx), arrayDelete(values, idx));
                
                if( keys.length==0 && parent!=null) {
                    tx.free(getPage());
                } else {
                    index.storeNode(tx, this, true);
                }
                
                return oldValue;
            }
        }
    }

    public Value put(Transaction tx, Key key, Value value) throws IOException {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }

        if( isBranch() ) {
            return getLeafNode(tx, this, key).put(tx, key, value);
        } else {
            int idx = Arrays.binarySearch(keys, key);
            
            Value oldValue=null;
            if (idx >= 0) {
                // Key was found... Overwrite
                oldValue = values[idx];
                values[idx] = value;
                setLeafData(keys, values);
            } else {
                // Key was not found, Insert it
                idx = -(idx + 1);
                setLeafData(arrayInsert(keys, key, idx), arrayInsert(values, value, idx));
            }
            
            try {
                index.storeNode(tx, this, allowOverflow());
            } catch ( Transaction.PageOverflowIOException e ) {
                // If we get an overflow 
                split(tx);
            }
            
            return oldValue;
        }
    }

    private void promoteValue(Transaction tx, Key key, long nodeId) throws IOException {

        int idx = Arrays.binarySearch(keys, key);
        idx = idx < 0 ? -(idx + 1) : idx + 1;
        setBranchData(arrayInsert(keys, key, idx), arrayInsert(children, nodeId, idx + 1));

        try {
            index.storeNode(tx, this, allowOverflow());
        } catch ( Transaction.PageOverflowIOException e ) {
            split(tx);
        }

    }

    /**
     * Internal to the BTreeNode method
     */
    private void split(Transaction tx) throws IOException {
        Key[] leftKeys;
        Key[] rightKeys;
        Value[] leftValues=null;
        Value[] rightValues=null;
        long[] leftChildren=null;
        long[] rightChildren=null;
        Key separator;

        int vc = keys.length;
        int pivot = vc / 2;

        // Split the node into two nodes
        if( isBranch() ) {

            leftKeys = createKeyArray(pivot);
            leftChildren = new long[leftKeys.length + 1];
            rightKeys = createKeyArray(vc - (pivot + 1));
            rightChildren = new long[rightKeys.length + 1];

            System.arraycopy(keys, 0, leftKeys, 0, leftKeys.length);
            System.arraycopy(children, 0, leftChildren, 0, leftChildren.length);
            System.arraycopy(keys, leftKeys.length + 1, rightKeys, 0, rightKeys.length);
            System.arraycopy(children, leftChildren.length, rightChildren, 0, rightChildren.length);

            // Is it a Simple Prefix BTree??
            Prefixer<Key> prefixer = index.getPrefixer();
            if(prefixer!=null) {
                separator = prefixer.getSimplePrefix(leftKeys[leftKeys.length - 1], rightKeys[0]);
            } else {
                separator = keys[leftKeys.length];
            }
                
            
        } else {

            leftKeys = createKeyArray(pivot);
            leftValues = createValueArray(leftKeys.length);
            rightKeys = createKeyArray(vc - pivot);
            rightValues = createValueArray(rightKeys.length);

            System.arraycopy(keys, 0, leftKeys, 0, leftKeys.length);
            System.arraycopy(values, 0, leftValues, 0, leftValues.length);
            System.arraycopy(keys, leftKeys.length, rightKeys, 0, rightKeys.length);
            System.arraycopy(values, leftValues.length, rightValues, 0, rightValues.length);

            // separator = getSeparator(leftVals[leftVals.length - 1],
            // rightVals[0]);
            separator = rightKeys[0];

        }

        // Promote the pivot to the parent branch
        if (parent == null) {
            
            // This can only happen if this is the root
            BTreeNode<Key,Value> rNode = this.index.createNode(tx, this);
            BTreeNode<Key,Value> lNode = this.index.createNode(tx, this);

            if( isBranch() ) {
                rNode.setBranchData(rightKeys, rightChildren);
                lNode.setBranchData(leftKeys, leftChildren);
            } else {
                rNode.setLeafData(rightKeys, rightValues);
                lNode.setLeafData(leftKeys, leftValues);
                lNode.setNext(rNode.getPageId());
            }

            Key[] v = createKeyArray(1);
            v[0]=separator;
            setBranchData(v, new long[] { lNode.getPageId(), rNode.getPageId() });

            index.storeNode(tx, this, true);
            index.storeNode(tx, rNode, true);
            index.storeNode(tx, lNode, true);
            
        } else {
            BTreeNode<Key,Value> rNode = this.index.createNode(tx, parent);
            
            if( isBranch() ) {
                setBranchData(leftKeys, leftChildren);
                rNode.setBranchData(rightKeys, rightChildren);
            } else {
                rNode.setNext(next);
                next = rNode.getPageId();
                setLeafData(leftKeys, leftValues);
                rNode.setLeafData(rightKeys, rightValues);
            }

            index.storeNode(tx, this, true);
            index.storeNode(tx, rNode, true);
            parent.promoteValue(tx, separator, rNode.getPageId());
        }
    }

    public void printStructure(Transaction tx, PrintWriter out, String prefix) throws IOException {
        if( prefix.length()>0 && parent == null ) {
            throw new IllegalStateException("Cycle back to root node detected.");
        }
        
        if( isBranch() ) {
            for(int i=0 ; i < children.length; i++) {
                BTreeNode<Key, Value> child = getChild(tx, i);
                if( i == children.length-1) {
                    out.println(prefix+"\\- "+child.getPageId()+(child.isBranch()?" ("+child.children.length+")":""));
                    child.printStructure(tx, out, prefix+"   ");
                } else {
                    out.println(prefix+"|- "+child.getPageId()+(child.isBranch()?" ("+child.children.length+")":"")+" : "+keys[i]);
                    child.printStructure(tx, out, prefix+"   ");
                }
            }
        }
    }
    
    
    public int getMinLeafDepth(Transaction tx, int depth) throws IOException {
        depth++;
        if( isBranch() ) {
            int min = Integer.MAX_VALUE;
            for(int i=0 ; i < children.length; i++) {
                min = Math.min(min, getChild(tx, i).getMinLeafDepth(tx, depth));
            }
            return min;
        } else {
//            print(depth*2, "- "+page.getPageId());
            return depth;
        }
    }

    public int getMaxLeafDepth(Transaction tx, int depth) throws IOException {
        depth++;
        if( isBranch() ) {
            int v = 0;
            for(int i=0 ; i < children.length; i++) {
                v = Math.max(v, getChild(tx, i).getMaxLeafDepth(tx, depth));
            }
            depth = v;
        } 
        return depth;
    }

    public Value get(Transaction tx, Key key) throws IOException {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        if( isBranch() ) {
            return getLeafNode(tx, this, key).get(tx, key);
        } else {
            int idx = Arrays.binarySearch(keys, key);
            if (idx < 0) {
                return null;
            } else {
                return values[idx];
            }
        }
    }
    
    public boolean isEmpty(final Transaction tx) throws IOException {
        return keys.length==0;
    }

    public void visit(Transaction tx, BTreeVisitor<Key, Value> visitor) throws IOException {
        if (visitor == null) {
            throw new IllegalArgumentException("Visitor cannot be null");
        }
        if( isBranch() ) {
            for(int i=0; i < this.children.length; i++) {
                Key key1 = null;
                if( i!=0 ) {
                    key1 = keys[i-1];
                }
                Key key2 = null;
                if( i!=this.children.length-1 ) {
                    key2 = keys[i];
                }
                if( visitor.isInterestedInKeysBetween(key1, key2) ) {
                    BTreeNode<Key, Value> child = getChild(tx, i);
                    child.visit(tx, visitor);
                }
            }
        } else {
            visitor.visit(Arrays.asList(keys), Arrays.asList(values));
        }
    }
    
    public Map.Entry<Key,Value> getFirst(Transaction tx) throws IOException {
        BTreeNode<Key, Value> node = this;
        while( node .isBranch() ) {
            node = node.getChild(tx, 0);
        }
        if( node.values.length>0 ) {
            return new KeyValueEntry(node.keys[0], node.values[0]);
        } else {
            return null;
        }
    }

    public Map.Entry<Key,Value> getLast(Transaction tx) throws IOException {
        BTreeNode<Key, Value> node = this;
        while( node.isBranch() ) {
            node = node.getChild(tx, node.children.length-1);
        }
        if( node.values.length>0 ) {
            int idx = node.values.length-1;
            return new KeyValueEntry(node.keys[idx], node.values[idx]);
        } else {
            return null;
        }
    }
    
    public BTreeNode<Key,Value> getFirstLeafNode(Transaction tx) throws IOException {
        BTreeNode<Key, Value> node = this;
        while( node .isBranch() ) {
            node = node.getChild(tx, 0);
        }
        return node;
    }
    
    public Iterator<Map.Entry iterator(final Transaction tx, Key startKey) throws IOException {
        if (startKey == null) {
            return iterator(tx);
        }
        if( isBranch() ) {
            return getLeafNode(tx, this, startKey).iterator(tx, startKey);
        } else {
            int idx = Arrays.binarySearch(keys, startKey);
            if (idx < 0) {
                idx = -(idx + 1);
            }
            return new BTreeIterator(tx, this, idx);
        }
    }

    public Iterator<Map.Entry iterator(final Transaction tx) throws IOException {
        return new BTreeIterator(tx, getFirstLeafNode(tx), 0);
    }
    
    public void clear(Transaction tx) throws IOException {
        if( isBranch() ) {
            for (int i = 0; i < children.length; i++) {
                BTreeNode<Key, Value> node = index.loadNode(tx, children[i], this);
                node.clear(tx);
                tx.free(node.getPage());
            }
        }
        // Reset the root node to be a leaf.
        if( parent == null ) {
            setLeafData(createKeyArray(0), createValueArray(0));
            next=-1;
            index.storeNode(tx, this, true);
        }
    }


    private static <Key,Value> BTreeNode getLeafNode(Transaction tx, final BTreeNode node, Key key) throws IOException {
        BTreeNode<Key, Value> current = node;
        while( true ) {
            if( current.isBranch() ) {
                int idx = Arrays.binarySearch(current.keys, key);
                idx = idx < 0 ? -(idx + 1) : idx + 1;
                BTreeNode<Key, Value> child = current.getChild(tx, idx);        

                // A little cycle detection for sanity's sake
                if( child == node ) {
                    throw new IOException("BTree corrupted: Cylce detected.");
                }
                
                current = child;
            } else {
                break;
            }
        }
        return current;
    }

    public boolean contains(Transaction tx, Key key) throws IOException {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }

        if( isBranch() ) {
            return getLeafNode(tx, this, key).contains(tx, key);
        } else {
            int idx = Arrays.binarySearch(keys, key);
            if (idx < 0) {
                return false;
            } else {
                return true;
            }
        }
    }

    ///////////////////////////////////////////////////////////////////
    // Implementation methods
    ///////////////////////////////////////////////////////////////////
 

    private boolean allowOverflow() {
        // Only allow page overflow if there are <= 3 keys in the node.  Otherwise a split will occur on overflow
        return this.keys.length<=3;
    }


    private void setLeafData(Key[] keys, Value[] values) {
        this.keys = keys;
        this.values = values;
        this.children = null;
    }
    
    private void setBranchData(Key[] keys, long[] nodeIds) {
        this.keys = keys;
        this.children = nodeIds;
        this.values = null;
    }

    @SuppressWarnings("unchecked")
    private Key[] createKeyArray(int size) {
        return (Key[])new Object[size];
    }

    @SuppressWarnings("unchecked")
    private Value[] createValueArray(int size) {
        return (Value[])new Object[size];
    }
    
    @SuppressWarnings("unchecked")
    static private <T> T[] arrayDelete(T[] vals, int idx) {
        T[] newVals = (T[])new Object[vals.length - 1];
        if (idx > 0) {
            System.arraycopy(vals, 0, newVals, 0, idx);
        }
        if (idx < newVals.length) {
            System.arraycopy(vals, idx + 1, newVals, idx, newVals.length - idx);
        }
        return newVals;
    }
    
    static private long[] arrayDelete(long[] vals, int idx) {
        long[] newVals = new long[vals.length - 1];
        if (idx > 0) {
            System.arraycopy(vals, 0, newVals, 0, idx);
        }
        if (idx < newVals.length) {
            System.arraycopy(vals, idx + 1, newVals, idx, newVals.length - idx);
        }
        return newVals;
    }

    @SuppressWarnings("unchecked")
    static private <T> T[] arrayInsert(T[] vals, T val, int idx) {
        T[] newVals = (T[])new Object[vals.length + 1];
        if (idx > 0) {
            System.arraycopy(vals, 0, newVals, 0, idx);
        }
        newVals[idx] = val;
        if (idx < vals.length) {
            System.arraycopy(vals, idx, newVals, idx + 1, vals.length - idx);
        }
        return newVals;
    }


    static private long[] arrayInsert(long[] vals, long val, int idx) {
        
        long[] newVals = new long[vals.length + 1];
        if (idx > 0) {
            System.arraycopy(vals, 0, newVals, 0, idx);
        }
        newVals[idx] = val;
        if (idx < vals.length) {
            System.arraycopy(vals, idx, newVals, idx + 1, vals.length - idx);
        }
        return newVals;
    }

    ///////////////////////////////////////////////////////////////////
    // Property Accessors
    ///////////////////////////////////////////////////////////////////
    private boolean isBranch() {
        return children!=null;
    }

    public long getPageId() {
        return page.getPageId();
    }

    public BTreeNode<Key, Value> getParent() {
        return parent;
    }

    public void setParent(BTreeNode<Key, Value> parent) {
        this.parent = parent;
    }

    public Page<BTreeNode getPage() {
        return page;
    }

    public void setPage(Page<BTreeNode page) {
        this.page = page;
    }

    public long getNext() {
        return next;
    }

    public void setNext(long next) {
        this.next = next;
    }
    
    @Override
    public String toString() {
        return "[BTreeNode "+(isBranch()?"branch":"leaf")+": "+Arrays.asList(keys)+"]";
    }

}