Merge pull request #16401 from MinaProtocol/georgeee/fix-16397-mini

Author: amc-ie

src/lib/network_pool/indexed_pool.ml

@@ -693,27 +693,27 @@ let revalidate :
     -> [ `Entire_pool | `Subset of Account_id.Set.t ]
     -> (Account_id.t -> Account.t)
     -> t * Transaction_hash.User_command_with_valid_signature.t Sequence.t =
- fun ({ config = { constraint_constants; _ }; _ } as t) ~logger scope f ->
+ fun ({ config = { constraint_constants; _ }; _ } as t_initial) ~logger scope f ->
   let requires_revalidation =
     match scope with
     | `Entire_pool ->
         Fn.const true
     | `Subset subset ->
         Set.mem subset
   in
-  Map.fold t.all_by_sender ~init:(t, Sequence.empty)
+  Map.fold t_initial.all_by_sender ~init:(t_initial, Sequence.empty)
     ~f:(fun
          ~key:sender
          ~data:(queue, currency_reserved)
-         ((t', dropped_acc) as acc)
+         ((t, dropped_acc) as acc)
        ->
       if not (requires_revalidation sender) then acc
       else
         let account : Account.t = f sender in
         let current_balance =
           Currency.Balance.to_amount
             (Account.liquid_balance_at_slot
-               ~global_slot:(global_slot_since_genesis t.config)
+               ~global_slot:(global_slot_since_genesis t_initial.config)
                account )
         in
         [%log debug]
@@ -739,14 +739,14 @@ let revalidate :
         then (
           [%log debug]
             "Account no longer has permission to send; dropping queue" ;
-          let dropped, t'' = remove_with_dependents_exn' t first_cmd in
-          (t'', Sequence.append dropped_acc dropped) )
+          let dropped, t_updated = remove_with_dependents_exn' t first_cmd in
+          (t_updated, Sequence.append dropped_acc dropped) )
         else if Account_nonce.(account.nonce < first_nonce) then (
           [%log debug]
             "Current account nonce precedes first nonce in queue; dropping \
              queue" ;
-          let dropped, t'' = remove_with_dependents_exn' t first_cmd in
-          (t'', Sequence.append dropped_acc dropped) )
+          let dropped, t_updated = remove_with_dependents_exn' t first_cmd in
+          (t_updated, Sequence.append dropped_acc dropped) )
         else
           (* current_nonce >= first_nonce *)
           let first_applicable_nonce_index =
@@ -763,64 +763,75 @@ let revalidate :
             "Current account nonce succeeds first nonce in queue; splitting \
              queue at $index"
             ~metadata:[ ("index", `Int first_applicable_nonce_index) ] ;
-          let drop_queue, keep_queue =
+          let dropped_for_nonce, retained_for_nonce =
             F_sequence.split_at queue first_applicable_nonce_index
           in
-          let currency_reserved' =
+          let currency_reserved_partially_updated =
             F_sequence.foldl
               (fun c cmd ->
                 Option.value_exn
                   Currency.Amount.(
                     c
                     - Option.value_exn
                         (currency_consumed ~constraint_constants cmd)) )
-              currency_reserved drop_queue
+              currency_reserved dropped_for_nonce
           in
-          let keep_queue', currency_reserved'', dropped_for_balance =
+          let keep_queue, currency_reserved_updated, dropped_for_balance =
             drop_until_sufficient_balance ~constraint_constants
-              (keep_queue, currency_reserved')
+              (retained_for_nonce, currency_reserved_partially_updated)
               current_balance
           in
           let to_drop =
-            Sequence.append (F_sequence.to_seq drop_queue) dropped_for_balance
+            Sequence.append
+              (F_sequence.to_seq dropped_for_nonce)
+              dropped_for_balance
           in
-          match Sequence.next to_drop with
-          | None ->
-              acc
-          | Some (head, tail) ->
-              let t'' =
-                Sequence.fold tail
-                  ~init:
-                    (remove_all_by_fee_and_hash_and_expiration_exn
-                       (remove_applicable_exn t' head)
-                       head )
-                  ~f:remove_all_by_fee_and_hash_and_expiration_exn
-              in
-              let t''' =
-                match F_sequence.uncons keep_queue' with
-                | None ->
-                    { t'' with
-                      all_by_sender = Map.remove t''.all_by_sender sender
-                    }
-                | Some (first_kept, _) ->
-                    let first_kept_unchecked =
-                      Transaction_hash.User_command_with_valid_signature.command
-                        first_kept
-                    in
-                    { t'' with
-                      all_by_sender =
-                        Map.set t''.all_by_sender ~key:sender
-                          ~data:(keep_queue', currency_reserved'')
-                    ; applicable_by_fee =
-                        Map_set.insert
-                          ( module Transaction_hash
-                                   .User_command_with_valid_signature )
-                          t''.applicable_by_fee
-                          (User_command.fee_per_wu first_kept_unchecked)
-                          first_kept
-                    }
-              in
-              (t''', Sequence.append dropped_acc to_drop) )
+          let keeping_prefix = F_sequence.is_empty dropped_for_nonce in
+          let keeping_suffix = Sequence.is_empty dropped_for_balance in
+          (* t with all_by_sender and applicable_by_fee fields updated *)
+          let t_partially_updated =
+            match F_sequence.uncons keep_queue with
+            | _ when keeping_prefix && keeping_suffix ->
+                (* Nothing dropped, nothing needs to be updated *)
+                t
+            | None ->
+                (* We drop the entire queue, first element needs to be removed from
+                   applicable_by_fee *)
+                let t' = remove_applicable_exn t first_cmd in
+                { t' with all_by_sender = Map.remove t'.all_by_sender sender }
+            | Some _ when keeping_prefix ->
+                (* We drop only transactions from the end of queue, keeping
+                   the head untouched, no need to update applicable_by_fee *)
+                { t with
+                  all_by_sender =
+                    Map.set t.all_by_sender ~key:sender
+                      ~data:(keep_queue, currency_reserved_updated)
+                }
+            | Some (first_kept, _) ->
+                (* We need to replace old queue head with the new queue head
+                   in applicable_by_fee *)
+                let first_kept_unchecked =
+                  Transaction_hash.User_command_with_valid_signature.command
+                    first_kept
+                in
+                let t' = remove_applicable_exn t first_cmd in
+                { t' with
+                  all_by_sender =
+                    Map.set t'.all_by_sender ~key:sender
+                      ~data:(keep_queue, currency_reserved_updated)
+                ; applicable_by_fee =
+                    Map_set.insert
+                      (module Transaction_hash.User_command_with_valid_signature)
+                      t'.applicable_by_fee
+                      (User_command.fee_per_wu first_kept_unchecked)
+                      first_kept
+                }
+          in
+          let t_updated =
+            Sequence.fold ~init:t_partially_updated
+              ~f:remove_all_by_fee_and_hash_and_expiration_exn to_drop
+          in
+          (t_updated, Sequence.append dropped_acc to_drop) )
 
 let expired_by_global_slot (t : t) :
     Transaction_hash.User_command_with_valid_signature.t Sequence.t =