/*
* JBoss, Home of Professional Open Source.
* See the COPYRIGHT.txt file distributed with this work for information
* regarding copyright ownership. Some portions may be licensed
* to Red Hat, Inc. under one or more contributor license agreements.
*
* This library is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* This library is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with this library; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
* 02110-1301 USA.
*/
package org.teiid.common.buffer;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantLock;
import org.teiid.client.BatchSerializer;
import org.teiid.common.buffer.LobManager.ReferenceMode;
import org.teiid.common.buffer.SPage.SearchResult;
import org.teiid.core.TeiidComponentException;
import org.teiid.core.TeiidProcessingException;
import org.teiid.core.TeiidRuntimeException;
import org.teiid.query.QueryPlugin;
import org.teiid.query.processor.relational.ListNestedSortComparator;
/**
* Self balancing search tree using skip list like logic
* This has similar performance similar to a B+/-Tree,
* but with fewer updates.
*/
@SuppressWarnings("unchecked")
public class STree implements Cloneable {
public enum InsertMode {ORDERED, NEW, UPDATE}
private static final Random seedGenerator = new Random(0);
protected int randomSeed;
private int mask = 1;
private int shift = 1;
protected HashMap<Long, SPage> pages = new HashMap<Long, SPage>();
protected volatile SPage[] header = new SPage[] {new SPage(this, true)};
protected BatchManager keyManager;
protected BatchManager leafManager;
protected ListNestedSortComparator comparator;
private int pageSize;
protected int leafSize;
protected int minPageSize;
protected int minStorageSize;
protected int keyLength;
protected boolean batchInsert;
protected SPage incompleteInsert;
protected LobManager lobManager;
protected ReentrantLock updateLock = new ReentrantLock();
private AtomicLong rowCount = new AtomicLong();
public STree(BatchManager manager,
BatchManager leafManager,
final ListNestedSortComparator comparator,
int pageSize,
int leafSize,
int keyLength,
LobManager lobManager) {
randomSeed = seedGenerator.nextInt() | 0x00000100; // ensure nonzero
this.keyManager = manager;
manager.setPrefersMemory(true);
this.leafManager = leafManager;
this.comparator = comparator;
this.pageSize = pageSize;
pageSize >>>= 3;
while (pageSize > 0) {
pageSize >>>= 1;
shift++;
mask <<= 1;
mask++;
}
this.leafSize = leafSize;
this.keyLength = keyLength;
this.lobManager = lobManager;
this.minPageSize = this.pageSize>>5;
this.minStorageSize = this.pageSize>>2;
}
public STree clone() {
updateLock.lock();
try {
STree clone = (STree) super.clone();
if (lobManager != null) {
clone.lobManager = lobManager.clone();
}
clone.updateLock = new ReentrantLock();
clone.rowCount = new AtomicLong(rowCount.get());
//clone the pages
clone.pages = new HashMap<Long, SPage>(pages);
for (Map.Entry<Long, SPage> entry : clone.pages.entrySet()) {
entry.setValue(entry.getValue().clone(clone));
}
//reset the pointers
for (Map.Entry<Long, SPage> entry : clone.pages.entrySet()) {
SPage clonePage = entry.getValue();
clonePage.next = clone.getPage(clonePage.next);
clonePage.prev = clone.getPage(clonePage.prev);
if (clonePage.children != null) {
for (int i = 0; i < clonePage.children.size(); i++) {
clonePage.children.set(i, clone.getPage(clonePage.children.get(i)));
}
}
}
clone.header = Arrays.copyOf(header, header.length);
for (int i = 0; i < header.length; i++) {
clone.header[i] = clone.pages.get(header[i].getId());
}
return clone;
} catch (CloneNotSupportedException e) {
throw new TeiidRuntimeException(QueryPlugin.Event.TEIID30039, e);
} finally {
updateLock.unlock();
}
}
private SPage getPage(SPage page) {
if (page == null) {
return page;
}
return pages.get(page.getId());
}
public void writeValuesTo(ObjectOutputStream oos) throws TeiidComponentException, IOException {
SPage page = header[0];
oos.writeLong(this.rowCount.get());
while (true) {
List<List<?>> batch = page.getValues();
BatchSerializer.writeBatch(oos, leafManager.getTypes(), batch);
if (page.next == null) {
break;
}
page = page.next;
}
}
public void setBatchInsert(boolean batchInsert) throws TeiidComponentException {
if (this.batchInsert == batchInsert) {
return;
}
this.batchInsert = batchInsert;
if (batchInsert || incompleteInsert == null) {
return;
}
SPage toFlush = incompleteInsert;
incompleteInsert = null;
if (toFlush.managedBatch != null) {
return;
}
toFlush.setValues(toFlush.getValues());
}
public void readValuesFrom(ObjectInputStream ois) throws IOException, ClassNotFoundException, TeiidComponentException {
long size = ois.readLong();
int sizeHint = this.getExpectedHeight(size);
batchInsert = true;
while (this.getRowCount() < size) {
List<List<Object>> batch = BatchSerializer.readBatch(ois, leafManager.getTypes());
for (List list : batch) {
this.insert(list, InsertMode.ORDERED, sizeHint);
}
}
batchInsert = false;
}
protected SPage findChildTail(SPage page) {
if (page == null) {
page = header[header.length - 1];
while (page.next != null) {
page = page.next;
}
return page;
}
if (page.children != null) {
page = page.children.get(page.children.size() - 1);
while (page.next != null) {
page = page.next;
}
}
return page;
}
/**
* Determine a new random level using an XOR rng.
*
* This uses the simplest of the generators described in George
* Marsaglia's "Xorshift RNGs" paper. This is not a high-quality
* generator but is acceptable here.
*
* See also the JSR-166 working group ConcurrentSkipListMap implementation.
*
* @return
*/
private int randomLevel() {
int x = randomSeed;
x ^= x << 13;
x ^= x >>> 17;
randomSeed = x ^= x << 5;
int level = 0;
while ((x & mask) == mask) {
++level;
x >>>= shift;
}
return level;
}
/**
* Search each level to find the pointer to the next level
* @param n
* @param places
* @return
* @throws IOException
* @throws TeiidComponentException
*/
List find(List n, List<SearchResult> places) throws TeiidComponentException {
SPage x = null;
for (int i = header.length - 1; i >= 0; i--) {
if (x == null) {
x = header[i];
}
SearchResult s = SPage.search(x, n, places);
if (places != null) {
places.add(s);
}
if ((s.index == -1 && s.page == header[i]) || s.values.isEmpty()) {
x = null;
continue; //start at the beginning of the next level
}
x = s.page;
int index = s.index;
boolean matched = true;
if (index < 0) {
matched = false;
index = Math.max(0, -index - 2);
}
if (i == 0) {
if (!matched) {
return null;
}
return s.values.get(index);
}
x = x.children.get(index);
}
return null;
}
public List find(List n) throws TeiidComponentException {
return find(n, new LinkedList<SPage.SearchResult>());
}
public List insert(List tuple, InsertMode mode, int sizeHint) throws TeiidComponentException {
if (tuple.size() != this.leafManager.getTypes().length) {
throw new AssertionError("Invalid tuple."); //$NON-NLS-1$
}
LinkedList<SearchResult> places = new LinkedList<SearchResult>();
List match = null;
if (this.lobManager != null) {
this.lobManager.updateReferences(tuple, ReferenceMode.CREATE);
}
if (mode == InsertMode.ORDERED) {
SPage last = null;
while (last == null || last.children != null) {
last = findChildTail(last);
//TODO: do this lazily
List<List<?>> batch = last.getValues();
places.add(new SearchResult(-batch.size() -1, last, batch));
}
} else {
match = find(tuple, places);
if (match != null) {
if (mode != InsertMode.UPDATE) {
return match;
}
SearchResult last = places.getLast();
SPage page = last.page;
last.values.set(last.index, tuple);
page.setValues(last.values);
if (this.lobManager != null) {
this.lobManager.updateReferences(tuple, ReferenceMode.REMOVE);
}
return match;
}
}
List key = extractKey(tuple);
int level = 0;
if (mode != InsertMode.ORDERED) {
if (sizeHint > -1) {
level = Math.min(sizeHint, randomLevel());
} else {
level = randomLevel();
}
} else if (!places.isEmpty() && places.getLast().values.size() == getPageSize(true)) {
long row = rowCount.get();
while (row != 0 && row%getPageSize(true) == 0) {
row = (row - getPageSize(true) + 1)/getPageSize(true);
level++;
}
}
assert header.length == places.size();
if (level >= header.length) {
header = Arrays.copyOf(header, level + 1);
}
rowCount.addAndGet(1);
SPage page = null;
for (int i = 0; i <= level; i++) {
if (places.isEmpty()) {
SPage newHead = new SPage(this, false);
List<List<?>> batch = newHead.getValues();
batch.add(key);
newHead.setValues(batch);
newHead.children.add(page);
header[i] = newHead;
page = newHead;
} else {
SearchResult result = places.removeLast();
Object value = (i == 0 ? tuple : page);
page = insert(key, result, places.peekLast(), value, mode == InsertMode.ORDERED);
}
}
return null;
}
public int getExpectedHeight(long sizeHint) {
if (sizeHint == 0) {
return 0;
}
int logSize = 1;
while (sizeHint > this.getPageSize(logSize==0)) {
logSize++;
sizeHint/=this.getPageSize(logSize==0);
}
return logSize;
}
List extractKey(List tuple) {
if (tuple.size() > keyLength) {
return new ArrayList(tuple.subList(0, keyLength));
}
return tuple;
}
SPage insert(List k, SearchResult result, SearchResult parent, Object value, boolean ordered) throws TeiidComponentException {
SPage page = result.page;
int index = -result.index - 1;
boolean leaf = !(value instanceof SPage);
if (result.values.size() == getPageSize(leaf)) {
SPage nextPage = new SPage(this, leaf);
List<List<?>> nextValues = nextPage.getValues();
nextPage.next = page.next;
nextPage.prev = page;
if (nextPage.next != null) {
nextPage.next.prev = nextPage;
}
page.next = nextPage;
boolean inNext = false;
if (!ordered) {
//split the values
nextValues.addAll(result.values.subList(getPageSize(leaf)/2, getPageSize(leaf)));
result.values.subList(getPageSize(leaf)/2, getPageSize(leaf)).clear();
if (!leaf) {
nextPage.children.addAll(page.children.subList(getPageSize(leaf)/2, getPageSize(false)));
page.children.subList(getPageSize(false)/2, getPageSize(false)).clear();
}
if (index <= getPageSize(leaf)/2) {
setValue(index, k, value, result.values, page);
} else {
inNext = true;
setValue(index - getPageSize(leaf)/2, k, value, nextValues, nextPage);
}
page.setValues(result.values);
if (parent != null) {
List min = nextPage.getValues().get(0);
SPage.correctParents(parent.page, min, page, nextPage);
}
} else {
inNext = true;
setValue(0, k, value, nextValues, nextPage);
}
nextPage.setValues(nextValues);
if (inNext) {
page = nextPage;
}
} else {
setValue(index, k, value, result.values, page);
page.setValues(result.values);
}
return page;
}
static void setValue(int index, List key, Object value, List<List<?>> values, SPage page) {
if (value instanceof SPage) {
values.add(index, key);
page.children.add(index, (SPage) value);
} else {
values.add(index, (List)value);
}
}
public List remove(List key) throws TeiidComponentException {
LinkedList<SearchResult> places = new LinkedList<SearchResult>();
List tuple = find(key, places);
if (tuple == null) {
return null;
}
rowCount.addAndGet(-1);
for (int i = 0; i < header.length; i++) {
SearchResult searchResult = places.removeLast();
if (searchResult.index < 0) {
continue;
}
searchResult.values.remove(searchResult.index);
boolean leaf = true;
if (searchResult.page.children != null) {
leaf = false;
searchResult.page.children.remove(searchResult.index);
}
int size = searchResult.values.size();
if (size == 0) {
if (header[i] != searchResult.page) {
searchResult.page.remove(false);
if (searchResult.page.next != null) {
searchResult.page.next.prev = searchResult.page.prev;
}
searchResult.page.prev.next = searchResult.page.next;
searchResult.page.next = header[i];
searchResult.page.prev = null;
continue;
}
header[i].remove(false);
if (header[i].next != null) {
header[i] = header[i].next;
header[i].prev = null;
} else {
if (i != 0) {
header = Arrays.copyOf(header, i);
break;
}
header[0] = new SPage(this, true);
}
continue;
} else if (size < getPageSize(leaf)/2) {
//check for merge
if (searchResult.page.next != null) {
List<List<?>> nextValues = searchResult.page.next.getValues();
if (nextValues.size() < getPageSize(leaf)/4) {
SPage.merge(places, nextValues, searchResult.page, searchResult.values);
continue;
}
}
if (searchResult.page.prev != null) {
List<List<?>> prevValues = searchResult.page.prev.getValues();
if (prevValues.size() < getPageSize(leaf)/4) {
SPage.merge(places, searchResult.values, searchResult.page.prev, prevValues);
continue;
}
}
}
searchResult.page.setValues(searchResult.values);
}
if (lobManager != null) {
lobManager.updateReferences(tuple, ReferenceMode.REMOVE);
}
return tuple;
}
public void remove() {
truncate(true);
this.keyManager.remove();
this.leafManager.remove();
}
public long getRowCount() {
return this.rowCount.get();
}
public long truncate(boolean force) {
long oldSize = rowCount.getAndSet(0);
for (int i = 0; i < header.length; i++) {
SPage page = header[i];
while (page != null) {
page.remove(force);
page = page.next;
}
}
header = new SPage[] {new SPage(this, true)};
return oldSize;
}
public int getHeight() {
return header.length;
}
@Override
public String toString() {
StringBuffer result = new StringBuffer();
for (int i = header.length -1; i >= 0; i--) {
SPage page = header[i];
result.append("Level ").append(i).append(" "); //$NON-NLS-1$ //$NON-NLS-2$
while (page != null) {
result.append(page);
result.append(", "); //$NON-NLS-1$
page = page.next;
}
result.append("\n"); //$NON-NLS-1$
}
return result.toString();
}
public int getKeyLength() {
return keyLength;
}
public void setPreferMemory(boolean preferMemory) {
this.leafManager.setPrefersMemory(preferMemory);
}
public boolean isPreferMemory() {
return this.leafManager.prefersMemory();
}
public ListNestedSortComparator getComparator() {
return comparator;
}
/**
* Quickly check if the index can be compacted
*/
public void compact() {
while (true) {
if (this.header.length == 1) {
return;
}
SPage child = this.header[header.length - 2];
if (child.next != null) {
//TODO: condense the page pointers
return;
}
//remove unneeded index level
this.header = Arrays.copyOf(this.header, header.length - 1);
}
}
public void removeRowIdFromKey() {
this.keyLength--;
int[] sortParameters = this.comparator.getSortParameters();
sortParameters = Arrays.copyOf(sortParameters, sortParameters.length - 1);
this.comparator.setSortParameters(sortParameters);
}
public void clearClonedFlags() {
for (SPage page : pages.values()) {
if (page.trackingObject != null) {
Long val = page.managedBatch;
if (val != null) {
SPage.REFERENCES.remove(val);
}
page.trackingObject = null;
//we don't really care about using synchronization or a volatile here
//since the worst case is that we'll just use gc cleanup
}
}
}
public int getPageSize(boolean leaf) {
if (leaf) {
return leafSize;
}
return pageSize;
}
BatchManager getBatchManager(boolean leaf) {
if (leaf) {
return leafManager;
}
return keyManager;
}
public TupleSource getTupleSource(final boolean destructive) {
return new TupleSource() {
SPage current = header[0];
List<List<?>> values;
int index = 0;
@Override
public List<?> nextTuple() throws TeiidComponentException,
TeiidProcessingException {
if (current == null) {
return null;
}
if (values == null) {
values = current.getValues();
}
if (index >= values.size()) {
if (destructive) {
current.remove(true);
}
values = null;
current = current.next;
if (current == null) {
return null;
}
values = current.getValues();
index = 0;
}
return values.get(index++);
}
@Override
public void closeSource() {
}
};
}
public void setMinStorageSize(int minStorageSize) {
this.minStorageSize = minStorageSize;
}
}