fix: update aex9 balance on any call and invalidate it on fork (#630)

Author: jyeshe

lib/ae_mdw/db/model.ex

@@ -220,14 +220,16 @@ defmodule AeMdw.Db.Model do
           )
 
   # AEX9 balance:
-  #     index = {contract_pk, account_pk}
-  #     amounts: float
+  #     index: {contract_pk, account_pk}
+  #     block_index: {kbi, mbi},
+  #     amount: float
   @type aex9_balance ::
           record(:aex9_balance,
             index: {Db.pubkey(), Db.pubkey()},
+            block_index: {Blocks.height(), Blocks.mbi()},
             amount: float()
           )
-  @aex9_balance_defaults [index: {<<>>, <<>>}, amount: nil]
+  @aex9_balance_defaults [index: {<<>>, <<>>}, block_index: {-1, -1}, amount: nil]
   defrecord :aex9_balance, @aex9_balance_defaults
 
   # AEX9 contract:

lib/ae_mdw/db/mutations/aex9_account_balance_mutation.ex

@@ -68,6 +68,7 @@ defmodule AeMdw.Db.ContractCallMutation do
     with true <- Contract.is_aex9?(contract_pk),
          {:ok, method_name, method_args} <- Contract.extract_successful_function(fun_arg_res),
          false <- Contract.is_non_stateful_aex9_function?(method_name) do
+      # writes already known presence
       update_aex9_presence(
         contract_pk,
         caller_pk,
@@ -77,18 +78,17 @@ defmodule AeMdw.Db.ContractCallMutation do
       )
     end
 
+    # update balance on any call
+    AsyncTasks.Producer.enqueue(:update_aex9_state, [contract_pk])
+
     :ok
   end
 
   #
   # Private functions
   #
   defp update_aex9_presence(contract_pk, caller_pk, txi, method_name, method_args) do
-    updated? = write_aex9_presence(method_name, method_args, contract_pk, caller_pk, txi)
-
-    if not updated? do
-      AsyncTasks.Producer.enqueue(:update_aex9_presence, [contract_pk])
-    end
+    :ok = write_aex9_presence(method_name, method_args, contract_pk, caller_pk, txi)
   end
 
   defp write_aex9_presence(
@@ -99,12 +99,10 @@ defmodule AeMdw.Db.ContractCallMutation do
          txi
        ) do
     DBContract.aex9_write_presence(contract_pk, txi, caller_pk)
-    true
   end
 
   defp write_aex9_presence("swap", [], contract_pk, caller_pk, txi) do
     DBContract.aex9_write_presence(contract_pk, txi, caller_pk)
-    true
   end
 
   defp write_aex9_presence(
@@ -119,18 +117,16 @@ defmodule AeMdw.Db.ContractCallMutation do
        ) do
     to_pk = Validate.id!(to_account_id)
     DBContract.aex9_write_presence(contract_pk, txi, to_pk)
-    true
   end
 
   defp write_aex9_presence(method_name, method_args, contract_pk, caller_pk, txi) do
     case Contract.get_aex9_transfer(caller_pk, method_name, method_args) do
       {from_pk, to_pk, _value} ->
         DBContract.aex9_write_presence(contract_pk, txi, from_pk)
         DBContract.aex9_write_presence(contract_pk, txi, to_pk)
-        true
 
       nil ->
-        false
+        :ok
     end
   end
 end

lib/ae_mdw/db/mutations/contract_call_mutation.ex

@@ -2,6 +2,7 @@ defmodule AeMdw.Db.Sync.Invalidate do
   # credo:disable-for-this-file
   alias AeMdw.Blocks
   alias AeMdw.Contract
+  alias AeMdw.Collection
   alias AeMdw.Database
   alias AeMdw.Db.Stream, as: DBS
   alias AeMdw.Db.Format
@@ -12,6 +13,7 @@ defmodule AeMdw.Db.Sync.Invalidate do
   alias AeMdw.Db.OracleInvalidationMutation
   alias AeMdw.Log
   alias AeMdw.Node, as: AE
+  alias AeMdw.Sync.AsyncTasks
   alias AeMdw.Validate
 
   require Model
@@ -39,6 +41,7 @@ defmodule AeMdw.Db.Sync.Invalidate do
     aex9_transfer_key_dels = aex9_transfer_key_dels(from_txi)
     aex9_account_presence_key_dels = aex9_account_presence_key_dels(from_txi)
     aex9_account_presence_key_writes = aex9_account_presence_key_writes(from_txi)
