fix: Fix schedules during replay (#12585)

Neeharika-Sompalli · Neeharika-Sompalli · commit 73ea9d64582b · 2024-04-11T00:12:53.000-05:00
Signed-off-by: Neeharika-Sompalli &lt;neeharika.sompalli@swirldslabs.com&gt;
diff --git a/hedera-node/hedera-app/src/main/java/com/hedera/node/app/statedumpers/scheduledtransactions/ScheduledTransactionsDumpUtils.java b/hedera-node/hedera-app/src/main/java/com/hedera/node/app/statedumpers/scheduledtransactions/ScheduledTransactionsDumpUtils.java
@@ -16,10 +16,12 @@
 
 package com.hedera.node.app.statedumpers.scheduledtransactions;
 
+import static com.hedera.node.app.service.mono.pbj.PbjConverter.fromPbj;
 import static com.hedera.node.app.service.mono.statedumpers.associations.BBMTokenAssociation.entityIdFrom;
 import static com.hedera.node.app.service.mono.statedumpers.scheduledtransactions.ScheduledTransactionsDumpUtils.reportOnScheduledTransactionsByEquality;
 import static com.hedera.node.app.service.mono.statedumpers.scheduledtransactions.ScheduledTransactionsDumpUtils.reportOnScheduledTransactionsByExpiry;
 import static com.hedera.node.app.service.mono.statedumpers.scheduledtransactions.ScheduledTransactionsDumpUtils.reportOnScheduledTransactionsById;
+import static com.hedera.node.app.service.schedule.impl.handlers.HandlerUtility.childAsOrdinary;
 import static com.swirlds.common.threading.manager.AdHocThreadManager.getStaticThreadManager;
 
 import com.hedera.hapi.node.base.ScheduleID;
@@ -28,7 +30,6 @@
 import com.hedera.hapi.node.state.schedule.Schedule;
 import com.hedera.hapi.node.state.schedule.ScheduleList;
 import com.hedera.node.app.service.mono.legacy.core.jproto.JKey;
-import com.hedera.node.app.service.mono.pbj.PbjConverter;
 import com.hedera.node.app.service.mono.state.adapters.VirtualMapLike;
 import com.hedera.node.app.service.mono.state.submerkle.RichInstant;
 import com.hedera.node.app.service.mono.statedumpers.DumpCheckpoint;
@@ -52,7 +53,6 @@
 import java.util.Optional;
 import java.util.TreeMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
-import java.util.concurrent.CopyOnWriteArrayList;
 
 public class ScheduledTransactionsDumpUtils {
 
@@ -63,38 +63,52 @@ public static void dumpModScheduledTransactions(
             @NonNull final VirtualMap<OnDiskKey<ProtoLong>, OnDiskValue<ScheduleList>> byExpiry,
             @NonNull final DumpCheckpoint checkpoint) {
         try (@NonNull final var writer = new Writer(path)) {
-            final var dumpableScheduledTransactionsById = gatherModScheduledTransactionsById(scheduledTransactions);
-            reportOnScheduledTransactionsById(writer, dumpableScheduledTransactionsById);
-            System.out.printf(
-                    "=== mod scheduled transactions report is %d bytes at checkpoint %s%n",
-                    writer.getSize(), checkpoint.name());
-            final var byEqualityDump = gatherModScheduledTransactionsByEquality(byEquality);
-            reportOnScheduledTransactionsByEquality(writer, byEqualityDump);
+            System.out.printf("=== Dumping schedule transactions %n ======");
+
+            final var byId = gatherModScheduledTransactionsById(scheduledTransactions);
+            reportOnScheduledTransactionsById(writer, byId);
+            System.out.println(
+                    "Size of byId in State : " + scheduledTransactions.size() + " and gathered : " + byId.size());
+
             // Not sure how to compare Equality Virtual map in mono and mod
             final var byExpiryDump = gatherModScheduledTransactionsByExpiry(byExpiry);
             reportOnScheduledTransactionsByExpiry(writer, byExpiryDump);
-            System.out.printf(
-                    "=== mod scheduled transactions by expiry report is %d bytes at checkpoint %s%n",
-                    writer.getSize(), checkpoint.name());
+            System.out.println(
+                    "Size of byExpiry in State : " + byExpiry.size() + " and gathered : " + byExpiryDump.size());
+
+            try {
+                final var byEqualityDump = gatherModScheduledTransactionsByEquality(byEquality);
+                reportOnScheduledTransactionsByEquality(writer, byEqualityDump);
+                System.out.println("Size of byEquality in State : " + byEquality.size() + " and gathered : "
+                        + byEqualityDump.size());
+            } catch (Exception e) {
+                e.printStackTrace();
+                System.out.println("Error in gathering byEqualityDump");
+            }
         }
     }
 
     private static List<BBMScheduledEqualityValue> gatherModScheduledTransactionsByEquality(
             final VirtualMap<OnDiskKey<ProtoBytes>, OnDiskValue<ScheduleList>> source) {
         final List<BBMScheduledEqualityValue> r = new ArrayList<>();
-        final var scheduledTransactions = new CopyOnWriteArrayList<BBMScheduledEqualityValue>();
+        final var scheduledTransactions = new ConcurrentLinkedQueue<BBMScheduledEqualityValue>();
 
         try {
             VirtualMapLike.from(source)
                     .extractVirtualMapDataC(
                             getStaticThreadManager(),
                             p -> scheduledTransactions.add(
                                     fromMod(p.key().getKey(), p.value().getValue())),
-                            8);
+                            1);
         } catch (final InterruptedException ex) {
             System.err.println("*** Traversal of scheduledTransactions by equality virtual map interrupted!");
             Thread.currentThread().interrupt();
         }
+
+        while (!scheduledTransactions.isEmpty()) {
+            final var mapping = scheduledTransactions.poll();
+            r.add(mapping);
+        }
         return r;
     }
 
@@ -132,11 +146,8 @@ private static Map<BBMScheduledId, BBMScheduledTransaction> gatherModScheduledTr
             VirtualMapLike.from(source)
                     .extractVirtualMapDataC(
                             getStaticThreadManager(),
-                            p -> {
-                                scheduledTransactions.add(Pair.of(
-                                        fromMod(p.key().getKey()),
-                                        fromMod(p.value().getValue())));
-                            },
+                            p -> scheduledTransactions.add(Pair.of(
+                                    fromMod(p.key().getKey()), fromMod(p.value().getValue()))),
                             8);
         } catch (final InterruptedException ex) {
             System.err.println("*** Traversal of scheduledTransactions virtual map interrupted!");
@@ -194,9 +205,9 @@ static BBMScheduledTransaction fromMod(@NonNull final Schedule value) {
                 RichInstant.fromJava(Instant.ofEpochSecond(value.calculatedExpirationSecond())),
                 RichInstant.fromJava(Instant.ofEpochSecond(
                         value.resolutionTime().seconds(), value.resolutionTime().nanos())),
-                PbjConverter.fromPbj(value.originalCreateTransaction()).toByteArray(),
-                PbjConverter.fromPbj(value.originalCreateTransaction()),
-                PbjConverter.fromPbj(value.scheduledTransaction()),
+                fromPbj(value.originalCreateTransaction()).toByteArray(),
+                fromPbj(childAsOrdinary(value)),
+                fromPbj(value.scheduledTransaction()),
                 value.signatories().stream().map(k -> toPrimitiveKey(k)).toList());
     }
 
