fix(api): consider expiration return addresses for ledger updates

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "chronicle"
-version = "1.0.0-rc.2"
+version = "1.0.0-rc.3"
 authors = ["IOTA Stiftung"]
 edition = "2021"
 description = "IOTA permanode implemented as an IOTA Node Extension (INX)."

documentation/api/api-explorer.yml

@@ -277,7 +277,7 @@ components:
           description: >-
             The total value held in unspent outputs owned by the given address 
             (includes funds held in storage deposit).
-        sigLockedBalance:
+        availableBalance:
           type: string
           description: >-
             The sum of value held in unspent outputs owned by the given address 
@@ -585,7 +585,7 @@ components:
     balance-example:
       value:
         totalBalance: 100000
-        sigLockedBalance: 99900
+        availableBalance: 99900
         ledgerIndex: 500000
     ledger-updates-address-example:
       value:

src/bin/inx-chronicle/api/explorer/responses.rs

@@ -78,7 +78,7 @@ impl From<LedgerUpdateByMilestoneRecord> for LedgerUpdateByMilestoneDto {
 #[serde(rename_all = "camelCase")]
 pub struct BalanceResponse {
     pub total_balance: String,
-    pub sig_locked_balance: String,
+    pub available_balance: String,
     pub ledger_index: MilestoneIndex,
 }
 

src/bin/inx-chronicle/api/explorer/routes.rs

@@ -155,22 +155,22 @@ async fn ledger_updates_by_milestone(
 }
 
 async fn balance(database: Extension<MongoDb>, Path(address): Path<String>) -> ApiResult<BalanceResponse> {
-    let ledger_index = database
+    let ledger_ms = database
         .collection::<MilestoneCollection>()
-        .get_ledger_index()
+        .get_newest_milestone()
         .await?
         .ok_or(MissingError::NoResults)?;
     let address = Address::from_str(&address).map_err(RequestError::from)?;
     let res = database
         .collection::<OutputCollection>()
-        .get_address_balance(address, ledger_index)
+        .get_address_balance(address, ledger_ms)
         .await?
         .ok_or(MissingError::NoResults)?;
 
     Ok(BalanceResponse {
         total_balance: res.total_balance,
-        sig_locked_balance: res.sig_locked_balance,
-        ledger_index,
+        available_balance: res.available_balance,
+        ledger_index: ledger_ms.milestone_index,
     })
 }
 

src/bin/inx-chronicle/migrations/migrate_2.rs

@@ -0,0 +1,122 @@
+// Copyright 2023 IOTA Stiftung
+// SPDX-License-Identifier: Apache-2.0
+
+use async_trait::async_trait;
+use chronicle::{
+    db::{
+        mongodb::collections::{LedgerUpdateCollection, OutputCollection},
+        MongoDb, MongoDbCollection, MongoDbCollectionExt,
+    },
+    model::{
+        ledger::{LedgerOutput, LedgerSpent, RentStructureBytes},
+        metadata::OutputMetadata,
+        utxo::{Output, OutputId},
+    },
+};
+use futures::{prelude::stream::TryStreamExt, StreamExt};
+use mongodb::bson::doc;
+use serde::Deserialize;
+use tokio::{task::JoinSet, try_join};
+
+use super::Migration;
+
+const INSERT_BATCH_SIZE: usize = 1000;
+
+pub struct Migrate;
+
+#[async_trait]
+impl Migration for Migrate {
+    const ID: usize = 2;
+    const APP_VERSION: &'static str = "1.0.0-rc.3";
+    const DATE: time::Date = time::macros::date!(2024 - 01 - 12);
+
+    async fn migrate(db: &MongoDb) -> eyre::Result<()> {
+        db.collection::<LedgerUpdateCollection>()
+            .collection()
+            .drop(None)
+            .await?;
+
+        let outputs_stream = db
+            .collection::<OutputCollection>()
+            .find::<OutputDocument>(doc! {}, None)
+            .await?;
+        let mut batched_stream = outputs_stream.try_chunks(INSERT_BATCH_SIZE);
+
+        let mut tasks = JoinSet::new();
+
+        while let Some(batch) = batched_stream.next().await {
+            let batch = batch?;
+            while tasks.len() >= 100 {
+                if let Some(res) = tasks.join_next().await {
+                    res??;
+                }
+            }
+            let db = db.clone();
+            tasks.spawn(async move {
+                let consumed = batch.iter().filter_map(Option::<LedgerSpent>::from).collect::<Vec<_>>();
+                let created = batch.into_iter().map(LedgerOutput::from).collect::<Vec<_>>();
+                try_join! {
+                    async {
+                        db.collection::<LedgerUpdateCollection>()
+                                .insert_unspent_ledger_updates(&created)
+                                .await
+                    },
+                    async {
+                        db.collection::<OutputCollection>().update_spent_outputs(&consumed).await
+                    },
+                    async {
+                        db.collection::<LedgerUpdateCollection>().insert_spent_ledger_updates(&consumed).await
+                    }
+                }
+                .and(Ok(()))
+            });
+        }
+
+        while let Some(res) = tasks.join_next().await {
+            res??;
+        }
+
+        Ok(())
+    }
+}
+
+#[derive(Deserialize)]
+pub struct OutputDocument {
+    #[serde(rename = "_id")]
+    output_id: OutputId,
+    output: Output,
+    metadata: OutputMetadata,
+    details: OutputDetails,
+}
+
+#[derive(Deserialize)]
+struct OutputDetails {
+    rent_structure: RentStructureBytes,
+}
+
+impl From<OutputDocument> for LedgerOutput {
+    fn from(value: OutputDocument) -> Self {
+        Self {
+            output_id: value.output_id,
+            block_id: value.metadata.block_id,
+            booked: value.metadata.booked,
+            output: value.output,
+            rent_structure: value.details.rent_structure,
+        }
+    }
+}
+
+impl From<&OutputDocument> for Option<LedgerSpent> {
+    fn from(value: &OutputDocument) -> Self {
+        value.metadata.spent_metadata.map(|spent_metadata| LedgerSpent {
+            spent_metadata,
+            output: LedgerOutput {
+                output_id: value.output_id,
+                block_id: value.metadata.block_id,
+                booked: value.metadata.booked,
+                output: value.output.clone(),
+                rent_structure: value.details.rent_structure,
+            },
+        })
+    }
+}

src/bin/inx-chronicle/migrations/mod.rs

@@ -12,15 +12,17 @@ use eyre::bail;
 
 pub mod migrate_0;
 pub mod migrate_1;
+pub mod migrate_2;
 
-pub type LatestMigration = migrate_1::Migrate;
+pub type LatestMigration = migrate_2::Migrate;
 
 /// The list of migrations, in order.
 const MIGRATIONS: &[&'static dyn DynMigration] = &[
     // In order to add a new migration, change the `LatestMigration` type above and add an entry at the bottom of this
     // list.
     &migrate_0::Migrate,
     &migrate_1::Migrate,
+    &migrate_2::Migrate,
 ];
 
 fn build_migrations(migrations: &[&'static dyn DynMigration]) -> HashMap<Option<usize>, &'static dyn DynMigration> {

src/db/mongodb/collections/ledger_update.rs

@@ -110,23 +110,18 @@ impl LedgerUpdateCollection {
         I: IntoIterator<Item = &'a LedgerSpent>,
         I::IntoIter: Send + Sync,
     {
-        let ledger_updates = outputs.into_iter().filter_map(
-            |LedgerSpent {
-                 output: LedgerOutput { output_id, output, .. },
-                 spent_metadata,
-             }| {
-                // Ledger updates
-                output.owning_address().map(|&address| LedgerUpdateDocument {
-                    _id: Id {
-                        milestone_index: spent_metadata.spent.milestone_index,
-                        output_id: *output_id,
-                        is_spent: true,
-                    },
-                    address,
-                    milestone_timestamp: spent_metadata.spent.milestone_timestamp,
-                })
-            },
-        );
+        let ledger_updates = outputs.into_iter().filter_map(|output| {
+            // Ledger updates
+            output.owning_address().map(|&address| LedgerUpdateDocument {
+                _id: Id {
+                    milestone_index: output.spent_metadata.spent.milestone_index,
+                    output_id: output.output_id(),
+                    is_spent: true,
+                },
+                address,
+                milestone_timestamp: output.spent_metadata.spent.milestone_timestamp,
+            })
+        });
         self.insert_many_ignore_duplicates(ledger_updates, InsertManyOptions::builder().ordered(false).build())
             .await?;
 
@@ -140,25 +135,18 @@ impl LedgerUpdateCollection {
         I: IntoIterator<Item = &'a LedgerOutput>,
         I::IntoIter: Send + Sync,
     {
-        let ledger_updates = outputs.into_iter().filter_map(
-            |LedgerOutput {
-                 output_id,
-                 booked,
-                 output,
-                 ..
-             }| {
-                // Ledger updates
-                output.owning_address().map(|&address| LedgerUpdateDocument {
-                    _id: Id {
-                        milestone_index: booked.milestone_index,
-                        output_id: *output_id,
-                        is_spent: false,
-                    },
-                    address,
-                    milestone_timestamp: booked.milestone_timestamp,
-                })
-            },
-        );
+        let ledger_updates = outputs.into_iter().filter_map(|output| {
+            // Ledger updates
+            output.owning_address().map(|&address| LedgerUpdateDocument {
+                _id: Id {
+                    milestone_index: output.booked.milestone_index,
+                    output_id: output.output_id,
+                    is_spent: false,
+                },
+                address,
+                milestone_timestamp: output.booked.milestone_timestamp,
+            })
+        });
         self.insert_many_ignore_duplicates(ledger_updates, InsertManyOptions::builder().ordered(false).build())
             .await?;
 

src/db/mongodb/collections/outputs/mod.rs

@@ -98,7 +98,7 @@ struct OutputDetails {
 
 impl From<&LedgerOutput> for OutputDocument {
     fn from(rec: &LedgerOutput) -> Self {
-        let address = rec.output.owning_address().copied();
+        let address = rec.owning_address().copied();
         let is_trivial_unlock = rec.output.is_trivial_unlock();
 
         Self {
@@ -141,6 +141,8 @@ impl From<&LedgerOutput> for OutputDocument {
 impl From<&LedgerSpent> for OutputDocument {
     fn from(rec: &LedgerSpent) -> Self {
         let mut res = Self::from(&rec.output);
+        // Update the address as the spending may have changed it
+        res.details.address = rec.owning_address().copied();
         res.metadata.spent_metadata.replace(rec.spent_metadata);
         res
     }
@@ -166,7 +168,7 @@ pub struct OutputWithMetadataResult {
 #[allow(missing_docs)]
 pub struct BalanceResult {
     pub total_balance: String,
-    pub sig_locked_balance: String,
+    pub available_balance: String,
 }
 
 #[derive(Clone, Debug, Default, Deserialize)]
@@ -420,27 +422,53 @@ impl OutputCollection {
     pub async fn get_address_balance(
         &self,
         address: Address,
-        ledger_index: MilestoneIndex,
+        ledger_ms: MilestoneIndexTimestamp,
     ) -> Result<Option<BalanceResult>, Error> {
         self
             .aggregate(
                 [
                     // Look at all (at ledger index o'clock) unspent output documents for the given address.
                     doc! { "$match": {
-                        "details.address": &address,
-                        "metadata.booked.milestone_index": { "$lte": ledger_index },
-                        "metadata.spent_metadata.spent.milestone_index": { "$not": { "$lte": ledger_index } }
+                        "$or": [
+                            { "details.address": &address },
+                            { "output.expiration_unlock_condition.return_address": &address }
+                        ],
+                        "metadata.booked.milestone_index": { "$lte": ledger_ms.milestone_index },
+                        "metadata.spent_metadata.spent.milestone_index": { "$not": { "$lte": ledger_ms.milestone_index } }
                     } },
                     doc! { "$group": {
                         "_id": null,
-                        "total_balance": { "$sum": { "$toDecimal": "$output.amount" } },
-                        "sig_locked_balance": { "$sum": { 
-                            "$cond": [ { "$eq": [ "$details.is_trivial_unlock", true] }, { "$toDecimal": "$output.amount" }, 0 ]
+                        "total_balance": { "$sum": {
+                            "$cond": [
+                                { "$or": [ 
+                                    { "$eq": [ "$details.address", &address ] },
+                                    { "$not": { "$lt": [ "$output.expiration_unlock_condition.timestamp", ledger_ms.milestone_timestamp ] } }
+                                ] },
+                                { "$toDecimal": "$output.amount" }, 0 
+                            ]
+                        } },
+                        "available_balance": { "$sum": {
+                            "$cond": [
+                                { "$or": [
+                                    { "$and": [
+                                        { "$eq": [ "$details.address", &address ] },
+                                        { "$or": [
+                                            { "$eq": [ "$details.is_trivial_unlock", true ] },
+                                            { "$not": { "$lt": [ "$output.timelock_unlock_condition.timestamp", ledger_ms.milestone_timestamp ] } }
+                                        ] }
+                                    ] },
+                                    { "$and": [
+                                        { "$eq": [ "$output.expiration_unlock_condition.return_address", &address ] },
+                                        { "$not": { "$lt": [ "$output.expiration_unlock_condition.timestamp", ledger_ms.milestone_timestamp ] } },
+                                    ] },
+                                ] },
+                                { "$toDecimal": "$output.amount" }, 0 
+                            ]
                         } },
                     } },
                     doc! { "$project": {
                         "total_balance": { "$toString": "$total_balance" },
-                        "sig_locked_balance": { "$toString": "$sig_locked_balance" },
+                        "available_balance": { "$toString": "$available_balance" },
                     } },
                 ],
                 None,

src/model/block/payload/transaction/output/ledger.rs

@@ -30,7 +30,7 @@ impl LedgerOutput {
     }
 
     pub fn owning_address(&self) -> Option<&Address> {
-        self.output.owning_address()
+        self.output.owning_address(None)
     }
 }
 
@@ -53,7 +53,9 @@ impl LedgerSpent {
     }
 
     pub fn owning_address(&self) -> Option<&Address> {
-        self.output.owning_address()
+        self.output
+            .output
+            .owning_address(self.spent_metadata.spent.milestone_timestamp)
     }
 }
 /// The different number of bytes that are used for computing the rent cost.