+    aex9_balance_key_dels = aex9_balance_key_dels(fork_height)
 
     int_contract_call_key_dels = int_contract_call_key_dels(from_txi)
     int_transfer_tx_key_dels = int_transfer_tx_key_dels(prev_kbi)
@@ -55,6 +58,7 @@ defmodule AeMdw.Db.Sync.Invalidate do
       DeleteKeysMutation.new(aex9_key_dels),
       DeleteKeysMutation.new(aex9_transfer_key_dels),
       DeleteKeysMutation.new(aex9_account_presence_key_dels),
+      DeleteKeysMutation.new(aex9_balance_key_dels),
       DeleteKeysMutation.new(contract_log_key_dels),
       DeleteKeysMutation.new(contract_call_key_dels),
       DeleteKeysMutation.new(int_contract_call_key_dels),
@@ -63,6 +67,15 @@ defmodule AeMdw.Db.Sync.Invalidate do
       UpdateIdsCountsMutation.new(fields_counts)
     ])
 
+    aex9_balance_key_dels
+    |> Enum.map(fn {contract_pk, _account_pk} -> contract_pk end)
+    |> Enum.uniq()
+    |> Enum.each(fn contract_pk ->
+      AsyncTasks.Producer.enqueue(:update_aex9_state, [contract_pk])
+    end)
+
+    AsyncTasks.Producer.commit_enqueued()
+
     # only ets writes
     do_writes(aex9_account_presence_key_writes, &AeMdw.Db.Contract.aex9_presence_cache_write/2)
   end
@@ -295,6 +308,19 @@ defmodule AeMdw.Db.Sync.Invalidate do
     }
   end
 
+  defp aex9_balance_key_dels(height) do
+    aex9_balance_keys =
+      Model.Aex9Balance
+      |> Collection.stream({<<>>, <<>>})
+      |> Stream.filter(fn key ->
+        Model.aex9_balance(block_index: {kbi, _mbi}) = Database.fetch!(Model.Aex9Balance, key)
+        kbi >= height
+      end)
+      |> Enum.to_list()
+
+    %{Model.Aex9Balance => aex9_balance_keys}
+  end
+
   # computes records to be written to cache, flushed at microblock boundary
   def aex9_account_presence_key_writes(from_txi) do
     bi = Model.tx(read_tx!(from_txi), :block_index)

lib/ae_mdw/db/sync/invalidate.ex

@@ -6,7 +6,6 @@ defmodule AeMdw.Db.Sync.Transaction do
   alias AeMdw.Blocks
   alias AeMdw.Contract
   alias AeMdw.Db.Model
-  alias AeMdw.Db.Aex9AccountBalanceMutation
   alias AeMdw.Db.Aex9CreateContractMutation
   alias AeMdw.Db.ContractCallMutation
   alias AeMdw.Db.ContractCreateMutation
@@ -187,20 +186,10 @@ defmodule AeMdw.Db.Sync.Transaction do
         create_txi
       )
 
