| // Copyright 2018 The Bazel Authors. All rights reserved. |
| // |
| // Licensed 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 com.google.devtools.build.lib.collect.nestedset; |
| |
| import com.google.common.base.Preconditions; |
| import com.google.devtools.build.lib.util.Fingerprint; |
| import com.google.devtools.build.lib.vfs.DigestHashFunction; |
| import com.google.devtools.build.lib.vfs.DigestHashFunction.DigestLength; |
| import java.util.concurrent.atomic.AtomicInteger; |
| import java.util.concurrent.atomic.AtomicReferenceArray; |
| import java.util.concurrent.locks.StampedLock; |
| |
| /** |
| * Map of key -> [digest bytes]. |
| * |
| * <p>This class uses a single array of keys and a big single block of bytes. To read/store digests |
| * we index straight into the byte array. This is more memory-efficient and uses less GC than a |
| * corresponding Map<Object, byte[]>. |
| * |
| * <p>Keys use reference equality. |
| * |
| * <p>Reading is lock free. During writes a read lock is taken. If we need to resize the table, a |
| * write lock is taken to flush all the readers and writers before the table is resized. |
| */ |
| final class DigestMap { |
| private static final Object INSERTION_IN_PROGRESS = new Object(); |
| private final DigestLength digestLength; |
| private final StampedLock readWriteLock = new StampedLock(); |
| |
| static class Table { |
| final int tableSize; |
| final int nextResize; |
| final AtomicReferenceArray<Object> keys; |
| final byte[] bytes; |
| |
| Table(int tableSize, int digestLength) { |
| this.tableSize = tableSize; |
| this.nextResize = getNextResize(tableSize); |
| this.keys = new AtomicReferenceArray<>(tableSize); |
| this.bytes = new byte[tableSize * digestLength]; |
| } |
| } |
| |
| private volatile Table table; |
| private final AtomicInteger allocatedSlots; |
| |
| DigestMap(DigestHashFunction digestHashFunction, int initialSize) { |
| Preconditions.checkArgument( |
| initialSize > 0 && (initialSize & (initialSize - 1)) == 0, |
| "initialSize must be a power of 2 greater than 0"); |
| this.digestLength = digestHashFunction.getDigestLength(); |
| this.table = new Table(initialSize, digestLength.getDigestMaximumLength()); |
| this.allocatedSlots = new AtomicInteger(); |
| } |
| |
| /** Finds the digest for the corresponding key and adds it to the passed fingerprint. */ |
| boolean readDigest(Object key, Fingerprint fingerprint) { |
| Table table = this.table; // Read once for duration of method |
| int index = findKey(table, key); |
| if (index >= 0) { |
| int offset = index * this.digestLength.getDigestMaximumLength(); |
| int digestLength = this.digestLength.getDigestLength(table.bytes, offset); |
| fingerprint.addBytes(table.bytes, offset, digestLength); |
| return true; |
| } |
| return false; |
| } |
| |
| private static int findKey(Table table, Object key) { |
| int hash = hash(key); |
| int index = hash & (table.tableSize - 1); |
| while (true) { |
| Object currentKey = table.keys.get(index); |
| if (currentKey == key) { |
| return index; |
| } else if (currentKey == null) { |
| return -1; |
| } |
| index = probe(index, table.tableSize); |
| } |
| } |
| |
| /** |
| * Inserts a digest for the corresponding key, then immediately reads it into another fingerprint. |
| * |
| * <p>This is equivalent to <code> |
| * digestMap.insertDigest(key, digest.digestAndReset(); digestMap.readDigest(key, readTo); </code> |
| * but it will be faster. |
| * |
| * @param key The key to insert. |
| * @param digest The fingerprint to insert. This will reset the fingerprint instance. |
| * @param readTo A fingerprint to read the just-added fingerprint into. |
| */ |
| void insertAndReadDigest(Object key, Fingerprint digest, Fingerprint readTo) { |
| // Check if we have to resize the table first and do that under write lock |
| // We assume that we are going to insert an item. If we do not do this, multiple |
| // threads could race and all think they do not need to resize, then some get stuck |
| // trying to insert the item. |
| Table table = this.table; |
| if (allocatedSlots.incrementAndGet() >= table.nextResize) { |
| long resizeLock = readWriteLock.writeLock(); |
| try { |
| // Guard against race to make sure only one thread resizes |
| if (table == this.table) { |
| resizeTableWriteLocked(); |
| } |
| } finally { |
| readWriteLock.unlockWrite(resizeLock); |
| } |
| } |
| final int index; |
| long stamp = readWriteLock.readLock(); |
| try { |
| table = this.table; // Grab the table again under read lock |
| index = insertKey(table, key, digest); |
| } finally { |
| readWriteLock.unlockRead(stamp); |
| } |
| // This can be done outside of the read lock since the slot is immutable once inserted |
| int offset = index * this.digestLength.getDigestMaximumLength(); |
| int digestLength = this.digestLength.getDigestLength(table.bytes, offset); |
| readTo.addBytes(table.bytes, offset, digestLength); |
| } |
| |
| // Inserts a key into the passed table and returns the index. |
| @SuppressWarnings("ThreadPriorityCheck") // We're not relying on thread scheduler for correctness |
| private int insertKey(Table table, Object key, Fingerprint digest) { |
| int hash = hash(key); |
| int index = hash & (table.tableSize - 1); |
| while (true) { |
| Object currentKey = table.keys.get(index); |
| if (currentKey == null) { |
| if (!table.keys.compareAndSet(index, null, INSERTION_IN_PROGRESS)) { |
| // We raced to insert a key in a free slot, retry this slot in case it's the same key. |
| // Failure to do so could lead to a double insertion. |
| continue; |
| } |
| digest.digestAndReset( |
| table.bytes, |
| index * digestLength.getDigestMaximumLength(), |
| digestLength.getDigestMaximumLength()); |
| table.keys.set(index, key); |
| return index; |
| } else if (currentKey == key) { |
| // Key is already present, give back the slot allocation |
| allocatedSlots.decrementAndGet(); |
| return index; |
| } else if (currentKey == INSERTION_IN_PROGRESS) { |
| // We are in the progress of inserting an item in this slot, but we don't yet know |
| // what the item is. Since it could be an insertion of ourselves we need to wait |
| // until done to avoid double insertion. We yield the thread in case the other |
| // thread is stuck between insertion and completion. |
| Thread.yield(); |
| continue; |
| } |
| index = probe(index, table.tableSize); |
| } |
| } |
| |
| private void resizeTableWriteLocked() { |
| int digestSize = this.digestLength.getDigestMaximumLength(); |
| Table oldTable = this.table; |
| Table newTable = new Table(oldTable.tableSize * 2, digestSize); |
| for (int i = 0; i < oldTable.tableSize; ++i) { |
| Object key = oldTable.keys.get(i); |
| if (key != null) { |
| int newIndex = firstFreeIndex(newTable.keys, newTable.tableSize, key); |
| newTable.keys.set(newIndex, key); |
| System.arraycopy( |
| oldTable.bytes, i * digestSize, newTable.bytes, newIndex * digestSize, digestSize); |
| } |
| } |
| this.table = newTable; |
| } |
| |
| private static int firstFreeIndex(AtomicReferenceArray<Object> keys, int tableSize, Object key) { |
| int hash = hash(key); |
| int index = hash & (tableSize - 1); |
| while (true) { |
| Object currentKey = keys.get(index); |
| if (currentKey == null) { |
| return index; |
| } |
| index = probe(index, tableSize); |
| } |
| } |
| |
| private static int hash(Object key) { |
| return smear(System.identityHashCode(key)); |
| } |
| |
| private static int probe(int index, int tableSize) { |
| return (index + 1) & (tableSize - 1); |
| } |
| |
| private static int getNextResize(int newTableSize) { |
| // 75% load |
| return (newTableSize * 3) / 4; |
| } |
| |
| /* |
| * This method was rewritten in Java from an intermediate step of the Murmur hash function in |
| * http://code.google.com/p/smhasher/source/browse/trunk/MurmurHash3.cpp, which contained the |
| * following header: |
| * |
| * MurmurHash3 was written by Austin Appleby, and is placed in the public domain. The author |
| * hereby disclaims copyright to this source code. |
| */ |
| private static int smear(int hashCode) { |
| return 0x1b873593 * Integer.rotateLeft(hashCode * 0xcc9e2d51, 15); |
| } |
| } |