refactor: add generic for legacy or horizon

Signed-off-by: Gustavo Inacio <[email protected]>

Author: gusinacio

Cargo.lock

@@ -39,7 +39,6 @@ tracing-subscriber.workspace = true
 tonic.workspace = true
 bigdecimal = { workspace = true, features = ["serde"] }
 graphql_client.workspace = true
-
 ruint = { version = "1.12.3", features = [
     "num-traits",
 ], default-features = false }
@@ -52,6 +51,7 @@ tap_aggregator.workspace = true
 futures = { version = "0.3.30", default-features = false }
 bon = "3.3"
 test-assets = { path = "../test-assets", optional = true }
+sealed = "0.6.0"
 
 [dev-dependencies]
 # Release-please breaks with cyclical dependencies if dev-dependencies

crates/tap-agent/Cargo.toml

@@ -39,6 +39,7 @@ use crate::{
     adaptative_concurrency::AdaptiveLimiter,
     agent::unaggregated_receipts::UnaggregatedReceipts,
     backoff::BackoffInfo,
+    tap::context::Legacy,
     tracker::{SenderFeeTracker, SimpleFeeTracker},
 };
 
@@ -210,11 +211,17 @@ impl SenderAccountConfig {
     }
 }
 
+pub enum AllocationType {
+    Legacy,
+    Horizon,
+}
+
 impl State {
     async fn create_sender_allocation(
         &self,
         sender_account_ref: ActorRef<SenderAccountMessage>,
         allocation_id: Address,
+        allocation_type: AllocationType,
     ) -> anyhow::Result<()> {
         tracing::trace!(
             %self.sender,
@@ -233,13 +240,18 @@ impl State {
             .config(AllocationConfig::from_sender_config(self.config))
             .build();
 
-        SenderAllocation::spawn_linked(
-            Some(self.format_sender_allocation(&allocation_id)),
-            SenderAllocation,
-            args,
-            sender_account_ref.get_cell(),
-        )
-        .await?;
+        match allocation_type {
+            AllocationType::Legacy => {
+                SenderAllocation::spawn_linked(
+                    Some(self.format_sender_allocation(&allocation_id)),
+                    SenderAllocation::<Legacy>::default(),
+                    args,
+                    sender_account_ref.get_cell(),
+                )
+                .await?;
+            }
+            AllocationType::Horizon => unimplemented!(),
+        }
         Ok(())
     }
     fn format_sender_allocation(&self, allocation_id: &Address) -> String {
@@ -646,7 +658,13 @@ impl Actor for SenderAccount {
 
         stream::iter(allocation_ids)
             // Create a sender allocation for each allocation
-            .map(|allocation_id| state.create_sender_allocation(myself.clone(), allocation_id))
+            .map(|allocation_id| {
+                state.create_sender_allocation(
+                    myself.clone(),
+                    allocation_id,
+                    AllocationType::Legacy,
+                )
+            })
             .buffer_unordered(10) // Limit concurrency to 10 allocations at a time
             .collect::<Vec<anyhow::Result<()>>>()
             .await
@@ -819,7 +837,11 @@ impl Actor for SenderAccount {
                 let mut new_allocation_ids = state.allocation_ids.clone();
                 for allocation_id in allocation_ids.difference(&state.allocation_ids) {
                     if let Err(error) = state
-                        .create_sender_allocation(myself.clone(), *allocation_id)
+                        .create_sender_allocation(
+                            myself.clone(),
+                            *allocation_id,
+                            AllocationType::Legacy,
+                        )
                         .await
                     {
                         tracing::error!(
@@ -870,7 +892,7 @@ impl Actor for SenderAccount {
             }
             SenderAccountMessage::NewAllocationId(allocation_id) => {
                 if let Err(error) = state
-                    .create_sender_allocation(myself.clone(), allocation_id)
+                    .create_sender_allocation(myself.clone(), allocation_id, AllocationType::Legacy)
                     .await
                 {
                     tracing::error!(
@@ -1013,7 +1035,7 @@ impl Actor for SenderAccount {
                 };
 
                 if let Err(error) = state
-                    .create_sender_allocation(myself.clone(), allocation_id)
+                    .create_sender_allocation(myself.clone(), allocation_id, AllocationType::Legacy)
                     .await
                 {
                     tracing::error!(

crates/tap-agent/src/agent/sender_account.rs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use std::{
+    marker::PhantomData,
     sync::Arc,
     time::{Duration, Instant},
 };
@@ -16,7 +17,7 @@ use tap_aggregator::grpc::{
     tap_aggregator_client::TapAggregatorClient, RavRequest as AggregatorRequest,
 };
 use tap_core::{
-    manager::adapters::RavRead,
+    manager::adapters::{RavRead, RavStore, ReceiptDelete, ReceiptRead},
     rav_request::RavRequest,
     receipt::{
         checks::{Check, CheckList},
@@ -43,7 +44,7 @@ use crate::{
     tap::{
         context::{
             checks::{AllocationId, Signature},
-            TapAgentContext,
+            ReceiptType, TapAgentContext,
         },
         signers_trimmed, TapReceipt,
     },
@@ -100,29 +101,28 @@ pub enum RavError {
     Other(#[from] anyhow::Error),
 }
 
-type TapManager = tap_core::manager::Manager<TapAgentContext, TapReceipt>;
-
-pub enum AllocationType {
-    Legacy,
-    Horizon,
-}
+type TapManager<T> = tap_core::manager::Manager<TapAgentContext<T>, TapReceipt>;
 
 /// Manages unaggregated fees and the TAP lifecyle for a specific (allocation, sender) pair.
-pub struct SenderAllocation;
+pub struct SenderAllocation<T>(PhantomData<T>);
+impl<T: ReceiptType> Default for SenderAllocation<T> {
+    fn default() -> Self {
+        Self(PhantomData)
+    }
+}
 
-pub struct SenderAllocationState {
+pub struct SenderAllocationState<T> {
     unaggregated_fees: UnaggregatedReceipts,
     invalid_receipts_fees: UnaggregatedReceipts,
     latest_rav: Option<SignedRav>,
     pgpool: PgPool,
-    tap_manager: TapManager,
+    tap_manager: TapManager<T>,
     allocation_id: Address,
     sender: Address,
     escrow_accounts: Receiver<EscrowAccounts>,
     domain_separator: Eip712Domain,
     sender_account_ref: ActorRef<SenderAccountMessage>,
     sender_aggregator: TapAggregatorClient<Channel>,
-    allocation_type: AllocationType,
     //config
     timestamp_buffer_ns: u64,
     rav_request_receipt_limit: u64,
@@ -157,8 +157,6 @@ pub struct SenderAllocationArgs {
     pub domain_separator: Eip712Domain,
     pub sender_account_ref: ActorRef<SenderAccountMessage>,
     pub sender_aggregator: TapAggregatorClient<Channel>,
-    #[builder(default = AllocationType::Legacy)]
-    pub allocation_type: AllocationType,
 
     //config
     pub config: AllocationConfig,
@@ -173,9 +171,16 @@ pub enum SenderAllocationMessage {
 }
 
 #[async_trait::async_trait]
-impl Actor for SenderAllocation {
+impl<T> Actor for SenderAllocation<T>
+where
+    T: ReceiptType + Send + Sync + 'static,
+    TapAgentContext<T>: RavRead<ReceiptAggregateVoucher>
+        + RavStore<ReceiptAggregateVoucher>
+        + ReceiptDelete
+        + ReceiptRead<TapReceipt>,
+{
     type Msg = SenderAllocationMessage;
-    type State = SenderAllocationState;
+    type State = SenderAllocationState<T>;
     type Arguments = SenderAllocationArgs;
 
     async fn pre_start(
@@ -350,10 +355,16 @@ impl Actor for SenderAllocation {
     }
 }
 
-impl SenderAllocationState {
+impl<T> SenderAllocationState<T>
+where
+    T: ReceiptType + Send + Sync,
+    TapAgentContext<T>: RavRead<ReceiptAggregateVoucher>
+        + RavStore<ReceiptAggregateVoucher>
+        + ReceiptDelete
+        + ReceiptRead<TapReceipt>,
+{
     async fn new(
         SenderAllocationArgs {
-            allocation_type,
             pgpool,
             allocation_id,
             sender,
@@ -397,7 +408,6 @@ impl SenderAllocationState {
         Ok(Self {
             pgpool,
             tap_manager,
-            allocation_type,
             allocation_id,
             sender,
             escrow_accounts,
@@ -1028,7 +1038,7 @@ pub mod tests {
             .rav_request_receipt_limit(rav_request_receipt_limit)
             .call()
             .await;
-        let actor = TestableActor::new(SenderAllocation);
+        let actor = TestableActor::new(SenderAllocation::default());
         let notify = actor.notify.clone();
 
         let (allocation_ref, _join_handle) = Actor::spawn(None, actor, args).await.unwrap();

crates/tap-agent/src/agent/sender_allocation.rs

@@ -1,6 +1,8 @@
 // Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
+use std::marker::PhantomData;
+
 use indexer_monitor::EscrowAccounts;
 use sqlx::PgPool;
 use thegraph_core::alloy::primitives::Address;
@@ -14,15 +16,28 @@ mod receipt;
 
 pub use error::AdapterError;
 
+#[sealed::sealed]
+pub trait ReceiptType {}
+
+pub enum Legacy {}
+pub enum Horizon {}
+
+#[sealed::sealed]
+impl ReceiptType for Legacy {}
+
+#[sealed::sealed]
+impl ReceiptType for Horizon {}
+
 #[derive(Clone)]
-pub struct TapAgentContext {
+pub struct TapAgentContext<T> {
     pgpool: PgPool,
     allocation_id: Address,
     sender: Address,
     escrow_accounts: Receiver<EscrowAccounts>,
+    _phantom: PhantomData<T>,
 }
 
-impl TapAgentContext {
+impl<T: ReceiptType> TapAgentContext<T> {
     pub fn new(
         pgpool: PgPool,
         allocation_id: Address,
@@ -34,6 +49,7 @@ impl TapAgentContext {
             allocation_id,
             sender,
             escrow_accounts,
+            _phantom: PhantomData,
         }
     }
 }

crates/tap-agent/src/tap/context.rs

@@ -5,7 +5,7 @@ use async_trait::async_trait;
 use tap_core::manager::adapters::SignatureChecker;
 use thegraph_core::alloy::primitives::Address;
 
-use super::{error::AdapterError, TapAgentContext};
+use super::{error::AdapterError, ReceiptType, TapAgentContext};
 
 // Conversion from eventuals::error::Closed to AdapterError::EscrowEventualError
 impl From<eventuals::error::Closed> for AdapterError {
@@ -22,7 +22,7 @@ impl From<eventuals::error::Closed> for AdapterError {
 // In any case, we don't want to fail a receipt because of this.
 // The receipt is fine, just the escrow account that is not.
 #[async_trait]
-impl SignatureChecker for TapAgentContext {
+impl<T: ReceiptType + Send + Sync> SignatureChecker for TapAgentContext<T> {
     type AdapterError = AdapterError;
 
     async fn verify_signer(&self, signer: Address) -> Result<bool, Self::AdapterError> {

crates/tap-agent/src/tap/context/escrow.rs

@@ -14,10 +14,10 @@ use tap_graph::{ReceiptAggregateVoucher, SignedRav};
 use thegraph_core::alloy::signers::Signature;
 use thegraph_core::alloy::{hex::ToHexExt, primitives::Address};
 
-use super::{error::AdapterError, TapAgentContext};
+use super::{error::AdapterError, Horizon, Legacy, TapAgentContext};
 
 #[async_trait::async_trait]
-impl RavRead<ReceiptAggregateVoucher> for TapAgentContext {
+impl RavRead<ReceiptAggregateVoucher> for TapAgentContext<Legacy> {
     type AdapterError = AdapterError;
 
     async fn last_rav(&self) -> Result<Option<SignedRav>, Self::AdapterError> {
@@ -87,7 +87,7 @@ impl RavRead<ReceiptAggregateVoucher> for TapAgentContext {
 }
 
 #[async_trait::async_trait]
-impl RavStore<ReceiptAggregateVoucher> for TapAgentContext {
+impl RavStore<ReceiptAggregateVoucher> for TapAgentContext<Legacy> {
     type AdapterError = AdapterError;
 
     async fn update_last_rav(&self, rav: SignedRav) -> Result<(), Self::AdapterError> {
@@ -129,6 +129,27 @@ impl RavStore<ReceiptAggregateVoucher> for TapAgentContext {
     }
 }
 
+#[async_trait::async_trait]
+impl RavRead<tap_graph::v2::ReceiptAggregateVoucher> for TapAgentContext<Horizon> {
+    type AdapterError = AdapterError;
+
+    async fn last_rav(&self) -> Result<Option<tap_graph::v2::SignedRav>, Self::AdapterError> {
+        unimplemented!()
+    }
+}
+
+#[async_trait::async_trait]
+impl RavStore<tap_graph::v2::ReceiptAggregateVoucher> for TapAgentContext<Horizon> {
+    type AdapterError = AdapterError;
+
+    async fn update_last_rav(
+        &self,
+        _rav: tap_graph::v2::SignedRav,
+    ) -> Result<(), Self::AdapterError> {
+        unimplemented!()
+    }
+}
+
 #[cfg(test)]
 mod test {
     use indexer_monitor::EscrowAccounts;