diff --git a/hedera-node/hedera-mono-service/src/main/java/com/hedera/node/app/service/mono/statedumpers/scheduledtransactions/ScheduledTransactionsDumpUtils.java b/hedera-node/hedera-mono-service/src/main/java/com/hedera/node/app/service/mono/statedumpers/scheduledtransactions/ScheduledTransactionsDumpUtils.java
@@ -49,18 +49,22 @@ public static void dumpMonoScheduledTransactions(
             final var byId = scheduledTransactions.byId();
             final var byEquality = scheduledTransactions.byEquality();
             final var byExpirationSecond = scheduledTransactions.byExpirationSecond();
-            System.out.printf("=== Dumping schedule transactions %s%n ======");
+
+            System.out.printf("=== Dumping schedule transactions %n ======");
+
             final var byIdDump = gatherMonoScheduledTransactionsByID(byId);
             reportOnScheduledTransactionsById(writer, byIdDump);
+            System.out.println("Size of byId in State : " + byId.size() + " and gathered : " + byIdDump.size());
 
             final var byEqualityDump = gatherMonoScheduledTransactionsByEquality(byEquality);
             reportOnScheduledTransactionsByEquality(writer, byEqualityDump);
-            // Not sure how to compare Equality Virtual map in mono and mod
+            System.out.println(
+                    "Size of byEquality in State : " + byEquality.size() + " and gathered : " + byEqualityDump.size());
+
             final var byExpiryDump = gatherMonoScheduledTransactionsByExpiry(byExpirationSecond);
             reportOnScheduledTransactionsByExpiry(writer, byExpiryDump);
-            System.out.printf(
-                    "=== mono scheduled transactions report by expiry is %d bytes at checkpoint %s%n",
-                    writer.getSize(), checkpoint.name());
+            System.out.println("Size of byExpiry in State : " + byExpirationSecond.size() + " and gathered : "
+                    + byExpiryDump.size());
         }
     }
 
