Skip to main content

tenzro_token/
fee_distribution.rs

1//! Fee processing and distribution
2//!
3//! This module processes incoming fees from settlements, transactions,
4//! and bridge operations, then distributes them according to the configured splits.
5
6use crate::error::{Result, TokenError};
7use crate::treasury::FeeDistributionConfig;
8use dashmap::DashMap;
9use serde::{Deserialize, Serialize};
10use std::collections::HashMap;
11use tenzro_types::asset::AssetId;
12use tenzro_types::primitives::Timestamp;
13use tracing::{debug, info};
14
15/// Fee source type
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
17pub enum FeeSource {
18    /// Transaction fee
19    Transaction,
20    /// Settlement fee
21    Settlement,
22    /// Bridge operation fee
23    Bridge,
24    /// Model inference fee
25    ModelInference,
26    /// Storage fee
27    Storage,
28    /// Other fee source
29    Other,
30}
31
32/// Fee record
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct FeeRecord {
35    /// Fee ID
36    pub fee_id: String,
37    /// Asset ID
38    pub asset_id: AssetId,
39    /// Total fee amount
40    pub amount: u128,
41    /// Fee source
42    pub source: FeeSource,
43    /// Treasury share
44    pub treasury_share: u128,
45    /// Burn share
46    pub burn_share: u128,
47    /// Staker share
48    pub staker_share: u128,
49    /// Timestamp
50    pub timestamp: Timestamp,
51}
52
53impl FeeRecord {
54    /// Creates a new fee record
55    pub fn new(
56        asset_id: AssetId,
57        amount: u128,
58        source: FeeSource,
59        treasury_share: u128,
60        burn_share: u128,
61        staker_share: u128,
62    ) -> Self {
63        Self {
64            fee_id: uuid::Uuid::new_v4().to_string(),
65            asset_id,
66            amount,
67            source,
68            treasury_share,
69            burn_share,
70            staker_share,
71            timestamp: Timestamp::now(),
72        }
73    }
74}
75
76/// Fee statistics
77#[derive(Debug, Clone, Serialize, Deserialize)]
78#[derive(Default)]
79pub struct FeeStats {
80    /// Total fees collected per asset
81    pub total_collected: HashMap<AssetId, u128>,
82    /// Total fees by source
83    pub by_source: HashMap<FeeSource, u128>,
84    /// Total to treasury
85    pub total_to_treasury: u128,
86    /// Total burned
87    pub total_burned: u128,
88    /// Total to stakers
89    pub total_to_stakers: u128,
90    /// Fee count
91    pub fee_count: u64,
92}
93
94
95/// Distribution history entry
96#[derive(Debug, Clone, Serialize, Deserialize)]
97pub struct DistributionHistory {
98    /// Time period start
99    pub period_start: Timestamp,
100    /// Time period end
101    pub period_end: Timestamp,
102    /// Total fees in period
103    pub total_fees: u128,
104    /// Fees per asset
105    pub fees_by_asset: HashMap<AssetId, u128>,
106    /// Treasury share
107    pub treasury_share: u128,
108    /// Burn share
109    pub burn_share: u128,
110    /// Staker share
111    pub staker_share: u128,
112}
113
114impl DistributionHistory {
115    /// Creates a new distribution history entry
116    pub fn new(period_start: Timestamp, period_end: Timestamp) -> Self {
117        Self {
118            period_start,
119            period_end,
120            total_fees: 0,
121            fees_by_asset: HashMap::new(),
122            treasury_share: 0,
123            burn_share: 0,
124            staker_share: 0,
125        }
126    }
127
128    /// Adds a fee record to the history
129    pub fn add_fee(&mut self, record: &FeeRecord) {
130        self.total_fees = self.total_fees.saturating_add(record.amount);
131
132        let asset_total = self.fees_by_asset.get(&record.asset_id).copied().unwrap_or(0);
133        self.fees_by_asset.insert(
134            record.asset_id.clone(),
135            asset_total.saturating_add(record.amount),
136        );
137
138        self.treasury_share = self.treasury_share.saturating_add(record.treasury_share);
139        self.burn_share = self.burn_share.saturating_add(record.burn_share);
140        self.staker_share = self.staker_share.saturating_add(record.staker_share);
141    }
142}
143
144/// Fee processor
145///
146/// Processes incoming fees and distributes them according to configuration.
147#[derive(Debug)]
148pub struct FeeProcessor {
149    /// Fee distribution configuration
150    config: parking_lot::RwLock<FeeDistributionConfig>,
151    /// Fee records (FeeId -> FeeRecord)
152    fee_records: DashMap<String, FeeRecord>,
153    /// Fee statistics
154    stats: parking_lot::RwLock<FeeStats>,
155    /// Distribution history (Period -> History)
156    history: DashMap<String, DistributionHistory>,
157    /// Current history period
158    current_period: parking_lot::RwLock<Option<String>>,
159}
160
161impl FeeProcessor {
162    /// Creates a new fee processor
163    pub fn new() -> Self {
164        Self {
165            config: parking_lot::RwLock::new(FeeDistributionConfig::default()),
166            fee_records: DashMap::new(),
167            stats: parking_lot::RwLock::new(FeeStats::default()),
168            history: DashMap::new(),
169            current_period: parking_lot::RwLock::new(None),
170        }
171    }
172
173    /// Processes a fee
174    ///
175    /// # Arguments
176    ///
177    /// * `asset_id` - Asset of the fee
178    /// * `amount` - Fee amount
179    /// * `source` - Source of the fee
180    ///
181    /// Returns the fee record
182    pub fn process_fee(&self, asset_id: AssetId, amount: u128, source: FeeSource) -> Result<FeeRecord> {
183        if amount == 0 {
184            return Err(TokenError::InvalidAmount("Fee amount must be greater than zero".to_string()));
185        }
186
187        // Calculate distribution
188        let config = self.config.read();
189        let distribution = config.calculate_distribution(amount);
190        drop(config);
191
192        // Create fee record
193        let record = FeeRecord::new(
194            asset_id.clone(),
195            amount,
196            source,
197            distribution.treasury_amount,
198            distribution.burn_amount,
199            distribution.staker_amount,
200        );
201
202        let fee_id = record.fee_id.clone();
203
204        // Update statistics
205        let mut stats = self.stats.write();
206
207        // Update total collected
208        let asset_total = stats.total_collected.get(&asset_id).copied().unwrap_or(0);
209        stats.total_collected.insert(
210            asset_id.clone(),
211            asset_total.checked_add(amount)
212                .ok_or_else(|| TokenError::ArithmeticOverflow {
213                    operation: "fee total collected".to_string(),
214                })?,
215        );
216
217        // Update by source
218        let source_total = stats.by_source.get(&source).copied().unwrap_or(0);
219        stats.by_source.insert(
220            source,
221            source_total.checked_add(amount)
222                .ok_or_else(|| TokenError::ArithmeticOverflow {
223                    operation: "fee by source".to_string(),
224                })?,
225        );
226
227        // Update distribution totals
228        stats.total_to_treasury = stats.total_to_treasury.checked_add(distribution.treasury_amount)
229            .ok_or_else(|| TokenError::ArithmeticOverflow {
230                operation: "fee total to treasury".to_string(),
231            })?;
232        stats.total_burned = stats.total_burned.checked_add(distribution.burn_amount)
233            .ok_or_else(|| TokenError::ArithmeticOverflow {
234                operation: "fee total burned".to_string(),
235            })?;
236        stats.total_to_stakers = stats.total_to_stakers.checked_add(distribution.staker_amount)
237            .ok_or_else(|| TokenError::ArithmeticOverflow {
238                operation: "fee total to stakers".to_string(),
239            })?;
240        stats.fee_count += 1;
241
242        drop(stats);
243
244        // Add to current history period
245        self.add_to_history(&record);
246
247        // Store fee record
248        self.fee_records.insert(fee_id.clone(), record.clone());
249
250        debug!(
251            "Processed fee: {} of {} (treasury: {}, burn: {}, stakers: {})",
252            amount, asset_id.as_str(),
253            distribution.treasury_amount,
254            distribution.burn_amount,
255            distribution.staker_amount
256        );
257
258        Ok(record)
259    }
260
261    /// Returns fee statistics
262    pub fn get_fee_stats(&self) -> FeeStats {
263        self.stats.read().clone()
264    }
265
266    /// Returns distribution history for a period
267    pub fn get_distribution_history(&self, period: &str) -> Option<DistributionHistory> {
268        self.history.get(period).map(|h| h.clone())
269    }
270
271    /// Returns all distribution history
272    pub fn get_all_history(&self) -> Vec<(String, DistributionHistory)> {
273        self.history
274            .iter()
275            .map(|entry| (entry.key().clone(), entry.value().clone()))
276            .collect()
277    }
278
279    /// Starts a new history period
280    pub fn start_new_period(&self, period_id: String, start_time: Timestamp) {
281        // Close current period if exists
282        if let Some(current) = self.current_period.read().as_ref()
283            && let Some(mut history) = self.history.get_mut(current)
284        {
285            history.period_end = Timestamp::now();
286        }
287
288        // Create new period
289        let history = DistributionHistory::new(start_time, Timestamp::now());
290        self.history.insert(period_id.clone(), history);
291        *self.current_period.write() = Some(period_id.clone());
292
293        info!("Started new distribution period: {}", period_id);
294    }
295
296    /// Updates the fee distribution configuration
297    pub fn update_config(&self, config: FeeDistributionConfig) -> Result<()> {
298        config.validate()?;
299        *self.config.write() = config;
300        info!("Updated fee distribution configuration");
301        Ok(())
302    }
303
304    /// Returns the current configuration
305    pub fn get_config(&self) -> FeeDistributionConfig {
306        self.config.read().clone()
307    }
308
309    /// Adds a fee record to the current history period
310    fn add_to_history(&self, record: &FeeRecord) {
311        let current_period = self.current_period.read().clone();
312
313        if let Some(period_id) = current_period {
314            if let Some(mut history) = self.history.get_mut(&period_id) {
315                history.add_fee(record);
316            }
317        } else {
318            // Auto-create a default period if none exists
319            drop(current_period);
320            let period_id = format!("period_{}", Timestamp::now().as_millis());
321            self.start_new_period(period_id.clone(), Timestamp::now());
322
323            if let Some(mut history) = self.history.get_mut(&period_id) {
324                history.add_fee(record);
325            }
326        }
327    }
328
329    /// Returns a fee record by ID
330    pub fn get_fee_record(&self, fee_id: &str) -> Option<FeeRecord> {
331        self.fee_records.get(fee_id).map(|r| r.clone())
332    }
333
334    /// Returns all fee records
335    pub fn get_all_fee_records(&self) -> Vec<FeeRecord> {
336        self.fee_records
337            .iter()
338            .map(|entry| entry.value().clone())
339            .collect()
340    }
341}
342
343impl Default for FeeProcessor {
344    fn default() -> Self {
345        Self::new()
346    }
347}
348
349#[cfg(test)]
350mod tests {
351    use super::*;
352
353    #[test]
354    fn test_process_fee() {
355        let processor = FeeProcessor::new();
356        let asset_id = AssetId::tnzo();
357
358        let record = processor.process_fee(asset_id.clone(), 10000, FeeSource::Transaction).unwrap();
359
360        assert_eq!(record.amount, 10000);
361        assert_eq!(record.treasury_share, 4000); // 40%
362        assert_eq!(record.burn_share, 3000); // 30%
363        assert_eq!(record.staker_share, 3000); // 30%
364
365        let stats = processor.get_fee_stats();
366        assert_eq!(stats.fee_count, 1);
367        assert_eq!(stats.total_to_treasury, 4000);
368    }
369
370    #[test]
371    fn test_fee_stats() {
372        let processor = FeeProcessor::new();
373        let asset_id = AssetId::tnzo();
374
375        processor.process_fee(asset_id.clone(), 10000, FeeSource::Transaction).unwrap();
376        processor.process_fee(asset_id.clone(), 5000, FeeSource::Settlement).unwrap();
377
378        let stats = processor.get_fee_stats();
379        assert_eq!(stats.fee_count, 2);
380        assert_eq!(*stats.total_collected.get(&asset_id).unwrap(), 15000);
381    }
382
383    #[test]
384    fn test_history_periods() {
385        let processor = FeeProcessor::new();
386        let asset_id = AssetId::tnzo();
387
388        processor.start_new_period("period1".to_string(), Timestamp::now());
389        processor.process_fee(asset_id.clone(), 10000, FeeSource::Transaction).unwrap();
390
391        let history = processor.get_distribution_history("period1").unwrap();
392        assert_eq!(history.total_fees, 10000);
393    }
394}