chore: warm token transfer (#11017)

Signed-off-by: Michael Heinrichs <[email protected]>

hedera-node/hedera-app-spi/src/main/java/com/hedera/node/app/spi/state/ReadableKVState.java

@@ -90,4 +90,13 @@ default boolean contains(@NonNull final K key) {
      * @return number of keys in the {@link ReadableKVState}.
      */
     long size();
+
+    /**
+     * Warms the system by preloading an entity into memory
+     *
+     * <p>The default implementation is empty because preloading data into memory is only used for some implementations.
+     *
+     * @param key the key of the entity
+     */
+    default void warm(@NonNull final K key) {}
 }

hedera-node/hedera-app-spi/src/main/java/com/hedera/node/app/spi/workflows/TransactionHandler.java

@@ -59,7 +59,7 @@ default void pureChecks(@NonNull final TransactionBody txn) throws PreCheckExcep
      * @param context the {@link PreHandleContext} which collects all information
      * @throws NullPointerException if {@code context} is {@code null}
      */
-    default void warmUp(@NonNull final PreHandleContext context) {}
+    default void warm(@NonNull final WarmupContext context) {}
 
     /**
      * Calculates the fees for a transaction

hedera-node/hedera-app-spi/src/main/java/com/hedera/node/app/spi/workflows/WarmupContext.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 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.hedera.node.app.spi.workflows;
+
+import com.hedera.hapi.node.transaction.TransactionBody;
+import edu.umd.cs.findbugs.annotations.NonNull;
+
+/**
+ * Represents the context of a single {@code warm()}-call.
+ */
+public interface WarmupContext {
+
+    /**
+     * Gets the {@link TransactionBody}
+     *
+     * @return the {@link TransactionBody} in this context
+     */
+    @NonNull
+    TransactionBody body();
+
+    /**
+     * Create a new store given the store's interface. This gives read-only access to the store.
+     *
+     * @param storeInterface The store interface to find and create a store for
+     * @param <C> Interface class for a Store
+     * @return An implementation of store interface provided, or null if the store
+     * @throws IllegalArgumentException if the storeInterface class provided is unknown to the app
+     * @throws NullPointerException if {@code storeInterface} is {@code null}
+     */
+    @NonNull
+    <C> C createStore(@NonNull final Class<C> storeInterface);
+}

hedera-node/hedera-app/src/main/java/com/hedera/node/app/state/merkle/disk/OnDiskReadableKVState.java

@@ -113,4 +113,10 @@ public long size() {
         logMapGetSize(getStateKey(), size);
         return size;
     }
+
+    @Override
+    public void warm(@NonNull final K key) {
+        final var k = new OnDiskKey<>(md, key);
+        virtualMap.warm(k);
+    }
 }

hedera-node/hedera-app/src/main/java/com/hedera/node/app/workflows/dispatcher/TransactionDispatcher.java

@@ -27,6 +27,7 @@
 import com.hedera.node.app.spi.workflows.PreCheckException;
 import com.hedera.node.app.spi.workflows.PreHandleContext;
 import com.hedera.node.app.spi.workflows.TransactionHandler;
+import com.hedera.node.app.spi.workflows.WarmupContext;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import javax.inject.Inject;
 import javax.inject.Singleton;
@@ -100,12 +101,12 @@ public void dispatchPreHandle(@NonNull final PreHandleContext context) throws Pr
      * @param context the context of the warmup workflow
      * @throws NullPointerException if {@code context} is {@code null}
      */