-    aex9_balance_mutation =
-      with :ok <- :aect_call.return_type(call_rec),
-           true <- Contract.is_aex9?(contract_pk),
-           {:ok, method_name, method_args} <- Contract.extract_successful_function(fun_arg_res) do
-        Aex9AccountBalanceMutation.new(method_name, method_args, contract_pk, caller_pk)
-      else
-        _error_or_false -> nil
-      end
-
     Enum.concat([
       child_mutations,
       events_mutations,
       [
-        aex9_balance_mutation,
         ContractCallMutation.new(
           contract_pk,
           caller_pk,

lib/ae_mdw/db/sync/transaction.ex

@@ -10,7 +10,7 @@ defmodule AeMdw.Sync.AsyncTasks.Consumer do
   alias AeMdw.Sync.AsyncTasks.Producer
   alias AeMdw.Sync.AsyncTasks.TaskSupervisor
   alias AeMdw.Sync.AsyncTasks.DeriveAex9Presence
-  alias AeMdw.Sync.AsyncTasks.UpdateAex9Presence
+  alias AeMdw.Sync.AsyncTasks.UpdateAex9State
 
   require Model
   require Logger
@@ -23,7 +23,7 @@ defmodule AeMdw.Sync.AsyncTasks.Consumer do
 
   @type_mod %{
     derive_aex9_presence: DeriveAex9Presence,
-    update_aex9_presence: UpdateAex9Presence
+    update_aex9_state: UpdateAex9State
   }
 
   # @max_retries 2

lib/ae_mdw/sync/async_tasks/consumer.ex

@@ -41,6 +41,7 @@ defmodule AeMdw.Sync.AsyncTasks.DeriveAex9Presence do
         m_balance =
           Model.aex9_balance(
             index: {contract_pk, account_pk},
+            block_index: {kbi, mbi},
             amount: amount
           )
 

lib/ae_mdw/sync/async_tasks/derive_aex9_presence.ex

@@ -1,6 +1,6 @@
-defmodule AeMdw.Sync.AsyncTasks.UpdateAex9Presence do
+defmodule AeMdw.Sync.AsyncTasks.UpdateAex9State do
   @moduledoc """
-  Async work to derive AEX9 presence from balance dry-running.
+  Async work to update AEX9 presence and balance through dry-run.
   """
   @behaviour AeMdw.Sync.AsyncTasks.Work
 
@@ -16,29 +16,34 @@ defmodule AeMdw.Sync.AsyncTasks.UpdateAex9Presence do
   require Logger
 
   @microsecs 1_000_000
+  @min_height 500_000
 
   @spec process(args :: list()) :: :ok
   def process([contract_pk]) do
-    Log.info("[update_aex9_presence] #{inspect(contract_pk)} ...")
+    Log.info("[update_aex9_state] #{inspect(contract_pk)} ...")
+    {_type, height, _hash} = height_hash = DBN.top_height_hash(false)
 
-    {time_delta, {balances, _last_block_tuple}} =
-      :timer.tc(fn -> DBN.aex9_balances(contract_pk) end)
+    if height > @min_height do
+      {time_delta, {balances, _height_hash}} =
+        :timer.tc(fn -> DBN.aex9_balances(contract_pk, height_hash) end)
 
-    Log.info("[update_aex9_presence] #{inspect(contract_pk)} after #{time_delta / @microsecs}s")
+      Log.info("[update_aex9_state] #{inspect(contract_pk)} after #{time_delta / @microsecs}s")
 
-    create_txi = Origin.tx_index!({:contract, contract_pk})
+      create_txi = Origin.tx_index!({:contract, contract_pk})
 
-    Enum.each(balances, fn {{:address, account_pk}, amount} ->
-      Contract.aex9_write_new_presence(contract_pk, create_txi, account_pk)
+      Enum.each(balances, fn {{:address, account_pk}, amount} ->
+        Contract.aex9_write_new_presence(contract_pk, create_txi, account_pk)
 
-      m_balance =
-        Model.aex9_balance(
-          index: {contract_pk, account_pk},
-          amount: amount
-        )
+        m_balance =
+          Model.aex9_balance(
+            index: {contract_pk, account_pk},
+            block_index: {height, -1},
+            amount: amount
+          )
 
-      Database.dirty_write(Model.Aex9Balance, m_balance)
-    end)
+        Database.dirty_write(Model.Aex9Balance, m_balance)
+      end)
+    end
 
     :ok
   end

lib/ae_mdw/sync/async_tasks/update_aex9_presence.ex

@@ -0,0 +1,38 @@
+defmodule AeMdw.Migrations.RecalculateAex9Balance do
+  @moduledoc """
+  Updates .
+  """
+  alias AeMdw.Database
+  alias AeMdw.Db.Model
+  alias AeMdw.Log
+  alias AeMdw.Sync.AsyncTasks
+
+  require Model
+
+  @spec run(boolean()) :: {:ok, {non_neg_integer(), pos_integer()}}
+  def run(_from_start?) do
+    begin = DateTime.utc_now()
+
+    all_keys = Database.all_keys(Model.Aex9Balance)
+
+    Database.commit([
+      AeMdw.Db.DeleteKeysMutation.new(%{Model.Aex9Balance => all_keys})
+    ])
+
+    indexing_count =
+      all_keys
+      |> Enum.map(fn {contract_pk, _account_pk} -> contract_pk end)
+      |> Enum.uniq()
+      |> Enum.map(fn contract_pk ->
+        AsyncTasks.Producer.enqueue(:update_aex9_state, [contract_pk])
+      end)
+      |> Enum.count()
+
+    AsyncTasks.Producer.commit_enqueued()
+
+    duration = DateTime.diff(DateTime.utc_now(), begin)
+    Log.info("Async indexation scheduled for #{indexing_count} contracts in #{duration}s")
+
+    {:ok, {indexing_count, duration}}
+  end
+end