Skip to content

Commit

Permalink
fix: Added asserts for threads in all modifying methods (#10584) (#10615
Browse files Browse the repository at this point in the history 
)

Signed-off-by: Ivan Malygin <ivan@swirldslabs.com>
  • Loading branch information
@imalygin
imalygin committed Dec 21, 2023
1 parent c151882 commit 3984a45
Showing 1 changed file with 102 additions and 78 deletions.
180 changes: 102 additions & 78 deletions ...rlds-virtualmap/src/main/java/com/swirlds/virtualmap/internal/merkle/VirtualRootNode.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -316,6 +316,12 @@ public static class ClassVersion {

private VirtualMapStatistics statistics;

/**
* This reference is used to assert that there is only one thread modifying the VM at a time.
* NOTE: This field is used *only* if assertions are enabled, otherwise it always has null value.
*/
private final AtomicReference<Thread> currentModifyingThreadRef = new AtomicReference<>(null);

/**
* Creates a new empty root node. This constructor is used for deserialization and
* reconnects, not for normal use.
Expand Down Expand Up @@ -791,9 +797,14 @@ public boolean containsKey(final K key) {
public V getForModify(final K key) {
throwIfImmutable();
Objects.requireNonNull(key, NO_NULL_KEYS_ALLOWED_MESSAGE);
final VirtualLeafRecord<K, V> rec = records.findLeafRecord(key, true);
statistics.countUpdatedEntities();
return rec == null ? null : rec.getValue();
assert currentModifyingThreadRef.compareAndSet(null, Thread.currentThread());
try {
final VirtualLeafRecord<K, V> rec = records.findLeafRecord(key, true);
statistics.countUpdatedEntities();
return rec == null ? null : rec.getValue();
} finally {
assert currentModifyingThreadRef.compareAndSet(Thread.currentThread(), null);
}
}

/**
Expand Down Expand Up @@ -826,20 +837,25 @@ public V get(final K key) {
public void put(final K key, final V value) {
throwIfImmutable();
assert !isHashed() : "Cannot modify already hashed node";
Objects.requireNonNull(key, NO_NULL_KEYS_ALLOWED_MESSAGE);
assert currentModifyingThreadRef.compareAndSet(null, Thread.currentThread());
try {
Objects.requireNonNull(key, NO_NULL_KEYS_ALLOWED_MESSAGE);

final long path = records.findKey(key);
if (path == INVALID_PATH) {
// The key is not stored. So add a new entry and return.
add(key, value);
statistics.countAddedEntities();
statistics.setSize(state.size());
return;
}

final long path = records.findKey(key);
if (path == INVALID_PATH) {
// The key is not stored. So add a new entry and return.
add(key, value);
statistics.countAddedEntities();
statistics.setSize(state.size());
return;
final VirtualLeafRecord<K, V> leaf = new VirtualLeafRecord<>(path, key, value);
cache.putLeaf(leaf);
statistics.countUpdatedEntities();
} finally {
assert currentModifyingThreadRef.compareAndSet(Thread.currentThread(), null);
}

final VirtualLeafRecord<K, V> leaf = new VirtualLeafRecord<>(path, key, value);
cache.putLeaf(leaf);
statistics.countUpdatedEntities();
}

/**
Expand All @@ -859,16 +875,20 @@ public void put(final K key, final V value) {
public V replace(final K key, final V value) {
throwIfImmutable();
Objects.requireNonNull(key, NO_NULL_KEYS_ALLOWED_MESSAGE);
assert currentModifyingThreadRef.compareAndSet(null, Thread.currentThread());
try {
// Attempt to replace the existing leaf
final boolean success = replaceImpl(key, value);
statistics.countUpdatedEntities();
if (success) {
return value;
}

// Attempt to replace the existing leaf
final boolean success = replaceImpl(key, value);
statistics.countUpdatedEntities();
if (success) {
return value;
// We failed to find an existing leaf (dirty or clean). So throw an ISE.
throw new IllegalStateException("Can not replace value that is not in the map");
} finally {
assert currentModifyingThreadRef.compareAndSet(Thread.currentThread(), null);
}

// We failed to find an existing leaf (dirty or clean). So throw an ISE.
throw new IllegalStateException("Can not replace value that is not in the map");
}

/**
Expand All @@ -882,67 +902,71 @@ public V replace(final K key, final V value) {
public V remove(final K key) {
throwIfImmutable();
Objects.requireNonNull(key);
assert currentModifyingThreadRef.compareAndSet(null, Thread.currentThread());
try {
// Verify whether the current leaf exists. If not, we can just return null.
VirtualLeafRecord<K, V> leafToDelete = records.findLeafRecord(key, true);
if (leafToDelete == null) {
return null;
}

// Verify whether the current leaf exists. If not, we can just return null.
VirtualLeafRecord<K, V> leafToDelete = records.findLeafRecord(key, true);
if (leafToDelete == null) {
return null;
}

// Mark the leaf as being deleted.
cache.deleteLeaf(leafToDelete);
statistics.countRemovedEntities();
// Mark the leaf as being deleted.
cache.deleteLeaf(leafToDelete);
statistics.countRemovedEntities();

// We're going to need these
final long lastLeafPath = state.getLastLeafPath();
final long firstLeafPath = state.getFirstLeafPath();
final long leafToDeletePath = leafToDelete.getPath();

// If the leaf was not the last leaf, then move the last leaf to take this spot
if (leafToDeletePath != lastLeafPath) {
final VirtualLeafRecord<K, V> lastLeaf = records.findLeafRecord(lastLeafPath, true);
assert lastLeaf != null;
cache.clearLeafPath(lastLeafPath);
lastLeaf.setPath(leafToDeletePath);
cache.putLeaf(lastLeaf);
// NOTE: at this point, if leafToDelete was in the cache at some "path" index, it isn't anymore!
// The lastLeaf has taken its place in the path index.
}
// We're going to need these
final long lastLeafPath = state.getLastLeafPath();
final long firstLeafPath = state.getFirstLeafPath();
final long leafToDeletePath = leafToDelete.getPath();

// If the leaf was not the last leaf, then move the last leaf to take this spot
if (leafToDeletePath != lastLeafPath) {
final VirtualLeafRecord<K, V> lastLeaf = records.findLeafRecord(lastLeafPath, true);
assert lastLeaf != null;
cache.clearLeafPath(lastLeafPath);
lastLeaf.setPath(leafToDeletePath);
cache.putLeaf(lastLeaf);
// NOTE: at this point, if leafToDelete was in the cache at some "path" index, it isn't anymore!
// The lastLeaf has taken its place in the path index.
}

// If the parent of the last leaf is root, then we can simply do some bookkeeping.
// Otherwise, we replace the parent of the last leaf with the sibling of the last leaf,
// and mark it dirty. This covers all cases.
final long lastLeafParent = getParentPath(lastLeafPath);
if (lastLeafParent == ROOT_PATH) {
if (firstLeafPath == lastLeafPath) {
// We just removed the very last leaf, so set these paths to be invalid
state.setFirstLeafPath(INVALID_PATH);
state.setLastLeafPath(INVALID_PATH);
// If the parent of the last leaf is root, then we can simply do some bookkeeping.
// Otherwise, we replace the parent of the last leaf with the sibling of the last leaf,
// and mark it dirty. This covers all cases.
final long lastLeafParent = getParentPath(lastLeafPath);
if (lastLeafParent == ROOT_PATH) {
if (firstLeafPath == lastLeafPath) {
// We just removed the very last leaf, so set these paths to be invalid
state.setFirstLeafPath(INVALID_PATH);
state.setLastLeafPath(INVALID_PATH);
} else {
// We removed the second to last leaf, so the first & last leaf paths are now the same.
state.setLastLeafPath(FIRST_LEFT_PATH);
}
} else {
// We removed the second to last leaf, so the first & last leaf paths are now the same.
state.setLastLeafPath(FIRST_LEFT_PATH);
final long lastLeafSibling = getSiblingPath(lastLeafPath);
final VirtualLeafRecord<K, V> sibling = records.findLeafRecord(lastLeafSibling, true);
assert sibling != null;
cache.clearLeafPath(lastLeafSibling);
cache.deleteHash(lastLeafParent);
sibling.setPath(lastLeafParent);
cache.putLeaf(sibling);

// Update the first & last leaf paths
state.setFirstLeafPath(lastLeafParent); // replaced by the sibling, it is now first
state.setLastLeafPath(lastLeafSibling - 1); // One left of the last leaf sibling
}
if (statistics != null) {
statistics.setSize(state.size());
}
} else {
final long lastLeafSibling = getSiblingPath(lastLeafPath);
final VirtualLeafRecord<K, V> sibling = records.findLeafRecord(lastLeafSibling, true);
assert sibling != null;
cache.clearLeafPath(lastLeafSibling);
cache.deleteHash(lastLeafParent);
sibling.setPath(lastLeafParent);
cache.putLeaf(sibling);

// Update the first & last leaf paths
state.setFirstLeafPath(lastLeafParent); // replaced by the sibling, it is now first
state.setLastLeafPath(lastLeafSibling - 1); // One left of the last leaf sibling
}
if (statistics != null) {
statistics.setSize(state.size());
}

// Get the value and return it (as read only).
final V value = leafToDelete.getValue();
//noinspection unchecked
return value == null ? null : (V) value.asReadOnly();
// Get the value and return it (as read only).
final V value = leafToDelete.getValue();
//noinspection unchecked
return value == null ? null : (V) value.asReadOnly();
} finally {
assert currentModifyingThreadRef.compareAndSet(Thread.currentThread(), null);
}
}

/*
Expand Down

0 comments on commit 3984a45

Please sign in to comment.