refactor: make extractor return a specific type

This commit improves how we deal with our types and communication with tycho-client.

The main objective here was to be able to detect any breaking change at compile time instead of runtime. We now deal with specific types to ensure serialization and deserialization are on the same Rust struct `dto::BlockChanges` and equivalent dto for `Command` and `Response`. We also removed the `NormalisedMessage` trait because it was not needed anymore.

Author: zizou0x

tycho-common/src/dto.rs

@@ -18,7 +18,7 @@ use utoipa::{IntoParams, ToSchema};
 use uuid::Uuid;
 
 use crate::{
-    models::{self, Address, ComponentId, StoreKey, StoreVal},
+    models::{self, blockchain::BlockAggregatedChanges, Address, ComponentId, StoreKey, StoreVal},
     serde_primitives::{
         hex_bytes, hex_bytes_option, hex_hashmap_key, hex_hashmap_key_value, hex_hashmap_value,
     },
@@ -145,7 +145,7 @@ pub enum Command {
 }
 
 /// A response sent from the server to the client
-#[derive(Deserialize, Serialize, Debug, PartialEq, Eq)]
+#[derive(Deserialize, Serialize, Debug, PartialEq, Eq, Clone)]
 #[serde(tag = "method", rename_all = "lowercase")]
 pub enum Response {
     NewSubscription { extractor_id: ExtractorIdentity, subscription_id: Uuid },
@@ -154,7 +154,7 @@ pub enum Response {
 
 /// A message sent from the server to the client
 #[allow(clippy::large_enum_variant)]
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Debug, Display, Clone)]
 #[serde(untagged)]
 pub enum WebSocketMessage {
     BlockChanges { subscription_id: Uuid, deltas: BlockChanges },
@@ -365,6 +365,122 @@ impl BlockChanges {
     pub fn n_changes(&self) -> usize {
         self.account_updates.len() + self.state_updates.len()
     }
+
+    pub fn drop_state(&self) -> Self {
+        Self {
+            extractor: self.extractor.clone(),
+            chain: self.chain,
+            block: self.block.clone(),
+            finalized_block_height: self.finalized_block_height,
+            revert: self.revert,
+            new_tokens: self.new_tokens.clone(),
+            account_updates: HashMap::new(),
+            state_updates: HashMap::new(),
+            new_protocol_components: self.new_protocol_components.clone(),
+            deleted_protocol_components: self.deleted_protocol_components.clone(),
+            component_balances: self.component_balances.clone(),
+            account_balances: self.account_balances.clone(),
+            component_tvl: self.component_tvl.clone(),
+            dci_update: self.dci_update.clone(),
+        }
+    }
+}
+
+impl From<models::blockchain::Block> for Block {
+    fn from(value: models::blockchain::Block) -> Self {
+        Self {
+            number: value.number,
+            hash: value.hash,
+            parent_hash: value.parent_hash,
+            chain: value.chain.into(),
+            ts: value.ts,
+        }
+    }
+}
+
+impl From<models::protocol::ComponentBalance> for ComponentBalance {
+    fn from(value: models::protocol::ComponentBalance) -> Self {
+        Self {
+            token: value.token,
+            balance: value.balance,
+            balance_float: value.balance_float,
+            modify_tx: value.modify_tx,
+            component_id: value.component_id,
+        }
+    }
+}
+
+impl From<models::contract::AccountBalance> for AccountBalance {
+    fn from(value: models::contract::AccountBalance) -> Self {
+        Self {
+            account: value.account,
+            token: value.token,
+            balance: value.balance,
+            modify_tx: value.modify_tx,
+        }
+    }
+}
+
+impl From<BlockAggregatedChanges> for BlockChanges {
+    fn from(value: BlockAggregatedChanges) -> Self {
+        Self {
+            extractor: value.extractor,
+            chain: value.chain.into(),
+            block: value.block.into(),
+            finalized_block_height: value.finalized_block_height,
+            revert: value.revert,
+            account_updates: value
+                .account_deltas
+                .into_iter()
+                .map(|(k, v)| (k, v.into()))
+                .collect(),
+            state_updates: value
+                .state_deltas
+                .into_iter()
+                .map(|(k, v)| (k, v.into()))
+                .collect(),
+            new_protocol_components: value
+                .new_protocol_components
+                .into_iter()
+                .map(|(k, v)| (k, v.into()))
+                .collect(),
+            deleted_protocol_components: value
+                .deleted_protocol_components
+                .into_iter()
+                .map(|(k, v)| (k, v.into()))
+                .collect(),
+            component_balances: value
+                .component_balances
+                .into_iter()
+                .map(|(component_id, v)| {
+                    let balances: HashMap<Bytes, ComponentBalance> = v
+                        .into_iter()
+                        .map(|(k, v)| (k, ComponentBalance::from(v)))
+                        .collect();
+                    (component_id, balances.into())
+                })
+                .collect(),
+            account_balances: value
+                .account_balances
+                .into_iter()
+                .map(|(k, v)| {
+                    (
+                        k,
+                        v.into_iter()
+                            .map(|(k, v)| (k, v.into()))
+                            .collect(),
+                    )
+                })
+                .collect(),
+            dci_update: value.dci_update.into(),
+            new_tokens: value
+                .new_tokens
+                .into_iter()
+                .map(|(k, v)| (k, v.into()))
+                .collect(),
+            component_tvl: value.component_tvl,
+        }
+    }
 }
 
 #[derive(PartialEq, Serialize, Deserialize, Clone, Debug, ToSchema)]
@@ -1338,6 +1454,42 @@ pub struct DCIUpdate {
     pub trace_results: HashMap<String, TracingResult>,
 }
 
+impl From<models::blockchain::DCIUpdate> for DCIUpdate {
+    fn from(value: models::blockchain::DCIUpdate) -> Self {
+        Self {
+            new_entrypoints: value
+                .new_entrypoints
+                .into_iter()
+                .map(|(k, v)| {
+                    (
+                        k,
+                        v.into_iter()
+                            .map(|v| v.into())
+                            .collect(),
+                    )
+                })
+                .collect(),
+            new_entrypoint_params: value
+                .new_entrypoint_params
+                .into_iter()
+                .map(|(k, v)| {
+                    (
+                        k,
+                        v.into_iter()
+                            .map(|(params, i)| (params.into(), i))
+                            .collect(),
+                    )
+                })
+                .collect(),
+            trace_results: value
+                .trace_results
+                .into_iter()
+                .map(|(k, v)| (k, v.into()))
+                .collect(),
+        }
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Default, PartialEq, ToSchema, Eq, Hash, Clone)]
 #[serde(deny_unknown_fields)]
 pub struct ComponentTvlRequestBody {

tycho-common/src/models/blockchain.rs

@@ -1,8 +1,4 @@
-use std::{
-    any::Any,
-    collections::{hash_map::Entry, HashMap, HashSet},
-    sync::Arc,
-};
+use std::collections::{hash_map::Entry, HashMap, HashSet};
 
 use chrono::NaiveDateTime;
 use serde::{ser::SerializeStruct, Deserialize, Serialize, Serializer};
@@ -16,8 +12,7 @@ use crate::{
             ComponentBalance, ProtocolChangesWithTx, ProtocolComponent, ProtocolComponentStateDelta,
         },
         token::Token,
-        Address, BlockHash, Chain, ComponentId, EntryPointId, ExtractorIdentity, MergeError,
-        NormalisedMessage, StoreKey,
+        Address, BlockHash, Chain, ComponentId, EntryPointId, MergeError, StoreKey,
     },
     Bytes,
 };
@@ -136,14 +131,9 @@ impl std::fmt::Display for BlockAggregatedChanges {
     }
 }
 