-    public void dispatchWarmup(@NonNull final PreHandleContext context) {
+    public void dispatchWarmup(@NonNull final WarmupContext context) {
         requireNonNull(context, "The supplied argument 'context' cannot be null!");
 
         try {
             final var handler = getHandler(context.body());
-            handler.warmUp(context);
+            handler.warm(context);
         } catch (UnsupportedOperationException ex) {
             // do nothing, the handler should have been used before we reach this point
         }

hedera-node/hedera-app/src/main/java/com/hedera/node/app/workflows/handle/CacheWarmer.java

@@ -0,0 +1,107 @@
+/*
+ * Copyright (C) 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.hedera.node.app.workflows.handle;
+
+import static java.util.Objects.requireNonNull;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.hedera.hapi.node.transaction.TransactionBody;
+import com.hedera.node.app.spi.workflows.PreCheckException;
+import com.hedera.node.app.spi.workflows.TransactionHandler;
+import com.hedera.node.app.state.HederaState;
+import com.hedera.node.app.workflows.TransactionChecker;
+import com.hedera.node.app.workflows.dispatcher.ReadableStoreFactory;
+import com.hedera.node.app.workflows.dispatcher.TransactionDispatcher;
+import com.hedera.node.app.workflows.prehandle.PreHandleResult;
+import com.hedera.pbj.runtime.io.buffer.Bytes;
+import com.swirlds.platform.system.Round;
+import com.swirlds.platform.system.events.ConsensusEvent;
+import com.swirlds.platform.system.transaction.Transaction;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import edu.umd.cs.findbugs.annotations.Nullable;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ForkJoinPool;
+import javax.inject.Inject;
+import javax.inject.Singleton;
+
+/**
+ * This class is used to warm up the cache. It is called at the beginning of a round with the current state
+ * and the round. It will start a background thread which iterates through all transactions and calls the
+ * {@link TransactionHandler#warm} method.
+ */
+@Singleton
+public class CacheWarmer {
+
+    private final TransactionChecker checker;
+    private final TransactionDispatcher dispatcher;
+    private final Executor executor;
+
+    @Inject
+    public CacheWarmer(@NonNull final TransactionChecker checker, @NonNull final TransactionDispatcher dispatcher) {
+        this(checker, dispatcher, ForkJoinPool.commonPool());
+    }
+
+    @VisibleForTesting
+    CacheWarmer(
+            @NonNull final TransactionChecker checker,
+            @NonNull final TransactionDispatcher dispatcher,
+            @NonNull final Executor executor) {
+        this.checker = checker;
+        this.dispatcher = requireNonNull(dispatcher);
+        this.executor = requireNonNull(executor);
+    }
+
+    /**
+     * Warms up the cache for the given round.
+     *
+     * @param state the current state
+     * @param round the current round
+     */
+    public void warm(@NonNull final HederaState state, @NonNull final Round round) {
+        executor.execute(() -> {
+            final ReadableStoreFactory storeFactory = new ReadableStoreFactory(state);
+            for (final ConsensusEvent event : round) {
+                event.forEachTransaction(platformTransaction -> executor.execute(() -> {
+                    final TransactionBody txBody = extractTransactionBody(platformTransaction);
+                    if (txBody != null) {
+                        final var context = new WarmupContextImpl(txBody, storeFactory);
+                        dispatcher.dispatchWarmup(context);
+                    }
+                }));
+            }
+        });
+    }
+
+    @Nullable
+    private TransactionBody extractTransactionBody(@NonNull final Transaction platformTransaction) {
+        // First we check if the transaction was already parsed during pre-handle (should be almost always the case)
+        final var metadata = platformTransaction.getMetadata();
+        if (metadata instanceof PreHandleResult result) {
+            return result.txInfo() == null ? null : result.txInfo().txBody();
+        }
+
+        // If not we parse it here using existing code. This is not ideal but should be rare.
+        // We can potentially optimize this by limiting the code to the bare minimum needed
+        // or keeping the result for later.
+        try {
+            final Bytes buffer = Bytes.wrap(platformTransaction.getContents());
+            return checker.parseAndCheck(buffer).txBody();
+        } catch (PreCheckException ex) {
+            return null;
+        }
+    }
+}

hedera-node/hedera-app/src/main/java/com/hedera/node/app/workflows/handle/HandleWorkflow.java

@@ -149,6 +149,7 @@ public class HandleWorkflow {
     private final Authorizer authorizer;
     private final NetworkUtilizationManager networkUtilizationManager;
     private final SynchronizedThrottleAccumulator synchronizedThrottleAccumulator;
+    private final CacheWarmer cacheWarmer;
 
     @Inject
     public HandleWorkflow(
@@ -172,7 +173,8 @@ public HandleWorkflow(
             @NonNull final Authorizer authorizer,
             @NonNull final NetworkUtilizationManager networkUtilizationManager,
             @NonNull final SynchronizedThrottleAccumulator synchronizedThrottleAccumulator,
-            @NonNull final ScheduleExpirationHook scheduleExpirationHook) {
+            @NonNull final ScheduleExpirationHook scheduleExpirationHook,
+            @NonNull final CacheWarmer cacheWarmer) {
         this.networkInfo = requireNonNull(networkInfo, "networkInfo must not be null");
         this.preHandleWorkflow = requireNonNull(preHandleWorkflow, "preHandleWorkflow must not be null");
         this.dispatcher = requireNonNull(dispatcher, "dispatcher must not be null");
@@ -199,6 +201,7 @@ public HandleWorkflow(
                 requireNonNull(synchronizedThrottleAccumulator, "synchronizedThrottleAccumulator must not be null");
         ;
         this.scheduleExpirationHook = requireNonNull(scheduleExpirationHook, "scheduleExpirationHook must not be null");
+        this.cacheWarmer = requireNonNull(cacheWarmer, "cacheWarmer must not be null");
     }
 
     /**
@@ -216,6 +219,9 @@ public void handleRound(
         // log start of round to transaction state log
         logStartRound(round);
 
+        // warm the cache
+        cacheWarmer.warm(state, round);
+
         // handle each event in the round
         for (final ConsensusEvent event : round) {
             final var creator = networkInfo.nodeInfo(event.getCreatorId().id());

hedera-node/hedera-app/src/main/java/com/hedera/node/app/workflows/handle/WarmupContextImpl.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright (C) 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.hedera.node.app.workflows.handle;
+
+import com.hedera.hapi.node.transaction.TransactionBody;
+import com.hedera.node.app.spi.workflows.WarmupContext;
+import com.hedera.node.app.workflows.TransactionInfo;
+import com.hedera.node.app.workflows.dispatcher.ReadableStoreFactory;
+import edu.umd.cs.findbugs.annotations.NonNull;
+
+/**
+ * The default implementation of {@link WarmupContext}.
+ */
+public class WarmupContextImpl implements WarmupContext {
+
+    @NonNull
+    private final TransactionBody txBody;
+
+    @NonNull
+    private final ReadableStoreFactory storeFactory;
+
+    /**
+     * Constructor of {@code WarmupContextImpl}
+     *
+     * @param txBody the {@link TransactionInfo} of the transaction
+     * @param storeFactory the {@link ReadableStoreFactory} to create stores
+     */
+    public WarmupContextImpl(@NonNull final TransactionBody txBody, @NonNull final ReadableStoreFactory storeFactory) {
+        this.txBody = txBody;
+        this.storeFactory = storeFactory;
+    }
+
+    @NonNull
+    @Override
+    public TransactionBody body() {
+        return txBody;
+    }
+
+    @NonNull
+    @Override
+    public <C> C createStore(@NonNull final Class<C> storeInterface) {
+        return storeFactory.getStore(storeInterface);
+    }
+}

hedera-node/hedera-app/src/main/java/com/hedera/node/app/workflows/prehandle/PreHandleWorkflowImpl.java

@@ -296,8 +296,6 @@ private PreHandleResult expandAndVerifySignatures(
             dispatcher.dispatchPureChecks(txBody);
             // Then gather the signatures from the transaction handler
             dispatcher.dispatchPreHandle(context);
-            // Finally, let the transaction handler do warm up of other state it may want to use later
-            dispatcher.dispatchWarmup(context);
         } catch (PreCheckException preCheck) {
             // It is quite possible those semantic checks and other tasks will fail and throw a PreCheckException.
             // In that case, the payer will end up paying for the transaction. So we still need to do the signature

hedera-node/hedera-app/src/test/java/com/hedera/node/app/state/merkle/disk/OnDiskReadableStateTest.java

@@ -18,6 +18,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.verify;
 
 import com.hedera.node.app.spi.fixtures.state.TestSchema;
 import com.hedera.node.app.spi.state.StateDefinition;
@@ -28,7 +29,11 @@
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
 
+@ExtendWith(MockitoExtension.class)
 class OnDiskReadableStateTest extends MerkleTestBase {
     private StateMetadata<String, String> md;
     private VirtualMap<OnDiskKey<String>, OnDiskValue<String>> virtualMap;
@@ -107,4 +112,12 @@ void get() {
             assertThat(state.get(G_KEY)).isNull();
         }
     }
+
+    @Test
+    @DisplayName("The method warm() calls the appropriate method on the virtual map")
+    void warm(@Mock VirtualMap<OnDiskKey<String>, OnDiskValue<String>> virtualMapMock) {
+        final var state = new OnDiskReadableKVState<>(md, virtualMapMock);
+        state.warm(A_KEY);
+        verify(virtualMapMock).warm(new OnDiskKey<>(md, A_KEY));
+    }
 }