com.swirlds.state.spi.WritableKVStateBase Maven / Gradle / Ivy
Show all versions of swirlds-state-api Show documentation
/*
* Copyright (C) 2023-2024 Hedera Hashgraph, LLC
*
* 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.swirlds.state.spi;
import static java.util.Objects.requireNonNull;
import edu.umd.cs.findbugs.annotations.NonNull;
import edu.umd.cs.findbugs.annotations.Nullable;
import java.util.*;
/**
* A base class for implementations of {@link WritableKVState}.
*
* @param The key type
* @param The value type
*/
public abstract class WritableKVStateBase extends ReadableKVStateBase implements WritableKVState {
/** A map of all modified values buffered in this mutable state */
private final Map modifications = new LinkedHashMap<>();
/**
* A list of listeners to be notified of changes to the state.
*/
private final List> listeners = new ArrayList<>();
/**
* Create a new StateBase.
*
* @param stateKey The state key. Cannot be null.
*/
protected WritableKVStateBase(@NonNull final String stateKey) {
super(stateKey);
}
/**
* Register a listener to be notified of changes to the state on {@link #commit()}. We do not support unregistering
* a listener, as the lifecycle of a {@link WritableKVState} is scoped to the set of mutations made to a state in a
* round; and there is no use case where an application would only want to be notified of a subset of those changes.
* @param listener the listener to register
*/
public void registerListener(@NonNull final KVChangeListener listener) {
requireNonNull(listener);
listeners.add(listener);
}
/**
* Flushes all changes into the underlying data store. This method should ONLY
* be called by the code that created the {@link WritableKVStateBase} instance or owns it. Don't
* cast and commit unless you own the instance!
*/
public void commit() {
for (final var entry : modifications.entrySet()) {
final var key = entry.getKey();
final var value = entry.getValue();
if (value == null) {
removeFromDataSource(key);
listeners.forEach(listener -> listener.mapDeleteChange(key));
} else {
putIntoDataSource(key, value);
listeners.forEach(listener -> listener.mapUpdateChange(key, value));
}
}
reset();
}
/**
* {@inheritDoc}
*
* Clears the set of modified keys and removed keys. Equivalent semantically to a "rollback"
* operation.
*/
@Override
public final void reset() {
super.reset();
modifications.clear();
}
/** {@inheritDoc} */
@Override
@Nullable
public final V get(@NonNull K key) {
// If there is a modification, then we've already done a "put" or "remove"
// and should return based on the modification
if (modifications.containsKey(key)) {
return modifications.get(key);
} else {
return super.get(key);
}
}
/** {@inheritDoc} */
@Nullable
@Override
public V getOriginalValue(@NonNull K key) {
return super.get(key);
}
/** {@inheritDoc} */
@Override
@Nullable
public final V getForModify(@NonNull final K key) {
Objects.requireNonNull(key);
// If there is a modification, then we've already done a "put" or "remove"
// and should return based on the modification
if (modifications.containsKey(key)) {
return modifications.get(key);
}
// If the modifications map does not contain an answer, but the read cache of the
// super class does, then it means we've looked this up before but never modified it.
// So we can just delegate to the super class.
if (hasBeenRead(key)) {
return super.get(key);
}
// We have not queried this key before, so let's look it up and store that we have
// read this key. And then return the value.
final var val = getForModifyFromDataSource(key);
markRead(key, val);
return val;
}
/** {@inheritDoc} */
@Override
public final void put(@NonNull final K key, @NonNull final V value) {
Objects.requireNonNull(key);
Objects.requireNonNull(value);
modifications.put(key, value);
}
/** {@inheritDoc} */
@Override
public final void remove(@NonNull final K key) {
Objects.requireNonNull(key);
modifications.put(key, null);
}
/**
* {@inheritDoc}
*
*
This is a bit of a pain, because we have to take into account modifications! If a key has
* been removed, then it must be omitted during iteration. If a key has been added, then it must
* be added into iteration.
*
* @return An iterator that iterates over all known keys.
*/
@NonNull
@Override
public Iterator keys() {
// Capture the set of keys that have been removed, and the set of keys that have been added.
final var removedKeys = new HashSet();
final var maybeAddedKeys = new HashSet();
for (final var mod : modifications.entrySet()) {
final var key = mod.getKey();
final var val = mod.getValue();
if (val == null) {
removedKeys.add(key);
} else {
maybeAddedKeys.add(key);
}
}
// Get the iterator from the backing store
final var backendItr = super.keys();
// Create and return a special iterator which will only include those keys that
// have NOT been removed, ARE in the backendItr, and includes keys that HAVE
// been added (i.e. are not in the backendItr).
return new KVStateKeyIterator<>(backendItr, removedKeys, maybeAddedKeys);
}
/** {@inheritDoc} */
@NonNull
@Override
public final Set modifiedKeys() {
return modifications.keySet();
}
/**
* {@inheritDoc}
* For the size of a {@link WritableKVState}, we need to take into account the size of the
* underlying data source, and the modifications that have been made to the state.
*
* - if the key is in backing store and is removed in modifications, then it is counted as removed
* - if the key is not in backing store and is added in modifications, then it is counted as addition
* - if the key is in backing store and is added in modifications, then it is not counted as the
* key already exists in state
* - if the key is not in backing store and is being tried to be removed in modifications,
* then it is not counted as the key does not exist in state.
*
* @return The size of the state.
*/
public long size() {
final var sizeOfBackingMap = sizeOfDataSource();
int numAdditions = 0;
int numRemovals = 0;
for (final var mod : modifications.entrySet()) {
boolean isPresentInBackingMap = readFromDataSource(mod.getKey()) != null;
boolean isRemovedInMod = mod.getValue() == null;
if (isPresentInBackingMap && isRemovedInMod) {
numRemovals++;
} else if (!isPresentInBackingMap && !isRemovedInMod) {
numAdditions++;
}
}
return sizeOfBackingMap + numAdditions - numRemovals;
}
/**
* Reads from the underlying data source in such a way as to cause any fast-copyable data
* structures underneath to make a fast copy.
*
* @param key key to read from state
* @return The value read from the underlying data source. May be null.
*/
protected abstract V getForModifyFromDataSource(@NonNull K key);
/**
* Puts the given key/value pair into the underlying data source.
*
* @param key key to update
* @param value value to put
*/
protected abstract void putIntoDataSource(@NonNull K key, @NonNull V value);
/**
* Removes the given key and implicit value from the underlying data source.
*
* @param key key to remove from the underlying data source
*/
protected abstract void removeFromDataSource(@NonNull K key);
/**
* Returns the size of the underlying data source. This can be a merkle map or a virtual map.
* @return size of the underlying data source.
*/
protected abstract long sizeOfDataSource();
/**
* A special iterator which includes all keys in the backend iterator, and all keys that have
* been added but are not part of the backend iterator, and excludes all keys that have been
* removed (even if they are in the backend iterator).
*
* For each key it gets back from the backing store, it must inspect that key to see if it is
* in the removedKeys or maybeAddedKeys. If it is in the removedKeys, then we don't return it to
* the caller, and pump another key from the backing store iterator to check instead. If it is
* in maybeAddedKeys, then remove it from maybeAddedKeys (since it is clearly not added) and
* then return it to the caller. At the very end, when the backing store iterator tells us it is
* out of keys, we start going through everything in maybeAddedKeys.
*
*
This iterator is not fail-fast.
*
* @param The type of key
*/
private static final class KVStateKeyIterator implements Iterator {
private final Iterator backendItr;
private final Set removedKeys;
private final Set maybeAddedKeys;
private Iterator addedItr;
private K next;
private KVStateKeyIterator(
@NonNull final Iterator backendItr,
@NonNull final Set removedKeys,
@NonNull final Set maybeAddedKeys) {
this.backendItr = backendItr;
this.removedKeys = removedKeys;
this.maybeAddedKeys = maybeAddedKeys;
}
@Override
public boolean hasNext() {
prepareNext();
return next != null;
}
@Override
public K next() {
prepareNext();
if (next == null) {
throw new NoSuchElementException();
}
final var ret = next;
next = null;
return ret;
}
private void prepareNext() {
while (next == null) {
if (backendItr.hasNext()) {
final var candidate = backendItr.next();
maybeAddedKeys.remove(candidate);
if (removedKeys.contains(candidate)) {
continue;
}
next = candidate;
return;
}
if (addedItr == null) {
addedItr = maybeAddedKeys.iterator();
}
if (addedItr.hasNext()) {
next = addedItr.next();
}
// If we get here, then there is nothing.
return;
}
}
}
}