-#[typetag::serde]
-impl NormalisedMessage for BlockAggregatedChanges {
-    fn source(&self) -> ExtractorIdentity {
-        ExtractorIdentity::new(self.chain, &self.extractor)
-    }
-
-    fn drop_state(&self) -> Arc<dyn NormalisedMessage> {
-        Arc::new(Self {
+impl BlockAggregatedChanges {
+    pub fn drop_state(&self) -> Self {
+        Self {
             extractor: self.extractor.clone(),
             chain: self.chain,
             block: self.block.clone(),
@@ -158,11 +148,7 @@ impl NormalisedMessage for BlockAggregatedChanges {
             account_balances: self.account_balances.clone(),
             component_tvl: self.component_tvl.clone(),
             dci_update: self.dci_update.clone(),
-        })
-    }
-
-    fn as_any(&self) -> &dyn Any {
-        self
+        }
     }
 }
 
@@ -176,6 +162,18 @@ impl BlockScoped for BlockAggregatedChanges {
     }
 }
 
+impl From<dto::Block> for Block {
+    fn from(value: dto::Block) -> Self {
+        Self {
+            number: value.number,
+            chain: value.chain.into(),
+            hash: value.hash,
+            parent_hash: value.parent_hash,
+            ts: value.ts,
+        }
+    }
+}
+
 #[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)]
 pub struct DCIUpdate {
     pub new_entrypoints: HashMap<ComponentId, HashSet<EntryPoint>>,

tycho-common/src/models/mod.rs

@@ -3,7 +3,7 @@ pub mod contract;
 pub mod protocol;
 pub mod token;
 
-use std::{collections::HashMap, fmt::Display, str::FromStr, sync::Arc};
+use std::{collections::HashMap, fmt::Display, str::FromStr};
 
 use serde::{Deserialize, Serialize};
 use strum_macros::{Display, EnumString};
@@ -173,6 +173,18 @@ impl std::fmt::Display for ExtractorIdentity {
     }
 }
 
+impl From<ExtractorIdentity> for dto::ExtractorIdentity {
+    fn from(value: ExtractorIdentity) -> Self {
+        dto::ExtractorIdentity { chain: value.chain.into(), name: value.name }
+    }
+}
+
+impl From<dto::ExtractorIdentity> for ExtractorIdentity {
+    fn from(value: dto::ExtractorIdentity) -> Self {
+        Self { chain: value.chain.into(), name: value.name }
+    }
+}
+
 #[derive(Debug, PartialEq, Clone)]
 pub struct ExtractionState {
     pub name: String,
@@ -200,18 +212,6 @@ impl ExtractionState {
     }
 }
 
-// TODO: replace with types from dto on extractor
-#[typetag::serde(tag = "type")]
-pub trait NormalisedMessage:
-    std::any::Any + std::fmt::Debug + std::fmt::Display + Send + Sync + 'static
-{
-    fn source(&self) -> ExtractorIdentity;
-
-    fn drop_state(&self) -> Arc<dyn NormalisedMessage>;
-
-    fn as_any(&self) -> &dyn std::any::Any;
-}
-
 #[derive(PartialEq, Debug, Clone, Default, Deserialize, Serialize)]
 pub enum ImplementationType {
     #[default]

tycho-indexer/src/extractor/mod.rs

@@ -6,10 +6,10 @@ use prost::DecodeError;
 use thiserror::Error;
 use tycho_common::{
     models::{
-        blockchain::{Block, BlockScoped},
+        blockchain::{Block, BlockAggregatedChanges, BlockScoped},
         contract::AccountBalance,
         protocol::ComponentBalance,
-        Address, BlockHash, ExtractorIdentity, MergeError, NormalisedMessage,
+        Address, BlockHash, ExtractorIdentity, MergeError,
     },
     storage::StorageError,
     Bytes,
@@ -77,7 +77,7 @@ pub enum RPCError {
     RequestError(String),
 }
 
-pub type ExtractorMsg = Arc<dyn NormalisedMessage>;
+pub type ExtractorMsg = Arc<BlockAggregatedChanges>;
 
 #[automock]
 #[async_trait]

tycho-indexer/src/extractor/protocol_extractor.rs

@@ -2876,10 +2876,6 @@ mod test_serial_db {
                 .unwrap();
 
 
-            let res = client_msg
-                .as_any()
-                .downcast_ref::<BlockAggregatedChanges>()
-                .expect("not good type");
             let base_ts = db_fixtures::yesterday_midnight().timestamp();
             let block_entity_changes_result = BlockAggregatedChanges {
                 extractor: "native_name".to_string(),
@@ -2959,8 +2955,8 @@ mod test_serial_db {
             };
 
             assert_eq!(
-                res,
-                &block_entity_changes_result
+                *client_msg,
+                block_entity_changes_result
             );
         })
             .await;
@@ -3057,10 +3053,6 @@ mod test_serial_db {
                 .unwrap()
                 .unwrap();
 
-            let res = client_msg
-                .as_any()
-                .downcast_ref::<BlockAggregatedChanges>()
-                .expect("not good type");
 
             let base_ts = db_fixtures::yesterday_midnight().timestamp();
             let account1 = Bytes::from_str("0000000000000000000000000000000000000001").unwrap();
@@ -3165,8 +3157,8 @@ mod test_serial_db {
             };
 
             assert_eq!(
-                res,
-                &block_account_expected
+                *client_msg,
+                block_account_expected
             );
         })
             .await;