Implement fees system v2 with improved architecture and concurrency

Co-authored-by: sharpy96 <sharpy96@gmail.com>

Author: cursoragent

fees.md

@@ -0,0 +1,397 @@
+# Fees System v2 - Comprehensive Analysis
+
+## Overview
+
+This document provides a comprehensive analysis of the fees system implementation in the fees-v2 branch, highlighting improvements over the previous version and identifying areas for further enhancement.
+
+## Major Improvements in v2
+
+### 1. **Better Architecture & Separation of Concerns**
+
+The v2 implementation introduces a clear three-phase architecture:
+
+```go
+// fees.go now has three distinct phases:
+LoadFees()         // Phase 1: Load fees from verifier into database
+HandleTransactions() // Phase 2: Process fee runs and create transactions  
+HandlePostTx()     // Phase 3: Monitor and finalize transactions
+```
+
+**Why this is better:** Each phase has clear responsibilities, making the system more maintainable and debuggable compared to the monolithic approach in v1.
+
+### 2. **Concurrency Control & Race Condition Prevention**
+
+```go
+type FeePlugin struct {
+    // ... other fields
+    transactingMutex sync.Mutex // when actual transactions are happening we cannot load fees
+}
+
+func (fp *FeePlugin) LoadFees(ctx context.Context, task *asynq.Task) error {
+    fp.transactingMutex.Lock()
+    defer fp.transactingMutex.Unlock()
+    // ...
+}
+```
+
+**Excellent improvement:** This addresses race condition concerns about concurrent fee collection that existed in v1.
+
+### 3. **Database Schema Improvements**
+
+```sql
+-- v2 uses tx_hash instead of tx_id
+tx_hash VARCHAR(66),  -- Much better than UUID reference
+
+-- Added proper fee run states
+CONSTRAINT fee_run_status_check CHECK (status IN ('draft', 'sent', 'completed', 'failed'))
+```
+
+**Major improvement:** Transaction tracking is now done via hash rather than UUID references, making it more reliable and easier to debug.
+
+### 4. **Smart Fee Loading Strategy**
+
+```go
+func (fp *FeePlugin) executeFeeLoading(ctx context.Context, feePolicy vtypes.PluginPolicy) error {
+    // Check if fee already exists
+    existingFee, err := fp.db.GetFees(ctx, fee.ID)
+    if len(existingFee) > 0 {
+        continue // Skip if already loaded
+    }
+    
+    // Look for existing draft run
+    run, err := fp.db.GetPendingFeeRun(ctx, feePolicy.ID)
+    if run == nil {
+        // Create new run
+    } else {
+        // Add to existing run
+    }
+}
+```
+
+**Much smarter approach:** Fees are batched intelligently and duplicate loading is prevented, improving efficiency and preventing data inconsistencies.
+
+### 5. **Proper Transaction Monitoring**
+
+```go
+// post_tx.go - New file for post-transaction logic
+func (fp *FeePlugin) updateStatus(ctx context.Context, run types.FeeRun, currentBlock uint64) error {
+    receipt, err := fp.ethClient.TransactionReceipt(ctx, hash)
+    if receipt.Status == 1 {
+        if receipt.BlockNumber.Uint64()+fp.config.SuccessConfirmations < currentBlock {
+            // Mark as successful with confirmation logic
+        }
+    }
+}
+```
+
+**Excellent addition:** Proper transaction monitoring with configurable confirmations ensures transactions are properly finalized.
+
+### 6. **Configuration Improvements**
+
+```go
+type FeeConfig struct {
+    // ... existing fields
+    EthProvider          string `mapstructure:"eth_provider"`         // Separate ETH provider
+    SuccessConfirmations uint64 `mapstructure:"success_confirmations"` // Configurable confirmations
+}
+```
+
+**Good improvement:** Separate RPC endpoints and configurable confirmation requirements provide better flexibility and reliability.
+
+### 7. **Enhanced Database Operations**
+
+The v2 implementation includes several new database methods that improve data management:
+
+- `GetPendingFeeRun()` - Efficiently finds existing draft runs
+- `GetFees()` - Checks for existing fees to prevent duplicates
+- `SetFeeRunSuccess()` - Proper status management
+- `GetAllFeeRuns()` with status filtering
+
+## Remaining Areas for Improvement
+
+### 1. **Error Handling & Recovery**
+
+```go
+// In HandlePostTx - TODO comments indicate missing logic
+if err == ethereum.NotFound {
+    // TODO rebroadcast logic
+    fp.logger.WithFields(logrus.Fields{"run_id": run.ID}).Info("tx not found on chain, rebroadcasting")
+    return nil
+}
+```
+
+**Still needs work:** Transaction rebroadcast logic is missing, which could lead to stuck transactions in high-congestion scenarios.
+
+### 2. **Hardcoded Concurrency Limits**
+
+```go
+// Good that semaphore was added, but could be improved
+sem := semaphore.NewWeighted(10)
+// Why 10? Should be configurable
+```
+
+**Minor improvement needed:** Make concurrency limits configurable rather than hardcoded to allow for different deployment scenarios.
+
+### 3. **Transaction Type Support**
+
+```go
+// helper.go has good transaction decoding, but could be more robust
+func decodeTx(rawHex string) (*erc20tx, error) {
+    // Check transaction type (EIP-1559 is 0x02)
+    if len(rawBytes) == 0 || rawBytes[0] != 0x02 {
+        return nil, fmt.Errorf("unsupported transaction type: 0x%02x", rawBytes[0])
+    }
+}
+```
+
+**Good addition but limited:** Only supports EIP-1559 transactions. Consider adding legacy transaction support for broader compatibility.
+
+### 4. **Database Query Optimization**
+
+```go
+// In GetAllFeeRuns - fetches ALL fees, could be expensive
+feesQuery := `select id, fee_run_id, amount from fee`
+feesRows, err := p.pool.Query(ctx, feesQuery)
+```
+
+**Performance concern:** This query gets ALL fees in the database rather than filtering by the runs being processed, which could become expensive at scale.
+
+### 5. **Failed Transaction Handling**
+
+```go
+} else {
+    // TODO failed tx logic
+    fp.logger.WithFields(logrus.Fields{"run_id": run.ID}).Info("tx failed, setting to failed")
+    return nil
+}
+```
+
+**Critical gap:** No logic to handle failed transactions, which could lead to lost fees or manual intervention requirements.
+
+## Security Improvements Made
+
+✅ **Address validation** is better with checksummed comparison  
+✅ **Concurrency control** prevents race conditions  
+✅ **Transaction validation** is more robust  
+✅ **Configuration validation** is comprehensive  
+✅ **Mutex protection** prevents concurrent state modification  
+✅ **Proper transaction decoding** with ABI validation  
+
+## Recommended Next Steps
+
+### 1. **Add Transaction Rebroadcast Logic**
+
+```go
+func (fp *FeePlugin) rebroadcastTransaction(ctx context.Context, run types.FeeRun) error {
+    // Implement logic to recreate and rebroadcast failed transactions
+    // Consider increasing gas price for rebroadcast
+    policy, err := fp.db.GetPluginPolicy(ctx, run.PolicyID)
+    if err != nil {
+        return err
+    }
+    
+    // Recreate transaction with higher gas price
+    newTxs, err := fp.proposeTransactions(ctx, *policy, run)
+    if err != nil {
+        return err
+    }
+    
+    // Update gas price by 10-20%
+    // Rebroadcast transaction
+    return fp.initSign(ctx, newTxs[0], *policy, run.ID)
+}
+```
+
+### 2. **Make Concurrency Configurable**
+
+```go
+type FeeConfig struct {
+    // Add to config
+    MaxConcurrentFeeLoads    uint64 `mapstructure:"max_concurrent_fee_loads"`
+    MaxConcurrentTxProcess   uint64 `mapstructure:"max_concurrent_tx_process"`
+    MaxConcurrentPostTx      uint64 `mapstructure:"max_concurrent_post_tx"`
+}
+```
+
+### 3. **Add Comprehensive Metrics**
+
+```go
+type FeeMetrics struct {
+    FeesLoaded           prometheus.Counter
+    TransactionsCreated  prometheus.Counter  
+    TransactionsSuccess  prometheus.Counter
+    TransactionsFailed   prometheus.Counter
+    ProcessingTime       prometheus.Histogram
+    FeeRunDuration       prometheus.Histogram
+}
+```
+
+### 4. **Optimize Database Queries**
+
+```go
+// In GetAllFeeRuns, only fetch fees for the specific runs
+func (p *PostgresBackend) GetAllFeeRuns(ctx context.Context, statuses ...types.FeeRunState) ([]types.FeeRun, error) {
+    // ... existing run query logic ...
+    
+    // Only fetch fees for the runs we actually retrieved
+    runIDs := make([]uuid.UUID, len(rm))
+    i := 0
+    for id := range rm {
+        runIDs[i] = id
+        i++
+    }
+    
+    feesQuery := `select id, fee_run_id, amount from fee where fee_run_id = ANY($1)`
+    feesRows, err := p.pool.Query(ctx, feesQuery, runIDs)
+    // ...
+}
+```
+
+### 5. **Add Failed Transaction Handling**
+
+```go
+func (fp *FeePlugin) handleFailedTransaction(ctx context.Context, run types.FeeRun) error {
+    // Mark run as failed
+    if err := fp.db.SetFeeRunFailed(ctx, run.ID); err != nil {
+        return err
+    }
+    
+    // Log detailed failure information
+    fp.logger.WithFields(logrus.Fields{
+        "run_id": run.ID,
+        "tx_hash": *run.TxHash,
+        "policy_id": run.PolicyID,
+    }).Error("Transaction failed, marking run as failed")
+    
+    // Consider notification/alerting logic here
+    return nil
+}
+```
+
+### 6. **Add Retry Logic with Exponential Backoff**
+
+```go
+func (fp *FeePlugin) executeWithRetry(ctx context.Context, operation func() error, maxRetries int) error {
+    for i := 0; i < maxRetries; i++ {
+        err := operation()
+        if err == nil {
+            return nil
+        }
+        
+        if !isRetryableError(err) {
+            return err
+        }
+        
+        backoff := time.Duration(i+1) * time.Second
+        time.Sleep(backoff)
+    }
+    return fmt.Errorf("max retries exceeded")
+}
+```
+
+### 7. **Add Health Check Endpoints**
+
+```go
+func (fp *FeePlugin) HealthCheck(ctx context.Context) error {
+    // Check database connectivity
+    if err := fp.db.Ping(ctx); err != nil {
+        return fmt.Errorf("database unhealthy: %w", err)
+    }
+    
+    // Check RPC connectivity
+    if _, err := fp.ethClient.BlockNumber(ctx); err != nil {
+        return fmt.Errorf("eth client unhealthy: %w", err)
+    }
+    
+    // Check verifier API
+    // ... additional health checks
+    
+    return nil
+}
+```
+
+## File Structure Analysis
+
+### Core Files
+- `fees.go` - Main plugin logic with three distinct phases
+- `transaction.go` - Transaction creation, validation, and signing
+- `post_tx.go` - Post-transaction monitoring and status updates
+- `policy.go` - Policy validation and recipe specifications
+- `config.go` - Configuration management with validation
+- `helper.go` - Transaction decoding utilities
+- `constraints.go` - Plugin constants and constraints
+
+### Database Layer
+- `storage/postgres/fees.go` - Fee-specific database operations
+- Database migrations with proper schema design
+
+### Configuration
+- Improved configuration with separate RPC endpoints
+- Validation logic for all required fields
+- Environment variable support
+
+## Performance Considerations
+
+### Current Optimizations
+- Semaphore-based concurrency control
+- Batching of fees into runs
+- Indexed database queries
+- Connection reuse for RPC calls
+
+### Areas for Improvement
+- Database query optimization (fetch only needed fees)
+- Configurable batch sizes
+- Connection pooling for RPC endpoints
+- Caching of frequently accessed data
+
+## Testing Recommendations
+
+### Unit Tests Needed
+- Fee loading logic with various edge cases
+- Transaction creation and validation
+- Database operations with mock data
+- Configuration validation
+
+### Integration Tests Needed
+- End-to-end fee collection flow
+- Transaction broadcasting and monitoring
+- Error scenarios and recovery
+- Concurrent operation testing
+
+### Load Tests Needed
+- High-volume fee processing
+- Concurrent fee run processing
+- Database performance under load
+- RPC endpoint stress testing
+
+## Overall Assessment
+
+The fees v2 implementation represents a **significant improvement** over the previous version:
+
+**V1 Assessment:** ~40% production ready
+**V2 Assessment:** ~80% production ready
+
+### ✅ **Major Improvements Made**
+- Concurrency control with mutex
+- Better database design with tx_hash
+- Separation of concerns with distinct phases
+- Smart fee batching to prevent duplicates
+- Transaction monitoring with confirmations
+- Improved error handling structure
+- Helper functions for transaction decoding
+
+### ⚠️ **Still Needs Work**
+- Transaction rebroadcast logic (critical)
+- Failed transaction handling (critical)
+- Database query optimization (performance)
+- Configurable concurrency limits (minor)
+- Comprehensive monitoring/metrics (operational)
+
+### 🎯 **Ready for Production After**
+1. Implementing transaction rebroadcast logic
+2. Adding failed transaction handling
+3. Completing TODO items in post_tx.go
+4. Adding comprehensive monitoring
+5. Load testing under realistic conditions
+
+The code is now much more production-ready and demonstrates good software engineering practices. The remaining work is primarily around edge case handling and operational concerns rather than fundamental architectural issues.