Remove old timestamps in scheduleld transaction manager cleanup

contracts/FlowTransactionSchedulerUtils.cdc

@@ -119,6 +119,7 @@ access(all) contract FlowTransactionSchedulerUtils {
         access(all) view fun getTransactionIDsByTimestamp(_ timestamp: UFix64): [UInt64]
         access(all) fun getTransactionIDsByTimestampRange(startTimestamp: UFix64, endTimestamp: UFix64): {UFix64: [UInt64]}
         access(all) view fun getTransactionStatus(id: UInt64): FlowTransactionScheduler.Status?
+        access(all) view fun getSortedTimestamps(): FlowTransactionScheduler.SortedTimestamps
     }
 
     /// Manager resource is meant to provide users and developers with a simple way
@@ -221,6 +222,7 @@ access(all) contract FlowTransactionSchedulerUtils {
 
             // Store the handler capability in our dictionary for later retrieval
             let id = scheduledTransaction.id
+            let actualTimestamp = scheduledTransaction.timestamp
             let handlerRef = handlerCap.borrow()
                 ?? panic("Invalid transaction handler: Could not borrow a reference to the transaction handler")
             let handlerTypeIdentifier = handlerRef.getType().identifier
@@ -250,14 +252,14 @@ access(all) contract FlowTransactionSchedulerUtils {
             self.scheduledTransactions[scheduledTransaction.id] <-! scheduledTransaction
 
             // Add the transaction to the sorted timestamps array
-            self.sortedTimestamps.add(timestamp: timestamp)
+            self.sortedTimestamps.add(timestamp: actualTimestamp)
 
             // Store the transaction in the ids by timestamp dictionary
-            if let ids = self.idsByTimestamp[timestamp] {
+            if let ids = self.idsByTimestamp[actualTimestamp] {
                 ids.append(id)
-                self.idsByTimestamp[timestamp] = ids
+                self.idsByTimestamp[actualTimestamp] = ids
             } else {
-                self.idsByTimestamp[timestamp] = [id]
+                self.idsByTimestamp[actualTimestamp] = [id]
             }
 
             return id
@@ -284,27 +286,26 @@ access(all) contract FlowTransactionSchedulerUtils {
         /// @param timestamp: The timestamp of the transaction to remove
         /// @param handlerTypeIdentifier: The type identifier of the handler of the transaction to remove
         access(self) fun removeID(id: UInt64, timestamp: UFix64, handlerTypeIdentifier: String) {
+            pre {
+                self.handlerInfos.containsKey(handlerTypeIdentifier): "Invalid handler type identifier: Handler with type identifier \(handlerTypeIdentifier) not found in manager"
+                self.idsByTimestamp.containsKey(timestamp): "Invalid timestamp: Timestamp \(timestamp) not found in manager for transaction with ID \(id)"
+            }
 
-            if let ids = self.idsByTimestamp[timestamp] {
-                let index = ids.firstIndex(of: id)
-                ids.remove(at: index!)
-                if ids.length == 0 {
-                    self.idsByTimestamp.remove(key: timestamp)
-                } else {
-                    self.idsByTimestamp[timestamp] = ids
-                }
+            let ids = &self.idsByTimestamp[timestamp]! as auth(Mutate) &[UInt64]
+            let index = ids.firstIndex(of: id)
+            ids.remove(at: index!)
+            if ids.length == 0 {
+                self.idsByTimestamp.remove(key: timestamp)
+                self.sortedTimestamps.remove(timestamp: timestamp)
             }
 
             let handlerUUID = self.handlerUUIDsByTransactionID.remove(key: id)
                 ?? panic("Invalid ID: Transaction with ID \(id) not found in manager")
 
             // Remove the transaction ID from the handler info array
-            if let handlers = self.handlerInfos[handlerTypeIdentifier] {
-                if let handlerInfo = handlers[handlerUUID] {
-                    handlerInfo.removeTransactionID(id: id)
-                    handlers[handlerUUID] = handlerInfo
-                }
-                self.handlerInfos[handlerTypeIdentifier] = handlers
+            let handlers = &self.handlerInfos[handlerTypeIdentifier]! as auth(Mutate) &{UInt64: HandlerInfo}
+            if let handlerInfo = handlers[handlerUUID] {
+                handlerInfo.removeTransactionID(id: id)
             }
         }
 
@@ -318,10 +319,21 @@ access(all) contract FlowTransactionSchedulerUtils {
             let pastTimestamps = self.sortedTimestamps.getBefore(current: currentTimestamp)
             for timestamp in pastTimestamps {
                 let ids = self.idsByTimestamp[timestamp] ?? []
+                if ids.length == 0 {
+                    self.sortedTimestamps.remove(timestamp: timestamp)
+                    continue
+                }
                 for id in ids {
                     let status = FlowTransactionScheduler.getStatus(id: id)
-                    if status == nil || status == FlowTransactionScheduler.Status.Unknown {
+                    if status == nil || status! != FlowTransactionScheduler.Status.Scheduled {
                         transactionsToRemove.append(id)
+                        // Need to temporarily limit the number of transactions to remove
+                        // because some managers on mainnet have already hit the limit and we need to batch them
+                        // to make sure they get cleaned up properly
+                        // This will be removed eventually
+                        if transactionsToRemove.length > 50 {
+                            break
+                        }
                     }
                 }
             }
@@ -541,6 +553,12 @@ access(all) contract FlowTransactionSchedulerUtils {
             }
             return FlowTransactionScheduler.Status.Unknown
         }
+
+        /// Gets the sorted timestamps struct
+        /// @return: The sorted timestamps struct
+        access(all) view fun getSortedTimestamps(): FlowTransactionScheduler.SortedTimestamps {
+            return self.sortedTimestamps
+        }
     }
 
     /// Create a new Manager instance

tests/scheduled_transaction_test_helpers.cdc

@@ -415,6 +415,15 @@ access(all) fun getManagedTxIDs(): [UInt64] {
     return result.returnValue as! [UInt64]
 }
 
+access(all) fun getManagerTimestamps(): [UFix64] {
+    var result = _executeScript(
+        "./scripts/get_manager_timestamps.cdc",
+        [admin.address]
+    )
+    Test.expect(result, Test.beSucceeded())
+    return result.returnValue as! [UFix64]
+}
+
 access(all) fun getHandlerTypeIdentifiers(): {String: [UInt64]} {
     var result = _executeScript(
         "../transactions/transactionScheduler/scripts/manager/get_handler_types.cdc",

tests/scripts/get_manager_timestamps.cdc

@@ -0,0 +1,8 @@
+import "FlowTransactionSchedulerUtils"
+import "FlowTransactionScheduler"
+
+access(all) fun main(address: Address): [UFix64] {
+    let managerRef = FlowTransactionSchedulerUtils.borrowManager(at: address)
+        ?? panic("Invalid address: Could not borrow a reference to the Scheduled Transaction Manager at address \(address)")
+    return managerRef.getSortedTimestamps().getBefore(current: getCurrentBlock().timestamp)
+}

tests/transactionScheduler_manager_test.cdc

@@ -473,6 +473,13 @@ access(all) fun testManagerExecuteAndCleanup() {
     // move time until after all the timestamps
     Test.moveTime(by: Fix64(futureDelta + 5.0))
 
+    // get the old timestamps
+    let oldTimestamps = getManagerTimestamps()
+    Test.assert(oldTimestamps.length == 3, message: "Should have 3 timestamps but got \(oldTimestamps.length)")
+    Test.assert(oldTimestamps.contains(timeInFuture), message: "Should contain timestamp \(timeInFuture)")
+    Test.assert(oldTimestamps.contains(timeInFuture + 1.0), message: "Should contain timestamp \(timeInFuture + 1.0)")
+    Test.assert(oldTimestamps.contains(timeInFuture + 2.0), message: "Should contain timestamp \(timeInFuture + 2.0)")
+
     // process the transactions
     processTransactions()
 
@@ -497,6 +504,10 @@ access(all) fun testManagerExecuteAndCleanup() {
         testName: "Test Manager Execute and Cleanup",
         failWithErr: nil)
 
+    // get the old timestamps, timeInFuture should have been removed
+    let timestamps = getManagerTimestamps()
+    Test.assert(timestamps.length == 0, message: "Should have 0 timestamps after cleanup but got \(timestamps.length)")
+
     // get status of id 1
     let status1 = getStatus(id: 1)
     Test.assert(status1 == statusUnknown, message: "Status for transaction 1 should be unknown but got \(status1!)")
@@ -515,13 +526,13 @@ access(all) fun testManagerExecuteAndCleanup() {
 
     // check that ID 5, 6, 8, and 9 are still in the manager and executed
     let status5 = getManagedTxStatus(5)
-    Test.assert(status5 == statusExecuted, message: "Status for transaction 5 should be executed but got \(status5!)")
+    Test.assert(status5 == statusUnknown, message: "Status for transaction 5 should be unknown but got \(status5!)")
     let status6 = getManagedTxStatus(6)
-    Test.assert(status6 == statusExecuted, message: "Status for transaction 6 should be executed but got \(status6!)")
+    Test.assert(status6 == statusUnknown, message: "Status for transaction 6 should be unknown but got \(status6!)")
     let status8 = getManagedTxStatus(8)
-    Test.assert(status8 == statusExecuted, message: "Status for transaction 8 should be executed but got \(status8!)")
+    Test.assert(status8 == statusUnknown, message: "Status for transaction 8 should be unknown but got \(status8!)")
     let status9 = getManagedTxStatus(9)
-    Test.assert(status9 == statusExecuted, message: "Status for transaction 9 should be executed but got \(status9!)")
+    Test.assert(status9 == statusUnknown, message: "Status for transaction 9 should be unknown but got \(status9!)")
 
     // test getting data for transactions which were executed and cleaned up
     let data2 = getManagedTxData(2)
@@ -537,28 +548,14 @@ access(all) fun testManagerExecuteAndCleanup() {
     let handlerType = "A.0000000000000007.TestFlowScheduledTransactionHandler.Handler"
     let txIds = getManagedTxIDsByHandler(handlerTypeIdentifier: handlerType, handlerUUID: nil)
 
-    // Should only have 5 transactions. the four executed and the one scheduled
-    Test.assert(txIds.length == 5, message: "Should have 4 transactions for the handler type but got \(txIds.length)")
-    var i: UInt64 = 1
-    while i <= 10 {
-        if i == 1 ||i == 2 || i == 3 || i == 4 || i == 7 {
-            Test.assert(!txIds.contains(i), message: "Should not contain transaction ID \(i)")
-        } else {
-            Test.assert(txIds.contains(i), message: "Should contain transaction ID \(i)")
-        }
-        i = i + 1
-    }
+    // Should only have 1 transaction, the one scheduled
+    Test.assert(txIds.length == 1, message: "Should have 1 transaction for the handler type but got \(txIds.length)")
+    Test.assert(txIds.contains(10), message: "Should contain transaction ID 10")
 
     // Test getting all transaction IDs
     let allTxIds = getManagedTxIDs()
-    Test.assert(allTxIds.length == 5, message: "Should have 5 total transactions")
-    i = 1
-    while i <= 10 {
-        if i == 1 || i == 2 || i == 3 || i == 4 || i == 7 {
-            Test.assert(!allTxIds.contains(i), message: "Should not contain transaction ID \(i)")
-        }
-        i = i + 1
-    }
+    Test.assert(allTxIds.length == 1, message: "Should have 1 total transaction")
+    Test.assert(allTxIds.contains(10), message: "Should contain transaction ID 10")
 
     // Test getting handler views from executed and cleaned up transaction ID
     let viewsFromTxId = getHandlerViewsFromTransactionID(3)

tests/transactionScheduler_misc_test.cdc

@@ -233,13 +233,13 @@ access(all) fun testScheduledTransactionDestroyHandler() {
 
     // Destroy the handler for both transactions
     let destroyHandlerCode = Test.readFile("./transactions/destroy_handler.cdc")
-    let executeTx = Test.Transaction(
+    let destroyHandlerTx = Test.Transaction(
         code: destroyHandlerCode,
         authorizers: [admin.address],
         signers: [admin],
         arguments: []
     )
-    var result = Test.executeTransaction(executeTx)
+    var result = Test.executeTransaction(destroyHandlerTx)
     Test.expect(result, Test.beSucceeded())
 
     // Cancel the second transaction