diff --git a/Cargo.lock b/Cargo.lock index 127202ed..903c8688 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5144,6 +5144,7 @@ dependencies = [ "tokio", "tokio-postgres", "torii", + "torii-common", "torii-erc1155", "torii-erc20", "torii-erc721", diff --git a/Cargo.toml b/Cargo.toml index 49d627ff..b8596143 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -196,6 +196,7 @@ tokio.workspace = true tonic-reflection.workspace = true tonic-web.workspace = true tonic.workspace = true +torii-common.workspace = true tower-http.workspace = true tracing.workspace = true prost.workspace = true diff --git a/bins/torii-arcade/src/main.rs b/bins/torii-arcade/src/main.rs index fe7d7b00..a9b611d7 100644 --- a/bins/torii-arcade/src/main.rs +++ b/bins/torii-arcade/src/main.rs @@ -29,17 +29,17 @@ use torii_ecs_sink::proto::world::world_server::WorldServer; use torii_ecs_sink::{EcsSink, FILE_DESCRIPTOR_SET as ECS_DESCRIPTOR_SET}; use torii_erc1155::proto::erc1155_server::Erc1155Server; use torii_erc1155::{ - Erc1155Decoder, Erc1155Service, Erc1155Sink, Erc1155Storage, + Erc1155Decoder, Erc1155MetadataCommandHandler, Erc1155Service, Erc1155Sink, Erc1155Storage, FILE_DESCRIPTOR_SET as ERC1155_DESCRIPTOR_SET, }; use torii_erc20::proto::erc20_server::Erc20Server; use torii_erc20::{ - Erc20Decoder, Erc20Service, Erc20Sink, Erc20Storage, + Erc20Decoder, Erc20MetadataCommandHandler, Erc20Service, Erc20Sink, Erc20Storage, FILE_DESCRIPTOR_SET as ERC20_DESCRIPTOR_SET, }; use torii_erc721::proto::erc721_server::Erc721Server; use torii_erc721::{ - Erc721Decoder, Erc721Service, Erc721Sink, Erc721Storage, + Erc721Decoder, Erc721MetadataCommandHandler, Erc721Service, Erc721Sink, Erc721Storage, FILE_DESCRIPTOR_SET as ERC721_DESCRIPTOR_SET, }; use torii_introspect_postgres_sink::processor::IntrospectPgDb; @@ -50,6 +50,11 @@ use torii_sqlite::{is_sqlite_memory_path, sqlite_connect_options}; type StarknetProvider = starknet::providers::jsonrpc::JsonRpcClient; +const TOKEN_COMMAND_QUEUE_SIZE: usize = 4096; +const TOKEN_METADATA_COMMAND_PARALLELISM: usize = 1; +const TOKEN_METADATA_MAX_RETRIES: u8 = 3; +const TOKEN_URI_FETCH_PARALLELISM: usize = 8; + struct ArcadeProjectionPipeline { introspect_sink: Box, arcade_sink: Box, @@ -336,6 +341,7 @@ async fn run_indexer(config: Config) -> Result<()> { let mut torii_config = torii::ToriiConfig::builder() .port(config.port) .database_root(&config.db_dir) + .command_bus_queue_size(TOKEN_COMMAND_QUEUE_SIZE) .etl_concurrency(EtlConcurrencyConfig { max_prefetch_batches: config.max_prefetch_batches, }) @@ -365,18 +371,29 @@ async fn run_indexer(config: Config) -> Result<()> { let mut erc20_grpc_service: Option = None; let mut erc721_grpc_service: Option = None; let mut erc1155_grpc_service: Option = None; + let mut token_uri_services = Vec::new(); if !erc20_addresses.is_empty() { let storage = Arc::new(Erc20Storage::new(&erc20_db_url).await?); let grpc_service = Erc20Service::new(storage.clone()); let sink = Box::new( - Erc20Sink::new(storage) + Erc20Sink::new(storage.clone()) .with_grpc_service(grpc_service.clone()) - .with_balance_tracking(provider.clone()), + .with_balance_tracking(provider.clone()) + .with_metadata_pipeline( + TOKEN_METADATA_COMMAND_PARALLELISM, + TOKEN_COMMAND_QUEUE_SIZE, + TOKEN_METADATA_MAX_RETRIES, + ), ); torii_config = torii_config .add_decoder(Arc::new(Erc20Decoder::new())) - .add_sink_boxed(sink); + .add_sink_boxed(sink) + .with_command_handler(Box::new(Erc20MetadataCommandHandler::new( + provider.clone(), + storage, + TOKEN_METADATA_MAX_RETRIES, + ))); reflection_builder = reflection_builder.register_encoded_file_descriptor_set(ERC20_DESCRIPTOR_SET); let decoder_id = DecoderId::new("erc20"); @@ -389,24 +406,29 @@ async fn run_indexer(config: Config) -> Result<()> { if !erc721_addresses.is_empty() { let storage = Arc::new(Erc721Storage::new(&erc721_db_url).await?); let grpc_service = Erc721Service::new(storage.clone()); - let mut sink = Erc721Sink::new(storage).with_grpc_service(grpc_service.clone()); + let mut sink = Erc721Sink::new(storage.clone()).with_grpc_service(grpc_service.clone()); if config.metadata_mode == MetadataMode::Inline { - let image_cache_dir = Path::new("./data").join("image-cache"); - let (token_uri_sender, _token_uri_service) = TokenUriService::spawn_with_image_cache( + let (token_uri_sender, token_uri_service) = TokenUriService::spawn_with_image_cache( Arc::new(MetadataFetcher::new(provider.clone())), - sink.storage().clone(), - 1024, - 8, - Some(image_cache_dir), - 8, + storage.clone(), + TOKEN_COMMAND_QUEUE_SIZE, + TOKEN_URI_FETCH_PARALLELISM, + Some(Path::new("./data").join("image-cache")), + 4, ); + token_uri_services.push(token_uri_service); sink = sink - .with_metadata_fetching(provider.clone()) + .with_metadata_commands() .with_token_uri_sender(token_uri_sender); } torii_config = torii_config .add_decoder(Arc::new(Erc721Decoder::new())) - .add_sink_boxed(Box::new(sink)); + .add_sink_boxed(Box::new(sink)) + .with_command_handler(Box::new(Erc721MetadataCommandHandler::new( + provider.clone(), + storage, + TOKEN_METADATA_MAX_RETRIES, + ))); reflection_builder = reflection_builder.register_encoded_file_descriptor_set(ERC721_DESCRIPTOR_SET); let decoder_id = DecoderId::new("erc721"); @@ -419,24 +441,30 @@ async fn run_indexer(config: Config) -> Result<()> { if !erc1155_addresses.is_empty() { let storage = Arc::new(Erc1155Storage::new(&erc1155_db_url).await?); let grpc_service = Erc1155Service::new(storage.clone()); - let mut sink = Erc1155Sink::new(storage) + let mut sink = Erc1155Sink::new(storage.clone()) .with_grpc_service(grpc_service.clone()) .with_balance_tracking(provider.clone()); if config.metadata_mode == MetadataMode::Inline { - let image_cache_dir = Path::new("./data").join("image-cache"); - let (token_uri_sender, _token_uri_service) = TokenUriService::spawn_with_image_cache( + let (token_uri_sender, token_uri_service) = TokenUriService::spawn_with_image_cache( Arc::new(MetadataFetcher::new(provider.clone())), - sink.storage().clone(), - 1024, - 8, - Some(image_cache_dir), - 8, + storage.clone(), + TOKEN_COMMAND_QUEUE_SIZE, + TOKEN_URI_FETCH_PARALLELISM, + Some(Path::new("./data").join("image-cache")), + 4, ); - sink = sink.with_token_uri_sender(token_uri_sender); + token_uri_services.push(token_uri_service); + sink = sink + .with_metadata_commands() + .with_token_uri_sender(token_uri_sender); } torii_config = torii_config .add_decoder(Arc::new(Erc1155Decoder::new())) - .add_sink_boxed(Box::new(sink)); + .add_sink_boxed(Box::new(sink)) + .with_command_handler(Box::new(Erc1155MetadataCommandHandler::new( + provider.clone(), + storage, + ))); reflection_builder = reflection_builder.register_encoded_file_descriptor_set(ERC1155_DESCRIPTOR_SET); let decoder_id = DecoderId::new("erc1155"); @@ -558,6 +586,8 @@ async fn run_indexer(config: Config) -> Result<()> { .await .map_err(|e| anyhow::anyhow!("Torii error: {e}"))?; + drop(token_uri_services); + Ok(()) } diff --git a/bins/torii-erc20-synth/src/main.rs b/bins/torii-erc20-synth/src/main.rs index 595146db..401d0eda 100644 --- a/bins/torii-erc20-synth/src/main.rs +++ b/bins/torii-erc20-synth/src/main.rs @@ -6,13 +6,59 @@ use std::fs; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use torii::command::{Command, CommandHandler}; use torii::etl::decoder::ContractFilter; use torii::etl::engine_db::{EngineDb, EngineDbConfig}; use torii::etl::extractor::{Extractor, SyntheticErc20Config, SyntheticErc20Extractor}; use torii::etl::sink::Sink; use torii::etl::{Decoder, DecoderContext}; +use torii_erc20::handlers::FetchErc20MetadataCommand; use torii_erc20::{Erc20Decoder, Erc20Sink, Erc20Storage}; -use torii_runtime_common::sink::{drop_postgres_schemas, initialize_sink}; +use torii_runtime_common::sink::{drop_postgres_schemas, initialize_sink_with_command_handlers}; + +const SYNTH_COMMAND_QUEUE_SIZE: usize = 4096; +const SYNTH_METADATA_COMMAND_PARALLELISM: usize = 1; +const SYNTH_METADATA_MAX_RETRIES: u8 = 1; + +struct SyntheticErc20MetadataCommandHandler { + storage: Arc, +} + +impl SyntheticErc20MetadataCommandHandler { + fn new(storage: Arc) -> Self { + Self { storage } + } + + fn suffix(token: impl std::fmt::LowerHex) -> String { + let token_hex = format!("{token:#x}"); + let trimmed = token_hex.trim_start_matches("0x"); + let start = trimmed.len().saturating_sub(6); + trimmed[start..].to_uppercase() + } +} + +#[torii::async_trait] +impl CommandHandler for SyntheticErc20MetadataCommandHandler { + fn supports(&self, command: &dyn Command) -> bool { + command.as_any().is::() + } + + async fn handle_command(&self, command: Box) -> Result<()> { + let command = command + .into_any() + .downcast::() + .map_err(|_| anyhow::anyhow!("synthetic ERC20 handler received unexpected command"))?; + let command = *command; + + let suffix = Self::suffix(command.token); + let name = format!("Synthetic ERC20 {suffix}"); + let symbol = format!("S{suffix}"); + + self.storage + .upsert_token_metadata(command.token, Some(&name), Some(&symbol), Some(18)) + .await + } +} #[derive(Parser, Debug, Clone)] #[command(name = "torii-erc20-synth")] @@ -202,8 +248,20 @@ async fn main() -> Result<()> { let storage = Arc::new(Erc20Storage::new(&cfg.db_url).await?); - let mut sink = Erc20Sink::new(storage.clone()); - initialize_sink(&mut sink, output_dir.clone()).await?; + let mut sink = Erc20Sink::new(storage.clone()).with_metadata_pipeline( + SYNTH_METADATA_COMMAND_PARALLELISM, + SYNTH_COMMAND_QUEUE_SIZE, + SYNTH_METADATA_MAX_RETRIES, + ); + let sink_runtime = initialize_sink_with_command_handlers( + &mut sink, + output_dir.clone(), + vec![Box::new(SyntheticErc20MetadataCommandHandler::new( + storage.clone(), + ))], + SYNTH_COMMAND_QUEUE_SIZE, + ) + .await?; let synthetic_cfg = SyntheticErc20Config { from_block: cfg.from_block, @@ -294,6 +352,8 @@ async fn main() -> Result<()> { #[cfg(feature = "profiling")] write_flamegraph(&output_dir, guard)?; + sink_runtime.shutdown().await; + tracing::info!( run_id = %run_id, output = %output_dir.display(), diff --git a/bins/torii-erc20/src/main.rs b/bins/torii-erc20/src/main.rs index 1cce7c97..ab7250b3 100644 --- a/bins/torii-erc20/src/main.rs +++ b/bins/torii-erc20/src/main.rs @@ -41,7 +41,14 @@ use torii_runtime_common::database::{backend_from_url_or_path, DatabaseBackend}; // Import from the library crate use torii_erc20::proto::erc20_server::Erc20Server; -use torii_erc20::{Erc20Decoder, Erc20Service, Erc20Sink, Erc20Storage, FILE_DESCRIPTOR_SET}; +use torii_erc20::{ + Erc20Decoder, Erc20MetadataCommandHandler, Erc20Service, Erc20Sink, Erc20Storage, + FILE_DESCRIPTOR_SET, +}; + +const ERC20_METADATA_COMMAND_PARALLELISM: usize = 1; +const ERC20_METADATA_COMMAND_QUEUE_SIZE: usize = 4096; +const ERC20_METADATA_MAX_RETRIES: u8 = 3; #[tokio::main] async fn main() -> Result<()> { @@ -110,7 +117,15 @@ async fn main() -> Result<()> { let grpc_service = Erc20Service::new(storage.clone()); // Create sink with gRPC service for dual publishing - let sink = Box::new(Erc20Sink::new(storage.clone()).with_grpc_service(grpc_service.clone())); + let sink = Box::new( + Erc20Sink::new(storage.clone()) + .with_grpc_service(grpc_service.clone()) + .with_metadata_pipeline( + ERC20_METADATA_COMMAND_PARALLELISM, + ERC20_METADATA_COMMAND_QUEUE_SIZE, + ERC20_METADATA_MAX_RETRIES, + ), + ); // Build gRPC router with Erc20 service and reflection let reflection = tonic_reflection::server::Builder::configure() @@ -126,10 +141,16 @@ async fn main() -> Result<()> { let mut torii_config = torii::ToriiConfig::builder() .port(config.port) .database_root(&db_setup.database_root) + .command_bus_queue_size(ERC20_METADATA_COMMAND_QUEUE_SIZE) .engine_database_url(db_setup.engine_url.clone()) .with_extractor(extractor) .add_decoder(decoder) .add_sink_boxed(sink) + .with_command_handler(Box::new(Erc20MetadataCommandHandler::new( + provider.clone(), + storage.clone(), + ERC20_METADATA_MAX_RETRIES, + ))) .with_grpc_router(grpc_router) .with_custom_reflection(true); // We've added our own reflection diff --git a/bins/torii-introspect-bin/src/main.rs b/bins/torii-introspect-bin/src/main.rs index ae3df9a4..4e16a1d7 100644 --- a/bins/torii-introspect-bin/src/main.rs +++ b/bins/torii-introspect-bin/src/main.rs @@ -16,6 +16,7 @@ use torii::etl::extractor::{ ContractEventConfig, EventExtractor, EventExtractorConfig, Extractor, RetryPolicy, }; use torii::{EtlConcurrencyConfig, ToriiConfigBuilder}; +use torii_common::{MetadataFetcher, TokenUriService}; use torii_config_common::apply_observability_env; use torii_dojo::decoder::DojoDecoder; use torii_dojo::store::postgres::PgStore; @@ -24,17 +25,17 @@ use torii_ecs_sink::proto::world::world_server::WorldServer; use torii_ecs_sink::{EcsSink, FILE_DESCRIPTOR_SET as ECS_DESCRIPTOR_SET}; use torii_erc1155::proto::erc1155_server::Erc1155Server; use torii_erc1155::{ - Erc1155Decoder, Erc1155Service, Erc1155Sink, Erc1155Storage, + Erc1155Decoder, Erc1155MetadataCommandHandler, Erc1155Service, Erc1155Sink, Erc1155Storage, FILE_DESCRIPTOR_SET as ERC1155_DESCRIPTOR_SET, }; use torii_erc20::proto::erc20_server::Erc20Server; use torii_erc20::{ - Erc20Decoder, Erc20Service, Erc20Sink, Erc20Storage, + Erc20Decoder, Erc20MetadataCommandHandler, Erc20Service, Erc20Sink, Erc20Storage, FILE_DESCRIPTOR_SET as ERC20_DESCRIPTOR_SET, }; use torii_erc721::proto::erc721_server::Erc721Server; use torii_erc721::{ - Erc721Decoder, Erc721Service, Erc721Sink, Erc721Storage, + Erc721Decoder, Erc721MetadataCommandHandler, Erc721Service, Erc721Sink, Erc721Storage, FILE_DESCRIPTOR_SET as ERC721_DESCRIPTOR_SET, }; use torii_introspect_postgres_sink::processor::IntrospectPgDb; @@ -46,6 +47,11 @@ type StarknetProvider = starknet::providers::jsonrpc::JsonRpcClient; type ReflectionBuilder = tonic_reflection::server::Builder<'static>; +const TOKEN_COMMAND_QUEUE_SIZE: usize = 4096; +const TOKEN_METADATA_COMMAND_PARALLELISM: usize = 1; +const TOKEN_METADATA_MAX_RETRIES: u8 = 3; +const TOKEN_URI_FETCH_PARALLELISM: usize = 8; + #[derive(Clone, Debug, Default)] struct TokenTargets { erc20: Vec, @@ -150,11 +156,22 @@ async fn configure_token_support( provider: Arc, mut torii_config: ToriiConfigBuilder, mut reflection_builder: ReflectionBuilder, -) -> Result<(ToriiConfigBuilder, ReflectionBuilder, TokenGrpcServices)> { +) -> Result<( + ToriiConfigBuilder, + ReflectionBuilder, + TokenGrpcServices, + Vec, +)> { let mut services = TokenGrpcServices::default(); + let mut token_uri_services = Vec::new(); if !token_targets.has_any() { - return Ok((torii_config, reflection_builder, services)); + return Ok(( + torii_config, + reflection_builder, + services, + token_uri_services, + )); } let db_setup = @@ -169,11 +186,22 @@ async fn configure_token_support( let grpc_service = Erc20Service::new(storage.clone()); let sink = Box::new( - Erc20Sink::new(storage) + Erc20Sink::new(storage.clone()) .with_grpc_service(grpc_service.clone()) - .with_balance_tracking(provider.clone()), + .with_balance_tracking(provider.clone()) + .with_metadata_pipeline( + TOKEN_METADATA_COMMAND_PARALLELISM, + TOKEN_COMMAND_QUEUE_SIZE, + TOKEN_METADATA_MAX_RETRIES, + ), ); - torii_config = torii_config.add_sink_boxed(sink); + torii_config = torii_config + .add_sink_boxed(sink) + .with_command_handler(Box::new(Erc20MetadataCommandHandler::new( + provider.clone(), + storage, + TOKEN_METADATA_MAX_RETRIES, + ))); services.erc20 = Some(grpc_service); reflection_builder = reflection_builder.register_encoded_file_descriptor_set(ERC20_DESCRIPTOR_SET); @@ -187,8 +215,28 @@ async fn configure_token_support( torii_config = torii_config.add_decoder(decoder); let grpc_service = Erc721Service::new(storage.clone()); - let sink = Box::new(Erc721Sink::new(storage).with_grpc_service(grpc_service.clone())); - torii_config = torii_config.add_sink_boxed(sink); + let (token_uri_sender, token_uri_service) = TokenUriService::spawn_with_image_cache( + Arc::new(MetadataFetcher::new(provider.clone())), + storage.clone(), + TOKEN_COMMAND_QUEUE_SIZE, + TOKEN_URI_FETCH_PARALLELISM, + Some(Path::new("./data").join("image-cache")), + 4, + ); + token_uri_services.push(token_uri_service); + let sink = Box::new( + Erc721Sink::new(storage.clone()) + .with_grpc_service(grpc_service.clone()) + .with_metadata_commands() + .with_token_uri_sender(token_uri_sender), + ); + torii_config = torii_config + .add_sink_boxed(sink) + .with_command_handler(Box::new(Erc721MetadataCommandHandler::new( + provider.clone(), + storage, + TOKEN_METADATA_MAX_RETRIES, + ))); services.erc721 = Some(grpc_service); reflection_builder = reflection_builder.register_encoded_file_descriptor_set(ERC721_DESCRIPTOR_SET); @@ -202,18 +250,39 @@ async fn configure_token_support( torii_config = torii_config.add_decoder(decoder); let grpc_service = Erc1155Service::new(storage.clone()); + let (token_uri_sender, token_uri_service) = TokenUriService::spawn_with_image_cache( + Arc::new(MetadataFetcher::new(provider.clone())), + storage.clone(), + TOKEN_COMMAND_QUEUE_SIZE, + TOKEN_URI_FETCH_PARALLELISM, + Some(Path::new("./data").join("image-cache")), + 4, + ); + token_uri_services.push(token_uri_service); let sink = Box::new( - Erc1155Sink::new(storage) + Erc1155Sink::new(storage.clone()) .with_grpc_service(grpc_service.clone()) - .with_balance_tracking(provider), + .with_balance_tracking(provider.clone()) + .with_metadata_commands() + .with_token_uri_sender(token_uri_sender), ); - torii_config = torii_config.add_sink_boxed(sink); + torii_config = torii_config + .add_sink_boxed(sink) + .with_command_handler(Box::new(Erc1155MetadataCommandHandler::new( + provider.clone(), + storage, + ))); services.erc1155 = Some(grpc_service); reflection_builder = reflection_builder.register_encoded_file_descriptor_set(ERC1155_DESCRIPTOR_SET); } - Ok((torii_config, reflection_builder, services)) + Ok(( + torii_config, + reflection_builder, + services, + token_uri_services, + )) } #[tokio::main] @@ -420,6 +489,7 @@ async fn run_with_postgres( let torii_config = torii::ToriiConfig::builder() .port(config.port) + .command_bus_queue_size(TOKEN_COMMAND_QUEUE_SIZE) .with_custom_reflection(true) .etl_concurrency(EtlConcurrencyConfig { max_prefetch_batches: config.max_prefetch_batches, @@ -430,14 +500,15 @@ async fn run_with_postgres( .add_sink_boxed(Box::new(ecs_sink)) .add_sink_boxed(Box::new(sink)); - let (torii_config, reflection_builder, token_services) = configure_token_support( - &token_targets, - token_db_setup.as_ref(), - token_provider, - torii_config, - reflection_builder, - ) - .await?; + let (torii_config, reflection_builder, token_services, _token_uri_services) = + configure_token_support( + &token_targets, + token_db_setup.as_ref(), + token_provider, + torii_config, + reflection_builder, + ) + .await?; let reflection = reflection_builder .build_v1() @@ -557,6 +628,7 @@ async fn run_with_sqlite( let torii_config = torii::ToriiConfig::builder() .port(config.port) + .command_bus_queue_size(TOKEN_COMMAND_QUEUE_SIZE) .with_custom_reflection(true) .etl_concurrency(EtlConcurrencyConfig { max_prefetch_batches: config.max_prefetch_batches, @@ -567,14 +639,15 @@ async fn run_with_sqlite( .add_sink_boxed(Box::new(ecs_sink)) .add_sink_boxed(Box::new(IntrospectSqliteDb::new(pool.clone(), ()))); - let (torii_config, reflection_builder, token_services) = configure_token_support( - &token_targets, - token_db_setup.as_ref(), - token_provider, - torii_config, - reflection_builder, - ) - .await?; + let (torii_config, reflection_builder, token_services, _token_uri_services) = + configure_token_support( + &token_targets, + token_db_setup.as_ref(), + token_provider, + torii_config, + reflection_builder, + ) + .await?; let reflection = reflection_builder .build_v1() diff --git a/bins/torii-introspect-synth/src/main.rs b/bins/torii-introspect-synth/src/main.rs index 1118277b..9078bd37 100644 --- a/bins/torii-introspect-synth/src/main.rs +++ b/bins/torii-introspect-synth/src/main.rs @@ -18,6 +18,7 @@ use std::fs; use std::path::PathBuf; use std::sync::Arc; use std::time::Instant; +use torii::command::CommandBus; use torii::etl::decoder::{ContractFilter, DecoderContext}; use torii::etl::engine_db::{EngineDb, EngineDbConfig}; use torii::etl::extractor::{BlockContext, ExtractionBatch, Extractor, TransactionContext}; @@ -425,10 +426,12 @@ async fn main() -> Result<()> { let started = Instant::now(); let mut sink = IntrospectPgDb::new(pool.clone(), TABLE_SCHEMA); + let command_bus = CommandBus::new(Vec::new(), 1)?; sink.initialize( Arc::new(EventBus::new(Arc::new(SubscriptionManager::new()))), &SinkContext { database_root: output_dir.clone(), + command_bus: command_bus.sender(), }, ) .await?; diff --git a/bins/torii-tokens-synth/Cargo.toml b/bins/torii-tokens-synth/Cargo.toml index a0678e88..00a651f9 100644 --- a/bins/torii-tokens-synth/Cargo.toml +++ b/bins/torii-tokens-synth/Cargo.toml @@ -10,6 +10,7 @@ path = "src/main.rs" [dependencies] torii = { path = "../../" } torii-runtime-common.workspace = true +torii-common.workspace = true torii-erc20 = { path = "../../crates/torii-erc20" } torii-erc721 = { path = "../../crates/torii-erc721" } torii-erc1155 = { path = "../../crates/torii-erc1155" } diff --git a/bins/torii-tokens-synth/src/main.rs b/bins/torii-tokens-synth/src/main.rs index 9af7c995..94c79e67 100644 --- a/bins/torii-tokens-synth/src/main.rs +++ b/bins/torii-tokens-synth/src/main.rs @@ -10,21 +10,226 @@ use std::fs; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use torii::command::{Command, CommandHandler}; use torii::etl::decoder::ContractFilter; use torii::etl::engine_db::{EngineDb, EngineDbConfig}; use torii::etl::extractor::SyntheticExtractor; use torii::etl::sink::Sink; use torii::etl::{Decoder, DecoderContext}; +use torii_common::{TokenUriResult, TokenUriStore}; +use torii_erc1155::handlers::{FetchErc1155MetadataCommand, RefreshErc1155TokenUriCommand}; use torii_erc1155::{ Erc1155Decoder, Erc1155Sink, Erc1155Storage, SyntheticErc1155Config, SyntheticErc1155Extractor, }; +use torii_erc20::handlers::FetchErc20MetadataCommand; use torii_erc20::{ Erc20Decoder, Erc20Sink, Erc20Storage, SyntheticErc20Config, SyntheticErc20Extractor, }; +use torii_erc721::handlers::{FetchErc721MetadataCommand, RefreshErc721TokenUriCommand}; use torii_erc721::{ Erc721Decoder, Erc721Sink, Erc721Storage, SyntheticErc721Config, SyntheticErc721Extractor, }; -use torii_runtime_common::sink::{drop_postgres_schemas, initialize_sink}; +use torii_runtime_common::sink::{drop_postgres_schemas, initialize_sink_with_command_handlers}; + +const SYNTH_COMMAND_QUEUE_SIZE: usize = 4096; +const SYNTH_METADATA_COMMAND_PARALLELISM: usize = 1; +const SYNTH_METADATA_MAX_RETRIES: u8 = 1; + +struct SyntheticErc20MetadataCommandHandler { + storage: Arc, +} + +impl SyntheticErc20MetadataCommandHandler { + fn new(storage: Arc) -> Self { + Self { storage } + } +} + +struct SyntheticErc721MetadataCommandHandler { + storage: Arc, +} + +impl SyntheticErc721MetadataCommandHandler { + fn new(storage: Arc) -> Self { + Self { storage } + } +} + +struct SyntheticErc721TokenUriCommandHandler { + storage: Arc, +} + +impl SyntheticErc721TokenUriCommandHandler { + fn new(storage: Arc) -> Self { + Self { storage } + } +} + +struct SyntheticErc1155MetadataCommandHandler { + storage: Arc, +} + +impl SyntheticErc1155MetadataCommandHandler { + fn new(storage: Arc) -> Self { + Self { storage } + } +} + +struct SyntheticErc1155TokenUriCommandHandler { + storage: Arc, +} + +impl SyntheticErc1155TokenUriCommandHandler { + fn new(storage: Arc) -> Self { + Self { storage } + } +} + +fn token_suffix(token: impl std::fmt::LowerHex) -> String { + let token_hex = format!("{token:#x}"); + let trimmed = token_hex.trim_start_matches("0x"); + let start = trimmed.len().saturating_sub(6); + trimmed[start..].to_uppercase() +} + +fn synthetic_token_uri_metadata(standard: &str, token_id: impl std::fmt::Display) -> String { + format!( + r#"{{"name":"Synthetic {standard} #{token_id}","description":"Synthetic metadata for command bus profiling","attributes":[{{"trait_type":"standard","value":"{standard}"}},{{"trait_type":"token_id","value":"{token_id}"}}]}}"# + ) +} + +#[torii::async_trait] +impl CommandHandler for SyntheticErc20MetadataCommandHandler { + fn supports(&self, command: &dyn Command) -> bool { + command.as_any().is::() + } + + async fn handle_command(&self, command: Box) -> Result<()> { + let command = command + .into_any() + .downcast::() + .map_err(|_| anyhow::anyhow!("synthetic ERC20 handler received unexpected command"))?; + let command = *command; + let suffix = token_suffix(command.token); + let name = format!("Synthetic ERC20 {suffix}"); + let symbol = format!("S{suffix}"); + self.storage + .upsert_token_metadata(command.token, Some(&name), Some(&symbol), Some(18)) + .await + } +} + +#[torii::async_trait] +impl CommandHandler for SyntheticErc721MetadataCommandHandler { + fn supports(&self, command: &dyn Command) -> bool { + command.as_any().is::() + } + + async fn handle_command(&self, command: Box) -> Result<()> { + let command = command + .into_any() + .downcast::() + .map_err(|_| anyhow::anyhow!("synthetic ERC721 handler received unexpected command"))?; + let command = *command; + let suffix = token_suffix(command.token); + let name = format!("Synthetic ERC721 {suffix}"); + let symbol = format!("N{suffix}"); + self.storage + .upsert_token_metadata( + command.token, + Some(&name), + Some(&symbol), + Some(10_000u64.into()), + ) + .await + } +} + +#[torii::async_trait] +impl CommandHandler for SyntheticErc721TokenUriCommandHandler { + fn supports(&self, command: &dyn Command) -> bool { + command.as_any().is::() + } + + async fn handle_command(&self, command: Box) -> Result<()> { + let command = command + .into_any() + .downcast::() + .map_err(|_| { + anyhow::anyhow!("synthetic ERC721 URI handler received unexpected command") + })?; + let command = *command; + + self.storage + .store_token_uri(&TokenUriResult { + contract: command.contract, + token_id: command.token_id, + uri: Some(format!( + "synthetic://erc721/{:#x}/{}", + command.contract, command.token_id + )), + metadata_json: Some(synthetic_token_uri_metadata("ERC721", command.token_id)), + }) + .await + } +} + +#[torii::async_trait] +impl CommandHandler for SyntheticErc1155MetadataCommandHandler { + fn supports(&self, command: &dyn Command) -> bool { + command.as_any().is::() + } + + async fn handle_command(&self, command: Box) -> Result<()> { + let command = command + .into_any() + .downcast::() + .map_err(|_| { + anyhow::anyhow!("synthetic ERC1155 handler received unexpected command") + })?; + let command = *command; + let suffix = token_suffix(command.token); + let name = format!("Synthetic ERC1155 {suffix}"); + let symbol = format!("M{suffix}"); + self.storage + .upsert_token_metadata( + command.token, + Some(&name), + Some(&symbol), + Some(10_000u64.into()), + ) + .await + } +} + +#[torii::async_trait] +impl CommandHandler for SyntheticErc1155TokenUriCommandHandler { + fn supports(&self, command: &dyn Command) -> bool { + command.as_any().is::() + } + + async fn handle_command(&self, command: Box) -> Result<()> { + let command = command + .into_any() + .downcast::() + .map_err(|_| { + anyhow::anyhow!("synthetic ERC1155 URI handler received unexpected command") + })?; + let command = *command; + + self.storage + .store_token_uri(&TokenUriResult { + contract: command.contract, + token_id: command.token_id, + uri: Some(format!( + "synthetic://erc1155/{:#x}/{}", + command.contract, command.token_id + )), + metadata_json: Some(synthetic_token_uri_metadata("ERC1155", command.token_id)), + }) + .await + } +} #[derive(Parser, Debug, Clone)] #[command(name = "torii-tokens-synth")] @@ -131,7 +336,13 @@ struct CommonConfig { output_root: PathBuf, /// Drop and recreate schemas before each run. - #[arg(long, default_value_t = true)] + #[arg( + long, + default_value_t = true, + num_args = 0..=1, + default_missing_value = "true", + action = clap::ArgAction::Set + )] reset_schema: bool, } @@ -298,8 +509,20 @@ async fn run_erc20_profile(config: Erc20ProfileConfig) -> Result<()> { let storage = Arc::new(Erc20Storage::new(&config.common.db_url).await?); - let mut sink = Erc20Sink::new(storage.clone()); - initialize_sink(&mut sink, output_dir.clone()).await?; + let mut sink = Erc20Sink::new(storage.clone()).with_metadata_pipeline( + SYNTH_METADATA_COMMAND_PARALLELISM, + SYNTH_COMMAND_QUEUE_SIZE, + SYNTH_METADATA_MAX_RETRIES, + ); + let sink_runtime = initialize_sink_with_command_handlers( + &mut sink, + output_dir.clone(), + vec![Box::new(SyntheticErc20MetadataCommandHandler::new( + storage.clone(), + ))], + SYNTH_COMMAND_QUEUE_SIZE, + ) + .await?; let synthetic_cfg = SyntheticErc20Config { from_block: config.common.from_block, @@ -383,6 +606,7 @@ async fn run_erc20_profile(config: Erc20ProfileConfig) -> Result<()> { ); write_artifacts(&output_dir, &report)?; + sink_runtime.shutdown().await; tracing::info!( run_id = %run_id, @@ -420,8 +644,19 @@ async fn run_erc721_profile(config: Erc721ProfileConfig) -> Result<()> { let storage = Arc::new(Erc721Storage::new(&config.common.db_url).await?); - let mut sink = Erc721Sink::new(storage.clone()); - initialize_sink(&mut sink, output_dir.clone()).await?; + let mut sink = Erc721Sink::new(storage.clone()) + .with_metadata_commands() + .with_token_uri_commands(); + let sink_runtime = initialize_sink_with_command_handlers( + &mut sink, + output_dir.clone(), + vec![ + Box::new(SyntheticErc721MetadataCommandHandler::new(storage.clone())), + Box::new(SyntheticErc721TokenUriCommandHandler::new(storage.clone())), + ], + SYNTH_COMMAND_QUEUE_SIZE, + ) + .await?; let synthetic_cfg = SyntheticErc721Config { from_block: config.common.from_block, @@ -509,6 +744,7 @@ async fn run_erc721_profile(config: Erc721ProfileConfig) -> Result<()> { ); write_artifacts(&output_dir, &report)?; + sink_runtime.shutdown().await; tracing::info!( run_id = %run_id, @@ -546,8 +782,19 @@ async fn run_erc1155_profile(config: Erc1155ProfileConfig) -> Result<()> { let storage = Arc::new(Erc1155Storage::new(&config.common.db_url).await?); - let mut sink = Erc1155Sink::new(storage.clone()); - initialize_sink(&mut sink, output_dir.clone()).await?; + let mut sink = Erc1155Sink::new(storage.clone()) + .with_metadata_commands() + .with_token_uri_commands(); + let sink_runtime = initialize_sink_with_command_handlers( + &mut sink, + output_dir.clone(), + vec![ + Box::new(SyntheticErc1155MetadataCommandHandler::new(storage.clone())), + Box::new(SyntheticErc1155TokenUriCommandHandler::new(storage.clone())), + ], + SYNTH_COMMAND_QUEUE_SIZE, + ) + .await?; let synthetic_cfg = SyntheticErc1155Config { from_block: config.common.from_block, @@ -637,6 +884,7 @@ async fn run_erc1155_profile(config: Erc1155ProfileConfig) -> Result<()> { ); write_artifacts(&output_dir, &report)?; + sink_runtime.shutdown().await; tracing::info!( run_id = %run_id, diff --git a/bins/torii-tokens/src/main.rs b/bins/torii-tokens/src/main.rs index b65bc53f..89a1ef4b 100644 --- a/bins/torii-tokens/src/main.rs +++ b/bins/torii-tokens/src/main.rs @@ -56,6 +56,7 @@ use torii::etl::extractor::{ }; use torii::etl::identification::ContractRegistry; use torii::EtlConcurrencyConfig; +use torii_common::{MetadataFetcher, TokenUriService}; use torii_config_common::apply_observability_env; use torii_runtime_common::database::resolve_token_db_setup; #[cfg(feature = "profiling")] @@ -64,24 +65,21 @@ use torii_runtime_common::database::DatabaseBackend; // Import from ERC20 library crate use torii_erc20::proto::erc20_server::Erc20Server; use torii_erc20::{ - Erc20Decoder, Erc20Rule, Erc20Service, Erc20Sink, Erc20Storage, + Erc20Decoder, Erc20MetadataCommandHandler, Erc20Rule, Erc20Service, Erc20Sink, Erc20Storage, FILE_DESCRIPTOR_SET as ERC20_DESCRIPTOR_SET, }; -// Import from ERC721 library crate use torii_erc721::proto::erc721_server::Erc721Server; use torii_erc721::{ - Erc721Decoder, Erc721Rule, Erc721Service, Erc721Sink, Erc721Storage, - FILE_DESCRIPTOR_SET as ERC721_DESCRIPTOR_SET, + Erc721Decoder, Erc721MetadataCommandHandler, Erc721Rule, Erc721Service, Erc721Sink, + Erc721Storage, FILE_DESCRIPTOR_SET as ERC721_DESCRIPTOR_SET, }; -use torii_common::{MetadataFetcher, TokenUriService}; - // Import from ERC1155 library crate use torii_erc1155::proto::erc1155_server::Erc1155Server; use torii_erc1155::{ - Erc1155Decoder, Erc1155Rule, Erc1155Service, Erc1155Sink, Erc1155Storage, - FILE_DESCRIPTOR_SET as ERC1155_DESCRIPTOR_SET, + Erc1155Decoder, Erc1155MetadataCommandHandler, Erc1155Rule, Erc1155Service, Erc1155Sink, + Erc1155Storage, FILE_DESCRIPTOR_SET as ERC1155_DESCRIPTOR_SET, }; async fn contracts_from_registry( @@ -518,6 +516,7 @@ async fn run_indexer(config: Config) -> Result<()> { .etl_concurrency(EtlConcurrencyConfig { max_prefetch_batches: config.max_prefetch_batches, }) + .command_bus_queue_size(config.metadata_queue_capacity) .engine_database_url(db_setup.engine_url.clone()) .with_extractor(extractor) .with_contract_identifier(registry); @@ -526,6 +525,7 @@ async fn run_indexer(config: Config) -> Result<()> { let mut erc20_grpc_service: Option = None; let mut erc721_grpc_service: Option = None; let mut erc1155_grpc_service: Option = None; + let mut token_uri_services = Vec::new(); // Global extraction modes create all token infra for runtime auto-discovery. let is_global_mode = @@ -544,6 +544,12 @@ async fn run_indexer(config: Config) -> Result<()> { torii_config = torii_config.add_decoder(decoder); let grpc_service = Erc20Service::new(storage.clone()); + torii_config = + torii_config.with_command_handler(Box::new(Erc20MetadataCommandHandler::new( + provider.clone(), + storage.clone(), + config.metadata_max_retries, + ))); let sink = Box::new( Erc20Sink::new(storage) .with_grpc_service(grpc_service.clone()) @@ -585,19 +591,25 @@ async fn run_indexer(config: Config) -> Result<()> { torii_config = torii_config.add_decoder(decoder); let grpc_service = Erc721Service::new(storage.clone()); + torii_config = + torii_config.with_command_handler(Box::new(Erc721MetadataCommandHandler::new( + provider.clone(), + storage.clone(), + config.metadata_max_retries, + ))); let mut sink = Erc721Sink::new(storage).with_grpc_service(grpc_service.clone()); if effective_metadata_mode == MetadataMode::Inline { - let image_cache_dir = Path::new("./data").join("image-cache"); - let (token_uri_sender, _token_uri_service) = TokenUriService::spawn_with_image_cache( + let (token_uri_sender, token_uri_service) = TokenUriService::spawn_with_image_cache( Arc::new(MetadataFetcher::new(provider.clone())), sink.storage().clone(), config.metadata_queue_capacity, - 8, // max concurrent metadata fetches - Some(image_cache_dir), - 8, // max concurrent image downloads + config.metadata_parallelism.max(1), + Some(Path::new("./data").join("image-cache")), + 4, ); + token_uri_services.push(token_uri_service); sink = sink - .with_metadata_fetching(provider.clone()) + .with_metadata_commands() .with_token_uri_sender(token_uri_sender); } else { tracing::info!("ERC721 metadata fetching disabled for throughput (deferred mode)"); @@ -634,22 +646,25 @@ async fn run_indexer(config: Config) -> Result<()> { torii_config = torii_config.add_decoder(decoder); let grpc_service = Erc1155Service::new(storage.clone()); + torii_config = torii_config.with_command_handler(Box::new( + Erc1155MetadataCommandHandler::new(provider.clone(), storage.clone()), + )); let mut sink = Erc1155Sink::new(storage) .with_grpc_service(grpc_service.clone()) .with_balance_tracking(provider.clone()); if effective_metadata_mode == MetadataMode::Inline { - let erc1155_fetcher = Arc::new(MetadataFetcher::new(provider.clone())); - let image_cache_dir = Path::new("./data").join("image-cache"); - let (erc1155_uri_sender, _erc1155_uri_service) = - TokenUriService::spawn_with_image_cache( - erc1155_fetcher, - sink.storage().clone(), - config.metadata_queue_capacity, - 8, - Some(image_cache_dir), - 8, - ); - sink = sink.with_token_uri_sender(erc1155_uri_sender); + let (token_uri_sender, token_uri_service) = TokenUriService::spawn_with_image_cache( + Arc::new(MetadataFetcher::new(provider.clone())), + sink.storage().clone(), + config.metadata_queue_capacity, + config.metadata_parallelism.max(1), + Some(Path::new("./data").join("image-cache")), + 4, + ); + token_uri_services.push(token_uri_service); + sink = sink + .with_metadata_commands() + .with_token_uri_sender(token_uri_sender); } else { tracing::info!("ERC1155 metadata fetching disabled for throughput (deferred mode)"); } @@ -749,6 +764,8 @@ async fn run_indexer(config: Config) -> Result<()> { .await .map_err(|e| anyhow::anyhow!("Torii error: {e}"))?; + drop(token_uri_services); + tracing::info!("Torii shutdown complete"); #[cfg(feature = "profiling")] diff --git a/crates/torii-common/src/lib.rs b/crates/torii-common/src/lib.rs index 4bbb5573..5a96b8dc 100644 --- a/crates/torii-common/src/lib.rs +++ b/crates/torii-common/src/lib.rs @@ -13,7 +13,8 @@ use starknet::core::types::{Felt, U256}; pub use metadata::{MetadataFetcher, TokenMetadata}; pub use token_uri::{ - TokenStandard, TokenUriRequest, TokenUriResult, TokenUriSender, TokenUriService, TokenUriStore, + process_token_uri_request, TokenStandard, TokenUriRequest, TokenUriResult, TokenUriSender, + TokenUriService, TokenUriStore, }; // ===== Felt conversions ===== diff --git a/crates/torii-common/src/metadata.rs b/crates/torii-common/src/metadata.rs index 681a1420..0fe1de28 100644 --- a/crates/torii-common/src/metadata.rs +++ b/crates/torii-common/src/metadata.rs @@ -5,11 +5,13 @@ //! selectors, felt-encoded strings and ByteArray returns. use starknet::core::codec::Decode; -use starknet::core::types::{BlockId, BlockTag, ByteArray, Felt, FunctionCall, U256}; +use starknet::core::types::{ + requests::CallRequest, BlockId, BlockTag, ByteArray, Felt, FunctionCall, U256, +}; use starknet::core::utils::parse_cairo_short_string; use starknet::macros::selector; use starknet::providers::jsonrpc::{HttpTransport, JsonRpcClient}; -use starknet::providers::Provider; +use starknet::providers::{Provider, ProviderRequestData, ProviderResponseData}; use std::sync::Arc; /// Token metadata (common fields for all ERC standards) @@ -169,6 +171,29 @@ impl MetadataFetcher { None } + /// Fetch `token_uri(token_id)` or `tokenURI(token_id)` for multiple ERC721 NFTs. + pub async fn fetch_token_uri_batch(&self, requests: &[(Felt, Felt)]) -> Vec> { + self.fetch_string_calls_batch( + requests, + &[ + (selector!("token_uri"), true), + (selector!("tokenURI"), true), + (selector!("token_uri"), false), + (selector!("tokenURI"), false), + ], + ) + .await + } + + /// Fetch `uri(token_id)` for multiple ERC1155 tokens. + pub async fn fetch_uri_batch(&self, requests: &[(Felt, Felt)]) -> Vec> { + self.fetch_string_calls_batch( + requests, + &[(selector!("uri"), true), (selector!("uri"), false)], + ) + .await + } + /// Fetch a string value (name or symbol) from a contract. /// /// Tries snake_case first, returns None on failure. @@ -325,6 +350,81 @@ impl MetadataFetcher { .ok() .filter(|s| !s.is_empty()) } + + async fn fetch_string_calls_batch( + &self, + requests: &[(Felt, Felt)], + attempts: &[(Felt, bool)], + ) -> Vec> { + if requests.is_empty() { + return Vec::new(); + } + + let mut results = vec![None; requests.len()]; + let mut unresolved = (0..requests.len()).collect::>(); + + for &(selector, use_u256) in attempts { + if unresolved.is_empty() { + break; + } + + let rpc_requests = unresolved + .iter() + .map(|&idx| { + let (contract, token_id) = requests[idx]; + let calldata = if use_u256 { + vec![token_id, Felt::ZERO] + } else { + vec![token_id] + }; + ProviderRequestData::Call(CallRequest { + request: FunctionCall { + contract_address: contract, + entry_point_selector: selector, + calldata, + }, + block_id: BlockId::Tag(BlockTag::Latest), + }) + }) + .collect::>(); + + let responses = match self.provider.batch_requests(&rpc_requests).await { + Ok(responses) => responses, + Err(error) => { + tracing::warn!( + target: "torii_common::metadata", + batch_size = unresolved.len(), + selector = %format!("{selector:#x}"), + use_u256, + error = %error, + "Failed to batch fetch string calls" + ); + break; + } + }; + + let mut still_unresolved = Vec::with_capacity(unresolved.len()); + for (response_idx, response) in responses.into_iter().enumerate() { + let request_idx = unresolved[response_idx]; + match response { + ProviderResponseData::Call(felts) => { + if let Some(value) = Self::decode_string_result(&felts) { + if !value.is_empty() { + results[request_idx] = Some(value); + continue; + } + } + still_unresolved.push(request_idx); + } + _ => still_unresolved.push(request_idx), + } + } + + unresolved = still_unresolved; + } + + results + } } #[cfg(test)] diff --git a/crates/torii-common/src/token_uri.rs b/crates/torii-common/src/token_uri.rs index d683aeea..71cdd249 100644 --- a/crates/torii-common/src/token_uri.rs +++ b/crates/torii-common/src/token_uri.rs @@ -18,13 +18,11 @@ //! - Raw JSON fallback for inline metadata use starknet::core::types::{Felt, U256}; -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::Duration; -use tokio::sync::OwnedSemaphorePermit; -use tokio::sync::{mpsc, Mutex}; +use tokio::sync::mpsc; use tokio::task::JoinHandle; use crate::MetadataFetcher; @@ -34,7 +32,9 @@ const INITIAL_BACKOFF: Duration = Duration::from_millis(100); const MAX_RETRIES: u32 = 5; const HTTP_TIMEOUT: Duration = Duration::from_secs(10); const IMAGE_FETCH_TIMEOUT: Duration = Duration::from_secs(20); -const DROP_WARN_INTERVAL_MS: u64 = 5_000; +const RPC_BATCH_SIZE_CAP: usize = 128; +const WRITE_BATCH_SIZE: usize = 128; +const WRITE_BATCH_DELAY: Duration = Duration::from_millis(25); /// Token standard hint for URI fetching #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -77,76 +77,100 @@ pub struct TokenUriResult { /// Callback trait for storing fetched token URI results. #[async_trait::async_trait] pub trait TokenUriStore: Send + Sync + 'static { - async fn store_token_uri(&self, result: &TokenUriResult) -> anyhow::Result<()>; + async fn store_token_uris_batch(&self, results: &[TokenUriResult]) -> anyhow::Result<()>; + + async fn store_token_uri(&self, result: &TokenUriResult) -> anyhow::Result<()> { + self.store_token_uris_batch(std::slice::from_ref(result)) + .await + } } /// Handle to send requests to the token URI service. #[derive(Clone)] pub struct TokenUriSender { - tx: mpsc::Sender, - dropped_count: Arc, - last_drop_warn_ms: Arc, + tx: mpsc::UnboundedSender, } impl TokenUriSender { - fn now_ms() -> u64 { - use std::time::{SystemTime, UNIX_EPOCH}; - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) - } - - fn maybe_flush_drop_warning(&self) { - let now_ms = Self::now_ms(); - let last = self.last_drop_warn_ms.load(Ordering::Relaxed); - if now_ms.saturating_sub(last) < DROP_WARN_INTERVAL_MS { - return; - } - if self - .last_drop_warn_ms - .compare_exchange(last, now_ms, Ordering::Relaxed, Ordering::Relaxed) - .is_err() - { - return; - } - let dropped = self.dropped_count.swap(0, Ordering::Relaxed); - if dropped > 0 { + /// Queue a token URI fetch request. + pub fn request_update(&self, request: TokenUriRequest) -> bool { + if let Err(_error) = self.tx.send(request) { tracing::warn!( target: "torii_common::token_uri", - dropped, - interval_ms = DROP_WARN_INTERVAL_MS, - "Dropping token URI requests: queue is full" + "Failed to send token URI request (channel closed)" ); - } - } - - /// Queue a token URI fetch request. - pub fn request_update(&self, request: TokenUriRequest) { - match self.tx.try_send(request) { - Ok(()) => {} - Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { - self.dropped_count.fetch_add(1, Ordering::Relaxed); - self.maybe_flush_drop_warning(); - } - Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { - tracing::warn!( - target: "torii_common::token_uri", - "Failed to send token URI request (channel closed)" - ); - } + false + } else { + true } } /// Queue updates for a batch of token IDs on the same contract. - pub fn request_batch(&self, contract: Felt, token_ids: &[U256], standard: TokenStandard) { + pub fn request_batch( + &self, + contract: Felt, + token_ids: &[U256], + standard: TokenStandard, + ) -> usize { + let mut accepted = 0; for &token_id in token_ids { - self.request_update(TokenUriRequest { + accepted += usize::from(self.request_update(TokenUriRequest { contract, token_id, standard, - }); + })); } + accepted + } +} + +struct RequestBacklog { + queued: HashMap, + ready: VecDeque, +} + +impl RequestBacklog { + fn with_capacity(capacity: usize) -> Self { + Self { + queued: HashMap::with_capacity(capacity), + ready: VecDeque::with_capacity(capacity), + } + } + + fn enqueue(&mut self, request: TokenUriRequest) { + let key = TaskKey { + contract: request.contract, + token_id: request.token_id, + }; + + if let Some(existing) = self.queued.get_mut(&key) { + *existing = request; + return; + } + + self.ready.push_back(key.clone()); + self.queued.insert(key, request); + } + + fn pop_batch(&mut self, batch_size: usize) -> Vec { + let mut batch = Vec::with_capacity(batch_size.min(self.ready.len())); + while batch.len() < batch_size { + let Some(key) = self.ready.pop_front() else { + break; + }; + if let Some(request) = self.queued.remove(&key) { + batch.push(request); + } + } + batch + } + + fn is_empty(&self) -> bool { + self.queued.is_empty() + } + + fn len(&self) -> usize { + self.queued.len() } } @@ -161,6 +185,12 @@ struct ImageCacheConfig { max_concurrent: usize, } +#[derive(Debug)] +struct BufferedTokenUriResult { + result: TokenUriResult, + image_uri: Option, +} + impl TokenUriService { /// Spawn the token URI service. /// @@ -187,163 +217,222 @@ impl TokenUriService { image_cache_dir: Option, image_max_concurrent: usize, ) -> (TokenUriSender, Self) { - let (tx, rx) = mpsc::channel(buffer_size); + let (tx, rx) = mpsc::unbounded_channel(); let image_cache = image_cache_dir.map(|dir| ImageCacheConfig { dir, max_concurrent: image_max_concurrent.max(1), }); - let handle = tokio::spawn(Self::run(rx, fetcher, store, max_concurrent, image_cache)); - let sender = TokenUriSender { - tx, - dropped_count: Arc::new(AtomicU64::new(0)), - last_drop_warn_ms: Arc::new(AtomicU64::new(0)), - }; + let handle = tokio::spawn(Self::run( + rx, + fetcher, + store, + buffer_size.max(1), + max_concurrent.max(1), + image_cache, + )); + let sender = TokenUriSender { tx }; (sender, Self { handle }) } /// Main processing loop. async fn run( - mut rx: mpsc::Receiver, + mut rx: mpsc::UnboundedReceiver, fetcher: Arc, store: Arc, + backlog_capacity: usize, max_concurrent: usize, image_cache: Option, ) { - let in_flight: Arc>>> = - Arc::new(Mutex::new(HashMap::new())); - let semaphore = Arc::new(tokio::sync::Semaphore::new(max_concurrent)); - let image_semaphore = image_cache - .as_ref() - .map(|cfg| Arc::new(tokio::sync::Semaphore::new(cfg.max_concurrent))); + let (result_tx, result_rx) = mpsc::channel::(WRITE_BATCH_SIZE * 2); + let writer = tokio::spawn(Self::run_writer(result_rx, store, image_cache.clone())); + let mut backlog = RequestBacklog::with_capacity(backlog_capacity); + let rpc_batch_size = backlog_capacity.clamp(1, RPC_BATCH_SIZE_CAP); + + loop { + if backlog.is_empty() { + match rx.recv().await { + Some(request) => backlog.enqueue(request), + None => break, + } + } - while let Some(request) = rx.recv().await { - let key = TaskKey { - contract: request.contract, - token_id: request.token_id, - }; + Self::drain_pending_requests(&mut rx, &mut backlog); + + while !backlog.is_empty() { + let batch = backlog.pop_batch(rpc_batch_size); + let resolved = + Self::resolve_token_uri_requests_batch(&fetcher, batch, max_concurrent).await; - let mut tasks = in_flight.lock().await; + for buffered in resolved { + if let Err(error) = result_tx.send(buffered).await { + tracing::warn!( + target: "torii_common::token_uri", + error = %error, + "Failed to enqueue resolved token URI result" + ); + break; + } + } - // Cancel previous task for same key - if let Some(old_handle) = tasks.remove(&key) { - old_handle.abort(); tracing::debug!( target: "torii_common::token_uri", - contract = %format!("{:#x}", key.contract), - token_id = %key.token_id, - "Cancelled previous fetch (superseded)" + queued = backlog.len(), + rpc_batch_size, + "Processed token URI queue batch" ); + + Self::drain_pending_requests(&mut rx, &mut backlog); } + } - let fetcher = fetcher.clone(); - let store = store.clone(); - let in_flight = in_flight.clone(); - let sem = semaphore.clone(); - let task_key = key.clone(); - let image_cache = image_cache.clone(); - let image_semaphore = image_semaphore.clone(); - - // Apply backpressure before spawning to keep task count bounded. - let permit: OwnedSemaphorePermit = match sem.acquire_owned().await { - Ok(p) => p, - Err(_) => break, - }; + drop(result_tx); + let _ = writer.await; - let handle = tokio::spawn(async move { - let _permit = permit; + tracing::info!( + target: "torii_common::token_uri", + "Token URI service shutting down" + ); + } - // Fetch the URI from chain - let uri = fetch_token_uri_with_retry( - &fetcher, - request.contract, - request.token_id, - request.standard, + async fn run_writer( + mut rx: mpsc::Receiver, + store: Arc, + image_cache: Option, + ) { + let image_semaphore = image_cache + .as_ref() + .map(|cfg| Arc::new(tokio::sync::Semaphore::new(cfg.max_concurrent))); + let mut pending: HashMap = HashMap::new(); + let mut closed = false; + + loop { + if closed { + if pending.is_empty() { + break; + } + Self::flush_results( + &store, + &mut pending, + image_cache.clone(), + image_semaphore.clone(), ) .await; + break; + } - // Apply ERC1155 {id} substitution - let uri = uri.map(|u| { - if request.standard == TokenStandard::Erc1155 { - let token_id_hex = format!("{:064x}", request.token_id); - u.replace("{id}", &token_id_hex) - } else { - u - } - }); - - // Resolve URI to JSON metadata - let metadata_json = if let Some(ref uri_str) = uri { - if uri_str.is_empty() { - None - } else { - resolve_metadata(uri_str).await - } - } else { - None - }; - - let result = TokenUriResult { - contract: request.contract, - token_id: request.token_id, - uri, - metadata_json, - }; - - if let Err(e) = store.store_token_uri(&result).await { - tracing::warn!( - target: "torii_common::token_uri", - contract = %format!("{:#x}", request.contract), - token_id = %request.token_id, - error = %e, - "Failed to store token URI result" + if pending.is_empty() { + if let Some(result) = rx.recv().await { + pending.insert( + TaskKey { + contract: result.result.contract, + token_id: result.result.token_id, + }, + result, ); } else { - tracing::debug!( - target: "torii_common::token_uri", - contract = %format!("{:#x}", request.contract), - token_id = %request.token_id, - uri = ?result.uri, - has_json = result.metadata_json.is_some(), - "Stored token URI" - ); + closed = true; + continue; } - - if let (Some(cfg), Some(meta)) = - (image_cache.clone(), result.metadata_json.as_ref()) - { - if let Some(image_uri) = extract_image_uri(meta) { - let contract = request.contract; - let token_id = request.token_id; - let _image_permit = match image_semaphore { - Some(sema) => sema.acquire_owned().await.ok(), - None => None, - }; - - if let Err(e) = - cache_image_locally(&cfg.dir, contract, token_id, &image_uri).await - { - tracing::debug!( - target: "torii_common::token_uri", - contract = %format!("{:#x}", contract), - token_id = %token_id, - image_uri = %image_uri, - error = %e, - "Failed to cache image locally" - ); + } else { + let delay = tokio::time::sleep(WRITE_BATCH_DELAY); + tokio::pin!(delay); + tokio::select! { + maybe_result = rx.recv() => { + match maybe_result { + Some(result) => { + pending.insert( + TaskKey { + contract: result.result.contract, + token_id: result.result.token_id, + }, + result, + ); + } + None => closed = true, } } + () = &mut delay => { + Self::flush_results(&store, &mut pending, image_cache.clone(), image_semaphore.clone()).await; + } } + } - in_flight.lock().await.remove(&task_key); - }); + if pending.len() >= WRITE_BATCH_SIZE { + Self::flush_results( + &store, + &mut pending, + image_cache.clone(), + image_semaphore.clone(), + ) + .await; + } + } + } - tasks.insert(key, handle); + async fn flush_results( + store: &Arc, + pending: &mut HashMap, + image_cache: Option, + image_semaphore: Option>, + ) { + if pending.is_empty() { + return; } - tracing::info!( - target: "torii_common::token_uri", - "Token URI service shutting down" - ); + let drained = pending + .drain() + .map(|(_, result)| result) + .collect::>(); + let results = drained + .iter() + .map(|buffered| buffered.result.clone()) + .collect::>(); + + if let Err(error) = store.store_token_uris_batch(&results).await { + tracing::warn!( + target: "torii_common::token_uri", + batch_size = results.len(), + error = %error, + "Failed to store token URI batch" + ); + return; + } + + for buffered in drained { + tracing::debug!( + target: "torii_common::token_uri", + contract = %format!("{:#x}", buffered.result.contract), + token_id = %buffered.result.token_id, + uri = ?buffered.result.uri, + has_json = buffered.result.metadata_json.is_some(), + "Stored token URI" + ); + + if let (Some(cfg), Some(image_uri)) = (image_cache.clone(), buffered.image_uri) { + let contract = buffered.result.contract; + let token_id = buffered.result.token_id; + let semaphore = image_semaphore.clone(); + tokio::spawn(async move { + let _image_permit = match semaphore { + Some(sema) => sema.acquire_owned().await.ok(), + None => None, + }; + + if let Err(error) = + cache_image_locally(&cfg.dir, contract, token_id, &image_uri).await + { + tracing::debug!( + target: "torii_common::token_uri", + contract = %format!("{:#x}", contract), + token_id = %token_id, + image_uri = %image_uri, + error = %error, + "Failed to cache image locally" + ); + } + }); + } + } } /// Wait for the service to finish. @@ -355,6 +444,189 @@ impl TokenUriService { pub fn abort(self) { self.handle.abort(); } + + fn drain_pending_requests( + rx: &mut mpsc::UnboundedReceiver, + backlog: &mut RequestBacklog, + ) { + while let Ok(request) = rx.try_recv() { + backlog.enqueue(request); + } + } + + async fn resolve_token_uri_requests_batch( + fetcher: &MetadataFetcher, + requests: Vec, + max_concurrent: usize, + ) -> Vec { + if requests.is_empty() { + return Vec::new(); + } + + let mut raw_uris = vec![None; requests.len()]; + let mut erc721_positions = Vec::new(); + let mut erc1155_positions = Vec::new(); + let mut erc721_requests = Vec::new(); + let mut erc1155_requests = Vec::new(); + + for (idx, request) in requests.iter().enumerate() { + let token_id = Felt::from(request.token_id.low()); + match request.standard { + TokenStandard::Erc721 => { + erc721_positions.push(idx); + erc721_requests.push((request.contract, token_id)); + } + TokenStandard::Erc1155 => { + erc1155_positions.push(idx); + erc1155_requests.push((request.contract, token_id)); + } + } + } + + let erc721_results = fetcher.fetch_token_uri_batch(&erc721_requests).await; + for (position, uri) in erc721_positions.into_iter().zip(erc721_results) { + raw_uris[position] = uri; + } + + let erc1155_results = fetcher.fetch_uri_batch(&erc1155_requests).await; + for (position, uri) in erc1155_positions.into_iter().zip(erc1155_results) { + raw_uris[position] = uri; + } + + let metadata_parallelism = max_concurrent.max(1); + let mut next_idx = 0usize; + let mut join_set = tokio::task::JoinSet::new(); + let mut resolved = (0..requests.len()).map(|_| None).collect::>(); + + while next_idx < requests.len() || !join_set.is_empty() { + while next_idx < requests.len() && join_set.len() < metadata_parallelism { + let request = requests[next_idx].clone(); + let raw_uri = raw_uris[next_idx].clone(); + join_set.spawn(async move { + let result = resolve_token_uri_from_uri(request, raw_uri).await; + let image_uri = result.metadata_json.as_deref().and_then(extract_image_uri); + (next_idx, BufferedTokenUriResult { result, image_uri }) + }); + next_idx += 1; + } + + let Some(joined) = join_set.join_next().await else { + break; + }; + match joined { + Ok((idx, buffered)) => resolved[idx] = Some(buffered), + Err(error) => { + tracing::warn!( + target: "torii_common::token_uri", + error = %error, + "Token URI metadata task failed" + ); + } + } + } + + resolved.into_iter().flatten().collect() + } +} + +pub async fn process_token_uri_request( + fetcher: &MetadataFetcher, + store: &S, + request: &TokenUriRequest, + image_cache_dir: Option<&Path>, +) -> anyhow::Result<()> { + let result = resolve_token_uri_request(fetcher, request).await; + + store.store_token_uri(&result).await?; + + if let (Some(root_dir), Some(meta)) = (image_cache_dir, result.metadata_json.as_ref()) { + if let Some(image_uri) = extract_image_uri(meta) { + if let Err(error) = + cache_image_locally(root_dir, request.contract, request.token_id, &image_uri).await + { + tracing::debug!( + target: "torii_common::token_uri", + contract = %format!("{:#x}", request.contract), + token_id = %request.token_id, + image_uri = %image_uri, + error = %error, + "Failed to cache image locally" + ); + } + } + } + + Ok(()) +} + +async fn resolve_token_uri_request( + fetcher: &MetadataFetcher, + request: &TokenUriRequest, +) -> TokenUriResult { + let uri = fetch_token_uri_with_retry( + fetcher, + request.contract, + request.token_id, + request.standard, + ) + .await; + + let uri = uri.map(|u| { + if request.standard == TokenStandard::Erc1155 { + let token_id_hex = format!("{:064x}", request.token_id); + u.replace("{id}", &token_id_hex) + } else { + u + } + }); + + let metadata_json = if let Some(ref uri_str) = uri { + if uri_str.is_empty() { + None + } else { + resolve_metadata(uri_str).await + } + } else { + None + }; + + TokenUriResult { + contract: request.contract, + token_id: request.token_id, + uri, + metadata_json, + } +} + +async fn resolve_token_uri_from_uri( + request: TokenUriRequest, + raw_uri: Option, +) -> TokenUriResult { + let uri = raw_uri.map(|value| { + if request.standard == TokenStandard::Erc1155 { + let token_id_hex = format!("{:064x}", request.token_id); + value.replace("{id}", &token_id_hex) + } else { + value + } + }); + + let metadata_json = if let Some(ref uri_str) = uri { + if uri_str.is_empty() { + None + } else { + resolve_metadata(uri_str).await + } + } else { + None + }; + + TokenUriResult { + contract: request.contract, + token_id: request.token_id, + uri, + metadata_json, + } } fn extract_image_uri(metadata_json: &str) -> Option { @@ -722,6 +994,24 @@ fn sanitize_json_string(s: &str) -> String { } // Inside a string + if c == '\n' { + result.push_str("\\n"); + backslash_count = 0; + continue; + } + + if c == '\r' { + result.push_str("\\r"); + backslash_count = 0; + continue; + } + + if c == '\t' { + result.push_str("\\t"); + backslash_count = 0; + continue; + } + if c == '\\' { backslash_count += 1; result.push('\\'); @@ -803,6 +1093,15 @@ mod tests { assert!(!sanitized.contains('\x02')); } + #[test] + fn test_sanitize_json_string_escapes_raw_newlines_inside_strings() { + let input = "{\"description\":\"line one\nline two\"}"; + let sanitized = sanitize_json_string(input); + assert_eq!(sanitized, "{\"description\":\"line one\\nline two\"}"); + let parsed: serde_json::Value = serde_json::from_str(&sanitized).unwrap(); + assert_eq!(parsed["description"], "line one\nline two"); + } + #[test] fn test_resolve_data_uri_base64() { let uri = "data:application/json;base64,eyJuYW1lIjoidGVzdCJ9"; @@ -833,4 +1132,49 @@ mod tests { let result = uri.replace("{id}", &token_id_hex); assert!(result.contains("000000000000000000000000000000000000000000000000000000000000002a")); } + + #[test] + fn request_backlog_deduplicates_same_token() { + let contract = Felt::from_hex_unchecked("0x123"); + let mut backlog = RequestBacklog::with_capacity(8); + backlog.enqueue(TokenUriRequest { + contract, + token_id: U256::from(1u64), + standard: TokenStandard::Erc721, + }); + backlog.enqueue(TokenUriRequest { + contract, + token_id: U256::from(1u64), + standard: TokenStandard::Erc1155, + }); + + assert_eq!(backlog.len(), 1); + let batch = backlog.pop_batch(8); + assert_eq!(batch.len(), 1); + assert_eq!(batch[0].standard, TokenStandard::Erc1155); + } + + #[test] + fn request_backlog_removes_items_when_popped() { + let contract = Felt::from_hex_unchecked("0x456"); + let mut backlog = RequestBacklog::with_capacity(8); + backlog.enqueue(TokenUriRequest { + contract, + token_id: U256::from(1u64), + standard: TokenStandard::Erc721, + }); + backlog.enqueue(TokenUriRequest { + contract, + token_id: U256::from(2u64), + standard: TokenStandard::Erc721, + }); + + let first = backlog.pop_batch(1); + assert_eq!(first.len(), 1); + assert_eq!(backlog.len(), 1); + + let second = backlog.pop_batch(8); + assert_eq!(second.len(), 1); + assert!(backlog.is_empty()); + } } diff --git a/crates/torii-erc1155/src/grpc_service.rs b/crates/torii-erc1155/src/grpc_service.rs index 9c28701a..d9b99c74 100644 --- a/crates/torii-erc1155/src/grpc_service.rs +++ b/crates/torii-erc1155/src/grpc_service.rs @@ -199,7 +199,7 @@ impl Erc1155Service { let uri_rows = self .storage - .get_token_uris_by_contract(contract) + .get_token_uris_batch(contract, &result.token_ids) .await .map_err(|e| Status::internal(format!("Query failed: {e}")))?; let by_token_id: HashMap, (Option, Option)> = uri_rows diff --git a/crates/torii-erc1155/src/handlers.rs b/crates/torii-erc1155/src/handlers.rs new file mode 100644 index 00000000..f82ae36a --- /dev/null +++ b/crates/torii-erc1155/src/handlers.rs @@ -0,0 +1,209 @@ +use anyhow::Result; +use async_trait::async_trait; +use prost::Message; +use prost_types::Any; +use starknet::core::types::Felt; +use starknet::providers::jsonrpc::{HttpTransport, JsonRpcClient}; +use std::collections::HashSet; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; +use torii::command::CommandHandler; +use torii::etl::sink::EventBus; +use torii::UpdateType; +use torii_common::{process_token_uri_request, u256_to_bytes, MetadataFetcher, TokenUriRequest}; + +use crate::proto; +use crate::storage::Erc1155Storage; + +#[derive(Debug, Clone)] +pub struct FetchErc1155MetadataCommand { + pub token: Felt, +} + +#[derive(Debug, Clone)] +pub struct RefreshErc1155TokenUriCommand { + pub contract: Felt, + pub token_id: starknet::core::types::U256, +} + +pub struct Erc1155MetadataCommandHandler { + fetcher: Arc, + storage: Arc, + event_bus: Mutex>>, + in_flight: Mutex>, +} + +impl Erc1155MetadataCommandHandler { + pub fn new(provider: Arc>, storage: Arc) -> Self { + Self { + fetcher: Arc::new(MetadataFetcher::new(provider)), + storage, + event_bus: Mutex::new(None), + in_flight: Mutex::new(HashSet::new()), + } + } + + fn matches_metadata_filters( + metadata: &proto::TokenMetadataEntry, + filters: &std::collections::HashMap, + ) -> bool { + if filters.is_empty() { + return true; + } + + if let Some(token_filter) = filters.get("token") { + let token_hex = format!("0x{}", hex::encode(&metadata.token)); + if !token_hex.eq_ignore_ascii_case(token_filter) { + return false; + } + } + + true + } +} + +#[async_trait] +impl CommandHandler for Erc1155MetadataCommandHandler { + fn supports(&self, command: &dyn torii::command::Command) -> bool { + command.as_any().is::() + } + + fn attach_event_bus(&self, event_bus: Arc) { + *self.event_bus.lock().unwrap() = Some(event_bus); + } + + async fn handle_command(&self, command: Box) -> Result<()> { + let command = command + .into_any() + .downcast::() + .map_err(|_| anyhow::anyhow!("erc1155 metadata handler received unexpected command"))?; + let command = *command; + + { + let mut in_flight = self.in_flight.lock().unwrap(); + if !in_flight.insert(command.token) { + tracing::debug!( + target: "torii_erc1155::handlers", + token = %format!("{:#x}", command.token), + "Skipping duplicate in-flight ERC1155 metadata command" + ); + return Ok(()); + } + } + + let meta = self.fetcher.fetch_erc1155_metadata(command.token).await; + let result: Result<()> = async { + self.storage + .upsert_token_metadata( + command.token, + meta.name.as_deref(), + meta.symbol.as_deref(), + meta.total_supply, + ) + .await?; + + let event_bus = self.event_bus.lock().unwrap().clone(); + if let Some(event_bus) = event_bus { + let meta_entry = proto::TokenMetadataEntry { + token: command.token.to_bytes_be().to_vec(), + name: meta.name, + symbol: meta.symbol, + total_supply: meta.total_supply.map(u256_to_bytes), + }; + + let mut buf = Vec::new(); + meta_entry.encode(&mut buf)?; + let any = Any { + type_url: "type.googleapis.com/torii.sinks.erc1155.TokenMetadataEntry" + .to_string(), + value: buf, + }; + + event_bus.publish_protobuf( + "erc1155.metadata", + "erc1155.metadata", + &any, + &meta_entry, + UpdateType::Created, + Self::matches_metadata_filters, + ); + } + + Ok(()) + } + .await; + + self.in_flight.lock().unwrap().remove(&command.token); + result + } +} + +pub struct Erc1155TokenUriCommandHandler { + fetcher: Arc, + storage: Arc, + image_cache_dir: Option, + in_flight: Mutex>, +} + +impl Erc1155TokenUriCommandHandler { + pub fn new( + provider: Arc>, + storage: Arc, + image_cache_dir: Option, + ) -> Self { + Self { + fetcher: Arc::new(MetadataFetcher::new(provider)), + storage, + image_cache_dir, + in_flight: Mutex::new(HashSet::new()), + } + } +} + +#[async_trait] +impl CommandHandler for Erc1155TokenUriCommandHandler { + fn supports(&self, command: &dyn torii::command::Command) -> bool { + command.as_any().is::() + } + + async fn handle_command(&self, command: Box) -> Result<()> { + let command = command + .into_any() + .downcast::() + .map_err(|_| { + anyhow::anyhow!("erc1155 token uri handler received unexpected command") + })?; + let command = *command; + + { + let mut in_flight = self.in_flight.lock().unwrap(); + if !in_flight.insert((command.contract, command.token_id)) { + tracing::debug!( + target: "torii_erc1155::handlers", + contract = %format!("{:#x}", command.contract), + token_id = %command.token_id, + "Skipping duplicate in-flight ERC1155 token URI command" + ); + return Ok(()); + } + } + + let result = process_token_uri_request( + self.fetcher.as_ref(), + self.storage.as_ref(), + &TokenUriRequest { + contract: command.contract, + token_id: command.token_id, + standard: torii_common::TokenStandard::Erc1155, + }, + self.image_cache_dir.as_deref(), + ) + .await; + + self.in_flight + .lock() + .unwrap() + .remove(&(command.contract, command.token_id)); + result + } +} diff --git a/crates/torii-erc1155/src/lib.rs b/crates/torii-erc1155/src/lib.rs index 6f2f5709..23b94621 100644 --- a/crates/torii-erc1155/src/lib.rs +++ b/crates/torii-erc1155/src/lib.rs @@ -39,6 +39,7 @@ pub mod balance_fetcher; pub mod decoder; pub mod grpc_service; +pub mod handlers; pub mod identification; pub mod sink; pub mod storage; @@ -56,6 +57,7 @@ pub const FILE_DESCRIPTOR_SET: &[u8] = include_bytes!("generated/erc1155_descrip pub use balance_fetcher::{Erc1155BalanceFetchRequest, Erc1155BalanceFetcher}; pub use decoder::{Erc1155Decoder, OperatorApproval, TransferBatch, TransferSingle, UriUpdate}; pub use grpc_service::Erc1155Service; +pub use handlers::{Erc1155MetadataCommandHandler, Erc1155TokenUriCommandHandler}; pub use identification::Erc1155Rule; pub use sink::Erc1155Sink; pub use storage::{ diff --git a/crates/torii-erc1155/src/sink.rs b/crates/torii-erc1155/src/sink.rs index a4195535..a3b1b1e8 100644 --- a/crates/torii-erc1155/src/sink.rs +++ b/crates/torii-erc1155/src/sink.rs @@ -17,6 +17,7 @@ use crate::decoder::{ TransferSingle as DecodedTransferSingle, UriUpdate as DecodedUriUpdate, }; use crate::grpc_service::Erc1155Service; +use crate::handlers::{FetchErc1155MetadataCommand, RefreshErc1155TokenUriCommand}; use crate::proto; use crate::storage::{Erc1155Storage, OperatorApprovalData, TokenTransferData, TokenUriData}; use anyhow::Result; @@ -28,13 +29,12 @@ use starknet::core::types::{Felt, U256}; use starknet::providers::jsonrpc::{HttpTransport, JsonRpcClient}; use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; +use torii::command::CommandBusSender; use torii::etl::sink::{EventBus, TopicInfo}; use torii::etl::{Envelope, ExtractionBatch, Sink, TypeId}; use torii::grpc::UpdateType; -use torii_common::{ - u256_to_bytes, MetadataFetcher, TokenStandard, TokenUriRequest, TokenUriSender, -}; +use torii_common::{u256_to_bytes, TokenStandard, TokenUriRequest, TokenUriSender}; /// Default threshold for "live" detection: 100 blocks from chain head. /// Events from blocks older than this won't be broadcast to real-time subscribers. @@ -56,12 +56,16 @@ pub struct Erc1155Sink { grpc_service: Option, /// Balance fetcher for RPC calls (None = balance tracking disabled) balance_fetcher: Option>, - /// Metadata fetcher for contract name/symbol - metadata_fetcher: Option>, - /// Token URI service sender for async URI fetching + /// Whether contract metadata commands should be dispatched. + metadata_commands_enabled: bool, + /// Whether token URI commands should be dispatched. + token_uri_commands_enabled: bool, + /// Command bus sender for background metadata and token URI work. + command_bus: Option, token_uri_sender: Option, - /// In-flight contract metadata fetches to avoid duplicate background jobs. - metadata_fetch_inflight: Arc>>, + /// Commands already queued but not yet observed in storage. + pending_metadata_commands: tokio::sync::Mutex>, + pending_token_uri_commands: tokio::sync::Mutex>, /// In-memory counters to avoid full-table COUNT(*) in the ingest hot path. total_transfers: AtomicU64, total_operator_approvals: AtomicU64, @@ -75,9 +79,12 @@ impl Erc1155Sink { event_bus: None, grpc_service: None, balance_fetcher: None, - metadata_fetcher: None, + metadata_commands_enabled: false, + token_uri_commands_enabled: false, + command_bus: None, token_uri_sender: None, - metadata_fetch_inflight: Arc::new(Mutex::new(HashSet::new())), + pending_metadata_commands: tokio::sync::Mutex::new(HashSet::new()), + pending_token_uri_commands: tokio::sync::Mutex::new(HashSet::new()), // Avoid startup full-table COUNT(*) scans on large datasets. total_transfers: AtomicU64::new(0), total_operator_approvals: AtomicU64::new(0), @@ -85,12 +92,6 @@ impl Erc1155Sink { } } - /// Enable async token URI fetching - pub fn with_token_uri_sender(mut self, sender: TokenUriSender) -> Self { - self.token_uri_sender = Some(sender); - self - } - /// Set the gRPC service for dual publishing pub fn with_grpc_service(mut self, service: Erc1155Service) -> Self { self.grpc_service = Some(service); @@ -105,8 +106,7 @@ impl Erc1155Sink { /// - Fetch actual balance from the chain and adjust /// - Record adjustments in an audit table pub fn with_balance_tracking(mut self, provider: Arc>) -> Self { - self.balance_fetcher = Some(Arc::new(Erc1155BalanceFetcher::new(provider.clone()))); - self.metadata_fetcher = Some(Arc::new(MetadataFetcher::new(provider))); + self.balance_fetcher = Some(Arc::new(Erc1155BalanceFetcher::new(provider))); self } @@ -162,28 +162,6 @@ impl Erc1155Sink { true } - /// Filter function for ERC1155 token metadata updates. - /// - /// Supports filters: - /// - "token": Filter by token contract address (hex string) - fn matches_metadata_filters( - metadata: &proto::TokenMetadataEntry, - filters: &HashMap, - ) -> bool { - if filters.is_empty() { - return true; - } - - if let Some(token_filter) = filters.get("token") { - let token_hex = format!("0x{}", hex::encode(&metadata.token)); - if !token_hex.eq_ignore_ascii_case(token_filter) { - return false; - } - } - - true - } - /// Filter function for ERC1155 URI updates. /// /// Supports filters: @@ -210,6 +188,32 @@ impl Erc1155Sink { true } + + fn enqueue_token_uri_request(&self, contract: Felt, token_id: U256) -> bool { + if let Some(sender) = &self.token_uri_sender { + return sender.request_update(TokenUriRequest { + contract, + token_id, + standard: TokenStandard::Erc1155, + }); + } + + if let Some(command_bus) = &self.command_bus { + if let Err(error) = + command_bus.dispatch(RefreshErc1155TokenUriCommand { contract, token_id }) + { + tracing::warn!( + target: "torii_erc1155::sink", + error = %error, + "Failed to dispatch ERC1155 token URI command" + ); + return false; + } + return true; + } + + false + } } #[async_trait] @@ -230,9 +234,10 @@ impl Sink for Erc1155Sink { async fn initialize( &mut self, event_bus: Arc, - _context: &torii::etl::sink::SinkContext, + context: &torii::etl::sink::SinkContext, ) -> Result<()> { self.event_bus = Some(event_bus); + self.command_bus = Some(context.command_bus.clone()); tracing::info!(target: "torii_erc1155::sink", "ERC1155 sink initialized"); Ok(()) } @@ -337,122 +342,106 @@ impl Sink for Erc1155Sink { } } - // Fetch metadata for any new token contracts - if let Some(ref fetcher) = self.metadata_fetcher { - let new_tokens: HashSet = transfers.iter().map(|t| t.token).collect(); - for token in new_tokens { - match self.storage.has_token_metadata(token).await { - Ok(exists) => { - if exists { - continue; - } + // Fetch metadata for any new token contracts. + if self.metadata_commands_enabled { + if let Some(ref command_bus) = self.command_bus { + let candidate_tokens = transfers + .iter() + .map(|transfer| transfer.token) + .collect::>() + .into_iter(); + let unchecked_tokens = { + let pending = self.pending_metadata_commands.lock().await; + candidate_tokens + .into_iter() + .filter(|token| !pending.contains(token)) + .collect::>() + }; - let scheduled = { - let mut inflight = self.metadata_fetch_inflight.lock().unwrap(); - inflight.insert(token) - }; - if !scheduled { - continue; + match self + .storage + .has_token_metadata_batch(&unchecked_tokens) + .await + { + Ok(existing_tokens) => { + let mut pending = self.pending_metadata_commands.lock().await; + for token in &existing_tokens { + pending.remove(token); } - let fetcher = fetcher.clone(); - let storage = self.storage.clone(); - let event_bus = self.event_bus.clone(); - let inflight = self.metadata_fetch_inflight.clone(); - tokio::spawn(async move { - let result = async { - let meta = fetcher.fetch_erc1155_metadata(token).await; - tracing::info!( - target: "torii_erc1155::sink", - token = %format!("{:#x}", token), - name = ?meta.name, - symbol = ?meta.symbol, - "Fetched token metadata" - ); - storage - .upsert_token_metadata( - token, - meta.name.as_deref(), - meta.symbol.as_deref(), - meta.total_supply, - ) - .await?; - - if let Some(event_bus) = &event_bus { - let meta_entry = proto::TokenMetadataEntry { - token: token.to_bytes_be().to_vec(), - name: meta.name, - symbol: meta.symbol, - total_supply: meta.total_supply.map(u256_to_bytes), - }; - - let mut buf = Vec::new(); - meta_entry.encode(&mut buf)?; - let any = Any { - type_url: - "type.googleapis.com/torii.sinks.erc1155.TokenMetadataEntry" - .to_string(), - value: buf, - }; - - event_bus.publish_protobuf( - "erc1155.metadata", - "erc1155.metadata", - &any, - &meta_entry, - UpdateType::Created, - Self::matches_metadata_filters, - ); - } - - Ok::<(), anyhow::Error>(()) + for token in unchecked_tokens + .into_iter() + .filter(|token| !existing_tokens.contains(token)) + { + if !pending.insert(token) { + continue; } - .await; - - if let Err(e) = result { + if let Err(error) = + command_bus.dispatch(FetchErc1155MetadataCommand { token }) + { + pending.remove(&token); tracing::warn!( target: "torii_erc1155::sink", token = %format!("{:#x}", token), - error = %e, - "Failed to fetch/store token metadata in background task" + error = %error, + "Failed to dispatch ERC1155 metadata command" ); } - - let mut inflight = inflight.lock().unwrap(); - inflight.remove(&token); - }); + } } - Err(e) => { - tracing::warn!(target: "torii_erc1155::sink", error = %e, "Failed to check token metadata"); + Err(error) => { + tracing::warn!( + target: "torii_erc1155::sink", + error = %error, + "Failed to batch-check token metadata" + ); } } } } - // Request token URI fetches for new token IDs - if let Some(ref sender) = self.token_uri_sender { - for transfer in &transfers { - match self - .storage - .has_token_uri(transfer.token, transfer.token_id) - .await - { - Ok(false) => { - sender.request_update(TokenUriRequest { - contract: transfer.token, - token_id: transfer.token_id, - standard: TokenStandard::Erc1155, - }); + // Request token URI fetches for new token IDs. + if self.token_uri_commands_enabled + && (self.token_uri_sender.is_some() || self.command_bus.is_some()) + { + let candidate_tokens = transfers + .iter() + .map(|transfer| (transfer.token, transfer.token_id)) + .collect::>() + .into_iter(); + let unchecked_tokens = { + let pending = self.pending_token_uri_commands.lock().await; + candidate_tokens + .filter(|token| !pending.contains(token)) + .collect::>() + }; + + match self.storage.has_token_uri_batch(&unchecked_tokens).await { + Ok(existing_tokens) => { + let mut pending = self.pending_token_uri_commands.lock().await; + for token in &existing_tokens { + pending.remove(token); } - Ok(true) => {} - Err(e) => { - tracing::warn!( - target: "torii_erc1155::sink", - error = %e, - "Failed to check token URI existence" - ); + + for (contract, token_id) in unchecked_tokens + .into_iter() + .filter(|token| !existing_tokens.contains(token)) + { + if !pending.insert((contract, token_id)) { + continue; + } + if !self.enqueue_token_uri_request(contract, token_id) { + pending.remove(&(contract, token_id)); + } } } + Err(error) => { + tracing::warn!( + target: "torii_erc1155::sink", + error = %error, + "Failed to batch-check token URI existence" + ); + } } } @@ -728,3 +717,25 @@ impl Sink for Erc1155Sink { Router::new() } } + +impl Erc1155Sink { + pub fn with_metadata_commands(mut self) -> Self { + self.metadata_commands_enabled = true; + self + } + + pub fn with_metadata_fetching(self, _provider: Arc>) -> Self { + self.with_metadata_commands() + } + + pub fn with_token_uri_commands(mut self) -> Self { + self.token_uri_commands_enabled = true; + self + } + + pub fn with_token_uri_sender(mut self, sender: TokenUriSender) -> Self { + self.token_uri_commands_enabled = true; + self.token_uri_sender = Some(sender); + self + } +} diff --git a/crates/torii-erc1155/src/storage.rs b/crates/torii-erc1155/src/storage.rs index 9ca469a5..e39345cd 100644 --- a/crates/torii-erc1155/src/storage.rs +++ b/crates/torii-erc1155/src/storage.rs @@ -7,9 +7,9 @@ //! - Records all adjustments in an audit table for debugging use anyhow::Result; -use rusqlite::{params, Connection}; +use rusqlite::{params, params_from_iter, Connection, ToSql}; use starknet::core::types::{Felt, U256}; -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::{Arc, Mutex}; use tokio_postgres::{types::ToSql as PgToSql, Client, NoTls}; use torii_common::{ @@ -18,6 +18,10 @@ use torii_common::{ use crate::balance_fetcher::Erc1155BalanceFetchRequest; +const SQLITE_MAX_BIND_VARS: usize = 900; +const SQLITE_TOKEN_BATCH_SIZE: usize = SQLITE_MAX_BIND_VARS; +const SQLITE_TOKEN_PAIR_BATCH_SIZE: usize = SQLITE_MAX_BIND_VARS / 2; + /// Maximum value for U256 (2^256 - 1) const U256_MAX: U256 = U256::from_words(u128::MAX, u128::MAX); @@ -121,6 +125,12 @@ pub struct TokenAttributeQueryResult { pub facets: Vec, } +#[derive(Debug, Clone)] +struct ResolvedFacetFilter { + key_id: i64, + value_ids: Vec, +} + /// Balance data for a (contract, wallet, token_id) tuple #[derive(Debug, Clone)] pub struct Erc1155BalanceData { @@ -227,6 +237,36 @@ impl Erc1155Storage { CREATE INDEX IF NOT EXISTS idx_token_attributes_token_key_value ON erc1155.token_attributes(token, key, value); CREATE INDEX IF NOT EXISTS idx_token_attributes_token_token_id ON erc1155.token_attributes(token, token_id); + CREATE TABLE IF NOT EXISTS erc1155.facet_keys ( + id BIGSERIAL PRIMARY KEY, + token BYTEA NOT NULL, + key_norm TEXT NOT NULL, + key_display TEXT NOT NULL, + UNIQUE(token, key_norm) + ); + CREATE INDEX IF NOT EXISTS idx_facet_keys_token_key_norm ON erc1155.facet_keys(token, key_norm); + + CREATE TABLE IF NOT EXISTS erc1155.facet_values ( + id BIGSERIAL PRIMARY KEY, + token BYTEA NOT NULL, + facet_key_id BIGINT NOT NULL REFERENCES erc1155.facet_keys(id) ON DELETE CASCADE, + value_norm TEXT NOT NULL, + value_display TEXT NOT NULL, + token_count BIGINT NOT NULL DEFAULT 0, + UNIQUE(token, facet_key_id, value_norm) + ); + CREATE INDEX IF NOT EXISTS idx_facet_values_token_key_value ON erc1155.facet_values(token, facet_key_id, value_norm); + + CREATE TABLE IF NOT EXISTS erc1155.facet_token_map ( + token BYTEA NOT NULL, + token_id BYTEA NOT NULL, + facet_key_id BIGINT NOT NULL REFERENCES erc1155.facet_keys(id) ON DELETE CASCADE, + facet_value_id BIGINT NOT NULL REFERENCES erc1155.facet_values(id) ON DELETE CASCADE, + PRIMARY KEY (token, token_id, facet_key_id) + ); + CREATE INDEX IF NOT EXISTS idx_facet_token_map_token_value_token_id ON erc1155.facet_token_map(token, facet_value_id, token_id); + CREATE INDEX IF NOT EXISTS idx_facet_token_map_token_token_id_key_value ON erc1155.facet_token_map(token, token_id, facet_key_id, facet_value_id); + CREATE TABLE IF NOT EXISTS erc1155.erc1155_balances ( id BIGSERIAL PRIMARY KEY, contract BYTEA NOT NULL, @@ -401,7 +441,40 @@ impl Erc1155Storage { CREATE INDEX IF NOT EXISTS idx_token_attributes_key ON token_attributes(key); CREATE INDEX IF NOT EXISTS idx_token_attributes_key_value ON token_attributes(key, value); CREATE INDEX IF NOT EXISTS idx_token_attributes_token_key_value ON token_attributes(token, key, value); - CREATE INDEX IF NOT EXISTS idx_token_attributes_token_token_id ON token_attributes(token, token_id);", + CREATE INDEX IF NOT EXISTS idx_token_attributes_token_token_id ON token_attributes(token, token_id); + + CREATE TABLE IF NOT EXISTS facet_keys ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + token BLOB NOT NULL, + key_norm TEXT NOT NULL, + key_display TEXT NOT NULL, + UNIQUE(token, key_norm) + ); + CREATE INDEX IF NOT EXISTS idx_facet_keys_token_key_norm ON facet_keys(token, key_norm); + + CREATE TABLE IF NOT EXISTS facet_values ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + token BLOB NOT NULL, + facet_key_id INTEGER NOT NULL, + value_norm TEXT NOT NULL, + value_display TEXT NOT NULL, + token_count INTEGER NOT NULL DEFAULT 0, + UNIQUE(token, facet_key_id, value_norm), + FOREIGN KEY (facet_key_id) REFERENCES facet_keys(id) ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS idx_facet_values_token_key_value ON facet_values(token, facet_key_id, value_norm); + + CREATE TABLE IF NOT EXISTS facet_token_map ( + token BLOB NOT NULL, + token_id BLOB NOT NULL, + facet_key_id INTEGER NOT NULL, + facet_value_id INTEGER NOT NULL, + PRIMARY KEY (token, token_id, facet_key_id), + FOREIGN KEY (facet_key_id) REFERENCES facet_keys(id) ON DELETE CASCADE, + FOREIGN KEY (facet_value_id) REFERENCES facet_values(id) ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS idx_facet_token_map_token_value_token_id ON facet_token_map(token, facet_value_id, token_id); + CREATE INDEX IF NOT EXISTS idx_facet_token_map_token_token_id_key_value ON facet_token_map(token, token_id, facet_key_id, facet_value_id);", )?; // Balance tracking tables @@ -827,10 +900,19 @@ impl Erc1155Storage { let facet_limit = facet_limit.clamp(1, 1000) as i64; if self.backend == StorageBackend::Postgres { + let resolved_filters = self.pg_resolve_facet_filters(token, &filters).await?; + let Some(resolved_filters) = resolved_filters else { + return Ok(TokenAttributeQueryResult { + token_ids: Vec::new(), + next_cursor_token_id: None, + total_hits: 0, + facets: Vec::new(), + }); + }; return self - .pg_query_token_ids_by_attributes( + .pg_query_token_ids_by_facets( token, - &filters, + &resolved_filters, cursor_token_id, page_fetch, include_facets, @@ -840,95 +922,25 @@ impl Erc1155Storage { } let conn = self.conn.lock().unwrap(); - - let mut query = String::from("SELECT u.token_id FROM token_uris u WHERE u.token = ?"); - let mut params: Vec> = vec![Box::new(felt_to_blob(token))]; - if let Some(cursor) = cursor_token_id { - query.push_str(" AND u.token_id > ?"); - params.push(Box::new(u256_to_blob(cursor))); - } - Self::sqlite_append_attribute_exists_filters(&mut query, &mut params, &filters, "u"); - query.push_str(" ORDER BY u.token_id ASC LIMIT ?"); - params.push(Box::new(page_fetch)); - - let mut stmt = conn.prepare(&query)?; - let params_refs: Vec<&dyn rusqlite::ToSql> = - params.iter().map(std::convert::AsRef::as_ref).collect(); - let rows = stmt.query_map(params_refs.as_slice(), |row| row.get::<_, Vec>(0))?; - let mut token_ids: Vec = rows - .collect::>, _>>()? - .into_iter() - .map(|b| blob_to_u256(&b)) - .collect(); - - let next_cursor_token_id = if token_ids.len() > page_limit as usize { - let next = token_ids[page_limit as usize]; - token_ids.truncate(page_limit as usize); - Some(next) - } else { - None - }; - - let mut count_query = String::from("SELECT COUNT(*) FROM token_uris u WHERE u.token = ?"); - let mut count_params: Vec> = vec![Box::new(felt_to_blob(token))]; - Self::sqlite_append_attribute_exists_filters( - &mut count_query, - &mut count_params, - &filters, - "u", - ); - let count_refs: Vec<&dyn rusqlite::ToSql> = count_params - .iter() - .map(std::convert::AsRef::as_ref) - .collect(); - let total_hits: i64 = - conn.query_row(&count_query, count_refs.as_slice(), |row| row.get(0))?; - - let facets = if include_facets { - let mut facet_query = String::from( - "SELECT a.key, a.value, COUNT(*) AS cnt - FROM token_attributes a - JOIN token_uris u ON u.token = a.token AND u.token_id = a.token_id - WHERE u.token = ?", - ); - let mut facet_params: Vec> = - vec![Box::new(felt_to_blob(token))]; - Self::sqlite_append_attribute_exists_filters( - &mut facet_query, - &mut facet_params, - &filters, - "u", - ); - facet_query.push_str( - " GROUP BY a.key, a.value - ORDER BY cnt DESC, a.key ASC, a.value ASC - LIMIT ?", - ); - facet_params.push(Box::new(facet_limit)); - - let facet_refs: Vec<&dyn rusqlite::ToSql> = facet_params - .iter() - .map(std::convert::AsRef::as_ref) - .collect(); - let mut stmt = conn.prepare(&facet_query)?; - let rows = stmt.query_map(facet_refs.as_slice(), |row| { - Ok(AttributeFacetCount { - key: row.get(0)?, - value: row.get(1)?, - count: row.get::<_, i64>(2)? as u64, - }) - })?; - rows.collect::, _>>()? - } else { - Vec::new() + let resolved_filters = self.sqlite_resolve_facet_filters(&conn, token, &filters)?; + let Some(resolved_filters) = resolved_filters else { + return Ok(TokenAttributeQueryResult { + token_ids: Vec::new(), + next_cursor_token_id: None, + total_hits: 0, + facets: Vec::new(), + }); }; - Ok(TokenAttributeQueryResult { - token_ids, - next_cursor_token_id, - total_hits: total_hits as u64, - facets, - }) + self.sqlite_query_token_ids_by_facets( + &conn, + token, + &resolved_filters, + cursor_token_id, + page_fetch, + include_facets, + facet_limit, + ) } /// Get transfer count @@ -1419,6 +1431,36 @@ impl Erc1155Storage { Ok(count > 0) } + pub async fn has_token_metadata_batch(&self, tokens: &[Felt]) -> Result> { + if tokens.is_empty() { + return Ok(HashSet::new()); + } + if self.backend == StorageBackend::Postgres { + return self.pg_has_token_metadata_batch(tokens).await; + } + + let conn = self.conn.lock().unwrap(); + let mut existing = HashSet::with_capacity(tokens.len()); + + for chunk in tokens.chunks(SQLITE_TOKEN_BATCH_SIZE) { + let placeholders = vec!["?"; chunk.len()].join(","); + let query = format!("SELECT token FROM token_metadata WHERE token IN ({placeholders})"); + let token_blobs: Vec> = + chunk.iter().map(|token| felt_to_blob(*token)).collect(); + let mut stmt = conn.prepare(&query)?; + let rows = stmt.query_map(params_from_iter(token_blobs.iter()), |row| { + let token_bytes: Vec = row.get(0)?; + Ok(blob_to_felt(&token_bytes)) + })?; + + for row in rows { + existing.insert(row?); + } + } + + Ok(existing) + } + /// Insert or update token metadata pub async fn upsert_token_metadata( &self, @@ -1586,6 +1628,45 @@ impl Erc1155Storage { Ok(count > 0) } + pub async fn has_token_uri_batch( + &self, + tokens: &[(Felt, U256)], + ) -> Result> { + if tokens.is_empty() { + return Ok(HashSet::new()); + } + if self.backend == StorageBackend::Postgres { + return self.pg_has_token_uri_batch(tokens).await; + } + + let conn = self.conn.lock().unwrap(); + let mut existing = HashSet::with_capacity(tokens.len()); + + for chunk in tokens.chunks(SQLITE_TOKEN_PAIR_BATCH_SIZE) { + let predicates = vec!["(token = ? AND token_id = ?)"; chunk.len()].join(" OR "); + let query = format!("SELECT token, token_id FROM token_uris WHERE {predicates}"); + let mut params_vec: Vec> = Vec::with_capacity(chunk.len() * 2); + for (token, token_id) in chunk { + params_vec.push(Box::new(felt_to_blob(*token))); + params_vec.push(Box::new(u256_to_blob(*token_id))); + } + let params_refs: Vec<&dyn ToSql> = + params_vec.iter().map(std::convert::AsRef::as_ref).collect(); + let mut stmt = conn.prepare(&query)?; + let rows = stmt.query_map(params_refs.as_slice(), |row| { + let token_bytes: Vec = row.get(0)?; + let token_id_bytes: Vec = row.get(1)?; + Ok((blob_to_felt(&token_bytes), blob_to_u256(&token_id_bytes))) + })?; + + for row in rows { + existing.insert(row?); + } + } + + Ok(existing) + } + /// Returns all token URI rows for a given contract. pub async fn get_token_uris_by_contract( &self, @@ -1611,6 +1692,48 @@ impl Erc1155Storage { rows.collect::, _>>().map_err(Into::into) } + pub async fn get_token_uris_batch( + &self, + token: Felt, + token_ids: &[U256], + ) -> Result, Option)>> { + if token_ids.is_empty() { + return Ok(Vec::new()); + } + if self.backend == StorageBackend::Postgres { + return self.pg_get_token_uris_batch(token, token_ids).await; + } + + let conn = self.conn.lock().unwrap(); + let token_blob = felt_to_blob(token); + let mut rows_out = Vec::new(); + for chunk in token_ids.chunks(SQLITE_TOKEN_BATCH_SIZE) { + let placeholders = chunk.iter().map(|_| "?").collect::>().join(","); + let query = format!( + "SELECT token_id, uri, metadata_json + FROM token_uris + WHERE token = ?1 AND token_id IN ({placeholders})" + ); + let mut params_vec: Vec> = Vec::with_capacity(chunk.len() + 1); + params_vec.push(Box::new(token_blob.clone())); + for token_id in chunk { + params_vec.push(Box::new(u256_to_blob(*token_id))); + } + let params_refs: Vec<&dyn ToSql> = + params_vec.iter().map(std::convert::AsRef::as_ref).collect(); + let mut stmt = conn.prepare(&query)?; + let rows = stmt.query_map(params_refs.as_slice(), |row| { + let token_id_bytes: Vec = row.get(0)?; + let uri: Option = row.get(1)?; + let metadata_json: Option = row.get(2)?; + Ok((blob_to_u256(&token_id_bytes), uri, metadata_json)) + })?; + rows_out.extend(rows.collect::, _>>()?); + } + + Ok(rows_out) + } + async fn pg_client(&self) -> Result> { let conn = self .pg_conn @@ -1656,54 +1779,267 @@ impl Erc1155Storage { .collect() } - fn sqlite_append_attribute_exists_filters( - query: &mut String, - params: &mut Vec>, + fn sqlite_resolve_facet_filters( + &self, + conn: &Connection, + token: Felt, filters: &[(String, Vec)], - parent_alias: &str, - ) { + ) -> Result>> { + let token_blob = felt_to_blob(token); + let mut resolved = Vec::with_capacity(filters.len()); + for (key, values) in filters { - query.push_str(&format!( - " AND EXISTS ( - SELECT 1 FROM token_attributes a - WHERE a.token = {parent_alias}.token - AND a.token_id = {parent_alias}.token_id - AND a.key = ? AND a.value IN (" - )); - params.push(Box::new(key.clone())); + let key_id: i64 = match conn.query_row( + "SELECT id FROM facet_keys WHERE token = ?1 AND key_norm = ?2", + params![&token_blob, key], + |row| row.get(0), + ) { + Ok(id) => id, + Err(rusqlite::Error::QueryReturnedNoRows) => return Ok(None), + Err(error) => return Err(error.into()), + }; + let placeholders = values.iter().map(|_| "?").collect::>().join(","); - query.push_str(&placeholders); - query.push_str("))"); + let query = format!( + "SELECT id FROM facet_values WHERE token = ? AND facet_key_id = ? AND value_norm IN ({placeholders})" + ); + let mut params_vec: Vec> = Vec::with_capacity(values.len() + 2); + params_vec.push(Box::new(token_blob.clone())); + params_vec.push(Box::new(key_id)); for value in values { - params.push(Box::new(value.clone())); + params_vec.push(Box::new(value.clone())); + } + let params_refs: Vec<&dyn ToSql> = + params_vec.iter().map(std::convert::AsRef::as_ref).collect(); + let mut stmt = conn.prepare(&query)?; + let rows = stmt.query_map(params_refs.as_slice(), |row| row.get::<_, i64>(0))?; + let value_ids = rows.collect::, _>>()?; + if value_ids.is_empty() { + return Ok(None); } + + resolved.push(ResolvedFacetFilter { key_id, value_ids }); } + + Ok(Some(resolved)) } - fn pg_append_attribute_exists_filters( - query: &mut String, - params: &mut Vec>, - filters: &[(String, Vec)], - parent_alias: &str, - ) { - for (key, values) in filters { - query.push_str(&format!( - " AND EXISTS ( - SELECT 1 FROM erc1155.token_attributes a - WHERE a.token = {parent_alias}.token - AND a.token_id = {parent_alias}.token_id - AND a.key = " - )); - query.push_str(&Self::pg_next_param(params, key.clone())); - query.push_str(" AND a.value IN ("); - let value_placeholders = values + fn sqlite_build_candidate_query( + token: Felt, + resolved_filters: &[ResolvedFacetFilter], + ) -> (String, Vec>) { + let token_blob = felt_to_blob(token); + if resolved_filters.is_empty() { + return ( + "SELECT token_id FROM token_uris WHERE token = ?".to_owned(), + vec![Box::new(token_blob)], + ); + } + + let mut query = String::new(); + let mut params: Vec> = Vec::new(); + for (idx, filter) in resolved_filters.iter().enumerate() { + if idx > 0 { + query.push_str(" INTERSECT "); + } + query.push_str( + "SELECT token_id FROM facet_token_map WHERE token = ? AND facet_key_id = ? AND facet_value_id IN (", + ); + params.push(Box::new(token_blob.clone())); + params.push(Box::new(filter.key_id)); + let placeholders = filter + .value_ids .iter() - .map(|value| Self::pg_next_param(params, value.clone())) + .map(|_| "?") .collect::>() .join(","); - query.push_str(&value_placeholders); - query.push_str("))"); + query.push_str(&placeholders); + query.push(')'); + for value_id in &filter.value_ids { + params.push(Box::new(*value_id)); + } + } + + (query, params) + } + + fn sqlite_query_token_ids_by_facets( + &self, + conn: &Connection, + token: Felt, + resolved_filters: &[ResolvedFacetFilter], + cursor_token_id: Option, + page_fetch: i64, + include_facets: bool, + facet_limit: i64, + ) -> Result { + let (candidate_query, candidate_params_for_page) = + Self::sqlite_build_candidate_query(token, resolved_filters); + let (_, candidate_params_for_count) = + Self::sqlite_build_candidate_query(token, resolved_filters); + + let mut page_query = format!( + "WITH candidate_tokens AS ({candidate_query}) SELECT token_id FROM candidate_tokens" + ); + let mut page_params = candidate_params_for_page; + if let Some(cursor) = cursor_token_id { + page_query.push_str(" WHERE token_id > ?"); + page_params.push(Box::new(u256_to_blob(cursor))); + } + page_query.push_str(" ORDER BY token_id ASC LIMIT ?"); + page_params.push(Box::new(page_fetch)); + + let page_refs: Vec<&dyn ToSql> = page_params + .iter() + .map(std::convert::AsRef::as_ref) + .collect(); + let mut stmt = conn.prepare(&page_query)?; + let rows = stmt.query_map(page_refs.as_slice(), |row| row.get::<_, Vec>(0))?; + let mut token_ids = rows + .collect::, _>>()? + .into_iter() + .map(|bytes| blob_to_u256(&bytes)) + .collect::>(); + + let page_limit = (page_fetch - 1) as usize; + let next_cursor_token_id = if token_ids.len() > page_limit { + let next = token_ids[page_limit]; + token_ids.truncate(page_limit); + Some(next) + } else { + None + }; + + let count_query = format!( + "WITH candidate_tokens AS ({candidate_query}) SELECT COUNT(*) FROM candidate_tokens" + ); + let count_refs: Vec<&dyn ToSql> = candidate_params_for_count + .iter() + .map(std::convert::AsRef::as_ref) + .collect(); + let total_hits: i64 = + conn.query_row(&count_query, count_refs.as_slice(), |row| row.get(0))?; + + let facets = if include_facets { + let (_, candidate_params_for_facets) = + Self::sqlite_build_candidate_query(token, resolved_filters); + let facet_query = format!( + "WITH candidate_tokens AS ({candidate_query}) + SELECT fk.key_display, fv.value_display, COUNT(*) AS cnt + FROM facet_token_map m + JOIN candidate_tokens c ON c.token_id = m.token_id + JOIN facet_keys fk ON fk.id = m.facet_key_id + JOIN facet_values fv ON fv.id = m.facet_value_id + WHERE m.token = ? + GROUP BY fk.key_display, fv.value_display + ORDER BY cnt DESC, fk.key_display ASC, fv.value_display ASC + LIMIT ?" + ); + let mut facet_params = candidate_params_for_facets; + facet_params.push(Box::new(felt_to_blob(token))); + facet_params.push(Box::new(facet_limit)); + let facet_refs: Vec<&dyn ToSql> = facet_params + .iter() + .map(std::convert::AsRef::as_ref) + .collect(); + let mut stmt = conn.prepare(&facet_query)?; + let rows = stmt.query_map(facet_refs.as_slice(), |row| { + Ok(AttributeFacetCount { + key: row.get(0)?, + value: row.get(1)?, + count: row.get::<_, i64>(2)? as u64, + }) + })?; + rows.collect::, _>>()? + } else { + Vec::new() + }; + + Ok(TokenAttributeQueryResult { + token_ids, + next_cursor_token_id, + total_hits: total_hits as u64, + facets, + }) + } + + async fn pg_resolve_facet_filters( + &self, + token: Felt, + filters: &[(String, Vec)], + ) -> Result>> { + let client = self.pg_client().await?; + let token_blob = felt_to_blob(token); + let mut resolved = Vec::with_capacity(filters.len()); + + for (key, values) in filters { + let key_row = client + .query_opt( + "SELECT id FROM erc1155.facet_keys WHERE token = $1 AND key_norm = $2", + &[&token_blob, key], + ) + .await?; + let Some(key_row) = key_row else { + return Ok(None); + }; + let key_id: i64 = key_row.get(0); + + let rows = client + .query( + "SELECT id FROM erc1155.facet_values WHERE token = $1 AND facet_key_id = $2 AND value_norm = ANY($3::text[])", + &[&token_blob, &key_id, values], + ) + .await?; + let value_ids = rows + .into_iter() + .map(|row| row.get::(0)) + .collect::>(); + if value_ids.is_empty() { + return Ok(None); + } + + resolved.push(ResolvedFacetFilter { key_id, value_ids }); + } + + Ok(Some(resolved)) + } + + fn pg_build_candidate_query( + token: Felt, + resolved_filters: &[ResolvedFacetFilter], + ) -> (String, Vec>) { + let mut candidate_query = String::new(); + let mut candidate_params: Vec> = Vec::new(); + if resolved_filters.is_empty() { + candidate_query.push_str("SELECT token_id FROM erc1155.token_uris WHERE token = "); + candidate_query.push_str(&Self::pg_next_param( + &mut candidate_params, + felt_to_blob(token), + )); + } else { + for (idx, filter) in resolved_filters.iter().enumerate() { + if idx > 0 { + candidate_query.push_str(" INTERSECT "); + } + candidate_query + .push_str("SELECT token_id FROM erc1155.facet_token_map WHERE token = "); + candidate_query.push_str(&Self::pg_next_param( + &mut candidate_params, + felt_to_blob(token), + )); + candidate_query.push_str(" AND facet_key_id = "); + candidate_query + .push_str(&Self::pg_next_param(&mut candidate_params, filter.key_id)); + candidate_query.push_str(" AND facet_value_id = ANY("); + candidate_query.push_str(&Self::pg_next_param( + &mut candidate_params, + filter.value_ids.clone(), + )); + candidate_query.push(')'); + } } + + (candidate_query, candidate_params) } async fn pg_insert_transfers_batch(&self, transfers: &[TokenTransferData]) -> Result { @@ -2023,27 +2359,30 @@ impl Erc1155Storage { Ok((transfers, next_cursor)) } - async fn pg_query_token_ids_by_attributes( + async fn pg_query_token_ids_by_facets( &self, token: Felt, - filters: &[(String, Vec)], + resolved_filters: &[ResolvedFacetFilter], cursor_token_id: Option, page_fetch: i64, include_facets: bool, facet_limit: i64, ) -> Result { let client = self.pg_client().await?; + let (candidate_query, candidate_params_for_page) = + Self::pg_build_candidate_query(token, resolved_filters); + let (_, candidate_params_for_count) = + Self::pg_build_candidate_query(token, resolved_filters); - let mut query = - String::from("SELECT u.token_id FROM erc1155.token_uris u WHERE u.token = "); - let mut params: Vec> = Vec::new(); - query.push_str(&Self::pg_next_param(&mut params, felt_to_blob(token))); + let mut query = format!( + "WITH candidate_tokens AS ({candidate_query}) SELECT token_id FROM candidate_tokens" + ); + let mut params = candidate_params_for_page; if let Some(cursor) = cursor_token_id { - query.push_str(" AND u.token_id > "); + query.push_str(" WHERE token_id > "); query.push_str(&Self::pg_next_param(&mut params, u256_to_blob(cursor))); } - Self::pg_append_attribute_exists_filters(&mut query, &mut params, filters, "u"); - query.push_str(" ORDER BY u.token_id ASC LIMIT "); + query.push_str(" ORDER BY token_id ASC LIMIT "); query.push_str(&Self::pg_next_param(&mut params, page_fetch)); let refs: Vec<&(dyn PgToSql + Sync)> = params @@ -2065,12 +2404,10 @@ impl Erc1155Storage { None }; - let mut count_query = - String::from("SELECT COUNT(*) FROM erc1155.token_uris u WHERE u.token = "); - let mut count_params: Vec> = Vec::new(); - count_query.push_str(&Self::pg_next_param(&mut count_params, felt_to_blob(token))); - Self::pg_append_attribute_exists_filters(&mut count_query, &mut count_params, filters, "u"); - let count_refs: Vec<&(dyn PgToSql + Sync)> = count_params + let count_query = format!( + "WITH candidate_tokens AS ({candidate_query}) SELECT COUNT(*) FROM candidate_tokens" + ); + let count_refs: Vec<&(dyn PgToSql + Sync)> = candidate_params_for_count .iter() .map(|p| p.as_ref() as &(dyn PgToSql + Sync)) .collect(); @@ -2078,23 +2415,22 @@ impl Erc1155Storage { let total_hits = total_row.get::(0) as u64; let facets = if include_facets { - let mut facet_query = String::from( - "SELECT a.key, a.value, COUNT(*) AS cnt - FROM erc1155.token_attributes a - JOIN erc1155.token_uris u ON u.token = a.token AND u.token_id = a.token_id - WHERE u.token = ", + let (_, candidate_params_for_facets) = + Self::pg_build_candidate_query(token, resolved_filters); + let mut facet_query = format!( + "WITH candidate_tokens AS ({candidate_query}) + SELECT fk.key_display, fv.value_display, COUNT(*) AS cnt + FROM erc1155.facet_token_map m + JOIN candidate_tokens c ON c.token_id = m.token_id + JOIN erc1155.facet_keys fk ON fk.id = m.facet_key_id + JOIN erc1155.facet_values fv ON fv.id = m.facet_value_id + WHERE m.token = " ); - let mut facet_params: Vec> = Vec::new(); + let mut facet_params = candidate_params_for_facets; facet_query.push_str(&Self::pg_next_param(&mut facet_params, felt_to_blob(token))); - Self::pg_append_attribute_exists_filters( - &mut facet_query, - &mut facet_params, - filters, - "u", - ); facet_query.push_str( - " GROUP BY a.key, a.value - ORDER BY cnt DESC, a.key ASC, a.value ASC + " GROUP BY fk.key_display, fv.value_display + ORDER BY cnt DESC, fk.key_display ASC, fv.value_display ASC LIMIT ", ); facet_query.push_str(&Self::pg_next_param(&mut facet_params, facet_limit)); @@ -2433,6 +2769,21 @@ impl Erc1155Storage { Ok(row.get::(0) > 0) } + async fn pg_has_token_metadata_batch(&self, tokens: &[Felt]) -> Result> { + let client = self.pg_client().await?; + let token_blobs: Vec> = tokens.iter().map(|token| felt_to_blob(*token)).collect(); + let rows = client + .query( + "SELECT token FROM erc1155.token_metadata WHERE token = ANY($1::bytea[])", + &[&token_blobs], + ) + .await?; + Ok(rows + .into_iter() + .map(|row| blob_to_felt(&row.get::>(0))) + .collect()) + } + async fn pg_upsert_token_metadata( &self, token: Felt, @@ -2564,6 +2915,39 @@ impl Erc1155Storage { Ok(row.get::(0) > 0) } + async fn pg_has_token_uri_batch( + &self, + tokens: &[(Felt, U256)], + ) -> Result> { + let client = self.pg_client().await?; + let token_blobs: Vec> = tokens + .iter() + .map(|(token, _)| felt_to_blob(*token)) + .collect(); + let token_id_blobs: Vec> = tokens + .iter() + .map(|(_, token_id)| u256_to_blob(*token_id)) + .collect(); + let rows = client + .query( + "SELECT DISTINCT u.token, u.token_id + FROM erc1155.token_uris u + JOIN unnest($1::bytea[], $2::bytea[]) AS i(token, token_id) + ON u.token = i.token AND u.token_id = i.token_id", + &[&token_blobs, &token_id_blobs], + ) + .await?; + Ok(rows + .into_iter() + .map(|row| { + ( + blob_to_felt(&row.get::>(0)), + blob_to_u256(&row.get::>(1)), + ) + }) + .collect()) + } + async fn pg_get_token_uris_by_contract( &self, token: Felt, @@ -2589,72 +2973,104 @@ impl Erc1155Storage { }) .collect()) } + + async fn pg_get_token_uris_batch( + &self, + token: Felt, + token_ids: &[U256], + ) -> Result, Option)>> { + let client = self.pg_client().await?; + let token_id_blobs: Vec> = token_ids + .iter() + .map(|token_id| u256_to_blob(*token_id)) + .collect(); + let rows = client + .query( + "SELECT token_id, uri, metadata_json + FROM erc1155.token_uris + WHERE token = $1 AND token_id = ANY($2::bytea[])", + &[&felt_to_blob(token), &token_id_blobs], + ) + .await?; + Ok(rows + .into_iter() + .map(|row| { + ( + blob_to_u256(&row.get::>(0)), + row.get(1), + row.get(2), + ) + }) + .collect()) + } } #[async_trait::async_trait] impl TokenUriStore for Erc1155Storage { - async fn store_token_uri(&self, result: &TokenUriResult) -> Result<()> { + async fn store_token_uris_batch(&self, results: &[TokenUriResult]) -> Result<()> { + if results.is_empty() { + return Ok(()); + } + + let expected_attributes = results + .iter() + .map(|result| extract_metadata_attributes(result.metadata_json.as_deref())) + .collect::>(); + if self.backend == StorageBackend::Postgres { let mut client = self.pg_client().await?; - let tx = client.transaction().await?; + let mut changed_indexes = Vec::new(); + for (idx, result) in results.iter().enumerate() { + if !pg_token_uri_state_matches(&client, result, &expected_attributes[idx]).await? { + changed_indexes.push(idx); + } + } - let token_blob = felt_to_blob(result.contract); - let token_id_blob = u256_to_blob(result.token_id); + if changed_indexes.is_empty() { + return Ok(()); + } - tx.execute( - "INSERT INTO erc1155.token_uris (token, token_id, uri, metadata_json, updated_at) - VALUES ($1, $2, $3, $4, EXTRACT(EPOCH FROM NOW())::BIGINT) - ON CONFLICT(token, token_id) DO UPDATE SET - uri = EXCLUDED.uri, - metadata_json = EXCLUDED.metadata_json, - updated_at = EXCLUDED.updated_at", - &[ - &token_blob, - &token_id_blob, - &result.uri.as_deref(), - &result.metadata_json.as_deref(), - ], - ) - .await?; + let tx = client.transaction().await?; - tx.execute( - "DELETE FROM erc1155.token_attributes WHERE token = $1 AND token_id = $2", - &[&token_blob, &token_id_blob], - ) - .await?; + for idx in changed_indexes { + let result = &results[idx]; + let attrs = &expected_attributes[idx]; + let token_blob = felt_to_blob(result.contract); + let token_id_blob = u256_to_blob(result.token_id); + + tx.execute( + "INSERT INTO erc1155.token_uris (token, token_id, uri, metadata_json, updated_at) + VALUES ($1, $2, $3, $4, EXTRACT(EPOCH FROM NOW())::BIGINT) + ON CONFLICT(token, token_id) DO UPDATE SET + uri = EXCLUDED.uri, + metadata_json = EXCLUDED.metadata_json, + updated_at = EXCLUDED.updated_at", + &[ + &token_blob, + &token_id_blob, + &result.uri.as_deref(), + &result.metadata_json.as_deref(), + ], + ) + .await?; - if let Some(metadata_json) = &result.metadata_json { - if let Ok(value) = serde_json::from_str::(metadata_json) { - if let Some(attrs) = value.get("attributes").and_then(|a| a.as_array()) { - for attr in attrs { - let key = attr - .get("trait_type") - .or_else(|| attr.get("key")) - .and_then(|v| v.as_str()); - let value = attr.get("value").and_then(|v| { - v.as_str().map(ToOwned::to_owned).or_else(|| { - if v.is_null() { - None - } else { - Some(v.to_string()) - } - }) - }); - - let key = key.and_then(sanitize_metadata_text); - let value = value.and_then(|v| sanitize_metadata_text(&v)); - - if let (Some(key), Some(value)) = (key, value) { - tx.execute( - "INSERT INTO erc1155.token_attributes (token, token_id, key, value) - VALUES ($1, $2, $3, $4) - ON CONFLICT (token, token_id, key) DO UPDATE SET value = EXCLUDED.value", - &[&token_blob, &token_id_blob, &key, &value], - ).await?; - } - } - } + tx.execute( + "DELETE FROM erc1155.token_attributes WHERE token = $1 AND token_id = $2", + &[&token_blob, &token_id_blob], + ) + .await?; + + for (key, value) in attrs { + tx.execute( + "INSERT INTO erc1155.token_attributes (token, token_id, key, value) + VALUES ($1, $2, $3, $4) + ON CONFLICT (token, token_id, key) DO UPDATE SET value = EXCLUDED.value", + &[&token_blob, &token_id_blob, key, value], + ) + .await?; } + + pg_sync_facets_for_token(&tx, result.contract, result.token_id, attrs).await?; } tx.commit().await?; @@ -2662,68 +3078,332 @@ impl TokenUriStore for Erc1155Storage { } let mut conn = self.conn.lock().unwrap(); + let mut changed_indexes = Vec::new(); + for (idx, result) in results.iter().enumerate() { + if !sqlite_token_uri_state_matches(&conn, result, &expected_attributes[idx])? { + changed_indexes.push(idx); + } + } + + if changed_indexes.is_empty() { + return Ok(()); + } + let tx = conn.transaction()?; + { + let mut upsert_stmt = tx.prepare_cached( + "INSERT INTO token_uris (token, token_id, uri, metadata_json, updated_at) + VALUES (?1, ?2, ?3, ?4, strftime('%s', 'now')) + ON CONFLICT(token, token_id) DO UPDATE SET + uri = excluded.uri, + metadata_json = excluded.metadata_json, + updated_at = excluded.updated_at", + )?; + let mut delete_attrs_stmt = tx.prepare_cached( + "DELETE FROM token_attributes WHERE token = ?1 AND token_id = ?2", + )?; + let mut insert_attr_stmt = tx.prepare_cached( + "INSERT OR REPLACE INTO token_attributes (token, token_id, key, value) + VALUES (?1, ?2, ?3, ?4)", + )?; + + for idx in changed_indexes { + let result = &results[idx]; + let attrs = &expected_attributes[idx]; + let token_blob = felt_to_blob(result.contract); + let token_id_blob = u256_to_blob(result.token_id); + + upsert_stmt.execute(params![ + &token_blob, + &token_id_blob, + result.uri.as_deref(), + result.metadata_json.as_deref() + ])?; + delete_attrs_stmt.execute(params![&token_blob, &token_id_blob])?; + + for (key, value) in attrs { + insert_attr_stmt.execute(params![&token_blob, &token_id_blob, key, value])?; + } + + sqlite_sync_facets_for_token(&tx, result.contract, result.token_id, attrs)?; + } + } + + tx.commit()?; + Ok(()) + } +} + +fn extract_metadata_attributes(metadata_json: Option<&str>) -> Vec<(String, String)> { + let Some(metadata_json) = metadata_json else { + return Vec::new(); + }; + let Ok(value) = serde_json::from_str::(metadata_json) else { + return Vec::new(); + }; + let Some(attrs) = value.get("attributes").and_then(|attrs| attrs.as_array()) else { + return Vec::new(); + }; + + let mut normalized = BTreeMap::new(); + for attr in attrs { + let key = attr + .get("trait_type") + .or_else(|| attr.get("key")) + .and_then(|value| value.as_str()) + .and_then(sanitize_metadata_text); + let value = attr + .get("value") + .and_then(|value| { + value.as_str().map(ToOwned::to_owned).or_else(|| { + if value.is_null() { + None + } else { + Some(value.to_string()) + } + }) + }) + .and_then(|value| sanitize_metadata_text(&value)); + + if let (Some(key), Some(value)) = (key, value) { + normalized.insert(key, value); + } + } + + normalized.into_iter().collect() +} - let token_blob = felt_to_blob(result.contract); - let token_id_blob = u256_to_blob(result.token_id); +fn sqlite_sync_facets_for_token( + tx: &rusqlite::Transaction<'_>, + token: Felt, + token_id: U256, + attributes: &[(String, String)], +) -> Result<()> { + let token_blob = felt_to_blob(token); + let token_id_blob = u256_to_blob(token_id); + + let existing_rows = { + let mut stmt = tx.prepare( + "SELECT facet_value_id FROM facet_token_map WHERE token = ?1 AND token_id = ?2", + )?; + let rows = stmt.query_map(params![&token_blob, &token_id_blob], |row| { + row.get::<_, i64>(0) + })?; + rows.collect::, _>>()? + }; + for facet_value_id in existing_rows { tx.execute( - "INSERT INTO token_uris (token, token_id, uri, metadata_json, updated_at) - VALUES (?1, ?2, ?3, ?4, strftime('%s', 'now')) - ON CONFLICT(token, token_id) DO UPDATE SET - uri = excluded.uri, - metadata_json = excluded.metadata_json, - updated_at = excluded.updated_at", - params![ - &token_blob, - &token_id_blob, - result.uri.as_deref(), - result.metadata_json.as_deref() - ], + "UPDATE facet_values + SET token_count = CASE WHEN token_count > 0 THEN token_count - 1 ELSE 0 END + WHERE id = ?1", + params![facet_value_id], )?; + } + tx.execute( + "DELETE FROM facet_token_map WHERE token = ?1 AND token_id = ?2", + params![&token_blob, &token_id_blob], + )?; + for (key, value) in attributes { tx.execute( - "DELETE FROM token_attributes WHERE token = ?1 AND token_id = ?2", - params![&token_blob, &token_id_blob], + "INSERT INTO facet_keys (token, key_norm, key_display) + VALUES (?1, ?2, ?3) + ON CONFLICT(token, key_norm) DO UPDATE SET key_display = excluded.key_display", + params![&token_blob, key, key], + )?; + let facet_key_id: i64 = tx.query_row( + "SELECT id FROM facet_keys WHERE token = ?1 AND key_norm = ?2", + params![&token_blob, key], + |row| row.get(0), )?; - if let Some(metadata_json) = &result.metadata_json { - if let Ok(value) = serde_json::from_str::(metadata_json) { - if let Some(attrs) = value.get("attributes").and_then(|a| a.as_array()) { - let mut attr_stmt = tx.prepare_cached( - "INSERT OR REPLACE INTO token_attributes (token, token_id, key, value) - VALUES (?1, ?2, ?3, ?4)", - )?; - - for attr in attrs { - let key = attr - .get("trait_type") - .or_else(|| attr.get("key")) - .and_then(|v| v.as_str()); - let value = attr.get("value").and_then(|v| { - v.as_str().map(ToOwned::to_owned).or_else(|| { - if v.is_null() { - None - } else { - Some(v.to_string()) - } - }) - }); + tx.execute( + "INSERT INTO facet_values (token, facet_key_id, value_norm, value_display) + VALUES (?1, ?2, ?3, ?4) + ON CONFLICT(token, facet_key_id, value_norm) DO UPDATE SET value_display = excluded.value_display", + params![&token_blob, facet_key_id, value, value], + )?; + let facet_value_id: i64 = tx.query_row( + "SELECT id FROM facet_values WHERE token = ?1 AND facet_key_id = ?2 AND value_norm = ?3", + params![&token_blob, facet_key_id, value], + |row| row.get(0), + )?; - let key = key.and_then(sanitize_metadata_text); - let value = value.and_then(|v| sanitize_metadata_text(&v)); + tx.execute( + "INSERT OR REPLACE INTO facet_token_map (token, token_id, facet_key_id, facet_value_id) + VALUES (?1, ?2, ?3, ?4)", + params![&token_blob, &token_id_blob, facet_key_id, facet_value_id], + )?; + tx.execute( + "UPDATE facet_values SET token_count = token_count + 1 WHERE id = ?1", + params![facet_value_id], + )?; + } - if let (Some(key), Some(value)) = (key, value) { - attr_stmt.execute(params![&token_blob, &token_id_blob, key, value])?; - } - } - } - } - } + Ok(()) +} - tx.commit()?; - Ok(()) +async fn pg_sync_facets_for_token( + tx: &tokio_postgres::Transaction<'_>, + token: Felt, + token_id: U256, + attributes: &[(String, String)], +) -> Result<()> { + let token_blob = felt_to_blob(token); + let token_id_blob = u256_to_blob(token_id); + + let existing_rows = tx + .query( + "SELECT facet_value_id FROM erc1155.facet_token_map WHERE token = $1 AND token_id = $2", + &[&token_blob, &token_id_blob], + ) + .await?; + for row in existing_rows { + let facet_value_id: i64 = row.get(0); + tx.execute( + "UPDATE erc1155.facet_values + SET token_count = GREATEST(token_count - 1, 0) + WHERE id = $1", + &[&facet_value_id], + ) + .await?; + } + tx.execute( + "DELETE FROM erc1155.facet_token_map WHERE token = $1 AND token_id = $2", + &[&token_blob, &token_id_blob], + ) + .await?; + + for (key, value) in attributes { + tx.execute( + "INSERT INTO erc1155.facet_keys (token, key_norm, key_display) + VALUES ($1, $2, $3) + ON CONFLICT(token, key_norm) DO UPDATE SET key_display = EXCLUDED.key_display", + &[&token_blob, key, key], + ) + .await?; + let facet_key_id: i64 = tx + .query_one( + "SELECT id FROM erc1155.facet_keys WHERE token = $1 AND key_norm = $2", + &[&token_blob, key], + ) + .await? + .get(0); + + tx.execute( + "INSERT INTO erc1155.facet_values (token, facet_key_id, value_norm, value_display) + VALUES ($1, $2, $3, $4) + ON CONFLICT(token, facet_key_id, value_norm) DO UPDATE SET value_display = EXCLUDED.value_display", + &[&token_blob, &facet_key_id, value, value], + ) + .await?; + let facet_value_id: i64 = tx + .query_one( + "SELECT id FROM erc1155.facet_values WHERE token = $1 AND facet_key_id = $2 AND value_norm = $3", + &[&token_blob, &facet_key_id, value], + ) + .await? + .get(0); + + tx.execute( + "INSERT INTO erc1155.facet_token_map (token, token_id, facet_key_id, facet_value_id) + VALUES ($1, $2, $3, $4) + ON CONFLICT(token, token_id, facet_key_id) DO UPDATE SET facet_value_id = EXCLUDED.facet_value_id", + &[&token_blob, &token_id_blob, &facet_key_id, &facet_value_id], + ) + .await?; + tx.execute( + "UPDATE erc1155.facet_values SET token_count = token_count + 1 WHERE id = $1", + &[&facet_value_id], + ) + .await?; } + + Ok(()) +} + +fn sqlite_token_uri_state_matches( + conn: &Connection, + result: &TokenUriResult, + expected_attributes: &[(String, String)], +) -> Result { + let token_blob = felt_to_blob(result.contract); + let token_id_blob = u256_to_blob(result.token_id); + let row = conn.query_row( + "SELECT uri, metadata_json FROM token_uris WHERE token = ?1 AND token_id = ?2", + params![&token_blob, &token_id_blob], + |row| { + Ok(( + row.get::>(0)?, + row.get::>(1)?, + )) + }, + ); + let (existing_uri, existing_metadata_json) = match row { + Ok(row) => row, + Err(rusqlite::Error::QueryReturnedNoRows) => return Ok(false), + Err(error) => return Err(error.into()), + }; + + if existing_uri.as_deref() != result.uri.as_deref() + || existing_metadata_json.as_deref() != result.metadata_json.as_deref() + { + return Ok(false); + } + + let mut stmt = conn.prepare( + "SELECT key, value FROM token_attributes + WHERE token = ?1 AND token_id = ?2 + ORDER BY key ASC, value ASC", + )?; + let rows = stmt.query_map(params![&token_blob, &token_id_blob], |row| { + Ok((row.get::(0)?, row.get::(1)?)) + })?; + let existing_attributes = rows.collect::, _>>()?; + + Ok(existing_attributes == expected_attributes) +} + +async fn pg_token_uri_state_matches( + client: &Client, + result: &TokenUriResult, + expected_attributes: &[(String, String)], +) -> Result { + let token_blob = felt_to_blob(result.contract); + let token_id_blob = u256_to_blob(result.token_id); + let row = client + .query_opt( + "SELECT uri, metadata_json FROM erc1155.token_uris WHERE token = $1 AND token_id = $2", + &[&token_blob, &token_id_blob], + ) + .await?; + let Some(row) = row else { + return Ok(false); + }; + + let existing_uri: Option = row.get(0); + let existing_metadata_json: Option = row.get(1); + if existing_uri.as_deref() != result.uri.as_deref() + || existing_metadata_json.as_deref() != result.metadata_json.as_deref() + { + return Ok(false); + } + + let rows = client + .query( + "SELECT key, value FROM erc1155.token_attributes + WHERE token = $1 AND token_id = $2 + ORDER BY key ASC, value ASC", + &[&token_blob, &token_id_blob], + ) + .await?; + let existing_attributes: Vec<(String, String)> = rows + .into_iter() + .map(|row| (row.get(0), row.get(1))) + .collect(); + + Ok(existing_attributes == expected_attributes) } fn sanitize_metadata_text(input: &str) -> Option { diff --git a/crates/torii-erc20/src/handlers.rs b/crates/torii-erc20/src/handlers.rs new file mode 100644 index 00000000..da4bfb52 --- /dev/null +++ b/crates/torii-erc20/src/handlers.rs @@ -0,0 +1,171 @@ +use anyhow::Result; +use async_trait::async_trait; +use prost::Message; +use prost_types::Any; +use starknet::core::types::Felt; +use starknet::providers::jsonrpc::{HttpTransport, JsonRpcClient}; +use std::collections::HashSet; +use std::sync::{Arc, Mutex}; +use torii::command::CommandHandler; +use torii::etl::sink::EventBus; +use torii::UpdateType; +use torii_common::MetadataFetcher; + +use crate::proto; +use crate::storage::Erc20Storage; + +#[derive(Debug, Clone)] +pub struct FetchErc20MetadataCommand { + pub token: Felt, +} + +pub struct Erc20MetadataCommandHandler { + fetcher: Arc, + storage: Arc, + event_bus: Mutex>>, + in_flight: Mutex>, + max_retries: u8, +} + +impl Erc20MetadataCommandHandler { + pub fn new( + provider: Arc>, + storage: Arc, + max_retries: u8, + ) -> Self { + Self { + fetcher: Arc::new(MetadataFetcher::new(provider)), + storage, + event_bus: Mutex::new(None), + in_flight: Mutex::new(HashSet::new()), + max_retries: max_retries.max(1), + } + } + + fn matches_metadata_filters( + metadata: &proto::TokenMetadataEntry, + filters: &std::collections::HashMap, + ) -> bool { + if filters.is_empty() { + return true; + } + + if let Some(token_filter) = filters.get("token") { + let token_hex = format!("0x{}", hex::encode(&metadata.token)); + if !token_hex.eq_ignore_ascii_case(token_filter) { + return false; + } + } + + true + } + + fn metadata_retry_delay(attempt: u8) -> std::time::Duration { + let shift = u32::from(attempt.saturating_sub(1).min(8)); + std::time::Duration::from_millis(250 * (1u64 << shift)) + } +} + +#[async_trait] +impl CommandHandler for Erc20MetadataCommandHandler { + fn supports(&self, command: &dyn torii::command::Command) -> bool { + command.as_any().is::() + } + + fn attach_event_bus(&self, event_bus: Arc) { + *self.event_bus.lock().unwrap() = Some(event_bus); + } + + async fn handle_command(&self, command: Box) -> Result<()> { + let command = command + .into_any() + .downcast::() + .map_err(|_| anyhow::anyhow!("erc20 metadata handler received unexpected command"))?; + let command = *command; + + { + let mut in_flight = self.in_flight.lock().unwrap(); + if !in_flight.insert(command.token) { + tracing::debug!( + target: "torii_erc20::handlers", + token = %format!("{:#x}", command.token), + "Skipping duplicate in-flight ERC20 metadata command" + ); + return Ok(()); + } + } + + let mut attempt = 1; + let result = loop { + let meta = self.fetcher.fetch_erc20_metadata(command.token).await; + let outcome: Result<()> = async { + self.storage + .upsert_token_metadata( + command.token, + meta.name.as_deref(), + meta.symbol.as_deref(), + meta.decimals, + ) + .await?; + + let event_bus = self.event_bus.lock().unwrap().clone(); + if let Some(event_bus) = event_bus { + let meta_entry = proto::TokenMetadataEntry { + token: command.token.to_bytes_be().to_vec(), + name: meta.name, + symbol: meta.symbol, + decimals: meta.decimals.map(|d| d as u32), + }; + + let mut buf = Vec::new(); + meta_entry.encode(&mut buf)?; + let any = Any { + type_url: "type.googleapis.com/torii.sinks.erc20.TokenMetadataEntry" + .to_string(), + value: buf, + }; + + event_bus.publish_protobuf( + "erc20.metadata", + "erc20.metadata", + &any, + &meta_entry, + UpdateType::Created, + Self::matches_metadata_filters, + ); + } + + Ok(()) + } + .await; + + match outcome { + Ok(()) => break Ok(()), + Err(error) if attempt < self.max_retries => { + let next_attempt = attempt + 1; + let delay = Self::metadata_retry_delay(next_attempt); + ::metrics::counter!("torii_erc20_metadata_retries_total").increment(1); + tracing::warn!( + target: "torii_erc20::handlers", + token = %format!("{:#x}", command.token), + attempt = next_attempt, + max_retries = self.max_retries, + delay_ms = delay.as_millis() as u64, + error = %error, + "ERC20 metadata command failed, retrying" + ); + tokio::time::sleep(delay).await; + attempt = next_attempt; + } + Err(error) => { + ::metrics::counter!("torii_erc20_metadata_terminal_failures_total") + .increment(1); + break Err(error); + } + } + }; + + self.in_flight.lock().unwrap().remove(&command.token); + result + } +} diff --git a/crates/torii-erc20/src/lib.rs b/crates/torii-erc20/src/lib.rs index 1eda63d4..70e2b54c 100644 --- a/crates/torii-erc20/src/lib.rs +++ b/crates/torii-erc20/src/lib.rs @@ -38,6 +38,7 @@ pub mod balance_fetcher; pub mod decoder; pub mod grpc_service; +pub mod handlers; pub mod identification; pub mod sink; pub mod storage; @@ -55,6 +56,7 @@ pub const FILE_DESCRIPTOR_SET: &[u8] = include_bytes!("generated/erc20_descripto pub use balance_fetcher::{BalanceFetchRequest, BalanceFetcher}; pub use decoder::{Approval, Erc20Decoder, Transfer}; pub use grpc_service::Erc20Service; +pub use handlers::Erc20MetadataCommandHandler; pub use identification::Erc20Rule; pub use sink::Erc20Sink; pub use storage::{ diff --git a/crates/torii-erc20/src/sink.rs b/crates/torii-erc20/src/sink.rs index a6686f5e..4a09e295 100644 --- a/crates/torii-erc20/src/sink.rs +++ b/crates/torii-erc20/src/sink.rs @@ -14,6 +14,7 @@ use crate::balance_fetcher::BalanceFetcher; use crate::decoder::{Approval as DecodedApproval, Transfer as DecodedTransfer}; use crate::grpc_service::Erc20Service; +use crate::handlers::FetchErc20MetadataCommand; use crate::proto; use crate::storage::{ApprovalData, Erc20Storage, TransferData}; use anyhow::Result; @@ -25,25 +26,16 @@ use starknet::core::types::{Felt, U256}; use starknet::providers::jsonrpc::{HttpTransport, JsonRpcClient}; use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{Arc, Mutex}; -use tokio::sync::{mpsc, Mutex as AsyncMutex}; +use std::sync::Arc; +use torii::command::CommandBusSender; use torii::etl::sink::{EventBus, TopicInfo}; use torii::etl::{Envelope, ExtractionBatch, Sink, TypeId}; use torii::grpc::UpdateType; -use torii_common::{u256_to_bytes, MetadataFetcher}; +use torii_common::u256_to_bytes; /// Default threshold for "live" detection: 100 blocks from chain head. /// Events from blocks older than this won't be broadcast to real-time subscribers. const LIVE_THRESHOLD_BLOCKS: u64 = 100; -const DEFAULT_METADATA_PARALLELISM: usize = 8; -const DEFAULT_METADATA_QUEUE_CAPACITY: usize = 2048; -const DEFAULT_METADATA_MAX_RETRIES: u8 = 5; - -#[derive(Debug, Clone, Copy)] -struct MetadataJob { - token: Felt, - attempt: u8, -} /// ERC20 transfer and approval sink /// @@ -61,18 +53,12 @@ pub struct Erc20Sink { grpc_service: Option, /// Balance fetcher for RPC calls (None = balance tracking disabled) balance_fetcher: Option>, - /// Metadata fetcher for token name/symbol/decimals - metadata_fetcher: Option>, - /// In-flight contract metadata fetches to avoid duplicate jobs. - metadata_fetch_inflight: Arc>>, - /// Background metadata worker job queue. - metadata_job_tx: Option>, - /// Max concurrent metadata workers. - metadata_parallelism: usize, - /// Metadata queue capacity. - metadata_queue_capacity: usize, - /// Maximum metadata retries with capped backoff. - metadata_max_retries: u8, + /// Whether contract metadata commands should be dispatched. + metadata_commands_enabled: bool, + /// Command bus sender for background metadata work. + command_bus: Option, + /// Commands already queued but not yet observed in storage. + pending_metadata_commands: tokio::sync::Mutex>, /// In-memory counters to avoid full-table COUNT(*) in the ingest hot path. total_transfers: AtomicU64, total_approvals: AtomicU64, @@ -85,31 +71,15 @@ impl Erc20Sink { event_bus: None, grpc_service: None, balance_fetcher: None, - metadata_fetcher: None, - metadata_fetch_inflight: Arc::new(Mutex::new(HashSet::new())), - metadata_job_tx: None, - metadata_parallelism: DEFAULT_METADATA_PARALLELISM, - metadata_queue_capacity: DEFAULT_METADATA_QUEUE_CAPACITY, - metadata_max_retries: DEFAULT_METADATA_MAX_RETRIES, + metadata_commands_enabled: false, + command_bus: None, + pending_metadata_commands: tokio::sync::Mutex::new(HashSet::new()), // Avoid startup full-table COUNT(*) scans on large datasets. total_transfers: AtomicU64::new(0), total_approvals: AtomicU64::new(0), } } - /// Configure the async metadata pipeline. - pub fn with_metadata_pipeline( - mut self, - parallelism: usize, - queue_capacity: usize, - max_retries: u8, - ) -> Self { - self.metadata_parallelism = parallelism.max(1); - self.metadata_queue_capacity = queue_capacity.max(1); - self.metadata_max_retries = max_retries.max(1); - self - } - /// Set the gRPC service for dual publishing pub fn with_grpc_service(mut self, service: Erc20Service) -> Self { self.grpc_service = Some(service); @@ -124,8 +94,7 @@ impl Erc20Sink { /// - Fetch actual balance from the chain and adjust /// - Record adjustments in an audit table pub fn with_balance_tracking(mut self, provider: Arc>) -> Self { - self.balance_fetcher = Some(Arc::new(BalanceFetcher::new(provider.clone()))); - self.metadata_fetcher = Some(Arc::new(MetadataFetcher::new(provider))); + self.balance_fetcher = Some(Arc::new(BalanceFetcher::new(provider))); self } @@ -239,34 +208,6 @@ impl Erc20Sink { true } - - /// Filter function for ERC20 token metadata updates. - /// - /// Supports filters: - /// - "token": Filter by token contract address (hex string) - fn matches_metadata_filters( - metadata: &proto::TokenMetadataEntry, - filters: &HashMap, - ) -> bool { - if filters.is_empty() { - return true; - } - - if let Some(token_filter) = filters.get("token") { - let token_hex = format!("0x{}", hex::encode(&metadata.token)); - if !token_hex.eq_ignore_ascii_case(token_filter) { - return false; - } - } - - true - } - - fn metadata_retry_delay(attempt: u8) -> std::time::Duration { - // attempt=1 => 250ms, then doubles up to capped retry count. - let shift = u32::from(attempt.saturating_sub(1).min(8)); - std::time::Duration::from_millis(250 * (1u64 << shift)) - } } #[async_trait] @@ -282,170 +223,10 @@ impl Sink for Erc20Sink { async fn initialize( &mut self, event_bus: Arc, - _context: &torii::etl::sink::SinkContext, + context: &torii::etl::sink::SinkContext, ) -> Result<()> { self.event_bus = Some(event_bus); - - if self.metadata_fetcher.is_some() && self.metadata_job_tx.is_none() { - let (metadata_tx, metadata_rx) = - mpsc::channel::(self.metadata_queue_capacity); - self.metadata_job_tx = Some(metadata_tx.clone()); - - let Some(fetcher) = self.metadata_fetcher.clone() else { - return Ok(()); - }; - let storage = self.storage.clone(); - let event_bus = self.event_bus.clone(); - let inflight = self.metadata_fetch_inflight.clone(); - let max_retries = self.metadata_max_retries; - let metadata_rx = Arc::new(AsyncMutex::new(metadata_rx)); - - let parallelism = self.metadata_parallelism; - for worker_id in 0..parallelism { - let worker_fetcher = fetcher.clone(); - let worker_storage = storage.clone(); - let worker_event_bus = event_bus.clone(); - let worker_inflight = inflight.clone(); - let worker_tx = metadata_tx.clone(); - let worker_rx = metadata_rx.clone(); - - tokio::spawn(async move { - loop { - let maybe_job = { - let mut rx = worker_rx.lock().await; - rx.recv().await - }; - let Some(job) = maybe_job else { - break; - }; - let token = job.token; - - let result: Result<()> = async { - let meta = worker_fetcher.fetch_erc20_metadata(token).await; - tracing::info!( - target: "torii_erc20::sink", - token = %format!("{:#x}", token), - name = ?meta.name, - symbol = ?meta.symbol, - decimals = ?meta.decimals, - "Fetched token metadata" - ); - - worker_storage - .upsert_token_metadata( - token, - meta.name.as_deref(), - meta.symbol.as_deref(), - meta.decimals, - ) - .await?; - - if let Some(event_bus) = &worker_event_bus { - let meta_entry = proto::TokenMetadataEntry { - token: token.to_bytes_be().to_vec(), - name: meta.name, - symbol: meta.symbol, - decimals: meta.decimals.map(|d| d as u32), - }; - - let mut buf = Vec::new(); - meta_entry.encode(&mut buf)?; - let any = Any { - type_url: - "type.googleapis.com/torii.sinks.erc20.TokenMetadataEntry" - .to_string(), - value: buf, - }; - - event_bus.publish_protobuf( - "erc20.metadata", - "erc20.metadata", - &any, - &meta_entry, - UpdateType::Created, - Self::matches_metadata_filters, - ); - } - - Ok(()) - } - .await; - - match result { - Ok(()) => { - let mut inflight = worker_inflight.lock().unwrap(); - inflight.remove(&token); - } - Err(error) => { - if job.attempt < max_retries { - let next_attempt = job.attempt + 1; - let delay = Self::metadata_retry_delay(next_attempt); - ::metrics::counter!("torii_erc20_metadata_retries_total") - .increment(1); - tracing::warn!( - target: "torii_erc20::sink", - token = %format!("{:#x}", token), - attempt = next_attempt, - max_retries, - delay_ms = delay.as_millis() as u64, - error = %error, - "Metadata fetch/store failed, scheduling retry" - ); - - tokio::time::sleep(delay).await; - if let Err(send_err) = worker_tx - .send(MetadataJob { - token, - attempt: next_attempt, - }) - .await - { - tracing::warn!( - target: "torii_erc20::sink", - token = %format!("{:#x}", token), - error = %send_err, - "Failed to enqueue metadata retry job" - ); - let mut inflight = worker_inflight.lock().unwrap(); - inflight.remove(&token); - } - } else { - ::metrics::counter!( - "torii_erc20_metadata_terminal_failures_total" - ) - .increment(1); - tracing::warn!( - target: "torii_erc20::sink", - token = %format!("{:#x}", token), - attempt = job.attempt, - max_retries, - error = %error, - "Metadata fetch/store failed after max retries" - ); - let mut inflight = worker_inflight.lock().unwrap(); - inflight.remove(&token); - } - } - } - } - - tracing::debug!( - target: "torii_erc20::sink", - worker_id, - "ERC20 metadata worker stopped" - ); - }); - } - - tracing::info!( - target: "torii_erc20::sink", - metadata_parallelism = self.metadata_parallelism, - metadata_queue_capacity = self.metadata_queue_capacity, - metadata_max_retries = self.metadata_max_retries, - "Started ERC20 async metadata workers" - ); - } - + self.command_bus = Some(context.command_bus.clone()); tracing::info!(target: "torii_erc20::sink", "ERC20 sink initialized"); Ok(()) } @@ -498,67 +279,62 @@ impl Sink for Erc20Sink { } } - // Enqueue metadata fetches for new token contracts (async workers). - if self.metadata_fetcher.is_some() { - let mut new_tokens: HashSet = HashSet::new(); - for transfer in &transfers { - new_tokens.insert(transfer.token); - } - for approval in &approvals { - new_tokens.insert(approval.token); - } - - let ordered_tokens: Vec = new_tokens.iter().copied().collect(); - let existing_metadata = - match self.storage.has_token_metadata_batch(&ordered_tokens).await { - Ok(existing) => existing, - Err(e) => { - tracing::warn!( - target: "torii_erc20::sink", - error = %e, - "Failed to batch-check token metadata" - ); - HashSet::new() - } - }; - - for token in ordered_tokens - .into_iter() - .filter(|t| !existing_metadata.contains(t)) - { - let scheduled = { - let mut inflight = self.metadata_fetch_inflight.lock().unwrap(); - inflight.insert(token) - }; - if !scheduled { - continue; + // Dispatch metadata fetches for new token contracts. + if self.metadata_commands_enabled { + if let Some(ref command_bus) = self.command_bus { + let mut new_tokens: HashSet = HashSet::new(); + for transfer in &transfers { + new_tokens.insert(transfer.token); + } + for approval in &approvals { + new_tokens.insert(approval.token); } - let Some(metadata_tx) = &self.metadata_job_tx else { - let mut inflight = self.metadata_fetch_inflight.lock().unwrap(); - inflight.remove(&token); - tracing::warn!( - target: "torii_erc20::sink", - token = %format!("{:#x}", token), - "Metadata workers not initialized, skipping async metadata enqueue" - ); - continue; + let candidate_tokens = new_tokens.into_iter(); + let unchecked_tokens = { + let pending = self.pending_metadata_commands.lock().await; + candidate_tokens + .into_iter() + .filter(|token| !pending.contains(token)) + .collect::>() }; - match metadata_tx.try_send(MetadataJob { token, attempt: 1 }) { - Ok(()) => { - ::metrics::counter!("torii_erc20_metadata_jobs_enqueued_total") - .increment(1); + match self + .storage + .has_token_metadata_batch(&unchecked_tokens) + .await + { + Ok(existing_metadata) => { + let mut pending = self.pending_metadata_commands.lock().await; + for token in &existing_metadata { + pending.remove(token); + } + + for token in unchecked_tokens + .into_iter() + .filter(|token| !existing_metadata.contains(token)) + { + if !pending.insert(token) { + continue; + } + if let Err(error) = + command_bus.dispatch(FetchErc20MetadataCommand { token }) + { + pending.remove(&token); + tracing::warn!( + target: "torii_erc20::sink", + token = %format!("{:#x}", token), + error = %error, + "Failed to dispatch ERC20 metadata command" + ); + } + } } - Err(e) => { - let mut inflight = self.metadata_fetch_inflight.lock().unwrap(); - inflight.remove(&token); - ::metrics::counter!("torii_erc20_metadata_jobs_dropped_total").increment(1); + Err(error) => { tracing::warn!( target: "torii_erc20::sink", - token = %format!("{:#x}", token), - error = %e, - "Failed to enqueue metadata job" + error = %error, + "Failed to batch-check token metadata" ); } } @@ -835,3 +611,16 @@ impl Sink for Erc20Sink { Router::new() } } + +impl Erc20Sink { + /// Enable background metadata commands. + pub fn with_metadata_pipeline( + mut self, + _parallelism: usize, + _queue_capacity: usize, + _max_retries: u8, + ) -> Self { + self.metadata_commands_enabled = true; + self + } +} diff --git a/crates/torii-erc721/src/grpc_service.rs b/crates/torii-erc721/src/grpc_service.rs index 39fbf49d..ae25479a 100644 --- a/crates/torii-erc721/src/grpc_service.rs +++ b/crates/torii-erc721/src/grpc_service.rs @@ -188,7 +188,7 @@ impl Erc721Service { let uri_rows = self .storage - .get_token_uris_by_contract(contract) + .get_token_uris_batch(contract, &result.token_ids) .await .map_err(|e| Status::internal(format!("Query failed: {e}")))?; let by_token_id: HashMap, (Option, Option)> = uri_rows diff --git a/crates/torii-erc721/src/handlers.rs b/crates/torii-erc721/src/handlers.rs new file mode 100644 index 00000000..e121cb00 --- /dev/null +++ b/crates/torii-erc721/src/handlers.rs @@ -0,0 +1,273 @@ +use anyhow::Result; +use async_trait::async_trait; +use prost::Message; +use prost_types::Any; +use starknet::core::types::Felt; +use starknet::providers::jsonrpc::{HttpTransport, JsonRpcClient}; +use std::collections::HashSet; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; +use torii::command::CommandHandler; +use torii::etl::sink::EventBus; +use torii::UpdateType; +use torii_common::{process_token_uri_request, u256_to_bytes, MetadataFetcher, TokenUriRequest}; + +use crate::proto; +use crate::storage::Erc721Storage; + +#[derive(Debug, Clone)] +pub struct FetchErc721MetadataCommand { + pub token: Felt, +} + +#[derive(Debug, Clone)] +pub struct RefreshErc721TokenUriCommand { + pub contract: Felt, + pub token_id: starknet::core::types::U256, +} + +pub struct Erc721MetadataCommandHandler { + fetcher: Arc, + storage: Arc, + event_bus: Mutex>>, + in_flight: Mutex>, + max_retries: u8, +} + +impl Erc721MetadataCommandHandler { + pub fn new( + provider: Arc>, + storage: Arc, + max_retries: u8, + ) -> Self { + Self { + fetcher: Arc::new(MetadataFetcher::new(provider)), + storage, + event_bus: Mutex::new(None), + in_flight: Mutex::new(HashSet::new()), + max_retries: max_retries.max(1), + } + } + + fn matches_metadata_filters( + metadata: &proto::TokenMetadataEntry, + filters: &std::collections::HashMap, + ) -> bool { + if filters.is_empty() { + return true; + } + + if let Some(token_filter) = filters.get("token") { + let token_hex = format!("0x{}", hex::encode(&metadata.token)); + if !token_hex.eq_ignore_ascii_case(token_filter) { + return false; + } + } + + true + } + + fn metadata_complete(meta: &torii_common::TokenMetadata) -> bool { + meta.name + .as_deref() + .is_some_and(|value| !value.trim().is_empty()) + && meta + .symbol + .as_deref() + .is_some_and(|value| !value.trim().is_empty()) + && meta.total_supply.is_some() + } + + fn metadata_retry_delay(attempt: u8) -> std::time::Duration { + let shift = u32::from(attempt.saturating_sub(1).min(8)); + std::time::Duration::from_millis(250 * (1u64 << shift)) + } +} + +#[async_trait] +impl CommandHandler for Erc721MetadataCommandHandler { + fn supports(&self, command: &dyn torii::command::Command) -> bool { + command.as_any().is::() + } + + fn attach_event_bus(&self, event_bus: Arc) { + *self.event_bus.lock().unwrap() = Some(event_bus); + } + + async fn handle_command(&self, command: Box) -> Result<()> { + let command = command + .into_any() + .downcast::() + .map_err(|_| anyhow::anyhow!("erc721 metadata handler received unexpected command"))?; + let command = *command; + + { + let mut in_flight = self.in_flight.lock().unwrap(); + if !in_flight.insert(command.token) { + tracing::debug!( + target: "torii_erc721::handlers", + token = %format!("{:#x}", command.token), + "Skipping duplicate in-flight ERC721 metadata command" + ); + return Ok(()); + } + } + + let mut attempt = 1; + let result = loop { + let meta = self.fetcher.fetch_erc721_metadata(command.token).await; + if !Self::metadata_complete(&meta) && attempt < self.max_retries { + let next_attempt = attempt + 1; + let delay = Self::metadata_retry_delay(next_attempt); + tracing::warn!( + target: "torii_erc721::handlers", + token = %format!("{:#x}", command.token), + attempt = next_attempt, + max_retries = self.max_retries, + delay_ms = delay.as_millis() as u64, + has_name = meta.name.as_deref().is_some_and(|value| !value.trim().is_empty()), + has_symbol = meta.symbol.as_deref().is_some_and(|value| !value.trim().is_empty()), + has_total_supply = meta.total_supply.is_some(), + "ERC721 metadata fetch returned incomplete data, retrying" + ); + tokio::time::sleep(delay).await; + attempt = next_attempt; + continue; + } + + let outcome: Result<()> = async { + self.storage + .upsert_token_metadata( + command.token, + meta.name.as_deref(), + meta.symbol.as_deref(), + meta.total_supply, + ) + .await?; + + let event_bus = self.event_bus.lock().unwrap().clone(); + if let Some(event_bus) = event_bus { + let meta_entry = proto::TokenMetadataEntry { + token: command.token.to_bytes_be().to_vec(), + name: meta.name, + symbol: meta.symbol, + total_supply: meta.total_supply.map(u256_to_bytes), + }; + + let mut buf = Vec::new(); + meta_entry.encode(&mut buf)?; + let any = Any { + type_url: "type.googleapis.com/torii.sinks.erc721.TokenMetadataEntry" + .to_string(), + value: buf, + }; + + event_bus.publish_protobuf( + "erc721.metadata", + "erc721.metadata", + &any, + &meta_entry, + UpdateType::Created, + Self::matches_metadata_filters, + ); + } + + Ok(()) + } + .await; + + match outcome { + Ok(()) => break Ok(()), + Err(error) if attempt < self.max_retries => { + let next_attempt = attempt + 1; + let delay = Self::metadata_retry_delay(next_attempt); + tracing::warn!( + target: "torii_erc721::handlers", + token = %format!("{:#x}", command.token), + attempt = next_attempt, + max_retries = self.max_retries, + delay_ms = delay.as_millis() as u64, + error = %error, + "ERC721 metadata command failed, retrying" + ); + tokio::time::sleep(delay).await; + attempt = next_attempt; + } + Err(error) => { + break Err(error); + } + } + }; + + self.in_flight.lock().unwrap().remove(&command.token); + result + } +} + +pub struct Erc721TokenUriCommandHandler { + fetcher: Arc, + storage: Arc, + image_cache_dir: Option, + in_flight: Mutex>, +} + +impl Erc721TokenUriCommandHandler { + pub fn new( + provider: Arc>, + storage: Arc, + image_cache_dir: Option, + ) -> Self { + Self { + fetcher: Arc::new(MetadataFetcher::new(provider)), + storage, + image_cache_dir, + in_flight: Mutex::new(HashSet::new()), + } + } +} + +#[async_trait] +impl CommandHandler for Erc721TokenUriCommandHandler { + fn supports(&self, command: &dyn torii::command::Command) -> bool { + command.as_any().is::() + } + + async fn handle_command(&self, command: Box) -> Result<()> { + let command = command + .into_any() + .downcast::() + .map_err(|_| anyhow::anyhow!("erc721 token uri handler received unexpected command"))?; + let command = *command; + + { + let mut in_flight = self.in_flight.lock().unwrap(); + if !in_flight.insert((command.contract, command.token_id)) { + tracing::debug!( + target: "torii_erc721::handlers", + contract = %format!("{:#x}", command.contract), + token_id = %command.token_id, + "Skipping duplicate in-flight ERC721 token URI command" + ); + return Ok(()); + } + } + + let result = process_token_uri_request( + self.fetcher.as_ref(), + self.storage.as_ref(), + &TokenUriRequest { + contract: command.contract, + token_id: command.token_id, + standard: torii_common::TokenStandard::Erc721, + }, + self.image_cache_dir.as_deref(), + ) + .await; + + self.in_flight + .lock() + .unwrap() + .remove(&(command.contract, command.token_id)); + result + } +} diff --git a/crates/torii-erc721/src/lib.rs b/crates/torii-erc721/src/lib.rs index 47ad159d..ec73d2d1 100644 --- a/crates/torii-erc721/src/lib.rs +++ b/crates/torii-erc721/src/lib.rs @@ -37,6 +37,7 @@ pub mod decoder; pub mod grpc_service; +pub mod handlers; pub mod identification; pub mod sink; pub mod storage; @@ -55,6 +56,7 @@ pub use decoder::{ BatchMetadataUpdate, Erc721Decoder, MetadataUpdate, NftApproval, NftTransfer, OperatorApproval, }; pub use grpc_service::Erc721Service; +pub use handlers::{Erc721MetadataCommandHandler, Erc721TokenUriCommandHandler}; pub use identification::Erc721Rule; pub use sink::Erc721Sink; pub use storage::{Erc721Storage, NftOwnershipData, NftTransferData, TransferCursor}; diff --git a/crates/torii-erc721/src/sink.rs b/crates/torii-erc721/src/sink.rs index a10a7156..cd8c8438 100644 --- a/crates/torii-erc721/src/sink.rs +++ b/crates/torii-erc721/src/sink.rs @@ -5,6 +5,7 @@ use crate::decoder::{ NftTransfer as DecodedNftTransfer, OperatorApproval as DecodedOperatorApproval, }; use crate::grpc_service::Erc721Service; +use crate::handlers::{FetchErc721MetadataCommand, RefreshErc721TokenUriCommand}; use crate::proto; use crate::storage::{Erc721Storage, NftTransferData, OperatorApprovalData}; use anyhow::Result; @@ -12,17 +13,16 @@ use async_trait::async_trait; use axum::Router; use prost::Message; use prost_types::Any; -use starknet::core::types::Felt; +use starknet::core::types::{Felt, U256}; use starknet::providers::jsonrpc::{HttpTransport, JsonRpcClient}; use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; +use torii::command::CommandBusSender; use torii::etl::sink::{EventBus, TopicInfo}; use torii::etl::{Envelope, ExtractionBatch, Sink, TypeId}; use torii::grpc::UpdateType; -use torii_common::{ - u256_to_bytes, MetadataFetcher, TokenStandard, TokenUriRequest, TokenUriSender, -}; +use torii_common::{u256_to_bytes, TokenStandard, TokenUriRequest, TokenUriSender}; /// Default threshold for "live" detection: 100 blocks from chain head. /// Events from blocks older than this won't be broadcast to real-time subscribers. @@ -41,12 +41,16 @@ pub struct Erc721Sink { storage: Arc, event_bus: Option>, grpc_service: Option, - /// Metadata fetcher for token name/symbol - metadata_fetcher: Option>, - /// Token URI service sender for async URI fetching + /// Whether contract metadata commands should be dispatched. + metadata_commands_enabled: bool, + /// Whether token URI commands should be dispatched. + token_uri_commands_enabled: bool, + /// Command bus sender for background metadata and token URI work. + command_bus: Option, token_uri_sender: Option, - /// In-flight contract metadata fetches to avoid duplicate background jobs. - metadata_fetch_inflight: Arc>>, + /// Commands already queued but not yet observed in storage. + pending_metadata_commands: tokio::sync::Mutex>, + pending_token_uri_commands: tokio::sync::Mutex>, /// In-memory counters to avoid full-table COUNT(*) in the ingest hot path. total_transfers: AtomicU64, total_operator_approvals: AtomicU64, @@ -58,27 +62,18 @@ impl Erc721Sink { storage, event_bus: None, grpc_service: None, - metadata_fetcher: None, + metadata_commands_enabled: false, + token_uri_commands_enabled: false, + command_bus: None, token_uri_sender: None, - metadata_fetch_inflight: Arc::new(Mutex::new(HashSet::new())), + pending_metadata_commands: tokio::sync::Mutex::new(HashSet::new()), + pending_token_uri_commands: tokio::sync::Mutex::new(HashSet::new()), // Avoid startup full-table COUNT(*) scans on large datasets. total_transfers: AtomicU64::new(0), total_operator_approvals: AtomicU64::new(0), } } - /// Enable metadata fetching with a provider - pub fn with_metadata_fetching(mut self, provider: Arc>) -> Self { - self.metadata_fetcher = Some(Arc::new(MetadataFetcher::new(provider))); - self - } - - /// Enable async token URI fetching - pub fn with_token_uri_sender(mut self, sender: TokenUriSender) -> Self { - self.token_uri_sender = Some(sender); - self - } - /// Set the gRPC service for dual publishing pub fn with_grpc_service(mut self, service: Erc721Service) -> Self { self.grpc_service = Some(service); @@ -137,26 +132,64 @@ impl Erc721Sink { true } - /// Filter function for ERC721 token metadata updates. - /// - /// Supports filters: - /// - "token": Filter by token contract address (hex string) - fn matches_metadata_filters( - metadata: &proto::TokenMetadataEntry, - filters: &HashMap, - ) -> bool { - if filters.is_empty() { - return true; + fn enqueue_token_uri_request(&self, contract: Felt, token_id: U256) -> bool { + if let Some(sender) = &self.token_uri_sender { + return sender.request_update(TokenUriRequest { + contract, + token_id, + standard: TokenStandard::Erc721, + }); } - if let Some(token_filter) = filters.get("token") { - let token_hex = format!("0x{}", hex::encode(&metadata.token)); - if !token_hex.eq_ignore_ascii_case(token_filter) { + if let Some(command_bus) = &self.command_bus { + if let Err(error) = + command_bus.dispatch(RefreshErc721TokenUriCommand { contract, token_id }) + { + tracing::warn!( + target: "torii_erc721::sink", + error = %error, + "Failed to dispatch ERC721 token URI refresh" + ); return false; } + return true; } - true + false + } + + async fn enqueue_token_uri_range_refresh( + &self, + contract: Felt, + from_token_id: U256, + to_token_id: U256, + ) { + let Ok(uris) = self.storage.get_token_uris_by_contract(contract).await else { + return; + }; + let token_ids = uris + .into_iter() + .map(|(token_id, _, _)| token_id) + .filter(|token_id| *token_id >= from_token_id && *token_id <= to_token_id) + .collect::>(); + + if let Some(sender) = &self.token_uri_sender { + let accepted = sender.request_batch(contract, &token_ids, TokenStandard::Erc721); + if accepted != token_ids.len() { + tracing::warn!( + target: "torii_erc721::sink", + contract = %format!("{:#x}", contract), + requested = token_ids.len(), + accepted, + "Failed to enqueue some ERC721 batch token URI refreshes" + ); + } + return; + } + + for token_id in token_ids { + self.enqueue_token_uri_request(contract, token_id); + } } } @@ -179,9 +212,10 @@ impl Sink for Erc721Sink { async fn initialize( &mut self, event_bus: Arc, - _context: &torii::etl::sink::SinkContext, + context: &torii::etl::sink::SinkContext, ) -> Result<()> { self.event_bus = Some(event_bus); + self.command_bus = Some(context.command_bus.clone()); tracing::info!(target: "torii_erc721::sink", "ERC721 sink initialized"); Ok(()) } @@ -244,12 +278,8 @@ impl Sink for Erc721Sink { .as_any() .downcast_ref::() { - if let Some(ref sender) = self.token_uri_sender { - sender.request_update(TokenUriRequest { - contract: update.token, - token_id: update.token_id, - standard: TokenStandard::Erc721, - }); + if self.token_uri_commands_enabled { + self.enqueue_token_uri_request(update.token, update.token_id); } } } @@ -260,24 +290,13 @@ impl Sink for Erc721Sink { .as_any() .downcast_ref::() { - if let Some(ref sender) = self.token_uri_sender { - // For batch updates, we need to know which token IDs exist in the range. - // Fetch them from storage and request URI updates for each. - if let Ok(uris) = - self.storage.get_token_uris_by_contract(update.token).await - { - for (token_id, _, _) in &uris { - if *token_id >= update.from_token_id - && *token_id <= update.to_token_id - { - sender.request_update(TokenUriRequest { - contract: update.token, - token_id: *token_id, - standard: TokenStandard::Erc721, - }); - } - } - } + if self.token_uri_commands_enabled { + self.enqueue_token_uri_range_refresh( + update.token, + update.from_token_id, + update.to_token_id, + ) + .await; } } } @@ -285,122 +304,106 @@ impl Sink for Erc721Sink { // but is less commonly needed for indexing purposes } - // Fetch metadata for any new token contracts - if let Some(ref fetcher) = self.metadata_fetcher { - let new_tokens: HashSet = transfers.iter().map(|t| t.token).collect(); - for token in new_tokens { - match self.storage.has_token_metadata(token).await { - Ok(exists) => { - if exists { - continue; - } + // Fetch metadata for any new token contracts. + if self.metadata_commands_enabled { + if let Some(ref command_bus) = self.command_bus { + let candidate_tokens = transfers + .iter() + .map(|transfer| transfer.token) + .collect::>() + .into_iter(); + let unchecked_tokens = { + let pending = self.pending_metadata_commands.lock().await; + candidate_tokens + .into_iter() + .filter(|token| !pending.contains(token)) + .collect::>() + }; - let scheduled = { - let mut inflight = self.metadata_fetch_inflight.lock().unwrap(); - inflight.insert(token) - }; - if !scheduled { - continue; + match self + .storage + .has_token_metadata_batch(&unchecked_tokens) + .await + { + Ok(existing_tokens) => { + let mut pending = self.pending_metadata_commands.lock().await; + for token in &existing_tokens { + pending.remove(token); } - let fetcher = fetcher.clone(); - let storage = self.storage.clone(); - let event_bus = self.event_bus.clone(); - let inflight = self.metadata_fetch_inflight.clone(); - tokio::spawn(async move { - let result = async { - let meta = fetcher.fetch_erc721_metadata(token).await; - tracing::info!( - target: "torii_erc721::sink", - token = %format!("{:#x}", token), - name = ?meta.name, - symbol = ?meta.symbol, - "Fetched token metadata" - ); - storage - .upsert_token_metadata( - token, - meta.name.as_deref(), - meta.symbol.as_deref(), - meta.total_supply, - ) - .await?; - - if let Some(event_bus) = &event_bus { - let meta_entry = proto::TokenMetadataEntry { - token: token.to_bytes_be().to_vec(), - name: meta.name, - symbol: meta.symbol, - total_supply: meta.total_supply.map(u256_to_bytes), - }; - - let mut buf = Vec::new(); - meta_entry.encode(&mut buf)?; - let any = Any { - type_url: - "type.googleapis.com/torii.sinks.erc721.TokenMetadataEntry" - .to_string(), - value: buf, - }; - - event_bus.publish_protobuf( - "erc721.metadata", - "erc721.metadata", - &any, - &meta_entry, - UpdateType::Created, - Self::matches_metadata_filters, - ); - } - - Ok::<(), anyhow::Error>(()) + for token in unchecked_tokens + .into_iter() + .filter(|token| !existing_tokens.contains(token)) + { + if !pending.insert(token) { + continue; } - .await; - - if let Err(e) = result { + if let Err(error) = + command_bus.dispatch(FetchErc721MetadataCommand { token }) + { + pending.remove(&token); tracing::warn!( target: "torii_erc721::sink", token = %format!("{:#x}", token), - error = %e, - "Failed to fetch/store token metadata in background task" + error = %error, + "Failed to dispatch ERC721 metadata command" ); } - - let mut inflight = inflight.lock().unwrap(); - inflight.remove(&token); - }); + } } - Err(e) => { - tracing::warn!(target: "torii_erc721::sink", error = %e, "Failed to check token metadata"); + Err(error) => { + tracing::warn!( + target: "torii_erc721::sink", + error = %error, + "Failed to batch-check token metadata" + ); } } } } - // Request token URI fetches for new token IDs - if let Some(ref sender) = self.token_uri_sender { - for transfer in &transfers { - match self - .storage - .has_token_uri(transfer.token, transfer.token_id) - .await - { - Ok(false) => { - sender.request_update(TokenUriRequest { - contract: transfer.token, - token_id: transfer.token_id, - standard: TokenStandard::Erc721, - }); + // Request token URI fetches for new token IDs. + if self.token_uri_commands_enabled + && (self.token_uri_sender.is_some() || self.command_bus.is_some()) + { + let candidate_tokens = transfers + .iter() + .map(|transfer| (transfer.token, transfer.token_id)) + .collect::>() + .into_iter(); + let unchecked_tokens = { + let pending = self.pending_token_uri_commands.lock().await; + candidate_tokens + .filter(|token| !pending.contains(token)) + .collect::>() + }; + + match self.storage.has_token_uri_batch(&unchecked_tokens).await { + Ok(existing_tokens) => { + let mut pending = self.pending_token_uri_commands.lock().await; + for token in &existing_tokens { + pending.remove(token); } - Ok(true) => {} - Err(e) => { - tracing::warn!( - target: "torii_erc721::sink", - error = %e, - "Failed to check token URI existence" - ); + + for (contract, token_id) in unchecked_tokens + .into_iter() + .filter(|token| !existing_tokens.contains(token)) + { + if !pending.insert((contract, token_id)) { + continue; + } + if !self.enqueue_token_uri_request(contract, token_id) { + pending.remove(&(contract, token_id)); + } } } + Err(error) => { + tracing::warn!( + target: "torii_erc721::sink", + error = %error, + "Failed to batch-check token URI existence" + ); + } } } @@ -544,3 +547,25 @@ impl Sink for Erc721Sink { Router::new() } } + +impl Erc721Sink { + pub fn with_metadata_commands(mut self) -> Self { + self.metadata_commands_enabled = true; + self + } + + pub fn with_metadata_fetching(self, _provider: Arc>) -> Self { + self.with_metadata_commands() + } + + pub fn with_token_uri_commands(mut self) -> Self { + self.token_uri_commands_enabled = true; + self + } + + pub fn with_token_uri_sender(mut self, sender: TokenUriSender) -> Self { + self.token_uri_commands_enabled = true; + self.token_uri_sender = Some(sender); + self + } +} diff --git a/crates/torii-erc721/src/storage.rs b/crates/torii-erc721/src/storage.rs index 06d11e17..0bfd23f5 100644 --- a/crates/torii-erc721/src/storage.rs +++ b/crates/torii-erc721/src/storage.rs @@ -4,14 +4,19 @@ //! Tracks current NFT ownership state (who owns each token). use anyhow::Result; -use rusqlite::{params, Connection}; +use rusqlite::{params, params_from_iter, Connection, ToSql}; use starknet::core::types::{Felt, U256}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::{Arc, Mutex}; use tokio_postgres::{types::ToSql as PgToSql, Client, NoTls}; use torii_common::{ blob_to_felt, blob_to_u256, felt_to_blob, u256_to_blob, TokenUriResult, TokenUriStore, }; +const SQLITE_MAX_BIND_VARS: usize = 900; +const SQLITE_TOKEN_BATCH_SIZE: usize = SQLITE_MAX_BIND_VARS; +const SQLITE_TOKEN_PAIR_BATCH_SIZE: usize = SQLITE_MAX_BIND_VARS / 2; + /// Storage for ERC721 NFT data pub struct Erc721Storage { backend: StorageBackend, @@ -99,6 +104,12 @@ pub struct TokenAttributeQueryResult { pub facets: Vec, } +#[derive(Debug, Clone)] +struct ResolvedFacetFilter { + key_id: i64, + value_ids: Vec, +} + impl Erc721Storage { /// Create or open the database pub async fn new(db_path: &str) -> Result { @@ -206,6 +217,36 @@ impl Erc721Storage { CREATE INDEX IF NOT EXISTS idx_token_attributes_key_value ON erc721.token_attributes(key, value); CREATE INDEX IF NOT EXISTS idx_token_attributes_token_key_value ON erc721.token_attributes(token, key, value); CREATE INDEX IF NOT EXISTS idx_token_attributes_token_token_id ON erc721.token_attributes(token, token_id); + + CREATE TABLE IF NOT EXISTS erc721.facet_keys ( + id BIGSERIAL PRIMARY KEY, + token BYTEA NOT NULL, + key_norm TEXT NOT NULL, + key_display TEXT NOT NULL, + UNIQUE(token, key_norm) + ); + CREATE INDEX IF NOT EXISTS idx_facet_keys_token_key_norm ON erc721.facet_keys(token, key_norm); + + CREATE TABLE IF NOT EXISTS erc721.facet_values ( + id BIGSERIAL PRIMARY KEY, + token BYTEA NOT NULL, + facet_key_id BIGINT NOT NULL REFERENCES erc721.facet_keys(id) ON DELETE CASCADE, + value_norm TEXT NOT NULL, + value_display TEXT NOT NULL, + token_count BIGINT NOT NULL DEFAULT 0, + UNIQUE(token, facet_key_id, value_norm) + ); + CREATE INDEX IF NOT EXISTS idx_facet_values_token_key_value ON erc721.facet_values(token, facet_key_id, value_norm); + + CREATE TABLE IF NOT EXISTS erc721.facet_token_map ( + token BYTEA NOT NULL, + token_id BYTEA NOT NULL, + facet_key_id BIGINT NOT NULL REFERENCES erc721.facet_keys(id) ON DELETE CASCADE, + facet_value_id BIGINT NOT NULL REFERENCES erc721.facet_values(id) ON DELETE CASCADE, + PRIMARY KEY (token, token_id, facet_key_id) + ); + CREATE INDEX IF NOT EXISTS idx_facet_token_map_token_value_token_id ON erc721.facet_token_map(token, facet_value_id, token_id); + CREATE INDEX IF NOT EXISTS idx_facet_token_map_token_token_id_key_value ON erc721.facet_token_map(token, token_id, facet_key_id, facet_value_id); ", ).await?; @@ -391,7 +432,40 @@ impl Erc721Storage { CREATE INDEX IF NOT EXISTS idx_token_attributes_key ON token_attributes(key); CREATE INDEX IF NOT EXISTS idx_token_attributes_key_value ON token_attributes(key, value); CREATE INDEX IF NOT EXISTS idx_token_attributes_token_key_value ON token_attributes(token, key, value); - CREATE INDEX IF NOT EXISTS idx_token_attributes_token_token_id ON token_attributes(token, token_id);", + CREATE INDEX IF NOT EXISTS idx_token_attributes_token_token_id ON token_attributes(token, token_id); + + CREATE TABLE IF NOT EXISTS facet_keys ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + token BLOB NOT NULL, + key_norm TEXT NOT NULL, + key_display TEXT NOT NULL, + UNIQUE(token, key_norm) + ); + CREATE INDEX IF NOT EXISTS idx_facet_keys_token_key_norm ON facet_keys(token, key_norm); + + CREATE TABLE IF NOT EXISTS facet_values ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + token BLOB NOT NULL, + facet_key_id INTEGER NOT NULL, + value_norm TEXT NOT NULL, + value_display TEXT NOT NULL, + token_count INTEGER NOT NULL DEFAULT 0, + UNIQUE(token, facet_key_id, value_norm), + FOREIGN KEY (facet_key_id) REFERENCES facet_keys(id) ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS idx_facet_values_token_key_value ON facet_values(token, facet_key_id, value_norm); + + CREATE TABLE IF NOT EXISTS facet_token_map ( + token BLOB NOT NULL, + token_id BLOB NOT NULL, + facet_key_id INTEGER NOT NULL, + facet_value_id INTEGER NOT NULL, + PRIMARY KEY (token, token_id, facet_key_id), + FOREIGN KEY (facet_key_id) REFERENCES facet_keys(id) ON DELETE CASCADE, + FOREIGN KEY (facet_value_id) REFERENCES facet_values(id) ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS idx_facet_token_map_token_value_token_id ON facet_token_map(token, facet_value_id, token_id); + CREATE INDEX IF NOT EXISTS idx_facet_token_map_token_token_id_key_value ON facet_token_map(token, token_id, facet_key_id, facet_value_id);", )?; tracing::info!(target: "torii_erc721::storage", db_path = %db_path, "ERC721 database initialized"); @@ -805,10 +879,19 @@ impl Erc721Storage { let facet_limit = facet_limit.clamp(1, 1000) as i64; if self.backend == StorageBackend::Postgres { + let resolved_filters = self.pg_resolve_facet_filters(token, &filters).await?; + let Some(resolved_filters) = resolved_filters else { + return Ok(TokenAttributeQueryResult { + token_ids: Vec::new(), + next_cursor_token_id: None, + total_hits: 0, + facets: Vec::new(), + }); + }; return self - .pg_query_token_ids_by_attributes( + .pg_query_token_ids_by_facets( token, - &filters, + &resolved_filters, cursor_token_id, page_fetch, include_facets, @@ -818,95 +901,25 @@ impl Erc721Storage { } let conn = self.conn.lock().unwrap(); - - let mut query = String::from("SELECT u.token_id FROM token_uris u WHERE u.token = ?"); - let mut params: Vec> = vec![Box::new(felt_to_blob(token))]; - if let Some(cursor) = cursor_token_id { - query.push_str(" AND u.token_id > ?"); - params.push(Box::new(u256_to_blob(cursor))); - } - Self::sqlite_append_attribute_exists_filters(&mut query, &mut params, &filters, "u"); - query.push_str(" ORDER BY u.token_id ASC LIMIT ?"); - params.push(Box::new(page_fetch)); - - let mut stmt = conn.prepare(&query)?; - let params_refs: Vec<&dyn rusqlite::ToSql> = - params.iter().map(std::convert::AsRef::as_ref).collect(); - let rows = stmt.query_map(params_refs.as_slice(), |row| row.get::<_, Vec>(0))?; - let mut token_ids: Vec = rows - .collect::>, _>>()? - .into_iter() - .map(|b| blob_to_u256(&b)) - .collect(); - - let next_cursor_token_id = if token_ids.len() > page_limit as usize { - let next = token_ids[page_limit as usize]; - token_ids.truncate(page_limit as usize); - Some(next) - } else { - None - }; - - let mut count_query = String::from("SELECT COUNT(*) FROM token_uris u WHERE u.token = ?"); - let mut count_params: Vec> = vec![Box::new(felt_to_blob(token))]; - Self::sqlite_append_attribute_exists_filters( - &mut count_query, - &mut count_params, - &filters, - "u", - ); - let count_refs: Vec<&dyn rusqlite::ToSql> = count_params - .iter() - .map(std::convert::AsRef::as_ref) - .collect(); - let total_hits: i64 = - conn.query_row(&count_query, count_refs.as_slice(), |row| row.get(0))?; - - let facets = if include_facets { - let mut facet_query = String::from( - "SELECT a.key, a.value, COUNT(*) AS cnt - FROM token_attributes a - JOIN token_uris u ON u.token = a.token AND u.token_id = a.token_id - WHERE u.token = ?", - ); - let mut facet_params: Vec> = - vec![Box::new(felt_to_blob(token))]; - Self::sqlite_append_attribute_exists_filters( - &mut facet_query, - &mut facet_params, - &filters, - "u", - ); - facet_query.push_str( - " GROUP BY a.key, a.value - ORDER BY cnt DESC, a.key ASC, a.value ASC - LIMIT ?", - ); - facet_params.push(Box::new(facet_limit)); - - let facet_refs: Vec<&dyn rusqlite::ToSql> = facet_params - .iter() - .map(std::convert::AsRef::as_ref) - .collect(); - let mut stmt = conn.prepare(&facet_query)?; - let rows = stmt.query_map(facet_refs.as_slice(), |row| { - Ok(AttributeFacetCount { - key: row.get(0)?, - value: row.get(1)?, - count: row.get::<_, i64>(2)? as u64, - }) - })?; - rows.collect::, _>>()? - } else { - Vec::new() + let resolved_filters = self.sqlite_resolve_facet_filters(&conn, token, &filters)?; + let Some(resolved_filters) = resolved_filters else { + return Ok(TokenAttributeQueryResult { + token_ids: Vec::new(), + next_cursor_token_id: None, + total_hits: 0, + facets: Vec::new(), + }); }; - Ok(TokenAttributeQueryResult { - token_ids, - next_cursor_token_id, - total_hits: total_hits as u64, - facets, - }) + self.sqlite_query_token_ids_by_facets( + &conn, + token, + &resolved_filters, + cursor_token_id, + page_fetch, + include_facets, + facet_limit, + ) } /// Get transfer count @@ -969,13 +982,57 @@ impl Erc721Storage { let conn = self.conn.lock().unwrap(); let token_blob = felt_to_blob(token); let count: i64 = conn.query_row( - "SELECT COUNT(*) FROM token_metadata WHERE token = ?", + "SELECT COUNT(*) FROM token_metadata + WHERE token = ? + AND name IS NOT NULL + AND TRIM(name) <> '' + AND symbol IS NOT NULL + AND TRIM(symbol) <> '' + AND total_supply IS NOT NULL", params![&token_blob], |row| row.get(0), )?; Ok(count > 0) } + pub async fn has_token_metadata_batch(&self, tokens: &[Felt]) -> Result> { + if tokens.is_empty() { + return Ok(HashSet::new()); + } + if self.backend == StorageBackend::Postgres { + return self.pg_has_token_metadata_batch(tokens).await; + } + + let conn = self.conn.lock().unwrap(); + let mut existing = HashSet::with_capacity(tokens.len()); + + for chunk in tokens.chunks(SQLITE_TOKEN_BATCH_SIZE) { + let placeholders = vec!["?"; chunk.len()].join(","); + let query = format!( + "SELECT token FROM token_metadata + WHERE token IN ({placeholders}) + AND name IS NOT NULL + AND TRIM(name) <> '' + AND symbol IS NOT NULL + AND TRIM(symbol) <> '' + AND total_supply IS NOT NULL" + ); + let token_blobs: Vec> = + chunk.iter().map(|token| felt_to_blob(*token)).collect(); + let mut stmt = conn.prepare(&query)?; + let rows = stmt.query_map(params_from_iter(token_blobs.iter()), |row| { + let token_bytes: Vec = row.get(0)?; + Ok(blob_to_felt(&token_bytes)) + })?; + + for row in rows { + existing.insert(row?); + } + } + + Ok(existing) + } + /// Insert or update token metadata pub async fn upsert_token_metadata( &self, @@ -1133,13 +1190,54 @@ impl Erc721Storage { let token_blob = felt_to_blob(token); let token_id_blob = u256_to_blob(token_id); let count: i64 = conn.query_row( - "SELECT COUNT(*) FROM token_uris WHERE token = ?1 AND token_id = ?2", + "SELECT COUNT(*) FROM token_uris + WHERE token = ?1 + AND token_id = ?2", params![&token_blob, &token_id_blob], |row| row.get(0), )?; Ok(count > 0) } + pub async fn has_token_uri_batch( + &self, + tokens: &[(Felt, U256)], + ) -> Result> { + if tokens.is_empty() { + return Ok(HashSet::new()); + } + if self.backend == StorageBackend::Postgres { + return self.pg_has_token_uri_batch(tokens).await; + } + + let conn = self.conn.lock().unwrap(); + let mut existing = HashSet::with_capacity(tokens.len()); + + for chunk in tokens.chunks(SQLITE_TOKEN_PAIR_BATCH_SIZE) { + let predicates = vec!["(token = ? AND token_id = ?)"; chunk.len()].join(" OR "); + let query = format!("SELECT token, token_id FROM token_uris WHERE {predicates}"); + let mut params_vec: Vec> = Vec::with_capacity(chunk.len() * 2); + for (token, token_id) in chunk { + params_vec.push(Box::new(felt_to_blob(*token))); + params_vec.push(Box::new(u256_to_blob(*token_id))); + } + let params_refs: Vec<&dyn ToSql> = + params_vec.iter().map(std::convert::AsRef::as_ref).collect(); + let mut stmt = conn.prepare(&query)?; + let rows = stmt.query_map(params_refs.as_slice(), |row| { + let token_bytes: Vec = row.get(0)?; + let token_id_bytes: Vec = row.get(1)?; + Ok((blob_to_felt(&token_bytes), blob_to_u256(&token_id_bytes))) + })?; + + for row in rows { + existing.insert(row?); + } + } + + Ok(existing) + } + /// Returns all token URI rows for a given contract. pub async fn get_token_uris_by_contract( &self, @@ -1165,6 +1263,48 @@ impl Erc721Storage { rows.collect::, _>>().map_err(Into::into) } + pub async fn get_token_uris_batch( + &self, + token: Felt, + token_ids: &[U256], + ) -> Result, Option)>> { + if token_ids.is_empty() { + return Ok(Vec::new()); + } + if self.backend == StorageBackend::Postgres { + return self.pg_get_token_uris_batch(token, token_ids).await; + } + + let conn = self.conn.lock().unwrap(); + let token_blob = felt_to_blob(token); + let mut rows_out = Vec::new(); + for chunk in token_ids.chunks(SQLITE_TOKEN_BATCH_SIZE) { + let placeholders = chunk.iter().map(|_| "?").collect::>().join(","); + let query = format!( + "SELECT token_id, uri, metadata_json + FROM token_uris + WHERE token = ?1 AND token_id IN ({placeholders})" + ); + let mut params_vec: Vec> = Vec::with_capacity(chunk.len() + 1); + params_vec.push(Box::new(token_blob.clone())); + for token_id in chunk { + params_vec.push(Box::new(u256_to_blob(*token_id))); + } + let params_refs: Vec<&dyn ToSql> = + params_vec.iter().map(std::convert::AsRef::as_ref).collect(); + let mut stmt = conn.prepare(&query)?; + let rows = stmt.query_map(params_refs.as_slice(), |row| { + let token_id_bytes: Vec = row.get(0)?; + let uri: Option = row.get(1)?; + let metadata_json: Option = row.get(2)?; + Ok((blob_to_u256(&token_id_bytes), uri, metadata_json)) + })?; + rows_out.extend(rows.collect::, _>>()?); + } + + Ok(rows_out) + } + async fn pg_client(&self) -> Result> { let conn = self .pg_conn @@ -1210,54 +1350,267 @@ impl Erc721Storage { .collect() } - fn sqlite_append_attribute_exists_filters( - query: &mut String, - params: &mut Vec>, + fn sqlite_resolve_facet_filters( + &self, + conn: &Connection, + token: Felt, filters: &[(String, Vec)], - parent_alias: &str, - ) { + ) -> Result>> { + let token_blob = felt_to_blob(token); + let mut resolved = Vec::with_capacity(filters.len()); + for (key, values) in filters { - query.push_str(&format!( - " AND EXISTS ( - SELECT 1 FROM token_attributes a - WHERE a.token = {parent_alias}.token - AND a.token_id = {parent_alias}.token_id - AND a.key = ? AND a.value IN (" - )); - params.push(Box::new(key.clone())); + let key_id: i64 = match conn.query_row( + "SELECT id FROM facet_keys WHERE token = ?1 AND key_norm = ?2", + params![&token_blob, key], + |row| row.get(0), + ) { + Ok(id) => id, + Err(rusqlite::Error::QueryReturnedNoRows) => return Ok(None), + Err(error) => return Err(error.into()), + }; + let placeholders = values.iter().map(|_| "?").collect::>().join(","); - query.push_str(&placeholders); - query.push_str("))"); + let query = format!( + "SELECT id FROM facet_values WHERE token = ? AND facet_key_id = ? AND value_norm IN ({placeholders})" + ); + let mut params_vec: Vec> = Vec::with_capacity(values.len() + 2); + params_vec.push(Box::new(token_blob.clone())); + params_vec.push(Box::new(key_id)); for value in values { - params.push(Box::new(value.clone())); + params_vec.push(Box::new(value.clone())); + } + let params_refs: Vec<&dyn ToSql> = + params_vec.iter().map(std::convert::AsRef::as_ref).collect(); + let mut stmt = conn.prepare(&query)?; + let rows = stmt.query_map(params_refs.as_slice(), |row| row.get::<_, i64>(0))?; + let value_ids = rows.collect::, _>>()?; + if value_ids.is_empty() { + return Ok(None); } + + resolved.push(ResolvedFacetFilter { key_id, value_ids }); } + + Ok(Some(resolved)) } - fn pg_append_attribute_exists_filters( - query: &mut String, - params: &mut Vec>, - filters: &[(String, Vec)], - parent_alias: &str, - ) { - for (key, values) in filters { - query.push_str(&format!( - " AND EXISTS ( - SELECT 1 FROM erc721.token_attributes a - WHERE a.token = {parent_alias}.token - AND a.token_id = {parent_alias}.token_id - AND a.key = " - )); - query.push_str(&Self::pg_next_param(params, key.clone())); - query.push_str(" AND a.value IN ("); - let value_placeholders = values + fn sqlite_build_candidate_query( + token: Felt, + resolved_filters: &[ResolvedFacetFilter], + ) -> (String, Vec>) { + let token_blob = felt_to_blob(token); + if resolved_filters.is_empty() { + return ( + "SELECT token_id FROM token_uris WHERE token = ?".to_owned(), + vec![Box::new(token_blob)], + ); + } + + let mut query = String::new(); + let mut params: Vec> = Vec::new(); + for (idx, filter) in resolved_filters.iter().enumerate() { + if idx > 0 { + query.push_str(" INTERSECT "); + } + query.push_str( + "SELECT token_id FROM facet_token_map WHERE token = ? AND facet_key_id = ? AND facet_value_id IN (", + ); + params.push(Box::new(token_blob.clone())); + params.push(Box::new(filter.key_id)); + let placeholders = filter + .value_ids .iter() - .map(|value| Self::pg_next_param(params, value.clone())) + .map(|_| "?") .collect::>() .join(","); - query.push_str(&value_placeholders); - query.push_str("))"); + query.push_str(&placeholders); + query.push(')'); + for value_id in &filter.value_ids { + params.push(Box::new(*value_id)); + } + } + + (query, params) + } + + fn sqlite_query_token_ids_by_facets( + &self, + conn: &Connection, + token: Felt, + resolved_filters: &[ResolvedFacetFilter], + cursor_token_id: Option, + page_fetch: i64, + include_facets: bool, + facet_limit: i64, + ) -> Result { + let (candidate_query, candidate_params_for_page) = + Self::sqlite_build_candidate_query(token, resolved_filters); + let (_, candidate_params_for_count) = + Self::sqlite_build_candidate_query(token, resolved_filters); + + let mut page_query = format!( + "WITH candidate_tokens AS ({candidate_query}) SELECT token_id FROM candidate_tokens" + ); + let mut page_params = candidate_params_for_page; + if let Some(cursor) = cursor_token_id { + page_query.push_str(" WHERE token_id > ?"); + page_params.push(Box::new(u256_to_blob(cursor))); + } + page_query.push_str(" ORDER BY token_id ASC LIMIT ?"); + page_params.push(Box::new(page_fetch)); + + let page_refs: Vec<&dyn ToSql> = page_params + .iter() + .map(std::convert::AsRef::as_ref) + .collect(); + let mut stmt = conn.prepare(&page_query)?; + let rows = stmt.query_map(page_refs.as_slice(), |row| row.get::<_, Vec>(0))?; + let mut token_ids = rows + .collect::, _>>()? + .into_iter() + .map(|bytes| blob_to_u256(&bytes)) + .collect::>(); + + let page_limit = (page_fetch - 1) as usize; + let next_cursor_token_id = if token_ids.len() > page_limit { + let next = token_ids[page_limit]; + token_ids.truncate(page_limit); + Some(next) + } else { + None + }; + + let count_query = format!( + "WITH candidate_tokens AS ({candidate_query}) SELECT COUNT(*) FROM candidate_tokens" + ); + let count_refs: Vec<&dyn ToSql> = candidate_params_for_count + .iter() + .map(std::convert::AsRef::as_ref) + .collect(); + let total_hits: i64 = + conn.query_row(&count_query, count_refs.as_slice(), |row| row.get(0))?; + + let facets = if include_facets { + let (_, candidate_params_for_facets) = + Self::sqlite_build_candidate_query(token, resolved_filters); + let facet_query = format!( + "WITH candidate_tokens AS ({candidate_query}) + SELECT fk.key_display, fv.value_display, COUNT(*) AS cnt + FROM facet_token_map m + JOIN candidate_tokens c ON c.token_id = m.token_id + JOIN facet_keys fk ON fk.id = m.facet_key_id + JOIN facet_values fv ON fv.id = m.facet_value_id + WHERE m.token = ? + GROUP BY fk.key_display, fv.value_display + ORDER BY cnt DESC, fk.key_display ASC, fv.value_display ASC + LIMIT ?" + ); + let mut facet_params = candidate_params_for_facets; + facet_params.push(Box::new(felt_to_blob(token))); + facet_params.push(Box::new(facet_limit)); + let facet_refs: Vec<&dyn ToSql> = facet_params + .iter() + .map(std::convert::AsRef::as_ref) + .collect(); + let mut stmt = conn.prepare(&facet_query)?; + let rows = stmt.query_map(facet_refs.as_slice(), |row| { + Ok(AttributeFacetCount { + key: row.get(0)?, + value: row.get(1)?, + count: row.get::<_, i64>(2)? as u64, + }) + })?; + rows.collect::, _>>()? + } else { + Vec::new() + }; + + Ok(TokenAttributeQueryResult { + token_ids, + next_cursor_token_id, + total_hits: total_hits as u64, + facets, + }) + } + + async fn pg_resolve_facet_filters( + &self, + token: Felt, + filters: &[(String, Vec)], + ) -> Result>> { + let client = self.pg_client().await?; + let token_blob = felt_to_blob(token); + let mut resolved = Vec::with_capacity(filters.len()); + + for (key, values) in filters { + let key_row = client + .query_opt( + "SELECT id FROM erc721.facet_keys WHERE token = $1 AND key_norm = $2", + &[&token_blob, key], + ) + .await?; + let Some(key_row) = key_row else { + return Ok(None); + }; + let key_id: i64 = key_row.get(0); + + let rows = client + .query( + "SELECT id FROM erc721.facet_values WHERE token = $1 AND facet_key_id = $2 AND value_norm = ANY($3::text[])", + &[&token_blob, &key_id, values], + ) + .await?; + let value_ids = rows + .into_iter() + .map(|row| row.get::(0)) + .collect::>(); + if value_ids.is_empty() { + return Ok(None); + } + + resolved.push(ResolvedFacetFilter { key_id, value_ids }); + } + + Ok(Some(resolved)) + } + + fn pg_build_candidate_query( + token: Felt, + resolved_filters: &[ResolvedFacetFilter], + ) -> (String, Vec>) { + let mut candidate_query = String::new(); + let mut candidate_params: Vec> = Vec::new(); + if resolved_filters.is_empty() { + candidate_query.push_str("SELECT token_id FROM erc721.token_uris WHERE token = "); + candidate_query.push_str(&Self::pg_next_param( + &mut candidate_params, + felt_to_blob(token), + )); + } else { + for (idx, filter) in resolved_filters.iter().enumerate() { + if idx > 0 { + candidate_query.push_str(" INTERSECT "); + } + candidate_query + .push_str("SELECT token_id FROM erc721.facet_token_map WHERE token = "); + candidate_query.push_str(&Self::pg_next_param( + &mut candidate_params, + felt_to_blob(token), + )); + candidate_query.push_str(" AND facet_key_id = "); + candidate_query + .push_str(&Self::pg_next_param(&mut candidate_params, filter.key_id)); + candidate_query.push_str(" AND facet_value_id = ANY("); + candidate_query.push_str(&Self::pg_next_param( + &mut candidate_params, + filter.value_ids.clone(), + )); + candidate_query.push(')'); + } } + + (candidate_query, candidate_params) } async fn pg_insert_transfers_batch(&self, transfers: &[NftTransferData]) -> Result { @@ -1593,26 +1946,30 @@ impl Erc721Storage { Ok((ownership, next_cursor)) } - async fn pg_query_token_ids_by_attributes( + async fn pg_query_token_ids_by_facets( &self, token: Felt, - filters: &[(String, Vec)], + resolved_filters: &[ResolvedFacetFilter], cursor_token_id: Option, page_fetch: i64, include_facets: bool, facet_limit: i64, ) -> Result { let client = self.pg_client().await?; + let (candidate_query, candidate_params_for_page) = + Self::pg_build_candidate_query(token, resolved_filters); + let (_, candidate_params_for_count) = + Self::pg_build_candidate_query(token, resolved_filters); - let mut query = String::from("SELECT u.token_id FROM erc721.token_uris u WHERE u.token = "); - let mut params: Vec> = Vec::new(); - query.push_str(&Self::pg_next_param(&mut params, felt_to_blob(token))); + let mut query = format!( + "WITH candidate_tokens AS ({candidate_query}) SELECT token_id FROM candidate_tokens" + ); + let mut params = candidate_params_for_page; if let Some(cursor) = cursor_token_id { - query.push_str(" AND u.token_id > "); + query.push_str(" WHERE token_id > "); query.push_str(&Self::pg_next_param(&mut params, u256_to_blob(cursor))); } - Self::pg_append_attribute_exists_filters(&mut query, &mut params, filters, "u"); - query.push_str(" ORDER BY u.token_id ASC LIMIT "); + query.push_str(" ORDER BY token_id ASC LIMIT "); query.push_str(&Self::pg_next_param(&mut params, page_fetch)); let refs: Vec<&(dyn PgToSql + Sync)> = params @@ -1634,12 +1991,10 @@ impl Erc721Storage { None }; - let mut count_query = - String::from("SELECT COUNT(*) FROM erc721.token_uris u WHERE u.token = "); - let mut count_params: Vec> = Vec::new(); - count_query.push_str(&Self::pg_next_param(&mut count_params, felt_to_blob(token))); - Self::pg_append_attribute_exists_filters(&mut count_query, &mut count_params, filters, "u"); - let count_refs: Vec<&(dyn PgToSql + Sync)> = count_params + let count_query = format!( + "WITH candidate_tokens AS ({candidate_query}) SELECT COUNT(*) FROM candidate_tokens" + ); + let count_refs: Vec<&(dyn PgToSql + Sync)> = candidate_params_for_count .iter() .map(|p| p.as_ref() as &(dyn PgToSql + Sync)) .collect(); @@ -1647,23 +2002,22 @@ impl Erc721Storage { let total_hits = total_row.get::(0) as u64; let facets = if include_facets { - let mut facet_query = String::from( - "SELECT a.key, a.value, COUNT(*) AS cnt - FROM erc721.token_attributes a - JOIN erc721.token_uris u ON u.token = a.token AND u.token_id = a.token_id - WHERE u.token = ", + let (_, candidate_params_for_facets) = + Self::pg_build_candidate_query(token, resolved_filters); + let mut facet_query = format!( + "WITH candidate_tokens AS ({candidate_query}) + SELECT fk.key_display, fv.value_display, COUNT(*) AS cnt + FROM erc721.facet_token_map m + JOIN candidate_tokens c ON c.token_id = m.token_id + JOIN erc721.facet_keys fk ON fk.id = m.facet_key_id + JOIN erc721.facet_values fv ON fv.id = m.facet_value_id + WHERE m.token = " ); - let mut facet_params: Vec> = Vec::new(); + let mut facet_params = candidate_params_for_facets; facet_query.push_str(&Self::pg_next_param(&mut facet_params, felt_to_blob(token))); - Self::pg_append_attribute_exists_filters( - &mut facet_query, - &mut facet_params, - filters, - "u", - ); facet_query.push_str( - " GROUP BY a.key, a.value - ORDER BY cnt DESC, a.key ASC, a.value ASC + " GROUP BY fk.key_display, fv.value_display + ORDER BY cnt DESC, fk.key_display ASC, fv.value_display ASC LIMIT ", ); facet_query.push_str(&Self::pg_next_param(&mut facet_params, facet_limit)); @@ -1731,13 +2085,40 @@ impl Erc721Storage { let client = self.pg_client().await?; let row = client .query_one( - "SELECT COUNT(*) FROM erc721.token_metadata WHERE token = $1", + "SELECT COUNT(*) FROM erc721.token_metadata + WHERE token = $1 + AND name IS NOT NULL + AND BTRIM(name) <> '' + AND symbol IS NOT NULL + AND BTRIM(symbol) <> '' + AND total_supply IS NOT NULL", &[&felt_to_blob(token)], ) .await?; Ok(row.get::(0) > 0) } + async fn pg_has_token_metadata_batch(&self, tokens: &[Felt]) -> Result> { + let client = self.pg_client().await?; + let token_blobs: Vec> = tokens.iter().map(|token| felt_to_blob(*token)).collect(); + let rows = client + .query( + "SELECT token FROM erc721.token_metadata + WHERE token = ANY($1::bytea[]) + AND name IS NOT NULL + AND BTRIM(name) <> '' + AND symbol IS NOT NULL + AND BTRIM(symbol) <> '' + AND total_supply IS NOT NULL", + &[&token_blobs], + ) + .await?; + Ok(rows + .into_iter() + .map(|row| blob_to_felt(&row.get::>(0))) + .collect()) + } + async fn pg_upsert_token_metadata( &self, token: Felt, @@ -1834,14 +2215,49 @@ impl Erc721Storage { let client = self.pg_client().await?; let row = client .query_one( - "SELECT COUNT(*) FROM erc721.token_uris WHERE token = $1 AND token_id = $2", + "SELECT COUNT(*) FROM erc721.token_uris + WHERE token = $1 + AND token_id = $2", &[&felt_to_blob(token), &u256_to_blob(token_id)], ) .await?; Ok(row.get::(0) > 0) } - async fn pg_get_token_uris_by_contract( + async fn pg_has_token_uri_batch( + &self, + tokens: &[(Felt, U256)], + ) -> Result> { + let client = self.pg_client().await?; + let token_blobs: Vec> = tokens + .iter() + .map(|(token, _)| felt_to_blob(*token)) + .collect(); + let token_id_blobs: Vec> = tokens + .iter() + .map(|(_, token_id)| u256_to_blob(*token_id)) + .collect(); + let rows = client + .query( + "SELECT DISTINCT u.token, u.token_id + FROM erc721.token_uris u + JOIN unnest($1::bytea[], $2::bytea[]) AS i(token, token_id) + ON u.token = i.token AND u.token_id = i.token_id", + &[&token_blobs, &token_id_blobs], + ) + .await?; + Ok(rows + .into_iter() + .map(|row| { + ( + blob_to_felt(&row.get::>(0)), + blob_to_u256(&row.get::>(1)), + ) + }) + .collect()) + } + + async fn pg_get_token_uris_by_contract( &self, token: Felt, ) -> Result, Option)>> { @@ -1866,143 +2282,729 @@ impl Erc721Storage { }) .collect()) } + + async fn pg_get_token_uris_batch( + &self, + token: Felt, + token_ids: &[U256], + ) -> Result, Option)>> { + let client = self.pg_client().await?; + let token_id_blobs: Vec> = token_ids + .iter() + .map(|token_id| u256_to_blob(*token_id)) + .collect(); + let rows = client + .query( + "SELECT token_id, uri, metadata_json + FROM erc721.token_uris + WHERE token = $1 AND token_id = ANY($2::bytea[])", + &[&felt_to_blob(token), &token_id_blobs], + ) + .await?; + Ok(rows + .into_iter() + .map(|row| { + ( + blob_to_u256(&row.get::>(0)), + row.get(1), + row.get(2), + ) + }) + .collect()) + } } #[async_trait::async_trait] impl TokenUriStore for Erc721Storage { - async fn store_token_uri(&self, result: &TokenUriResult) -> Result<()> { + async fn store_token_uris_batch(&self, results: &[TokenUriResult]) -> Result<()> { + if results.is_empty() { + return Ok(()); + } + + let expected_attributes = results + .iter() + .map(|result| extract_metadata_attributes(result.metadata_json.as_deref())) + .collect::>(); + let token_blobs = results + .iter() + .map(|result| felt_to_blob(result.contract)) + .collect::>(); + let token_id_blobs = results + .iter() + .map(|result| u256_to_blob(result.token_id)) + .collect::>(); + if self.backend == StorageBackend::Postgres { let mut client = self.pg_client().await?; - let tx = client.transaction().await?; + let mut changed_indexes = Vec::new(); + for (idx, result) in results.iter().enumerate() { + if !pg_token_uri_state_matches(&client, result, &expected_attributes[idx]).await? { + changed_indexes.push(idx); + } + } - let token_blob = felt_to_blob(result.contract); - let token_id_blob = u256_to_blob(result.token_id); + if changed_indexes.is_empty() { + let facet_inputs = results + .iter() + .enumerate() + .map(|(idx, _)| FacetSyncInput { + token_blob: token_blobs[idx].clone(), + token_id_blob: token_id_blobs[idx].clone(), + attributes: &expected_attributes[idx], + }) + .collect::>(); + let tx = client.transaction().await?; + pg_sync_facets_for_tokens(&tx, &facet_inputs).await?; + tx.commit().await?; + return Ok(()); + } - tx.execute( - "INSERT INTO erc721.token_uris (token, token_id, uri, metadata_json, updated_at) - VALUES ($1, $2, $3, $4, EXTRACT(EPOCH FROM NOW())::BIGINT) - ON CONFLICT(token, token_id) DO UPDATE SET - uri = EXCLUDED.uri, - metadata_json = EXCLUDED.metadata_json, - updated_at = EXCLUDED.updated_at", - &[ - &token_blob, - &token_id_blob, - &result.uri.as_deref(), - &result.metadata_json.as_deref(), - ], - ) - .await?; + let tx = client.transaction().await?; - tx.execute( - "DELETE FROM erc721.token_attributes WHERE token = $1 AND token_id = $2", - &[&token_blob, &token_id_blob], - ) - .await?; + for idx in changed_indexes { + let result = &results[idx]; + let attrs = &expected_attributes[idx]; + let token_blob = &token_blobs[idx]; + let token_id_blob = &token_id_blobs[idx]; + + tx.execute( + "INSERT INTO erc721.token_uris (token, token_id, uri, metadata_json, updated_at) + VALUES ($1, $2, $3, $4, EXTRACT(EPOCH FROM NOW())::BIGINT) + ON CONFLICT(token, token_id) DO UPDATE SET + uri = EXCLUDED.uri, + metadata_json = EXCLUDED.metadata_json, + updated_at = EXCLUDED.updated_at", + &[ + &token_blob, + &token_id_blob, + &result.uri.as_deref(), + &result.metadata_json.as_deref(), + ], + ) + .await?; - if let Some(metadata_json) = &result.metadata_json { - if let Ok(value) = serde_json::from_str::(metadata_json) { - if let Some(attrs) = value.get("attributes").and_then(|a| a.as_array()) { - for attr in attrs { - let key = attr - .get("trait_type") - .or_else(|| attr.get("key")) - .and_then(|v| v.as_str()); - let value = attr.get("value").and_then(|v| { - v.as_str().map(ToOwned::to_owned).or_else(|| { - if v.is_null() { - None - } else { - Some(v.to_string()) - } - }) - }); - - let key = key.and_then(sanitize_metadata_text); - let value = value.and_then(|v| sanitize_metadata_text(&v)); - - if let (Some(key), Some(value)) = (key, value) { - tx.execute( - "INSERT INTO erc721.token_attributes (token, token_id, key, value) - VALUES ($1, $2, $3, $4) - ON CONFLICT (token, token_id, key) DO UPDATE SET value = EXCLUDED.value", - &[&token_blob, &token_id_blob, &key, &value], - ).await?; - } - } - } + tx.execute( + "DELETE FROM erc721.token_attributes WHERE token = $1 AND token_id = $2", + &[&token_blob, &token_id_blob], + ) + .await?; + + for (key, value) in attrs { + tx.execute( + "INSERT INTO erc721.token_attributes (token, token_id, key, value) + VALUES ($1, $2, $3, $4) + ON CONFLICT (token, token_id, key) DO UPDATE SET value = EXCLUDED.value", + &[&token_blob, &token_id_blob, key, value], + ) + .await?; } } + let facet_inputs = results + .iter() + .enumerate() + .map(|(idx, _)| FacetSyncInput { + token_blob: token_blobs[idx].clone(), + token_id_blob: token_id_blobs[idx].clone(), + attributes: &expected_attributes[idx], + }) + .collect::>(); + pg_sync_facets_for_tokens(&tx, &facet_inputs).await?; tx.commit().await?; return Ok(()); } let mut conn = self.conn.lock().unwrap(); - let tx = conn.transaction()?; + let mut changed_indexes = Vec::new(); + for (idx, result) in results.iter().enumerate() { + if !sqlite_token_uri_state_matches(&conn, result, &expected_attributes[idx])? { + changed_indexes.push(idx); + } + } - let token_blob = felt_to_blob(result.contract); - let token_id_blob = u256_to_blob(result.token_id); + if changed_indexes.is_empty() { + let facet_inputs = results + .iter() + .enumerate() + .map(|(idx, _)| FacetSyncInput { + token_blob: token_blobs[idx].clone(), + token_id_blob: token_id_blobs[idx].clone(), + attributes: &expected_attributes[idx], + }) + .collect::>(); + let tx = conn.transaction()?; + sqlite_sync_facets_for_tokens(&tx, &facet_inputs)?; + tx.commit()?; + return Ok(()); + } - tx.execute( - "INSERT INTO token_uris (token, token_id, uri, metadata_json, updated_at) - VALUES (?1, ?2, ?3, ?4, strftime('%s', 'now')) - ON CONFLICT(token, token_id) DO UPDATE SET - uri = excluded.uri, - metadata_json = excluded.metadata_json, - updated_at = excluded.updated_at", - params![ - &token_blob, - &token_id_blob, - result.uri.as_deref(), - result.metadata_json.as_deref() - ], - )?; + let tx = conn.transaction()?; + { + let mut upsert_stmt = tx.prepare_cached( + "INSERT INTO token_uris (token, token_id, uri, metadata_json, updated_at) + VALUES (?1, ?2, ?3, ?4, strftime('%s', 'now')) + ON CONFLICT(token, token_id) DO UPDATE SET + uri = excluded.uri, + metadata_json = excluded.metadata_json, + updated_at = excluded.updated_at", + )?; + let mut delete_attrs_stmt = tx.prepare_cached( + "DELETE FROM token_attributes WHERE token = ?1 AND token_id = ?2", + )?; + let mut insert_attr_stmt = tx.prepare_cached( + "INSERT OR REPLACE INTO token_attributes (token, token_id, key, value) + VALUES (?1, ?2, ?3, ?4)", + )?; - tx.execute( - "DELETE FROM token_attributes WHERE token = ?1 AND token_id = ?2", - params![&token_blob, &token_id_blob], - )?; + for idx in changed_indexes { + let result = &results[idx]; + let attrs = &expected_attributes[idx]; + let token_blob = &token_blobs[idx]; + let token_id_blob = &token_id_blobs[idx]; + + upsert_stmt.execute(params![ + token_blob, + token_id_blob, + result.uri.as_deref(), + result.metadata_json.as_deref() + ])?; + delete_attrs_stmt.execute(params![token_blob, token_id_blob])?; - if let Some(metadata_json) = &result.metadata_json { - if let Ok(value) = serde_json::from_str::(metadata_json) { - if let Some(attrs) = value.get("attributes").and_then(|a| a.as_array()) { - let mut attr_stmt = tx.prepare_cached( - "INSERT OR REPLACE INTO token_attributes (token, token_id, key, value) - VALUES (?1, ?2, ?3, ?4)", - )?; - - for attr in attrs { - let key = attr - .get("trait_type") - .or_else(|| attr.get("key")) - .and_then(|v| v.as_str()); - let value = attr.get("value").and_then(|v| { - v.as_str().map(ToOwned::to_owned).or_else(|| { - if v.is_null() { - None - } else { - Some(v.to_string()) - } - }) - }); - - let key = key.and_then(sanitize_metadata_text); - let value = value.and_then(|v| sanitize_metadata_text(&v)); - - if let (Some(key), Some(value)) = (key, value) { - attr_stmt.execute(params![&token_blob, &token_id_blob, key, value])?; - } - } + for (key, value) in attrs { + insert_attr_stmt.execute(params![token_blob, token_id_blob, key, value])?; } } } + let facet_inputs = results + .iter() + .enumerate() + .map(|(idx, _)| FacetSyncInput { + token_blob: token_blobs[idx].clone(), + token_id_blob: token_id_blobs[idx].clone(), + attributes: &expected_attributes[idx], + }) + .collect::>(); + sqlite_sync_facets_for_tokens(&tx, &facet_inputs)?; tx.commit()?; Ok(()) } } +fn extract_metadata_attributes(metadata_json: Option<&str>) -> Vec<(String, String)> { + let Some(metadata_json) = metadata_json else { + return Vec::new(); + }; + let Ok(value) = serde_json::from_str::(metadata_json) else { + return Vec::new(); + }; + let Some(attrs) = value.get("attributes").and_then(|attrs| attrs.as_array()) else { + return Vec::new(); + }; + + let mut normalized = BTreeMap::new(); + for attr in attrs { + let key = attr + .get("trait_type") + .or_else(|| attr.get("key")) + .and_then(|value| value.as_str()) + .and_then(sanitize_metadata_text); + let value = attr + .get("value") + .and_then(|value| { + value.as_str().map(ToOwned::to_owned).or_else(|| { + if value.is_null() { + None + } else { + Some(value.to_string()) + } + }) + }) + .and_then(|value| sanitize_metadata_text(&value)); + + if let (Some(key), Some(value)) = (key, value) { + normalized.insert(key, value); + } + } + + normalized.into_iter().collect() +} + +type TokenFacetKey = (Vec, Vec); +type FacetAssignmentMap = BTreeMap; + +struct FacetSyncInput<'a> { + token_blob: Vec, + token_id_blob: Vec, + attributes: &'a [(String, String)], +} + +#[derive(Default)] +struct FacetSyncPlan { + deletes: Vec<(Vec, Vec, i64)>, + inserts: Vec<(Vec, Vec, i64, i64)>, + count_deltas: HashMap, +} + +fn build_facet_sync_plan( + entries: &[FacetSyncInput<'_>], + existing: &HashMap, + desired: &HashMap, +) -> FacetSyncPlan { + let mut plan = FacetSyncPlan::default(); + + for entry in entries { + let token_key = (entry.token_blob.clone(), entry.token_id_blob.clone()); + let existing_map = existing.get(&token_key).cloned().unwrap_or_default(); + let desired_map = desired.get(&token_key).cloned().unwrap_or_default(); + + if existing_map == desired_map { + continue; + } + + for (facet_key_id, existing_value_id) in &existing_map { + if desired_map.get(facet_key_id) != Some(existing_value_id) { + plan.deletes.push(( + entry.token_blob.clone(), + entry.token_id_blob.clone(), + *facet_key_id, + )); + *plan.count_deltas.entry(*existing_value_id).or_insert(0) -= 1; + } + } + + for (facet_key_id, desired_value_id) in &desired_map { + if existing_map.get(facet_key_id) != Some(desired_value_id) { + plan.inserts.push(( + entry.token_blob.clone(), + entry.token_id_blob.clone(), + *facet_key_id, + *desired_value_id, + )); + *plan.count_deltas.entry(*desired_value_id).or_insert(0) += 1; + } + } + } + + plan.count_deltas.retain(|_, delta| *delta != 0); + plan +} + +fn sqlite_sync_facets_for_tokens( + tx: &rusqlite::Transaction<'_>, + entries: &[FacetSyncInput<'_>], +) -> Result<()> { + if entries.is_empty() { + return Ok(()); + } + + let existing = sqlite_load_existing_facet_assignments(tx, entries)?; + let desired = sqlite_resolve_desired_facet_assignments(tx, entries)?; + let plan = build_facet_sync_plan(entries, &existing, &desired); + sqlite_apply_facet_sync_plan(tx, &plan) +} + +fn sqlite_load_existing_facet_assignments( + tx: &rusqlite::Transaction<'_>, + entries: &[FacetSyncInput<'_>], +) -> Result> { + let mut assignments: HashMap = HashMap::new(); + + for chunk in entries.chunks(SQLITE_TOKEN_PAIR_BATCH_SIZE) { + let conditions = vec!["(token = ? AND token_id = ?)"; chunk.len()].join(" OR "); + let query = format!( + "SELECT token, token_id, facet_key_id, facet_value_id + FROM facet_token_map + WHERE {conditions}" + ); + + let mut params: Vec<&dyn ToSql> = Vec::with_capacity(chunk.len() * 2); + for entry in chunk { + params.push(&entry.token_blob); + params.push(&entry.token_id_blob); + } + + let mut stmt = tx.prepare(&query)?; + let rows = stmt.query_map(params.as_slice(), |row| { + Ok(( + row.get::>(0)?, + row.get::>(1)?, + row.get::(2)?, + row.get::(3)?, + )) + })?; + + for row in rows { + let (token_blob, token_id_blob, facet_key_id, facet_value_id) = row?; + assignments + .entry((token_blob, token_id_blob)) + .or_default() + .insert(facet_key_id, facet_value_id); + } + } + + Ok(assignments) +} + +fn sqlite_resolve_desired_facet_assignments( + tx: &rusqlite::Transaction<'_>, + entries: &[FacetSyncInput<'_>], +) -> Result> { + let mut desired: HashMap = HashMap::new(); + let mut key_cache: HashMap<(Vec, String), i64> = HashMap::new(); + let mut value_cache: HashMap<(Vec, i64, String), i64> = HashMap::new(); + + let mut insert_key_stmt = tx.prepare_cached( + "INSERT INTO facet_keys (token, key_norm, key_display) + VALUES (?1, ?2, ?3) + ON CONFLICT(token, key_norm) DO UPDATE SET key_display = excluded.key_display", + )?; + let mut select_key_stmt = + tx.prepare_cached("SELECT id FROM facet_keys WHERE token = ?1 AND key_norm = ?2")?; + let mut insert_value_stmt = tx.prepare_cached( + "INSERT INTO facet_values (token, facet_key_id, value_norm, value_display) + VALUES (?1, ?2, ?3, ?4) + ON CONFLICT(token, facet_key_id, value_norm) DO UPDATE SET value_display = excluded.value_display", + )?; + let mut select_value_stmt = tx.prepare_cached( + "SELECT id FROM facet_values WHERE token = ?1 AND facet_key_id = ?2 AND value_norm = ?3", + )?; + + for entry in entries { + let mut assignments = BTreeMap::new(); + + for (key, value) in entry.attributes { + let key_cache_key = (entry.token_blob.clone(), key.clone()); + let facet_key_id = if let Some(facet_key_id) = key_cache.get(&key_cache_key) { + *facet_key_id + } else { + insert_key_stmt.execute(params![&entry.token_blob, key, key])?; + let facet_key_id = select_key_stmt + .query_row(params![&entry.token_blob, key], |row| { + row.get::(0) + })?; + key_cache.insert(key_cache_key, facet_key_id); + facet_key_id + }; + + let value_cache_key = (entry.token_blob.clone(), facet_key_id, value.clone()); + let facet_value_id = if let Some(facet_value_id) = value_cache.get(&value_cache_key) { + *facet_value_id + } else { + insert_value_stmt.execute(params![ + &entry.token_blob, + facet_key_id, + value, + value + ])?; + let facet_value_id = select_value_stmt + .query_row(params![&entry.token_blob, facet_key_id, value], |row| { + row.get::(0) + })?; + value_cache.insert(value_cache_key, facet_value_id); + facet_value_id + }; + + assignments.insert(facet_key_id, facet_value_id); + } + + desired.insert( + (entry.token_blob.clone(), entry.token_id_blob.clone()), + assignments, + ); + } + + Ok(desired) +} + +fn sqlite_apply_facet_sync_plan( + tx: &rusqlite::Transaction<'_>, + plan: &FacetSyncPlan, +) -> Result<()> { + if plan.deletes.is_empty() && plan.inserts.is_empty() && plan.count_deltas.is_empty() { + return Ok(()); + } + + let mut delete_stmt = tx.prepare_cached( + "DELETE FROM facet_token_map WHERE token = ?1 AND token_id = ?2 AND facet_key_id = ?3", + )?; + let mut insert_stmt = tx.prepare_cached( + "INSERT OR REPLACE INTO facet_token_map (token, token_id, facet_key_id, facet_value_id) + VALUES (?1, ?2, ?3, ?4)", + )?; + let mut update_count_stmt = tx.prepare_cached( + "UPDATE facet_values + SET token_count = MAX(token_count + ?2, 0) + WHERE id = ?1", + )?; + + for (token_blob, token_id_blob, facet_key_id) in &plan.deletes { + delete_stmt.execute(params![token_blob, token_id_blob, facet_key_id])?; + } + + for (token_blob, token_id_blob, facet_key_id, facet_value_id) in &plan.inserts { + insert_stmt.execute(params![ + token_blob, + token_id_blob, + facet_key_id, + facet_value_id + ])?; + } + + for (facet_value_id, delta) in &plan.count_deltas { + update_count_stmt.execute(params![facet_value_id, delta])?; + } + + Ok(()) +} + +async fn pg_sync_facets_for_tokens( + tx: &tokio_postgres::Transaction<'_>, + entries: &[FacetSyncInput<'_>], +) -> Result<()> { + if entries.is_empty() { + return Ok(()); + } + + let existing = pg_load_existing_facet_assignments(tx, entries).await?; + let desired = pg_resolve_desired_facet_assignments(tx, entries).await?; + let plan = build_facet_sync_plan(entries, &existing, &desired); + pg_apply_facet_sync_plan(tx, &plan).await +} + +async fn pg_load_existing_facet_assignments( + tx: &tokio_postgres::Transaction<'_>, + entries: &[FacetSyncInput<'_>], +) -> Result> { + let mut assignments: HashMap = HashMap::new(); + + for chunk in entries.chunks(SQLITE_TOKEN_PAIR_BATCH_SIZE) { + let mut query = String::from( + "SELECT token, token_id, facet_key_id, facet_value_id + FROM erc721.facet_token_map + WHERE ", + ); + let mut params: Vec<&(dyn PgToSql + Sync)> = Vec::with_capacity(chunk.len() * 2); + + for (idx, entry) in chunk.iter().enumerate() { + if idx > 0 { + query.push_str(" OR "); + } + let token_param = params.len() + 1; + let token_id_param = params.len() + 2; + query.push_str(&format!( + "(token = ${token_param} AND token_id = ${token_id_param})" + )); + params.push(&entry.token_blob); + params.push(&entry.token_id_blob); + } + + let rows = tx.query(&query, ¶ms).await?; + for row in rows { + let token_blob: Vec = row.get(0); + let token_id_blob: Vec = row.get(1); + let facet_key_id: i64 = row.get(2); + let facet_value_id: i64 = row.get(3); + assignments + .entry((token_blob, token_id_blob)) + .or_default() + .insert(facet_key_id, facet_value_id); + } + } + + Ok(assignments) +} + +async fn pg_resolve_desired_facet_assignments( + tx: &tokio_postgres::Transaction<'_>, + entries: &[FacetSyncInput<'_>], +) -> Result> { + let mut desired: HashMap = HashMap::new(); + let mut key_cache: HashMap<(Vec, String), i64> = HashMap::new(); + let mut value_cache: HashMap<(Vec, i64, String), i64> = HashMap::new(); + + for entry in entries { + let mut assignments = BTreeMap::new(); + + for (key, value) in entry.attributes { + let key_cache_key = (entry.token_blob.clone(), key.clone()); + let facet_key_id = if let Some(facet_key_id) = key_cache.get(&key_cache_key) { + *facet_key_id + } else { + tx.execute( + "INSERT INTO erc721.facet_keys (token, key_norm, key_display) + VALUES ($1, $2, $3) + ON CONFLICT(token, key_norm) DO UPDATE SET key_display = EXCLUDED.key_display", + &[&entry.token_blob, key, key], + ) + .await?; + let facet_key_id = tx + .query_one( + "SELECT id FROM erc721.facet_keys WHERE token = $1 AND key_norm = $2", + &[&entry.token_blob, key], + ) + .await? + .get(0); + key_cache.insert(key_cache_key, facet_key_id); + facet_key_id + }; + + let value_cache_key = (entry.token_blob.clone(), facet_key_id, value.clone()); + let facet_value_id = if let Some(facet_value_id) = value_cache.get(&value_cache_key) { + *facet_value_id + } else { + tx.execute( + "INSERT INTO erc721.facet_values (token, facet_key_id, value_norm, value_display) + VALUES ($1, $2, $3, $4) + ON CONFLICT(token, facet_key_id, value_norm) DO UPDATE SET value_display = EXCLUDED.value_display", + &[&entry.token_blob, &facet_key_id, value, value], + ) + .await?; + let facet_value_id = tx + .query_one( + "SELECT id FROM erc721.facet_values WHERE token = $1 AND facet_key_id = $2 AND value_norm = $3", + &[&entry.token_blob, &facet_key_id, value], + ) + .await? + .get(0); + value_cache.insert(value_cache_key, facet_value_id); + facet_value_id + }; + + assignments.insert(facet_key_id, facet_value_id); + } + + desired.insert( + (entry.token_blob.clone(), entry.token_id_blob.clone()), + assignments, + ); + } + + Ok(desired) +} + +async fn pg_apply_facet_sync_plan( + tx: &tokio_postgres::Transaction<'_>, + plan: &FacetSyncPlan, +) -> Result<()> { + if plan.deletes.is_empty() && plan.inserts.is_empty() && plan.count_deltas.is_empty() { + return Ok(()); + } + + for (token_blob, token_id_blob, facet_key_id) in &plan.deletes { + tx.execute( + "DELETE FROM erc721.facet_token_map WHERE token = $1 AND token_id = $2 AND facet_key_id = $3", + &[token_blob, token_id_blob, facet_key_id], + ) + .await?; + } + + for (token_blob, token_id_blob, facet_key_id, facet_value_id) in &plan.inserts { + tx.execute( + "INSERT INTO erc721.facet_token_map (token, token_id, facet_key_id, facet_value_id) + VALUES ($1, $2, $3, $4) + ON CONFLICT(token, token_id, facet_key_id) DO UPDATE SET facet_value_id = EXCLUDED.facet_value_id", + &[token_blob, token_id_blob, facet_key_id, facet_value_id], + ) + .await?; + } + + for (facet_value_id, delta) in &plan.count_deltas { + tx.execute( + "UPDATE erc721.facet_values + SET token_count = GREATEST(token_count + $2, 0) + WHERE id = $1", + &[facet_value_id, delta], + ) + .await?; + } + + Ok(()) +} + +fn sqlite_token_uri_state_matches( + conn: &Connection, + result: &TokenUriResult, + expected_attributes: &[(String, String)], +) -> Result { + let token_blob = felt_to_blob(result.contract); + let token_id_blob = u256_to_blob(result.token_id); + let row = conn.query_row( + "SELECT uri, metadata_json FROM token_uris WHERE token = ?1 AND token_id = ?2", + params![&token_blob, &token_id_blob], + |row| { + Ok(( + row.get::>(0)?, + row.get::>(1)?, + )) + }, + ); + let (existing_uri, existing_metadata_json) = match row { + Ok(row) => row, + Err(rusqlite::Error::QueryReturnedNoRows) => return Ok(false), + Err(error) => return Err(error.into()), + }; + + if existing_uri.as_deref() != result.uri.as_deref() + || existing_metadata_json.as_deref() != result.metadata_json.as_deref() + { + return Ok(false); + } + + let mut stmt = conn.prepare( + "SELECT key, value FROM token_attributes + WHERE token = ?1 AND token_id = ?2 + ORDER BY key ASC, value ASC", + )?; + let rows = stmt.query_map(params![&token_blob, &token_id_blob], |row| { + Ok((row.get::(0)?, row.get::(1)?)) + })?; + let existing_attributes = rows.collect::, _>>()?; + + Ok(existing_attributes == expected_attributes) +} + +async fn pg_token_uri_state_matches( + client: &Client, + result: &TokenUriResult, + expected_attributes: &[(String, String)], +) -> Result { + let token_blob = felt_to_blob(result.contract); + let token_id_blob = u256_to_blob(result.token_id); + let row = client + .query_opt( + "SELECT uri, metadata_json FROM erc721.token_uris WHERE token = $1 AND token_id = $2", + &[&token_blob, &token_id_blob], + ) + .await?; + let Some(row) = row else { + return Ok(false); + }; + + let existing_uri: Option = row.get(0); + let existing_metadata_json: Option = row.get(1); + if existing_uri.as_deref() != result.uri.as_deref() + || existing_metadata_json.as_deref() != result.metadata_json.as_deref() + { + return Ok(false); + } + + let rows = client + .query( + "SELECT key, value FROM erc721.token_attributes + WHERE token = $1 AND token_id = $2 + ORDER BY key ASC, value ASC", + &[&token_blob, &token_id_blob], + ) + .await?; + let existing_attributes: Vec<(String, String)> = rows + .into_iter() + .map(|row| (row.get(0), row.get(1))) + .collect(); + + Ok(existing_attributes == expected_attributes) +} + fn sanitize_metadata_text(input: &str) -> Option { let sanitized = input.replace('\0', "").trim().to_owned(); if sanitized.is_empty() { @@ -2011,3 +3013,296 @@ fn sanitize_metadata_text(input: &str) -> Option { Some(sanitized) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn temp_db_path(test_name: &str) -> String { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("system time before unix epoch") + .as_nanos(); + std::env::temp_dir() + .join(format!("torii-erc721-{test_name}-{nanos}.db")) + .to_string_lossy() + .to_string() + } + + #[tokio::test] + async fn has_token_metadata_requires_complete_erc721_row() { + let db_path = temp_db_path("complete-metadata"); + let storage = Erc721Storage::new(&db_path).await.expect("create storage"); + let token = Felt::from_hex_unchecked( + "0x046da8955829adf2bda310099a0063451923f02e648cf25a1203aac6335cf0e4", + ); + + storage + .upsert_token_metadata(token, None, None, Some(U256::from(301u64))) + .await + .expect("insert partial metadata"); + assert!(!storage + .has_token_metadata(token) + .await + .expect("check partial metadata")); + + storage + .upsert_token_metadata(token, Some("Beasts"), Some("BEAST"), None) + .await + .expect("complete metadata"); + assert!(storage + .has_token_metadata(token) + .await + .expect("check complete metadata")); + + let _ = std::fs::remove_file(db_path); + } + + #[tokio::test] + async fn store_token_uri_skips_unchanged_rewrite() { + let db_path = temp_db_path("token-uri-noop"); + let storage = Erc721Storage::new(&db_path).await.expect("create storage"); + let result = TokenUriResult { + contract: Felt::from_hex_unchecked( + "0x046da8955829adf2bda310099a0063451923f02e648cf25a1203aac6335cf0e4", + ), + token_id: U256::from(1u64), + uri: Some("ipfs://beasts/1".to_owned()), + metadata_json: Some( + r#"{"name":"Beast #1","attributes":[{"trait_type":"Class","value":"Wolf"}]}"# + .to_owned(), + ), + }; + + storage + .store_token_uri(&result) + .await + .expect("insert token uri"); + + { + let conn = storage.conn.lock().unwrap(); + conn.execute( + "UPDATE token_uris SET updated_at = 123 WHERE token = ?1 AND token_id = ?2", + params![felt_to_blob(result.contract), u256_to_blob(result.token_id)], + ) + .expect("set sentinel updated_at"); + } + + storage + .store_token_uri(&result) + .await + .expect("store unchanged token uri"); + + let (updated_at, attr_count): (i64, i64) = { + let conn = storage.conn.lock().unwrap(); + let updated_at: i64 = conn + .query_row( + "SELECT updated_at FROM token_uris WHERE token = ?1 AND token_id = ?2", + params![felt_to_blob(result.contract), u256_to_blob(result.token_id)], + |row| row.get(0), + ) + .expect("read updated_at"); + let attr_count: i64 = conn + .query_row( + "SELECT COUNT(*) FROM token_attributes WHERE token = ?1 AND token_id = ?2", + params![felt_to_blob(result.contract), u256_to_blob(result.token_id)], + |row| row.get(0), + ) + .expect("count attributes"); + (updated_at, attr_count) + }; + + assert_eq!(updated_at, 123); + assert_eq!(attr_count, 1); + + let _ = std::fs::remove_file(db_path); + } + + #[tokio::test] + async fn store_token_uri_repairs_missing_attributes_for_same_metadata() { + let db_path = temp_db_path("token-uri-backfill"); + let storage = Erc721Storage::new(&db_path).await.expect("create storage"); + let result = TokenUriResult { + contract: Felt::from_hex_unchecked( + "0x046da8955829adf2bda310099a0063451923f02e648cf25a1203aac6335cf0e4", + ), + token_id: U256::from(2u64), + uri: Some("ipfs://beasts/2".to_owned()), + metadata_json: Some( + r#"{"name":"Beast #2","attributes":[{"trait_type":"Class","value":"Bear"}]}"# + .to_owned(), + ), + }; + + storage + .store_token_uri(&result) + .await + .expect("insert token uri"); + + { + let conn = storage.conn.lock().unwrap(); + conn.execute( + "DELETE FROM token_attributes WHERE token = ?1 AND token_id = ?2", + params![felt_to_blob(result.contract), u256_to_blob(result.token_id)], + ) + .expect("delete attributes"); + } + + storage + .store_token_uri(&result) + .await + .expect("repair missing attributes"); + + let attr_count: i64 = { + let conn = storage.conn.lock().unwrap(); + conn.query_row( + "SELECT COUNT(*) FROM token_attributes WHERE token = ?1 AND token_id = ?2", + params![felt_to_blob(result.contract), u256_to_blob(result.token_id)], + |row| row.get(0), + ) + .expect("count attributes") + }; + + assert_eq!(attr_count, 1); + + let _ = std::fs::remove_file(db_path); + } + + #[tokio::test] + async fn store_token_uri_repairs_missing_facets_for_same_metadata() { + let db_path = temp_db_path("token-uri-facet-backfill"); + let storage = Erc721Storage::new(&db_path).await.expect("create storage"); + let result = TokenUriResult { + contract: Felt::from_hex_unchecked( + "0x046da8955829adf2bda310099a0063451923f02e648cf25a1203aac6335cf0e4", + ), + token_id: U256::from(3u64), + uri: Some("ipfs://beasts/3".to_owned()), + metadata_json: Some( + r#"{"name":"Beast #3","attributes":[{"trait_type":"Class","value":"Wolf"},{"trait_type":"Color","value":"Red"}]}"# + .to_owned(), + ), + }; + + storage + .store_token_uri(&result) + .await + .expect("insert token uri"); + + { + let conn = storage.conn.lock().unwrap(); + conn.execute( + "UPDATE facet_values SET token_count = 0 WHERE id IN ( + SELECT facet_value_id FROM facet_token_map WHERE token = ?1 AND token_id = ?2 + )", + params![felt_to_blob(result.contract), u256_to_blob(result.token_id)], + ) + .expect("zero facet counts"); + conn.execute( + "DELETE FROM facet_token_map WHERE token = ?1 AND token_id = ?2", + params![felt_to_blob(result.contract), u256_to_blob(result.token_id)], + ) + .expect("delete facet mappings"); + } + + storage + .store_token_uri(&result) + .await + .expect("repair missing facets"); + + let (facet_map_count, positive_counts): (i64, i64) = { + let conn = storage.conn.lock().unwrap(); + let facet_map_count: i64 = conn + .query_row( + "SELECT COUNT(*) FROM facet_token_map WHERE token = ?1 AND token_id = ?2", + params![felt_to_blob(result.contract), u256_to_blob(result.token_id)], + |row| row.get(0), + ) + .expect("count facet mappings"); + let positive_counts: i64 = conn + .query_row( + "SELECT COUNT(*) FROM facet_values WHERE token_count > 0", + [], + |row| row.get(0), + ) + .expect("count positive facet values"); + (facet_map_count, positive_counts) + }; + + assert_eq!(facet_map_count, 2); + assert_eq!(positive_counts, 2); + + let _ = std::fs::remove_file(db_path); + } + + #[tokio::test] + async fn query_token_ids_by_attributes_uses_facet_index() { + let db_path = temp_db_path("facet-query"); + let storage = Erc721Storage::new(&db_path).await.expect("create storage"); + let contract = Felt::from_hex_unchecked( + "0x046da8955829adf2bda310099a0063451923f02e648cf25a1203aac6335cf0e4", + ); + + let fixtures = [ + ( + 1u64, + r#"{"attributes":[{"trait_type":"Color","value":"Red"},{"trait_type":"Class","value":"Wolf"}]}"#, + ), + ( + 2u64, + r#"{"attributes":[{"trait_type":"Color","value":"Blue"},{"trait_type":"Class","value":"Wolf"}]}"#, + ), + ( + 3u64, + r#"{"attributes":[{"trait_type":"Color","value":"Red"},{"trait_type":"Class","value":"Bear"}]}"#, + ), + ]; + + for (token_id, metadata_json) in fixtures { + storage + .store_token_uri(&TokenUriResult { + contract, + token_id: U256::from(token_id), + uri: Some(format!("ipfs://beasts/{token_id}")), + metadata_json: Some(metadata_json.to_owned()), + }) + .await + .expect("store token uri fixture"); + } + + let result = storage + .query_token_ids_by_attributes( + contract, + &[ + ( + "Color".to_owned(), + vec!["Red".to_owned(), "Blue".to_owned()], + ), + ("Class".to_owned(), vec!["Wolf".to_owned()]), + ], + None, + 10, + true, + 10, + ) + .await + .expect("query by attributes"); + + assert_eq!(result.token_ids, vec![U256::from(1u64), U256::from(2u64)]); + assert_eq!(result.total_hits, 2); + assert!(result + .facets + .iter() + .any(|facet| { facet.key == "Color" && facet.value == "Red" && facet.count == 1 })); + assert!(result + .facets + .iter() + .any(|facet| { facet.key == "Color" && facet.value == "Blue" && facet.count == 1 })); + assert!(result + .facets + .iter() + .any(|facet| { facet.key == "Class" && facet.value == "Wolf" && facet.count == 2 })); + + let _ = std::fs::remove_file(db_path); + } +} diff --git a/crates/torii-runtime-common/src/sink.rs b/crates/torii-runtime-common/src/sink.rs index 50d573fc..23a09dff 100644 --- a/crates/torii-runtime-common/src/sink.rs +++ b/crates/torii-runtime-common/src/sink.rs @@ -1,13 +1,46 @@ use anyhow::{Context, Result}; use std::sync::Arc; use tokio_postgres::NoTls; +use torii::command::{CommandBus, CommandHandler}; use torii::etl::sink::{EventBus, Sink, SinkContext}; use torii::grpc::SubscriptionManager; +pub struct InitializedSinkRuntime { + command_bus: CommandBus, +} + +impl InitializedSinkRuntime { + pub async fn shutdown(self) { + self.command_bus.shutdown().await; + } +} + pub async fn initialize_sink(sink: &mut dyn Sink, database_root: std::path::PathBuf) -> Result<()> { + let runtime = initialize_sink_with_command_handlers(sink, database_root, Vec::new(), 1).await?; + std::mem::drop(runtime); + Ok(()) +} + +pub async fn initialize_sink_with_command_handlers( + sink: &mut dyn Sink, + database_root: std::path::PathBuf, + handlers: Vec>, + queue_size: usize, +) -> Result { let event_bus = Arc::new(EventBus::new(Arc::new(SubscriptionManager::new()))); - sink.initialize(event_bus, &SinkContext { database_root }) - .await + for handler in &handlers { + handler.attach_event_bus(event_bus.clone()); + } + let command_bus = CommandBus::new(handlers, queue_size)?; + sink.initialize( + event_bus, + &SinkContext { + database_root, + command_bus: command_bus.sender(), + }, + ) + .await?; + Ok(InitializedSinkRuntime { command_bus }) } pub async fn drop_postgres_schemas( diff --git a/src/command.rs b/src/command.rs new file mode 100644 index 00000000..cc70c530 --- /dev/null +++ b/src/command.rs @@ -0,0 +1,273 @@ +use anyhow::Result; +use async_trait::async_trait; +use std::any::Any; +use std::fmt; +use std::sync::Arc; +use tokio::sync::mpsc; +use tokio::task::{JoinError, JoinHandle, JoinSet}; + +use crate::etl::sink::EventBus; + +pub trait Command: Send + Sync + 'static { + fn as_any(&self) -> &dyn Any; + + fn into_any(self: Box) -> Box; + + fn name(&self) -> &'static str { + std::any::type_name::() + } +} + +impl Command for T +where + T: Send + Sync + 'static, +{ + fn as_any(&self) -> &dyn Any { + self + } + + fn into_any(self: Box) -> Box { + self + } +} + +#[async_trait] +pub trait CommandHandler: Send + Sync { + fn supports(&self, command: &dyn Command) -> bool; + + fn attach_event_bus(&self, _event_bus: Arc) {} + + async fn handle_command(&self, command: Box) -> Result<()>; +} + +struct RoutedCommand { + handler_index: usize, + command: Box, + command_name: &'static str, +} + +#[derive(Debug)] +pub enum CommandDispatchError { + Unsupported(&'static str), + Ambiguous(&'static str), + QueueFull(&'static str), + Closed(&'static str), +} + +impl fmt::Display for CommandDispatchError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Unsupported(name) => write!(f, "no command handler is registered for {name}"), + Self::Ambiguous(name) => write!(f, "multiple command handlers matched {name}"), + Self::QueueFull(name) => write!(f, "command queue is full for {name}"), + Self::Closed(name) => write!(f, "command queue is closed for {name}"), + } + } +} + +impl std::error::Error for CommandDispatchError {} + +#[derive(Clone)] +pub struct CommandBusSender { + tx: mpsc::Sender, + handlers: Arc>>, +} + +impl CommandBusSender { + pub fn dispatch(&self, command: C) -> std::result::Result<(), CommandDispatchError> + where + C: Command, + { + let command: Box = Box::new(command); + let command_name = command.name(); + + let mut matches = self + .handlers + .iter() + .enumerate() + .filter(|(_, handler)| handler.supports(command.as_ref())); + + let Some((handler_index, _)) = matches.next() else { + ::metrics::counter!( + "torii_command_dispatch_total", + "command" => command_name, + "status" => "unsupported" + ) + .increment(1); + return Err(CommandDispatchError::Unsupported(command_name)); + }; + + if matches.next().is_some() { + ::metrics::counter!( + "torii_command_dispatch_total", + "command" => command_name, + "status" => "ambiguous" + ) + .increment(1); + return Err(CommandDispatchError::Ambiguous(command_name)); + } + + let routed = RoutedCommand { + handler_index, + command, + command_name, + }; + + match self.tx.try_send(routed) { + Ok(()) => { + ::metrics::counter!( + "torii_command_dispatch_total", + "command" => command_name, + "status" => "enqueued" + ) + .increment(1); + Ok(()) + } + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + ::metrics::counter!( + "torii_command_dispatch_total", + "command" => command_name, + "status" => "dropped_full" + ) + .increment(1); + Err(CommandDispatchError::QueueFull(command_name)) + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { + ::metrics::counter!( + "torii_command_dispatch_total", + "command" => command_name, + "status" => "closed" + ) + .increment(1); + Err(CommandDispatchError::Closed(command_name)) + } + } + } +} + +pub struct CommandBus { + sender: CommandBusSender, + receiver_handle: JoinHandle<()>, +} + +impl CommandBus { + pub fn new(handlers: Vec>, queue_size: usize) -> Result { + let handlers: Arc>> = + Arc::new(handlers.into_iter().map(Arc::from).collect()); + + let (tx, mut rx) = mpsc::channel::(queue_size.max(1)); + let sender = CommandBusSender { + tx, + handlers: handlers.clone(), + }; + + let receiver_handle = tokio::spawn(async move { + let mut tasks = JoinSet::new(); + let mut receiver_closed = false; + + loop { + tokio::select! { + Some(result) = tasks.join_next(), if !tasks.is_empty() => { + if let Err(error) = result { + log_join_error(error); + } + } + maybe_command = rx.recv(), if !receiver_closed => { + match maybe_command { + Some(routed) => { + let Some(handler) = handlers.get(routed.handler_index).cloned() else { + tracing::error!( + target: "torii::command_bus", + command = routed.command_name, + handler_index = routed.handler_index, + "Resolved command handler is missing" + ); + continue; + }; + + tasks.spawn(async move { + let start = std::time::Instant::now(); + let result = handler.handle_command(routed.command).await; + ::metrics::histogram!( + "torii_command_handle_duration_seconds", + "command" => routed.command_name + ) + .record(start.elapsed().as_secs_f64()); + + match result { + Ok(()) => { + ::metrics::counter!( + "torii_command_handle_total", + "command" => routed.command_name, + "status" => "ok" + ) + .increment(1); + } + Err(error) => { + ::metrics::counter!( + "torii_command_handle_total", + "command" => routed.command_name, + "status" => "error" + ) + .increment(1); + tracing::warn!( + target: "torii::command_bus", + command = routed.command_name, + error = %error, + "Command handler failed" + ); + } + } + }); + } + None => receiver_closed = true, + } + } + else => break, + } + + if receiver_closed && tasks.is_empty() { + break; + } + } + }); + + Ok(Self { + sender, + receiver_handle, + }) + } + + pub fn sender(&self) -> CommandBusSender { + self.sender.clone() + } + + pub async fn shutdown(self) { + self.receiver_handle.abort(); + match self.receiver_handle.await { + Ok(()) => {} + Err(error) if error.is_cancelled() => { + tracing::info!(target: "torii::command_bus", "Command bus stopped"); + } + Err(error) => { + tracing::warn!( + target: "torii::command_bus", + error = %error, + "Command bus shutdown join failed" + ); + } + } + } +} + +fn log_join_error(error: JoinError) { + if error.is_cancelled() { + tracing::debug!(target: "torii::command_bus", "Command task cancelled"); + } else { + tracing::warn!( + target: "torii::command_bus", + error = %error, + "Command task join failed" + ); + } +} diff --git a/src/etl/sink/mod.rs b/src/etl/sink/mod.rs index 539b7ce2..308e4368 100644 --- a/src/etl/sink/mod.rs +++ b/src/etl/sink/mod.rs @@ -7,6 +7,7 @@ use std::path::PathBuf; use std::sync::Arc; use super::envelope::{Envelope, TypeId}; +use crate::command::CommandBusSender; use crate::grpc::SubscriptionManager; pub use multi::MultiSink; @@ -25,6 +26,8 @@ pub struct SinkContext { /// Sinks can create databases at `{database_root}/{sink_name}.db` /// to co-locate all Torii databases for easy backup and management. pub database_root: PathBuf, + /// Command bus sender for background fire-and-forget work. + pub command_bus: CommandBusSender, // Future extensions: // pub metrics: Arc, // pub config: Arc, diff --git a/src/lib.rs b/src/lib.rs index 581e18b0..7cb6ba9b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ //! This library aims at providing a modular and high-performance blockchain indexer. //! The current implementation is still WIP, but gives a good idea of the architecture and the capabilities. +pub mod command; pub mod etl; pub mod grpc; pub mod http; @@ -34,6 +35,7 @@ use tokio_util::sync::CancellationToken; use tonic::transport::Server; use tower_http::cors::{Any as CorsAny, CorsLayer}; +use command::{CommandBus, CommandHandler}; use etl::decoder::{ContractFilter, DecoderId}; use etl::extractor::{Extractor, SyntheticExtractor, SyntheticExtractorAdapter}; use etl::identification::{ContractIdentifier, IdentificationRule}; @@ -180,6 +182,12 @@ pub struct ToriiConfig { /// ETL concurrency controls. pub etl_concurrency: EtlConcurrencyConfig, + + /// Background command handlers registered for the command bus. + pub command_handlers: Vec>, + + /// Command bus queue size. + pub command_bus_queue_size: usize, } impl ToriiConfig { @@ -215,6 +223,8 @@ pub struct ToriiConfigBuilder { contract_identifier: Option>, shutdown_timeout: Option, etl_concurrency: Option, + command_handlers: Vec>, + command_bus_queue_size: Option, } impl ToriiConfigBuilder { @@ -514,6 +524,21 @@ impl ToriiConfigBuilder { self } + pub fn with_command_handler(mut self, handler: Box) -> Self { + self.command_handlers.push(handler); + self + } + + pub fn with_command_handlers(mut self, handlers: Vec>) -> Self { + self.command_handlers.extend(handlers); + self + } + + pub fn command_bus_queue_size(mut self, size: usize) -> Self { + self.command_bus_queue_size = Some(size.max(1)); + self + } + /// Builds the Torii configuration. /// /// # Panics @@ -547,6 +572,8 @@ impl ToriiConfigBuilder { contract_identifier: self.contract_identifier, shutdown_timeout: self.shutdown_timeout.unwrap_or(30), etl_concurrency: self.etl_concurrency.unwrap_or_default(), + command_handlers: self.command_handlers, + command_bus_queue_size: self.command_bus_queue_size.unwrap_or(4096), } } } @@ -576,10 +603,15 @@ pub async fn run(config: ToriiConfig) -> Result<(), Box> let subscription_manager = Arc::new(SubscriptionManager::new()); let event_bus = Arc::new(EventBus::new(subscription_manager.clone())); + for handler in &config.command_handlers { + handler.attach_event_bus(event_bus.clone()); + } + let command_bus = CommandBus::new(config.command_handlers, config.command_bus_queue_size)?; // Create SinkContext for initialization let sink_context = etl::sink::SinkContext { database_root: config.database_root.clone(), + command_bus: command_bus.sender(), }; let mut initialized_sinks: Vec> = Vec::new(); @@ -1163,6 +1195,7 @@ pub async fn run(config: ToriiConfig) -> Result<(), Box> } tracing::info!(target: "torii::main", "Torii shutdown complete"); + command_bus.shutdown().await; Ok(()) }