diff --git a/hedera-node/hedera-schedule-service-impl/src/main/java/com/hedera/node/app/service/schedule/impl/ScheduleStoreUtility.java b/hedera-node/hedera-schedule-service-impl/src/main/java/com/hedera/node/app/service/schedule/impl/ScheduleStoreUtility.java
@@ -19,11 +19,16 @@
 import com.google.common.hash.Hasher;
 import com.google.common.hash.Hashing;
 import com.hedera.hapi.node.base.Key;
+import com.hedera.hapi.node.base.ScheduleID;
 import com.hedera.hapi.node.scheduled.SchedulableTransactionBody;
 import com.hedera.hapi.node.state.schedule.Schedule;
+import com.hedera.hapi.node.state.schedule.ScheduleList;
 import com.hedera.pbj.runtime.io.buffer.Bytes;
 import edu.umd.cs.findbugs.annotations.NonNull;
+import edu.umd.cs.findbugs.annotations.Nullable;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Objects;
 
 public final class ScheduleStoreUtility {
@@ -66,4 +71,29 @@ private static void addToHash(final Hasher hasher, final SchedulableTransactionB
         hasher.putInt(bytes.length);
         hasher.putBytes(bytes);
     }
+
+    private static boolean isScheduleInList(final ScheduleID scheduleId, final ScheduleList scheduleList) {
+        return scheduleList.schedulesOrElse(Collections.emptyList()).stream()
+                .anyMatch(s -> s.scheduleIdOrThrow().equals(scheduleId));
+    }
+
+    static ScheduleList addOrReplace(final Schedule schedule, @Nullable final ScheduleList scheduleList) {
+        if (scheduleList == null) {
+            return new ScheduleList(Collections.singletonList(schedule));
+        }
+        final var newScheduleList = scheduleList.copyBuilder();
+        final var scheduleId = schedule.scheduleIdOrThrow();
+        final var schedules = new ArrayList<>(scheduleList.schedulesOrElse(Collections.emptyList()));
+        if (!isScheduleInList(scheduleId, scheduleList)) {
+            schedules.add(schedule);
+        } else {
+            for (int i = 0; i < schedules.size(); i++) {
+                final var existingSchedule = schedules.get(i);
+                if (existingSchedule.scheduleIdOrThrow().equals(scheduleId)) {
+                    schedules.set(i, schedule);
+                }
+            }
+        }
+        return newScheduleList.schedules(schedules).build();
+    }
 }
diff --git a/hedera-node/hedera-schedule-service-impl/src/main/java/com/hedera/node/app/service/schedule/impl/WritableScheduleStoreImpl.java b/hedera-node/hedera-schedule-service-impl/src/main/java/com/hedera/node/app/service/schedule/impl/WritableScheduleStoreImpl.java
@@ -16,6 +16,8 @@
 
 package com.hedera.node.app.service.schedule.impl;
 
+import static com.hedera.node.app.service.schedule.impl.ScheduleStoreUtility.addOrReplace;
+
 import com.hedera.hapi.node.base.ScheduleID;
 import com.hedera.hapi.node.base.Timestamp;
 import com.hedera.hapi.node.state.primitives.ProtoBytes;
@@ -32,9 +34,6 @@
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
 import java.time.Instant;
-import java.util.ArrayList;
-import java.util.LinkedList;
-import java.util.List;
 import java.util.Objects;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -119,25 +118,18 @@ public Schedule getForModify(@Nullable final ScheduleID idToFind) {
     @Override
     public void put(@NonNull final Schedule scheduleToAdd) {
         schedulesByIdMutable.put(scheduleToAdd.scheduleIdOrThrow(), scheduleToAdd);
+
         final ProtoBytes newHash = new ProtoBytes(ScheduleStoreUtility.calculateBytesHash(scheduleToAdd));
         final ScheduleList inStateEquality = schedulesByEqualityMutable.get(newHash);
-        List<Schedule> byEquality = inStateEquality != null ? new LinkedList<>(inStateEquality.schedules()) : null;
-        if (byEquality == null) {
-            byEquality = new LinkedList<>();
-        }
-        byEquality.add(scheduleToAdd);
-        schedulesByEqualityMutable.put(newHash, new ScheduleList(byEquality));
+        final var newEqualityScheduleList = addOrReplace(scheduleToAdd, inStateEquality);
+        schedulesByEqualityMutable.put(newHash, newEqualityScheduleList);
+
         // calculated expiration time is never null...
         final ProtoLong expirationSecond = new ProtoLong(scheduleToAdd.calculatedExpirationSecond());
         final ScheduleList inStateExpiration = schedulesByExpirationMutable.get(expirationSecond);
         // we should not be modifying the schedules list directly. This could cause ISS
-        List<Schedule> byExpiration = inStateExpiration != null ? new ArrayList<>(inStateExpiration.schedules()) : null;
-        if (byExpiration == null) {
-            byExpiration = new LinkedList<>();
-        }
-        byExpiration.add(scheduleToAdd);
-        final var newScheduleList = new ScheduleList(byExpiration);
-        schedulesByExpirationMutable.put(expirationSecond, newScheduleList);
+        final var newExpiryScheduleList = addOrReplace(scheduleToAdd, inStateExpiration);
+        schedulesByExpirationMutable.put(expirationSecond, newExpiryScheduleList);
     }
 
     @NonNull
diff --git a/hedera-node/hedera-schedule-service-impl/src/test/java/com/hedera/node/app/service/schedule/impl/WritableScheduleStoreImplTest.java b/hedera-node/hedera-schedule-service-impl/src/test/java/com/hedera/node/app/service/schedule/impl/WritableScheduleStoreImplTest.java
@@ -21,6 +21,8 @@
 import com.hedera.hapi.node.base.Key;
 import com.hedera.hapi.node.base.ScheduleID;
 import com.hedera.hapi.node.base.Timestamp;
+import com.hedera.hapi.node.state.primitives.ProtoBytes;
+import com.hedera.hapi.node.state.primitives.ProtoLong;
 import com.hedera.hapi.node.state.schedule.Schedule;
 import com.hedera.node.app.spi.workflows.PreCheckException;
 import edu.umd.cs.findbugs.annotations.NonNull;
@@ -69,6 +71,40 @@ void verifyPutModifiesState() {
         assertThat(actual.signatories()).containsExactlyInAnyOrderElementsOf(modifiedSignatories);
     }
 
+    @Test
+    void verifyPutDoesDedupliction() {
+        final ScheduleID idToDelete = scheduleInState.scheduleId();
+        Schedule actual = writableById.getForModify(idToDelete);
+        assertThat(actual).isNotNull();
+        assertThat(actual.signatories()).containsExactlyInAnyOrderElementsOf(scheduleInState.signatories());
+        final Set<Key> modifiedSignatories = Set.of(schedulerKey, payerKey);
+        final Schedule modified = replaceSignatoriesAndMarkExecuted(actual, modifiedSignatories, testConsensusTime);
+        final var hash = new ProtoBytes(ScheduleStoreUtility.calculateBytesHash(actual));
+
+        final var equalityList = writableByEquality.get(hash);
+        assertThat(equalityList.schedules().size()).isEqualTo(1);
+
+        final var expiryList = writableByExpiration.get(new ProtoLong(actual.calculatedExpirationSecond()));
+        assertThat(expiryList.schedules().size()).isEqualTo(1);
+
+        writableSchedules.put(modified);
+        writableSchedules.put(modified);
+        writableSchedules.put(modified);
+
+        actual = scheduleStore.get(idToDelete);
+        assertThat(actual).isNotNull();
+        assertThat(actual.executed()).isTrue();
+        assertThat(actual.resolutionTime()).isNotNull().isEqualTo(asTimestamp(testConsensusTime));
+        assertThat(actual.signatories()).containsExactlyInAnyOrderElementsOf(modifiedSignatories);
+
+        // size doesn't increase when same element is put multiple times
+        final var equalityListAfter = writableByEquality.get(hash);
+        assertThat(equalityListAfter.schedules().size()).isEqualTo(1);
+
+        final var expiryListAfter = writableByExpiration.get(new ProtoLong(actual.calculatedExpirationSecond()));
+        assertThat(expiryListAfter.schedules().size()).isEqualTo(1);
+    }
+
     @Test
     void purgesExpiredSchedules() {
         final ScheduleID idToDelete = scheduleInState.scheduleId();
